barely working http proxy

testing/scenarios/dokodemo_test.go

@@ -24,6 +24,7 @@ func TestDokodemoTCP(t *testing.T) {
 	}
 	_, err := tcpServer.Start()
 	assert.Error(err).IsNil()
+	defer tcpServer.Close()
 
 	assert.Error(InitializeServerSetOnce("test_2")).IsNil()
 

testing/scenarios/router_test.go

@@ -24,6 +24,7 @@ func TestRouter(t *testing.T) {
 	}
 	_, err := tcpServer.Start()
 	assert.Error(err).IsNil()
+	defer tcpServer.Close()
 
 	tcpServer2Accessed := false
 	tcpServer2 := &tcp.Server{
@@ -35,6 +36,7 @@ func TestRouter(t *testing.T) {
 	}
 	_, err = tcpServer2.Start()
 	assert.Error(err).IsNil()
+	defer tcpServer2.Close()
 
 	assert.Error(InitializeServerSetOnce("test_3")).IsNil()
 

testing/scenarios/socks_end_test.go

@@ -31,12 +31,13 @@ func TestTCPConnection(t *testing.T) {
 	}
 	_, err := tcpServer.Start()
 	assert.Error(err).IsNil()
+	defer tcpServer.Close()
 
 	assert.Error(InitializeServerSetOnce("test_1")).IsNil()
 
 	socksPort := v2net.Port(50000)
 
-	for i := 0; i < 100; i++ {
+	for i := 0; i < 10; i++ {
 		conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
 			IP:   []byte{127, 0, 0, 1},
 			Port: int(socksPort),
@@ -98,6 +99,7 @@ func TestTCPBind(t *testing.T) {
 	}
 	_, err := tcpServer.Start()
 	assert.Error(err).IsNil()
+	defer tcpServer.Close()
 
 	assert.Error(InitializeServerSetOnce("test_1")).IsNil()
 
@@ -146,6 +148,7 @@ func TestUDPAssociate(t *testing.T) {
 	}
 	_, err := udpServer.Start()
 	assert.Error(err).IsNil()
+	defer udpServer.Close()
 
 	assert.Error(InitializeServerSetOnce("test_1")).IsNil()
 

testing/servers/tcp/tcp.go

@@ -2,7 +2,6 @@ package tcp
 
 import (
 	"fmt"
-	"io/ioutil"
 	"net"
 
 	v2net "github.com/v2ray/v2ray-core/common/net"
@@ -11,6 +10,7 @@ import (
 type Server struct {
 	Port         v2net.Port
 	MsgProcessor func(msg []byte) []byte
+	accepting    bool
 }
 
 func (server *Server) Start() (v2net.Address, error) {
@@ -28,7 +28,9 @@ func (server *Server) Start() (v2net.Address, error) {
 }
 
 func (server *Server) acceptConnections(listener *net.TCPListener) {
-	for {
+	server.accepting = true
+	defer listener.Close()
+	for server.accepting {
 		conn, err := listener.Accept()
 		if err != nil {
 			fmt.Printf("Failed accept TCP connection: %v", err)
@@ -40,12 +42,17 @@ func (server *Server) acceptConnections(listener *net.TCPListener) {
 }
 
 func (server *Server) handleConnection(conn net.Conn) {
-	request, err := ioutil.ReadAll(conn)
-	if err != nil {
-		fmt.Printf("Failed to read request: %v", err)
-		return
+	for true {
+		request, err := v2net.ReadFrom(conn, nil)
+		if err != nil {
+			break
+		}
+		response := server.MsgProcessor(request.Value)
+		conn.Write(response)
 	}
-	response := server.MsgProcessor(request)
-	conn.Write(response)
 	conn.Close()
 }
+
+func (this *Server) Close() {
+	this.accepting = true
+}

testing/servers/udp/udp.go

@@ -10,6 +10,7 @@ import (
 type Server struct {
 	Port         v2net.Port
 	MsgProcessor func(msg []byte) []byte
+	accepting    bool
 }
 
 func (server *Server) Start() (v2net.Address, error) {
@@ -27,6 +28,8 @@ func (server *Server) Start() (v2net.Address, error) {
 }
 
 func (server *Server) handleConnection(conn *net.UDPConn) {
+	server.accepting = true
+	defer conn.Close()
 	for {
 		buffer := make([]byte, 2*1024)
 		nBytes, addr, err := conn.ReadFromUDP(buffer)
@@ -39,3 +42,7 @@ func (server *Server) handleConnection(conn *net.UDPConn) {
 		conn.WriteToUDP(response, addr)
 	}
 }
+
+func (server *Server) Close() {
+	server.accepting = false
+}