tcphub now takes tls config as construtor parameter

rxiaoch · May 2, 2016 · 7d43952 · 7d43952
1 parent f594f5b
commit 7d43952
Show file tree

Hide file tree

Showing 6 changed files with 24 additions and 13 deletions.
diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go
@@ -107,7 +107,7 @@ func (this *DokodemoDoor) handleUDPPackets(payload *alloc.Buffer, dest v2net.Des
 }
 
 func (this *DokodemoDoor) ListenTCP(port v2net.Port) error {
-	tcpListener, err := hub.ListenTCP(port, this.HandleTCPConnection)
+	tcpListener, err := hub.ListenTCP(port, this.HandleTCPConnection, nil)
 	if err != nil {
 		log.Error("Dokodemo: Failed to listen on port ", port, ": ", err)
 		return err

diff --git a/proxy/http/http.go b/proxy/http/http.go
@@ -60,7 +60,7 @@ func (this *HttpProxyServer) Listen(port v2net.Port) error {
 	}
 	this.listeningPort = port
 
-	tcpListener, err := hub.ListenTCP(port, this.handleConnection)
+	tcpListener, err := hub.ListenTCP(port, this.handleConnection, nil)
 	if err != nil {
 		log.Error("Http: Failed listen on port ", port, ": ", err)
 		return err

diff --git a/proxy/shadowsocks/shadowsocks.go b/proxy/shadowsocks/shadowsocks.go
@@ -65,7 +65,7 @@ func (this *Shadowsocks) Listen(port v2net.Port) error {
 		}
 	}
 
-	tcpHub, err := hub.ListenTCP(port, this.handleConnection)
+	tcpHub, err := hub.ListenTCP(port, this.handleConnection, nil)
 	if err != nil {
 		log.Error("Shadowsocks: Failed to listen TCP on port ", port, ": ", err)
 		return err

diff --git a/proxy/socks/server.go b/proxy/socks/server.go
@@ -77,7 +77,7 @@ func (this *Server) Listen(port v2net.Port) error {
 	}
 	this.listeningPort = port
 
-	listener, err := hub.ListenTCP(port, this.handleConnection)
+	listener, err := hub.ListenTCP(port, this.handleConnection, nil)
 	if err != nil {
 		log.Error("Socks: failed to listen on port ", port, ": ", err)
 		return err

diff --git a/proxy/vmess/inbound/inbound.go b/proxy/vmess/inbound/inbound.go
@@ -103,7 +103,7 @@ func (this *VMessInboundHandler) Listen(port v2net.Port) error {
 	}
 	this.listeningPort = port
 
-	tcpListener, err := hub.ListenTCP(port, this.HandleConnection)
+	tcpListener, err := hub.ListenTCP(port, this.HandleConnection, nil)
 	if err != nil {
 		log.Error("Unable to listen tcp port ", port, ": ", err)
 		return err

diff --git a/transport/hub/tcp.go b/transport/hub/tcp.go
@@ -1,6 +1,7 @@
 package hub
 
 import (
+	"crypto/tls"
 	"errors"
 	"net"
 
@@ -13,12 +14,12 @@ var (
 )
 
 type TCPHub struct {
-	listener     *net.TCPListener
+	listener     net.Listener
 	connCallback ConnectionHandler
 	accepting    bool
 }
 
-func ListenTCP(port v2net.Port, callback ConnectionHandler) (*TCPHub, error) {
+func ListenTCP(port v2net.Port, callback ConnectionHandler, tlsConfig *tls.Config) (*TCPHub, error) {
 	listener, err := net.ListenTCP("tcp", &net.TCPAddr{
 		IP:   []byte{0, 0, 0, 0},
 		Port: int(port),
@@ -27,12 +28,22 @@ func ListenTCP(port v2net.Port, callback ConnectionHandler) (*TCPHub, error) {
 	if err != nil {
 		return nil, err
 	}
-	tcpListener := &TCPHub{
-		listener:     listener,
-		connCallback: callback,
+	var hub *TCPHub
+	if tlsConfig != nil {
+		tlsListener := tls.NewListener(listener, tlsConfig)
+		hub = &TCPHub{
+			listener:     tlsListener,
+			connCallback: callback,
+		}
+	} else {
+		hub = &TCPHub{
+			listener:     listener,
+			connCallback: callback,
+		}
 	}
-	go tcpListener.start()
-	return tcpListener, nil
+
+	go hub.start()
+	return hub, nil
 }
 
 func (this *TCPHub) Close() {
@@ -44,7 +55,7 @@ func (this *TCPHub) Close() {
 func (this *TCPHub) start() {
 	this.accepting = true
 	for this.accepting {
-		conn, err := this.listener.AcceptTCP()
+		conn, err := this.listener.Accept()
 		if err != nil {
 			if this.accepting {
 				log.Warning("Listener: Failed to accept new TCP connection: ", err)