Custom connections

v8/client/network.go

@@ -74,7 +74,13 @@ func (cl *Client) sendKDCUDP(realm string, b []byte) ([]byte, error) {
 	if err != nil {
 		return r, err
 	}
-	r, err = dialSendUDP(kdcs, b)
+
+	if cl.Config.LibDefaults.UDPConnGen == nil {
+		r, err = dialSendUDP(kdcs, b)
+	} else {
+		r, err = dialSendUDPCustom(kdcs, b, cl.Config.LibDefaults.UDPConnGen)
+	}
+
 	if err != nil {
 		return r, err
 	}
@@ -105,6 +111,38 @@ func dialSendUDP(kdcs map[int]string, b []byte) ([]byte, error) {
 	return nil, fmt.Errorf("error sending to a KDC: %s", strings.Join(errs, "; "))
 }
 
+// dialSendUDPCustom establishes a custom UDP connection to a KDC.
+func dialSendUDPCustom(kdcs map[int]string, b []byte, udpConnGen func(string) (*net.UDPConn, error)) ([]byte, error) {
+	var errs []string
+	var err error
+	var conn net.Conn
+
+	for i := 1; i <= len(kdcs); i++ {
+		if udpConnGen == nil {
+			conn, err = net.DialTimeout("udp", kdcs[i], 5*time.Second)
+		} else {
+			conn, err = udpConnGen(kdcs[i])
+		}
+
+		if err != nil {
+			errs = append(errs, fmt.Sprintf("error establishing connection to %s: %v", kdcs[i], err))
+			continue
+		}
+		if err := conn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
+			errs = append(errs, fmt.Sprintf("error setting deadline on connection to %s: %v", kdcs[i], err))
+			continue
+		}
+		// conn is guaranteed to be a UDPConn
+		rb, err := sendUDP(conn.(*net.UDPConn), b)
+		if err != nil {
+			errs = append(errs, fmt.Sprintf("error sneding to %s: %v", kdcs[i], err))
+			continue
+		}
+		return rb, nil
+	}
+	return nil, fmt.Errorf("error sending to a KDC: %s", strings.Join(errs, "; "))
+}
+
 // sendUDP sends bytes to connection over UDP.
 func sendUDP(conn *net.UDPConn, b []byte) ([]byte, error) {
 	var r []byte
@@ -132,14 +170,19 @@ func (cl *Client) sendKDCTCP(realm string, b []byte) ([]byte, error) {
 	if err != nil {
 		return r, err
 	}
-	r, err = dialSendTCP(kdcs, b)
+	if cl.Config.LibDefaults.TCPConnGen == nil {
+		r, err = dialSendTCP(kdcs, b)
+	} else {
+		r, err = dialSendTCPCustom(kdcs, b, cl.Config.LibDefaults.TCPConnGen)
+	}
+
 	if err != nil {
 		return r, err
 	}
 	return checkForKRBError(r)
 }
 
-// dialKDCTCP establishes a TCP connection to a KDC.
+// dialSendTCP establishes a TCP connection to a KDC.
 func dialSendTCP(kdcs map[int]string, b []byte) ([]byte, error) {
 	var errs []string
 	for i := 1; i <= len(kdcs); i++ {
@@ -163,6 +206,38 @@ func dialSendTCP(kdcs map[int]string, b []byte) ([]byte, error) {
 	return nil, fmt.Errorf("error sending to a KDC: %s", strings.Join(errs, "; "))
 }
 
+// dialSendTCPCustom establishes a custom TCP connection to a KDC.
+func dialSendTCPCustom(kdcs map[int]string, b []byte, tcpConnGen func(string) (*net.TCPConn, error)) ([]byte, error) {
+	var errs []string
+	var err error
+	var conn net.Conn
+
+	for i := 1; i <= len(kdcs); i++ {
+		if tcpConnGen == nil {
+			conn, err = net.DialTimeout("tcp", kdcs[i], 5*time.Second)
+		} else {
+			conn, err = tcpConnGen(kdcs[i])
+		}
+
+		if err != nil {
+			errs = append(errs, fmt.Sprintf("error establishing connection to %s: %v", kdcs[i], err))
+			continue
+		}
+		if err := conn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
+			errs = append(errs, fmt.Sprintf("error setting deadline on connection to %s: %v", kdcs[i], err))
+			continue
+		}
+		// conn is guaranteed to be a TCPConn
+		rb, err := sendTCP(conn.(*net.TCPConn), b)
+		if err != nil {
+			errs = append(errs, fmt.Sprintf("error sneding to %s: %v", kdcs[i], err))
+			continue
+		}
+		return rb, nil
+	}
+	return nil, fmt.Errorf("error sending to a KDC: %s", strings.Join(errs, "; "))
+}
+
 // sendTCP sends bytes to connection over TCP.
 func sendTCP(conn *net.TCPConn, b []byte) ([]byte, error) {
 	defer conn.Close()

v8/config/krb5conf.go

@@ -81,6 +81,8 @@ type LibDefaults struct {
 	TicketLifetime        time.Duration //default 1 day
 	UDPPreferenceLimit    int           // 1 means to always use tcp. MIT krb5 has a default value of 1465, and it prevents user setting more than 32700.
 	VerifyAPReqNofail     bool          //default false
+	TCPConnGen            func(string) (*net.TCPConn, error)
+	UDPConnGen            func(string) (*net.UDPConn, error)
 }
 
 // Create a new LibDefaults struct.