瀏覽代碼

lib/connections: Allow "tcp4" and "tcp6" addresses

Audrius Butkevicius 9 年之前
父節點
當前提交
690837dbe5
共有 1 個文件被更改,包括 62 次插入55 次删除
  1. 62 55
      lib/connections/connections_tcp.go

+ 62 - 55
lib/connections/connections_tcp.go

@@ -17,75 +17,82 @@ import (
 )
 
 func init() {
-	dialers["tcp"] = tcpDialer
-	listeners["tcp"] = tcpListener
-}
-
-func tcpDialer(uri *url.URL, tlsCfg *tls.Config) (*tls.Conn, error) {
-	// Check that there is a port number in uri.Host, otherwise add one.
-	host, port, err := net.SplitHostPort(uri.Host)
-	if err != nil && strings.HasPrefix(err.Error(), "missing port") {
-		// addr is on the form "1.2.3.4"
-		uri.Host = net.JoinHostPort(uri.Host, "22000")
-	} else if err == nil && port == "" {
-		// addr is on the form "1.2.3.4:"
-		uri.Host = net.JoinHostPort(host, "22000")
-	}
-
-	// Don't try to resolve the address before dialing. The dialer may be a
-	// proxy, and we should let the proxy do the resolving in that case.
-	conn, err := dialer.Dial("tcp", uri.Host)
-	if err != nil {
-		l.Debugln(err)
-		return nil, err
-	}
-
-	tc := tls.Client(conn, tlsCfg)
-	err = tc.Handshake()
-	if err != nil {
-		tc.Close()
-		return nil, err
+	for _, network := range []string{"tcp", "tcp4", "tcp6"} {
+		dialers[network] = makeTcpDialer(network)
+		listeners[network] = makeTcpListener(network)
 	}
 
-	return tc, nil
 }
 
-func tcpListener(uri *url.URL, tlsCfg *tls.Config, conns chan<- model.IntermediateConnection) {
-	tcaddr, err := net.ResolveTCPAddr("tcp", uri.Host)
-	if err != nil {
-		l.Fatalln("listen (BEP/tcp):", err)
-		return
-	}
-	listener, err := net.ListenTCP("tcp", tcaddr)
-	if err != nil {
-		l.Fatalln("listen (BEP/tcp):", err)
-		return
-	}
-
-	for {
-		conn, err := listener.Accept()
-		if err != nil {
-			l.Warnln("Accepting connection (BEP/tcp):", err)
-			continue
+func makeTcpDialer(network string) DialerFactory {
+	return func(uri *url.URL, tlsCfg *tls.Config) (*tls.Conn, error) {
+		// Check that there is a port number in uri.Host, otherwise add one.
+		host, port, err := net.SplitHostPort(uri.Host)
+		if err != nil && strings.HasPrefix(err.Error(), "missing port") {
+			// addr is on the form "1.2.3.4"
+			uri.Host = net.JoinHostPort(uri.Host, "22000")
+		} else if err == nil && port == "" {
+			// addr is on the form "1.2.3.4:"
+			uri.Host = net.JoinHostPort(host, "22000")
 		}
 
-		l.Debugln("connect from", conn.RemoteAddr())
-
-		err = dialer.SetTCPOptions(conn.(*net.TCPConn))
+		// Don't try to resolve the address before dialing. The dialer may be a
+		// proxy, and we should let the proxy do the resolving in that case.
+		conn, err := dialer.Dial(network, uri.Host)
 		if err != nil {
-			l.Infoln(err)
+			l.Debugln(err)
+			return nil, err
 		}
 
-		tc := tls.Server(conn, tlsCfg)
+		tc := tls.Client(conn, tlsCfg)
 		err = tc.Handshake()
 		if err != nil {
-			l.Infoln("TLS handshake (BEP/tcp):", err)
 			tc.Close()
-			continue
+			return nil, err
 		}
 
-		conns <- model.IntermediateConnection{
-			tc, model.ConnectionTypeDirectAccept,
+		return tc, nil
+	}
+}
+
+func makeTcpListener(network string) ListenerFactory {
+	return func(uri *url.URL, tlsCfg *tls.Config, conns chan<- model.IntermediateConnection) {
+		tcaddr, err := net.ResolveTCPAddr(network, uri.Host)
+		if err != nil {
+			l.Fatalln("listen (BEP/tcp):", err)
+			return
+		}
+		listener, err := net.ListenTCP(network, tcaddr)
+		if err != nil {
+			l.Fatalln("listen (BEP/tcp):", err)
+			return
+		}
+
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				l.Warnln("Accepting connection (BEP/tcp):", err)
+				continue
+			}
+
+			l.Debugln("connect from", conn.RemoteAddr())
+
+			err = dialer.SetTCPOptions(conn.(*net.TCPConn))
+			if err != nil {
+				l.Infoln(err)
+			}
+
+			tc := tls.Server(conn, tlsCfg)
+			err = tc.Handshake()
+			if err != nil {
+				l.Infoln("TLS handshake (BEP/tcp):", err)
+				tc.Close()
+				continue
+			}
+
+			conns <- model.IntermediateConnection{
+				tc, model.ConnectionTypeDirectAccept,
+			}
 		}
 	}
 }