refine socks udp handling

proxy/socks/server_udp.go

@@ -5,7 +5,6 @@ import (
 	"v2ray.com/core/common/log"
 	v2net "v2ray.com/core/common/net"
 	"v2ray.com/core/proxy"
-	"v2ray.com/core/proxy/socks/protocol"
 	"v2ray.com/core/transport/internet/udp"
 )
 
@@ -24,39 +23,33 @@ func (v *Server) listenUDP() error {
 }
 
 func (v *Server) handleUDPPayload(payload *buf.Buffer, session *proxy.SessionInfo) {
+	defer payload.Release()
+
 	source := session.Source
 	log.Info("Socks: Client UDP connection from ", source)
-	request, err := protocol.ReadUDPRequest(payload.Bytes())
-	payload.Release()
+	request, data, err := DecodeUDPPacket(payload.Bytes())
 
 	if err != nil {
-		log.Error("Socks: Failed to parse UDP request: ", err)
-		return
-	}
-	if request.Data.Len() == 0 {
-		request.Data.Release()
-		return
-	}
-	if request.Fragment != 0 {
-		log.Warning("Socks: Dropping fragmented UDP packets.")
-		// TODO handle fragments
-		request.Data.Release()
+		log.Error("Socks|Server: Failed to parse UDP request: ", err)
 		return
 	}
 
-	log.Info("Socks: Send packet to ", request.Destination(), " with ", request.Data.Len(), " bytes")
+	if len(data) == 0 {
+		return
+	}
+
+	log.Info("Socks: Send packet to ", request.Destination(), " with ", len(data), " bytes")
 	log.Access(source, request.Destination, log.AccessAccepted, "")
-	v.udpServer.Dispatch(&proxy.SessionInfo{Source: source, Destination: request.Destination(), Inbound: v.meta}, request.Data, func(destination v2net.Destination, payload *buf.Buffer) {
-		response := &protocol.Socks5UDPRequest{
-			Fragment: 0,
-			Address:  request.Destination().Address,
-			Port:     request.Destination().Port,
-			Data:     payload,
-		}
+
+	dataBuf := buf.NewSmall()
+	dataBuf.Append(data)
+	v.udpServer.Dispatch(&proxy.SessionInfo{Source: source, Destination: request.Destination(), Inbound: v.meta}, dataBuf, func(destination v2net.Destination, payload *buf.Buffer) {
+		defer payload.Release()
+
 		log.Info("Socks: Writing back UDP response with ", payload.Len(), " bytes to ", destination)
 
-		udpMessage := buf.NewLocal(2048)
-		response.Write(udpMessage)
+		udpMessage := EncodeUDPPacket(request, payload.Bytes())
+		defer udpMessage.Release()
 
 		v.udpMutex.RLock()
 		if !v.accepting {
@@ -65,10 +58,9 @@ func (v *Server) handleUDPPayload(payload *buf.Buffer, session *proxy.SessionInf
 		}
 		nBytes, err := v.udpHub.WriteTo(udpMessage.Bytes(), destination)
 		v.udpMutex.RUnlock()
-		udpMessage.Release()
-		response.Data.Release()
+
 		if err != nil {
-			log.Error("Socks: failed to write UDP message (", nBytes, " bytes) to ", destination, ": ", err)
+			log.Warning("Socks: failed to write UDP message (", nBytes, " bytes) to ", destination, ": ", err)
 		}
 	})
 }