Move mocked connection handlers to proxy/testing/mocks

proxy/freedom/freedom_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/v2ray/v2ray-core/proxy/common/connhandler"
 	_ "github.com/v2ray/v2ray-core/proxy/socks"
 	"github.com/v2ray/v2ray-core/proxy/socks/config/json"
+	proxymocks "github.com/v2ray/v2ray-core/proxy/testing/mocks"
 	"github.com/v2ray/v2ray-core/testing/mocks"
 	"github.com/v2ray/v2ray-core/testing/servers/tcp"
 	"github.com/v2ray/v2ray-core/testing/servers/udp"
@@ -38,9 +39,10 @@ func TestUDPSend(t *testing.T) {
 	udpServerAddr, err := udpServer.Start()
 	assert.Error(err).IsNil()
 
-	ich := &mocks.InboundConnectionHandler{
-		Data2Send:    []byte("Not Used"),
-		DataReturned: bytes.NewBuffer(make([]byte, 0, 1024)),
+	connOutput := bytes.NewBuffer(make([]byte, 0, 1024))
+	ich := &proxymocks.InboundConnectionHandler{
+		ConnInput:  bytes.NewReader([]byte("Not Used")),
+		ConnOutput: connOutput,
 	}
 
 	connhandler.RegisterInboundConnectionHandlerFactory("mock_ich", ich)
@@ -64,12 +66,11 @@ func TestUDPSend(t *testing.T) {
 	err = point.Start()
 	assert.Error(err).IsNil()
 
-	data2SendBuffer := alloc.NewBuffer()
-	data2SendBuffer.Clear()
+	data2SendBuffer := alloc.NewBuffer().Clear()
 	data2SendBuffer.Append([]byte(data2Send))
 	dest := v2net.NewUDPDestination(udpServerAddr)
 	ich.Communicate(v2net.NewPacket(dest, data2SendBuffer, false))
-	assert.Bytes(ich.DataReturned.Bytes()).Equals([]byte("Processed: Data to be sent to remote"))
+	assert.Bytes(connOutput.Bytes()).Equals([]byte("Processed: Data to be sent to remote"))
 }
 
 func TestSocksTcpConnect(t *testing.T) {

proxy/socks/socks_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/v2ray/v2ray-core/app/point"
 	"github.com/v2ray/v2ray-core/proxy/common/connhandler"
 	"github.com/v2ray/v2ray-core/proxy/socks/config/json"
+	proxymocks "github.com/v2ray/v2ray-core/proxy/testing/mocks"
 	"github.com/v2ray/v2ray-core/testing/mocks"
 	"github.com/v2ray/v2ray-core/testing/unit"
 )
@@ -19,9 +20,11 @@ func TestSocksTcpConnect(t *testing.T) {
 	assert := unit.Assert(t)
 	port := uint16(12385)
 
-	och := &mocks.OutboundConnectionHandler{
-		Data2Send:   bytes.NewBuffer(make([]byte, 0, 1024)),
-		Data2Return: []byte("The data to be returned to socks server."),
+	connInput := []byte("The data to be returned to socks server.")
+	connOutput := bytes.NewBuffer(make([]byte, 0, 1024))
+	och := &proxymocks.OutboundConnectionHandler{
+		ConnOutput: connOutput,
+		ConnInput:  bytes.NewReader(connInput),
 	}
 
 	connhandler.RegisterOutboundConnectionHandlerFactory("mock_och", och)
@@ -63,8 +66,8 @@ func TestSocksTcpConnect(t *testing.T) {
 	assert.Error(err).IsNil()
 	conn.Close()
 
-	assert.Bytes([]byte(data2Send)).Equals(och.Data2Send.Bytes())
-	assert.Bytes(dataReturned).Equals(och.Data2Return)
+	assert.Bytes([]byte(data2Send)).Equals(connOutput.Bytes())
+	assert.Bytes(dataReturned).Equals(connInput)
 	assert.String(targetServer).Equals(och.Destination.Address().String())
 }
 
@@ -72,9 +75,11 @@ func TestSocksTcpConnectWithUserPass(t *testing.T) {
 	assert := unit.Assert(t)
 	port := uint16(12386)
 
-	och := &mocks.OutboundConnectionHandler{
-		Data2Send:   bytes.NewBuffer(make([]byte, 0, 1024)),
-		Data2Return: []byte("The data to be returned to socks server."),
+	connInput := []byte("The data to be returned to socks server.")
+	connOutput := bytes.NewBuffer(make([]byte, 0, 1024))
+	och := &proxymocks.OutboundConnectionHandler{
+		ConnInput:  bytes.NewReader(connInput),
+		ConnOutput: connOutput,
 	}
 
 	connhandler.RegisterOutboundConnectionHandlerFactory("mock_och", och)
@@ -122,8 +127,8 @@ func TestSocksTcpConnectWithUserPass(t *testing.T) {
 	assert.Error(err).IsNil()
 	conn.Close()
 
-	assert.Bytes([]byte(data2Send)).Equals(och.Data2Send.Bytes())
-	assert.Bytes(dataReturned).Equals(och.Data2Return)
+	assert.Bytes([]byte(data2Send)).Equals(connOutput.Bytes())
+	assert.Bytes(dataReturned).Equals(connInput)
 	assert.String(targetServer).Equals(och.Destination.Address().String())
 }
 
@@ -131,9 +136,11 @@ func TestSocksTcpConnectWithWrongUserPass(t *testing.T) {
 	assert := unit.Assert(t)
 	port := uint16(12389)
 
-	och := &mocks.OutboundConnectionHandler{
-		Data2Send:   bytes.NewBuffer(make([]byte, 0, 1024)),
-		Data2Return: []byte("The data to be returned to socks server."),
+	connInput := []byte("The data to be returned to socks server.")
+	connOutput := bytes.NewBuffer(make([]byte, 0, 1024))
+	och := &proxymocks.OutboundConnectionHandler{
+		ConnInput:  bytes.NewReader(connInput),
+		ConnOutput: connOutput,
 	}
 
 	connhandler.RegisterOutboundConnectionHandlerFactory("mock_och", och)
@@ -176,9 +183,11 @@ func TestSocksTcpConnectWithWrongAuthMethod(t *testing.T) {
 	assert := unit.Assert(t)
 	port := uint16(38405)
 
-	och := &mocks.OutboundConnectionHandler{
-		Data2Send:   bytes.NewBuffer(make([]byte, 0, 1024)),
-		Data2Return: []byte("The data to be returned to socks server."),
+	connInput := []byte("The data to be returned to socks server.")
+	connOutput := bytes.NewBuffer(make([]byte, 0, 1024))
+	och := &proxymocks.OutboundConnectionHandler{
+		ConnInput:  bytes.NewReader(connInput),
+		ConnOutput: connOutput,
 	}
 
 	connhandler.RegisterOutboundConnectionHandlerFactory("mock_och", och)
@@ -221,9 +230,11 @@ func TestSocksUdpSend(t *testing.T) {
 	assert := unit.Assert(t)
 	port := uint16(12372)
 
-	och := &mocks.OutboundConnectionHandler{
-		Data2Send:   bytes.NewBuffer(make([]byte, 0, 1024)),
-		Data2Return: []byte("The data to be returned to socks server."),
+	connInput := []byte("The data to be returned to socks server.")
+	connOutput := bytes.NewBuffer(make([]byte, 0, 1024))
+	och := &proxymocks.OutboundConnectionHandler{
+		ConnInput:  bytes.NewReader(connInput),
+		ConnOutput: connOutput,
 	}
 
 	connhandler.RegisterOutboundConnectionHandlerFactory("mock_och", och)
@@ -270,7 +281,7 @@ func TestSocksUdpSend(t *testing.T) {
 	nBytes, err := conn.Read(response)
 
 	assert.Error(err).IsNil()
-	assert.Bytes(response[10:nBytes]).Equals(och.Data2Return)
-	assert.Bytes(data2Send).Equals(och.Data2Send.Bytes())
+	assert.Bytes(response[10:nBytes]).Equals(connInput)
+	assert.Bytes(data2Send).Equals(connOutput.Bytes())
 	assert.String(och.Destination.String()).Equals("udp:8.8.4.4:53")
 }

proxy/testing/mocks/inboundhandler.go

@@ -0,0 +1,55 @@
+package mocks
+
+import (
+	"io"
+	"sync"
+
+	"github.com/v2ray/v2ray-core/app"
+	v2net "github.com/v2ray/v2ray-core/common/net"
+	"github.com/v2ray/v2ray-core/proxy/common/connhandler"
+)
+
+type InboundConnectionHandler struct {
+	Port       uint16
+	Dispatcher app.PacketDispatcher
+	ConnInput  io.Reader
+	ConnOutput io.Writer
+}
+
+func (this *InboundConnectionHandler) Listen(port uint16) error {
+	this.Port = port
+	return nil
+}
+
+func (this *InboundConnectionHandler) Communicate(packet v2net.Packet) error {
+	ray := this.Dispatcher.DispatchToOutbound(packet)
+
+	input := ray.InboundInput()
+	output := ray.InboundOutput()
+
+	readFinish := &sync.Mutex{}
+	writeFinish := &sync.Mutex{}
+
+	readFinish.Lock()
+	writeFinish.Lock()
+
+	go func() {
+		v2net.ReaderToChan(input, this.ConnInput)
+		close(input)
+		readFinish.Unlock()
+	}()
+
+	go func() {
+		v2net.ChanToWriter(this.ConnOutput, output)
+		writeFinish.Unlock()
+	}()
+
+	readFinish.Lock()
+	writeFinish.Lock()
+	return nil
+}
+
+func (this *InboundConnectionHandler) Create(dispatcher app.PacketDispatcher, config interface{}) (connhandler.InboundConnectionHandler, error) {
+	this.Dispatcher = dispatcher
+	return this, nil
+}

proxy/testing/mocks/outboundhandler.go

@@ -0,0 +1,49 @@
+package mocks
+
+import (
+	"io"
+	"sync"
+
+	v2net "github.com/v2ray/v2ray-core/common/net"
+	"github.com/v2ray/v2ray-core/proxy/common/connhandler"
+	"github.com/v2ray/v2ray-core/transport/ray"
+)
+
+type OutboundConnectionHandler struct {
+	Destination v2net.Destination
+	ConnInput   io.Reader
+	ConnOutput  io.Writer
+}
+
+func (this *OutboundConnectionHandler) Dispatch(packet v2net.Packet, ray ray.OutboundRay) error {
+	input := ray.OutboundInput()
+	output := ray.OutboundOutput()
+
+	this.Destination = packet.Destination()
+	if packet.Chunk() != nil {
+		this.ConnOutput.Write(packet.Chunk().Value)
+		packet.Chunk().Release()
+	}
+
+	if packet.MoreChunks() {
+		writeFinish := &sync.Mutex{}
+
+		writeFinish.Lock()
+
+		go func() {
+			v2net.ChanToWriter(this.ConnOutput, input)
+			writeFinish.Unlock()
+		}()
+
+		writeFinish.Lock()
+	}
+
+	v2net.ReaderToChan(output, this.ConnInput)
+	close(output)
+
+	return nil
+}
+
+func (this *OutboundConnectionHandler) Create(config interface{}) (connhandler.OutboundConnectionHandler, error) {
+	return this, nil
+}

proxy/vmess/vmess_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/v2ray/v2ray-core/common/alloc"
 	v2net "github.com/v2ray/v2ray-core/common/net"
 	"github.com/v2ray/v2ray-core/proxy/common/connhandler"
+	proxymocks "github.com/v2ray/v2ray-core/proxy/testing/mocks"
 	"github.com/v2ray/v2ray-core/proxy/vmess/config"
 	"github.com/v2ray/v2ray-core/proxy/vmess/config/json"
 	"github.com/v2ray/v2ray-core/testing/mocks"
@@ -17,14 +18,15 @@ import (
 func TestVMessInAndOut(t *testing.T) {
 	assert := unit.Assert(t)
 
-	data2Send := "The data to be send to outbound server."
 	testAccount, err := config.NewID("ad937d9d-6e23-4a5a-ba23-bce5092a7c51")
 	assert.Error(err).IsNil()
 
 	portA := uint16(17392)
-	ich := &mocks.InboundConnectionHandler{
-		Data2Send:    []byte(data2Send),
-		DataReturned: bytes.NewBuffer(make([]byte, 0, 1024)),
+	ichConnInput := []byte("The data to be send to outbound server.")
+	ichConnOutput := bytes.NewBuffer(make([]byte, 0, 1024))
+	ich := &proxymocks.InboundConnectionHandler{
+		ConnInput:  bytes.NewReader(ichConnInput),
+		ConnOutput: ichConnOutput,
 	}
 
 	connhandler.RegisterInboundConnectionHandlerFactory("mock_ich", ich)
@@ -59,9 +61,11 @@ func TestVMessInAndOut(t *testing.T) {
 
 	portB := uint16(13829)
 
-	och := &mocks.OutboundConnectionHandler{
-		Data2Send:   bytes.NewBuffer(make([]byte, 0, 1024)),
-		Data2Return: []byte("The data to be returned to inbound server."),
+	ochConnInput := []byte("The data to be returned to inbound server.")
+	ochConnOutput := bytes.NewBuffer(make([]byte, 0, 1024))
+	och := &proxymocks.OutboundConnectionHandler{
+		ConnInput:  bytes.NewReader(ochConnInput),
+		ConnOutput: ochConnOutput,
 	}
 
 	connhandler.RegisterOutboundConnectionHandlerFactory("mock_och", och)
@@ -90,8 +94,8 @@ func TestVMessInAndOut(t *testing.T) {
 
 	dest := v2net.NewTCPDestination(v2net.IPAddress([]byte{1, 2, 3, 4}, 80))
 	ich.Communicate(v2net.NewPacket(dest, nil, true))
-	assert.Bytes([]byte(data2Send)).Equals(och.Data2Send.Bytes())
-	assert.Bytes(ich.DataReturned.Bytes()).Equals(och.Data2Return)
+	assert.Bytes(ichConnInput).Equals(ochConnOutput.Bytes())
+	assert.Bytes(ichConnOutput.Bytes()).Equals(ochConnInput)
 }
 
 func TestVMessInAndOutUDP(t *testing.T) {
@@ -102,9 +106,11 @@ func TestVMessInAndOutUDP(t *testing.T) {
 	assert.Error(err).IsNil()
 
 	portA := uint16(17394)
-	ich := &mocks.InboundConnectionHandler{
-		Data2Send:    []byte(data2Send),
-		DataReturned: bytes.NewBuffer(make([]byte, 0, 1024)),
+	ichConnInput := []byte("The data to be send to outbound server.")
+	ichConnOutput := bytes.NewBuffer(make([]byte, 0, 1024))
+	ich := &proxymocks.InboundConnectionHandler{
+		ConnInput:  bytes.NewReader(ichConnInput),
+		ConnOutput: ichConnOutput,
 	}
 
 	connhandler.RegisterInboundConnectionHandlerFactory("mock_ich", ich)
@@ -139,9 +145,11 @@ func TestVMessInAndOutUDP(t *testing.T) {
 
 	portB := uint16(13841)
 
-	och := &mocks.OutboundConnectionHandler{
-		Data2Send:   bytes.NewBuffer(make([]byte, 0, 1024)),
-		Data2Return: []byte("The data to be returned to inbound server."),
+	ochConnInput := []byte("The data to be returned to inbound server.")
+	ochConnOutput := bytes.NewBuffer(make([]byte, 0, 1024))
+	och := &proxymocks.OutboundConnectionHandler{
+		ConnInput:  bytes.NewReader(ochConnInput),
+		ConnOutput: ochConnOutput,
 	}
 
 	connhandler.RegisterOutboundConnectionHandlerFactory("mock_och", och)
@@ -174,6 +182,7 @@ func TestVMessInAndOutUDP(t *testing.T) {
 	data2SendBuffer.Append([]byte(data2Send))
 	dest := v2net.NewUDPDestination(v2net.IPAddress([]byte{1, 2, 3, 4}, 80))
 	ich.Communicate(v2net.NewPacket(dest, data2SendBuffer, false))
-	assert.Bytes([]byte(data2Send)).Equals(och.Data2Send.Bytes())
-	assert.Bytes(ich.DataReturned.Bytes()).Equals(och.Data2Return)
+
+	assert.Bytes(ichConnInput).Equals(ochConnOutput.Bytes())
+	assert.Bytes(ichConnOutput.Bytes()).Equals(ochConnInput)
 }

testing/mocks/inboundhandler.go

@@ -1,43 +0,0 @@
-package mocks
-
-import (
-	"bytes"
-
-	"github.com/v2ray/v2ray-core/app"
-	"github.com/v2ray/v2ray-core/common/alloc"
-	v2net "github.com/v2ray/v2ray-core/common/net"
-	"github.com/v2ray/v2ray-core/proxy/common/connhandler"
-)
-
-type InboundConnectionHandler struct {
-	Data2Send    []byte
-	DataReturned *bytes.Buffer
-	Port         uint16
-	Dispatcher   app.PacketDispatcher
-}
-
-func (handler *InboundConnectionHandler) Listen(port uint16) error {
-	handler.Port = port
-	return nil
-}
-
-func (handler *InboundConnectionHandler) Communicate(packet v2net.Packet) error {
-	ray := handler.Dispatcher.DispatchToOutbound(packet)
-
-	input := ray.InboundInput()
-	output := ray.InboundOutput()
-
-	buffer := alloc.NewBuffer()
-	buffer.Clear()
-	buffer.Append(handler.Data2Send)
-	input <- buffer
-	close(input)
-
-	v2net.ChanToWriter(handler.DataReturned, output)
-	return nil
-}
-
-func (handler *InboundConnectionHandler) Create(dispatcher app.PacketDispatcher, config interface{}) (connhandler.InboundConnectionHandler, error) {
-	handler.Dispatcher = dispatcher
-	return handler, nil
-}

testing/mocks/outboundhandler.go

@@ -1,48 +0,0 @@
-package mocks
-
-import (
-	"bytes"
-
-	"github.com/v2ray/v2ray-core/common/alloc"
-	v2net "github.com/v2ray/v2ray-core/common/net"
-	"github.com/v2ray/v2ray-core/proxy/common/connhandler"
-	"github.com/v2ray/v2ray-core/transport/ray"
-)
-
-type OutboundConnectionHandler struct {
-	Data2Send   *bytes.Buffer
-	Data2Return []byte
-	Destination v2net.Destination
-}
-
-func (handler *OutboundConnectionHandler) Dispatch(packet v2net.Packet, ray ray.OutboundRay) error {
-	input := ray.OutboundInput()
-	output := ray.OutboundOutput()
-
-	handler.Destination = packet.Destination()
-	if packet.Chunk() != nil {
-		handler.Data2Send.Write(packet.Chunk().Value)
-	}
-
-	go func() {
-		for {
-			data, open := <-input
-			if !open {
-				break
-			}
-			handler.Data2Send.Write(data.Value)
-			data.Release()
-		}
-		response := alloc.NewBuffer()
-		response.Clear()
-		response.Append(handler.Data2Return)
-		output <- response
-		close(output)
-	}()
-
-	return nil
-}
-
-func (handler *OutboundConnectionHandler) Create(config interface{}) (connhandler.OutboundConnectionHandler, error) {
-	return handler, nil
-}