TLS dialer

dnshostprovider.go

@@ -18,6 +18,22 @@ type DNSHostProvider struct {
 	lookupHost func(string) ([]string, error) // Override of net.LookupHost, for testing.
 }
 
+func addrsByHostname(server string) ([]string, error) {
+	res := []string{}
+	host, port, err := net.SplitHostPort(server)
+	if err != nil {
+		return nil, err
+	}
+	addrs, err := net.LookupHost(host)
+	if err != nil {
+		return nil, err
+	}
+	for _, addr := range addrs {
+		res = append(res, net.JoinHostPort(addr, port))
+	}
+	return res, nil
+}
+
 // Init is called first, with the servers specified in the connection
 // string. It uses DNS to look up addresses for each server, then
 // shuffles them all together.
@@ -32,16 +48,12 @@ func (hp *DNSHostProvider) Init(servers []string) error {
 
 	found := []string{}
 	for _, server := range servers {
-		host, port, err := net.SplitHostPort(server)
-		if err != nil {
-			return err
-		}
-		addrs, err := lookupHost(host)
+		addrs, err := addrsByHostname(server)
 		if err != nil {
 			return err
 		}
 		for _, addr := range addrs {
-			found = append(found, net.JoinHostPort(addr, port))
+			found = append(found, addr)
 		}
 	}
 

tls_dialer.go

@@ -0,0 +1,60 @@
+package zk
+
+import (
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
+	"net"
+	"os"
+	"time"
+)
+
+func CreateTLSConfig(rootCAFile, certFile, keyFile string) (*tls.Config, error) {
+	rootCABytes, err := os.ReadFile(rootCAFile)
+	if err != nil {
+		return nil, err
+	}
+
+	rootCA := x509.NewCertPool()
+	ok := rootCA.AppendCertsFromPEM(rootCABytes)
+	if !ok {
+		return nil, err
+	}
+
+	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
+	if err != nil {
+		return nil, err
+	}
+
+	return &tls.Config{
+		Certificates: []tls.Certificate{cert},
+		RootCAs:      rootCA,
+	}, nil
+}
+
+func GetTLSDialer(servers []string, dialer *net.Dialer, tlsConfig *tls.Config) (Dialer, error) {
+	if len(servers) == 0 {
+		return nil, errors.New("zk: server list must not be empty")
+	}
+	srvs := FormatServers(servers)
+
+	addrToHostname := map[string]string{}
+
+	for _, server := range srvs {
+		ips, err := addrsByHostname(server)
+		if err != nil {
+			return nil, err
+		}
+		for _, ip := range ips {
+			addrToHostname[ip] = server
+		}
+	}
+
+	return func(network, address string, _ time.Duration) (net.Conn, error) {
+		server, ok := addrToHostname[address]
+		if !ok {
+			server = address
+		}
+		return tls.DialWithDialer(dialer, network, server, tlsConfig)
+	}, nil
+}