UDP support

proxy/socks/config/json/config.go

@@ -13,6 +13,7 @@ type SocksConfig struct {
 	AuthMethod string `json:"auth"`
 	Username   string `json:"user"`
 	Password   string `json:"pass"`
+	UDPEnabled bool   `json:"udp"`
 }
 
 func (config SocksConfig) IsNoAuth() bool {

proxy/vmess/config.go

@@ -26,6 +26,7 @@ func (u *VMessUser) ToUser() (user.User, error) {
 // VMessInboundConfig is
 type VMessInboundConfig struct {
 	AllowedClients []VMessUser `json:"clients"`
+	UDPEnabled     bool        `json:"udp"`
 }
 
 func loadInboundConfig(rawConfig []byte) (VMessInboundConfig, error) {

proxy/vmess/vmess_test.go

@@ -73,3 +73,67 @@ func TestVMessInAndOut(t *testing.T) {
 	assert.Bytes([]byte(data2Send)).Equals(och.Data2Send.Bytes())
 	assert.Bytes(ich.DataReturned.Bytes()).Equals(och.Data2Return)
 }
+
+func TestVMessInAndOutUDP(t *testing.T) {
+	assert := unit.Assert(t)
+
+	data2Send := "The data to be send to outbound server."
+
+	portA := uint16(17394)
+	ich := &mocks.InboundConnectionHandler{
+		Data2Send:    []byte(data2Send),
+		DataReturned: bytes.NewBuffer(make([]byte, 0, 1024)),
+	}
+
+	core.RegisterInboundConnectionHandlerFactory("mock_ich", ich)
+
+	configA := mocks.Config{
+		PortValue: portA,
+		InboundConfigValue: &mocks.ConnectionConfig{
+			ProtocolValue: "mock_ich",
+			ContentValue:  nil,
+		},
+		OutboundConfigValue: &mocks.ConnectionConfig{
+			ProtocolValue: "vmess",
+			ContentValue:  []byte("{\"vnext\":[{\"address\": \"127.0.0.1\", \"network\": \"udp\", \"port\": 13841, \"users\":[{\"id\": \"ad937d9d-6e23-4a5a-ba23-bce5092a7c51\"}]}]}"),
+		},
+	}
+
+	pointA, err := core.NewPoint(&configA)
+	assert.Error(err).IsNil()
+
+	err = pointA.Start()
+	assert.Error(err).IsNil()
+
+	portB := uint16(13841)
+
+	och := &mocks.OutboundConnectionHandler{
+		Data2Send:   bytes.NewBuffer(make([]byte, 0, 1024)),
+		Data2Return: []byte("The data to be returned to inbound server."),
+	}
+
+	core.RegisterOutboundConnectionHandlerFactory("mock_och", och)
+
+	configB := mocks.Config{
+		PortValue: portB,
+		InboundConfigValue: &mocks.ConnectionConfig{
+			ProtocolValue: "vmess",
+			ContentValue:  []byte("{\"clients\": [{\"id\": \"ad937d9d-6e23-4a5a-ba23-bce5092a7c51\"}], \"udp\": true}"),
+		},
+		OutboundConfigValue: &mocks.ConnectionConfig{
+			ProtocolValue: "mock_och",
+			ContentValue:  nil,
+		},
+	}
+
+	pointB, err := core.NewPoint(&configB)
+	assert.Error(err).IsNil()
+
+	err = pointB.Start()
+	assert.Error(err).IsNil()
+
+	dest := v2net.NewUDPDestination(v2net.IPAddress([]byte{1, 2, 3, 4}, 80))
+	ich.Communicate(v2net.NewPacket(dest, []byte(data2Send), false))
+	assert.Bytes([]byte(data2Send)).Equals(och.Data2Send.Bytes())
+	assert.Bytes(ich.DataReturned.Bytes()).Equals(och.Data2Return)
+}

proxy/vmess/vmessin.go

@@ -25,15 +25,17 @@ var (
 )
 
 type VMessInboundHandler struct {
-	vPoint    *core.Point
-	clients   user.UserSet
-	accepting bool
+	vPoint     *core.Point
+	clients    user.UserSet
+	accepting  bool
+	udpEnabled bool
 }
 
-func NewVMessInboundHandler(vp *core.Point, clients user.UserSet) *VMessInboundHandler {
+func NewVMessInboundHandler(vp *core.Point, clients user.UserSet, udpEnabled bool) *VMessInboundHandler {
 	return &VMessInboundHandler{
-		vPoint:  vp,
-		clients: clients,
+		vPoint:     vp,
+		clients:    clients,
+		udpEnabled: udpEnabled,
 	}
 }
 
@@ -45,6 +47,10 @@ func (handler *VMessInboundHandler) Listen(port uint16) error {
 	handler.accepting = true
 	go handler.AcceptConnections(listener)
 
+	if handler.udpEnabled {
+		handler.ListenUDP(port)
+	}
+
 	return nil
 }
 
@@ -143,7 +149,8 @@ func (factory *VMessInboundHandlerFactory) Create(vp *core.Point, rawConfig []by
 		}
 		allowedClients.AddUser(user)
 	}
-	return NewVMessInboundHandler(vp, allowedClients), nil
+
+	return NewVMessInboundHandler(vp, allowedClients, config.UDPEnabled), nil
 }
 
 func init() {

proxy/vmess/vmessin_udp.go

@@ -31,13 +31,13 @@ func (handler *VMessInboundHandler) ListenUDP(port uint16) error {
 	return nil
 }
 
-func (handler *VMessInboundHandler) AcceptPackets(conn *net.UDPConn) error {
+func (handler *VMessInboundHandler) AcceptPackets(conn *net.UDPConn) {
 	for {
-		buffer := make([]byte, 0, bufferSize)
+		buffer := make([]byte, bufferSize)
 		nBytes, addr, err := conn.ReadFromUDP(buffer)
 		if err != nil {
 			log.Error("VMessIn failed to read UDP packets: %v", err)
-			return err
+			continue
 		}
 
 		reader := bytes.NewReader(buffer[:nBytes])
@@ -46,23 +46,23 @@ func (handler *VMessInboundHandler) AcceptPackets(conn *net.UDPConn) error {
 		request, err := requestReader.Read(reader)
 		if err != nil {
 			log.Warning("VMessIn: Invalid request from (%s): %v", addr.String(), err)
-			return err
+			continue
 		}
 
 		cryptReader, err := v2io.NewAesDecryptReader(request.RequestKey[:], request.RequestIV[:], reader)
 		if err != nil {
 			log.Error("VMessIn: Failed to create decrypt reader: %v", err)
-			return err
+			continue
 		}
 
 		data := make([]byte, bufferSize)
 		nBytes, err = cryptReader.Read(data)
 		if err != nil {
 			log.Warning("VMessIn: Unable to decrypt data: %v", err)
-			return err
+			continue
 		}
 
-		packet := v2net.NewPacket(request.Destination(), data, false)
+		packet := v2net.NewPacket(request.Destination(), data[:nBytes], false)
 		go handler.handlePacket(conn, request, packet, addr)
 	}
 }

proxy/vmess/vmessout.go

@@ -35,9 +35,10 @@ type VMessOutboundHandler struct {
 
 func NewVMessOutboundHandler(vp *core.Point, vNextList, vNextListUDP []VNextServer, firstPacket v2net.Packet) *VMessOutboundHandler {
 	return &VMessOutboundHandler{
-		vPoint:    vp,
-		packet:    firstPacket,
-		vNextList: vNextList,
+		vPoint:       vp,
+		packet:       firstPacket,
+		vNextList:    vNextList,
+		vNextListUDP: vNextListUDP,
 	}
 }
 
@@ -65,7 +66,11 @@ func pickVNext(serverList []VNextServer) (v2net.Destination, user.User) {
 }
 
 func (handler *VMessOutboundHandler) Start(ray core.OutboundRay) error {
-	vNextAddress, vNextUser := pickVNext(handler.vNextList)
+	vNextList := handler.vNextList
+	if handler.packet.Destination().IsUDP() {
+		vNextList = handler.vNextListUDP
+	}
+	vNextAddress, vNextUser := pickVNext(vNextList)
 
 	command := protocol.CmdTCP
 	if handler.packet.Destination().IsUDP() {
@@ -86,9 +91,9 @@ func (handler *VMessOutboundHandler) Start(ray core.OutboundRay) error {
 }
 
 func startCommunicate(request *protocol.VMessRequest, dest v2net.Destination, ray core.OutboundRay, firstPacket v2net.Packet) error {
-	conn, err := net.DialTCP(dest.Network(), nil, &net.TCPAddr{dest.Address().IP(), int(dest.Address().Port()), ""})
+	conn, err := net.Dial(dest.Network(), dest.Address().String())
 	if err != nil {
-		log.Error("Failed to open tcp (%s): %v", dest.String(), err)
+		log.Error("Failed to open %s: %v", dest.String(), err)
 		if ray != nil {
 			close(ray.OutboundOutput())
 		}
@@ -105,15 +110,17 @@ func startCommunicate(request *protocol.VMessRequest, dest v2net.Destination, ra
 	responseFinish.Lock()
 
 	go handleRequest(conn, request, firstPacket, input, &requestFinish)
-	go handleResponse(conn, request, output, &responseFinish)
+	go handleResponse(conn, request, output, &responseFinish, dest.IsUDP())
 
 	requestFinish.Lock()
-	conn.CloseWrite()
+	if tcpConn, ok := conn.(*net.TCPConn); ok {
+		tcpConn.CloseWrite()
+	}
 	responseFinish.Lock()
 	return nil
 }
 
-func handleRequest(conn *net.TCPConn, request *protocol.VMessRequest, firstPacket v2net.Packet, input <-chan []byte, finish *sync.Mutex) {
+func handleRequest(conn net.Conn, request *protocol.VMessRequest, firstPacket v2net.Packet, input <-chan []byte, finish *sync.Mutex) {
 	defer finish.Unlock()
 	encryptRequestWriter, err := v2io.NewAesEncryptWriter(request.RequestKey[:], request.RequestIV[:], conn)
 	if err != nil {
@@ -153,7 +160,7 @@ func handleRequest(conn *net.TCPConn, request *protocol.VMessRequest, firstPacke
 	return
 }
 
-func handleResponse(conn *net.TCPConn, request *protocol.VMessRequest, output chan<- []byte, finish *sync.Mutex) {
+func handleResponse(conn net.Conn, request *protocol.VMessRequest, output chan<- []byte, finish *sync.Mutex, isUDP bool) {
 	defer finish.Unlock()
 	defer close(output)
 	responseKey := md5.Sum(request.RequestKey[:])
@@ -165,18 +172,24 @@ func handleResponse(conn *net.TCPConn, request *protocol.VMessRequest, output ch
 		return
 	}
 
-	response := protocol.VMessResponse{}
-	_, err = decryptResponseReader.Read(response[:])
+	buffer := make([]byte, 2*1024)
+
+	nBytes, err := decryptResponseReader.Read(buffer)
 	if err != nil {
 		//log.Error("VMessOut: Failed to read VMess response (%d bytes): %v", nBytes, err)
 		return
 	}
-	if !bytes.Equal(response[:], request.ResponseHeader[:]) {
+	if !bytes.Equal(buffer[:4], request.ResponseHeader[:]) {
 		log.Warning("VMessOut: unexepcted response header. The connection is probably hijacked.")
 		return
 	}
 
-	v2net.ReaderToChan(output, decryptResponseReader)
+	output <- buffer[4:nBytes]
+
+	if !isUDP {
+		v2net.ReaderToChan(output, decryptResponseReader)
+	}
+
 	return
 }
 

release/config/in_socks.json

@@ -1,3 +1,4 @@
 {
-  "auth": "noauth"
+  "auth": "noauth",
+  "udp": false
 }

release/config/in_vmess.json

@@ -1,5 +1,6 @@
 {
   "clients": [
     {"id": "ad937d9d-6e23-4a5a-ba23-bce5092a7c51"}
-  ]
+  ],
+  "udp": false
 }

testing/mocks/outboundhandler.go

@@ -25,7 +25,9 @@ func (handler *OutboundConnectionHandler) Start(ray core.OutboundRay) error {
 			}
 			handler.Data2Send.Write(data)
 		}
-		output <- handler.Data2Return
+		dataCopy := make([]byte, len(handler.Data2Return))
+		copy(dataCopy, handler.Data2Return)
+		output <- dataCopy
 		close(output)
 	}()
 
@@ -38,5 +40,9 @@ func (handler *OutboundConnectionHandler) Initialize(config []byte) error {
 
 func (handler *OutboundConnectionHandler) Create(point *core.Point, packet v2net.Packet) (core.OutboundConnectionHandler, error) {
 	handler.Destination = packet.Destination()
+	if packet.Chunk() != nil {
+		handler.Data2Send.Write(packet.Chunk())
+	}
+
 	return handler, nil
 }