test case for UDP in freedom

proxy/freedom/freedom.go

@@ -40,9 +40,6 @@ func (vconn *FreedomConnection) Start(ray core.OutboundRay) error {
 	}
 
 	if !vconn.packet.MoreChunks() {
-		if ray != nil {
-			close(ray.OutboundOutput())
-		}
 		writeMutex.Unlock()
 	} else {
 		go dumpInput(conn, input, &writeMutex)

proxy/freedom/freedom_test.go

@@ -1,53 +1,88 @@
 package freedom
 
 import (
+	"bytes"
 	"io/ioutil"
 	"net"
-	"sync"
 	"testing"
 
 	"golang.org/x/net/proxy"
 
 	"github.com/v2ray/v2ray-core"
+	v2net "github.com/v2ray/v2ray-core/common/net"
 	_ "github.com/v2ray/v2ray-core/proxy/socks"
 	"github.com/v2ray/v2ray-core/testing/mocks"
+	"github.com/v2ray/v2ray-core/testing/servers/tcp"
+	"github.com/v2ray/v2ray-core/testing/servers/udp"
 	"github.com/v2ray/v2ray-core/testing/unit"
 )
 
-func TestSocksTcpConnect(t *testing.T) {
+func TestUDPSend(t *testing.T) {
 	assert := unit.Assert(t)
-	port := 48274
 
 	data2Send := "Data to be sent to remote"
-	data2Return := "Data to be returned to local"
 
-	var serverReady sync.Mutex
-	serverReady.Lock()
+	udpServer := &udp.Server{
+		Port: 0,
+		MsgProcessor: func(data []byte) []byte {
+			buffer := make([]byte, 0, 2048)
+			buffer = append(buffer, []byte("Processed: ")...)
+			buffer = append(buffer, data...)
+			return buffer
+		},
+	}
 
-	go func() {
-		listener, err := net.ListenTCP("tcp", &net.TCPAddr{
-			IP:   []byte{0, 0, 0, 0},
-			Port: port,
-			Zone: "",
-		})
-		assert.Error(err).IsNil()
+	udpServerAddr, err := udpServer.Start()
+	assert.Error(err).IsNil()
 
-		serverReady.Unlock()
-		conn, err := listener.Accept()
-		assert.Error(err).IsNil()
+	ich := &mocks.InboundConnectionHandler{
+		Data2Send:    []byte("Not Used"),
+		DataReturned: bytes.NewBuffer(make([]byte, 0, 1024)),
+	}
 
-		buffer := make([]byte, 1024)
-		nBytes, err := conn.Read(buffer)
-		assert.Error(err).IsNil()
+	core.RegisterInboundConnectionHandlerFactory("mock_ich", ich)
 
-		if string(buffer[:nBytes]) == data2Send {
-			_, err = conn.Write([]byte(data2Return))
-			assert.Error(err).IsNil()
-		}
+	pointPort := uint16(38724)
+	config := mocks.Config{
+		PortValue: pointPort,
+		InboundConfigValue: &mocks.ConnectionConfig{
+			ProtocolValue: "mock_ich",
+			ContentValue:  nil,
+		},
+		OutboundConfigValue: &mocks.ConnectionConfig{
+			ProtocolValue: "freedom",
+			ContentValue:  nil,
+		},
+	}
 
-		conn.Close()
-		listener.Close()
-	}()
+	point, err := core.NewPoint(&config)
+	assert.Error(err).IsNil()
+
+	err = point.Start()
+	assert.Error(err).IsNil()
+
+	dest := v2net.NewUDPDestination(udpServerAddr)
+	ich.Communicate(v2net.NewPacket(dest, []byte(data2Send), false))
+	assert.Bytes(ich.DataReturned.Bytes()).Equals([]byte("Processed: Data to be sent to remote"))
+}
+
+func TestSocksTcpConnect(t *testing.T) {
+	assert := unit.Assert(t)
+	port := uint16(38293)
+
+	data2Send := "Data to be sent to remote"
+
+	tcpServer := &tcp.Server{
+		Port: port,
+		MsgProcessor: func(data []byte) []byte {
+			buffer := make([]byte, 0, 2048)
+			buffer = append(buffer, []byte("Processed: ")...)
+			buffer = append(buffer, data...)
+			return buffer
+		},
+	}
+	_, err := tcpServer.Start()
+	assert.Error(err).IsNil()
 
 	pointPort := uint16(38724)
 	config := mocks.Config{
@@ -71,9 +106,7 @@ func TestSocksTcpConnect(t *testing.T) {
 	socks5Client, err := proxy.SOCKS5("tcp", "127.0.0.1:38724", nil, proxy.Direct)
 	assert.Error(err).IsNil()
 
-	serverReady.Lock()
-
-	targetServer := "127.0.0.1:48274"
+	targetServer := "127.0.0.1:38293"
 	conn, err := socks5Client.Dial("tcp", targetServer)
 	assert.Error(err).IsNil()
 
@@ -86,5 +119,5 @@ func TestSocksTcpConnect(t *testing.T) {
 	assert.Error(err).IsNil()
 	conn.Close()
 
-	assert.Bytes(dataReturned).Equals([]byte(data2Return))
+	assert.Bytes(dataReturned).Equals([]byte("Processed: Data to be sent to remote"))
 }

testing/servers/tcp/tcp.go

@@ -0,0 +1,51 @@
+package tcp
+
+import (
+	"fmt"
+	"io/ioutil"
+	"net"
+
+	v2net "github.com/v2ray/v2ray-core/common/net"
+)
+
+type Server struct {
+	Port         uint16
+	MsgProcessor func(msg []byte) []byte
+}
+
+func (server *Server) Start() (v2net.Address, error) {
+	listener, err := net.ListenTCP("tcp", &net.TCPAddr{
+		IP:   []byte{0, 0, 0, 0},
+		Port: int(server.Port),
+		Zone: "",
+	})
+	if err != nil {
+		return nil, err
+	}
+	go server.acceptConnections(listener)
+	localAddr := listener.Addr().(*net.TCPAddr)
+	return v2net.IPAddress(localAddr.IP, uint16(localAddr.Port)), nil
+}
+
+func (server *Server) acceptConnections(listener *net.TCPListener) {
+	for {
+		conn, err := listener.Accept()
+		if err != nil {
+			fmt.Printf("Failed accept TCP connection: %v", err)
+			continue
+		}
+
+		go server.handleConnection(conn)
+	}
+}
+
+func (server *Server) handleConnection(conn net.Conn) {
+	request, err := ioutil.ReadAll(conn)
+	if err != nil {
+		fmt.Printf("Failed to read request: %v", err)
+		return
+	}
+	response := server.MsgProcessor(request)
+	conn.Write(response)
+	conn.Close()
+}

testing/servers/udp/udp.go

@@ -0,0 +1,41 @@
+package udp
+
+import (
+	"fmt"
+	"net"
+
+	v2net "github.com/v2ray/v2ray-core/common/net"
+)
+
+type Server struct {
+	Port         uint16
+	MsgProcessor func(msg []byte) []byte
+}
+
+func (server *Server) Start() (v2net.Address, error) {
+	conn, err := net.ListenUDP("udp", &net.UDPAddr{
+		IP:   []byte{0, 0, 0, 0},
+		Port: int(server.Port),
+		Zone: "",
+	})
+	if err != nil {
+		return nil, err
+	}
+	go server.handleConnection(conn)
+	localAddr := conn.LocalAddr().(*net.UDPAddr)
+	return v2net.IPAddress(localAddr.IP, uint16(localAddr.Port)), nil
+}
+
+func (server *Server) handleConnection(conn *net.UDPConn) {
+	for {
+		buffer := make([]byte, 2*1024)
+		nBytes, addr, err := conn.ReadFromUDP(buffer)
+		if err != nil {
+			fmt.Printf("Failed to read from UDP: %v\n", err)
+			continue
+		}
+
+		response := server.MsgProcessor(buffer[:nBytes])
+		conn.WriteToUDP(response, addr)
+	}
+}