From 7d43952690a26a4a43faae18ea5b8dc583ddb784 Mon Sep 17 00:00:00 2001
From: v2ray <admin@v2ray.com>
Date: Tue, 3 May 2016 00:16:07 +0200
Subject: [PATCH] tcphub now takes tls config as construtor parameter

---
 proxy/dokodemo/dokodemo.go       |  2 +-
 proxy/http/http.go               |  2 +-
 proxy/shadowsocks/shadowsocks.go |  2 +-
 proxy/socks/server.go            |  2 +-
 proxy/vmess/inbound/inbound.go   |  2 +-
 transport/hub/tcp.go             | 27 +++++++++++++++++++--------
 6 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go
index 524eda03..57b33f3d 100644
--- 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
index 17754061..8e24d014 100644
--- 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
index 54b3ec38..410ebe9d 100644
--- 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
index 864efac1..a4998a71 100644
--- 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
index 073b6d31..bf848304 100644
--- 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
index eef64183..1487cffe 100644
--- 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)