apply udpHub in socks proxy

proxy/socks/socks.go

@@ -3,7 +3,6 @@ package socks
 import (
 	"errors"
 	"io"
-	"net"
 	"sync"
 	"time"
 
@@ -30,7 +29,7 @@ type SocksServer struct {
 	packetDispatcher dispatcher.PacketDispatcher
 	config           *Config
 	tcpListener      *hub.TCPHub
-	udpConn          *net.UDPConn
+	udpHub           *hub.UDPHub
 	udpAddress       v2net.Destination
 	udpServer        *hub.UDPServer
 	listeningPort    v2net.Port
@@ -55,10 +54,10 @@ func (this *SocksServer) Close() {
 		this.tcpListener = nil
 		this.tcpMutex.Unlock()
 	}
-	if this.udpConn != nil {
-		this.udpConn.Close()
+	if this.udpHub != nil {
 		this.udpMutex.Lock()
-		this.udpConn = nil
+		this.udpHub.Close()
+		this.udpHub = nil
 		this.udpMutex.Unlock()
 	}
 }

proxy/socks/udp.go

@@ -1,8 +1,6 @@
 package socks
 
 import (
-	"net"
-
 	"github.com/v2ray/v2ray-core/common/alloc"
 	"github.com/v2ray/v2ray-core/common/log"
 	v2net "github.com/v2ray/v2ray-core/common/net"
@@ -11,90 +9,64 @@ import (
 )
 
 func (this *SocksServer) ListenUDP(port v2net.Port) error {
-	addr := &net.UDPAddr{
-		IP:   net.IP{0, 0, 0, 0},
-		Port: int(port),
-		Zone: "",
-	}
-	conn, err := net.ListenUDP("udp", addr)
+	this.udpServer = hub.NewUDPServer(this.packetDispatcher)
+	udpHub, err := hub.ListenUDP(port, this.handleUDPPayload)
 	if err != nil {
-		log.Error("Socks: failed to listen UDP on port ", port, ": ", err)
+		log.Error("Socks: Failed to listen on udp port ", port)
 		return err
 	}
 	this.udpMutex.Lock()
 	this.udpAddress = v2net.UDPDestination(this.config.Address, port)
-	this.udpConn = conn
-	this.udpServer = hub.NewUDPServer(this.packetDispatcher)
+	this.udpHub = udpHub
 	this.udpMutex.Unlock()
-
-	go this.AcceptPackets()
 	return nil
 }
 
-func (this *SocksServer) AcceptPackets() error {
-	for this.accepting {
-		buffer := alloc.NewBuffer()
+func (this *SocksServer) handleUDPPayload(payload *alloc.Buffer, source v2net.Destination) {
+	log.Info("Socks: Client UDP connection from ", source)
+	request, err := protocol.ReadUDPRequest(payload.Value)
+	payload.Release()
+
+	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()
+		return
+	}
+
+	udpPacket := v2net.NewPacket(request.Destination(), request.Data, false)
+	log.Info("Socks: Send packet to ", udpPacket.Destination(), " with ", request.Data.Len(), " bytes")
+	this.udpServer.Dispatch(source, udpPacket, func(packet v2net.Packet) {
+		response := &protocol.Socks5UDPRequest{
+			Fragment: 0,
+			Address:  udpPacket.Destination().Address(),
+			Port:     udpPacket.Destination().Port(),
+			Data:     packet.Chunk(),
+		}
+		log.Info("Socks: Writing back UDP response with ", response.Data.Len(), " bytes to ", packet.Destination())
+
+		udpMessage := alloc.NewSmallBuffer().Clear()
+		response.Write(udpMessage)
+
 		this.udpMutex.RLock()
 		if !this.accepting {
 			this.udpMutex.RUnlock()
-			return nil
+			return
 		}
-		nBytes, addr, err := this.udpConn.ReadFromUDP(buffer.Value)
+		nBytes, err := this.udpHub.WriteTo(udpMessage.Value, packet.Destination())
 		this.udpMutex.RUnlock()
+		udpMessage.Release()
+		response.Data.Release()
 		if err != nil {
-			log.Error("Socks: failed to read UDP packets: ", err)
-			buffer.Release()
-			continue
-		}
-		log.Info("Socks: Client UDP connection from ", addr)
-		request, err := protocol.ReadUDPRequest(buffer.Value[:nBytes])
-		buffer.Release()
-		if err != nil {
-			log.Error("Socks: failed to parse UDP request: ", err)
-			continue
-		}
-		if request.Data == nil || request.Data.Len() == 0 {
-			continue
-		}
-		if request.Fragment != 0 {
-			log.Warning("Socks: Dropping fragmented UDP packets.")
-			// TODO handle fragments
-			request.Data.Release()
-			continue
+			log.Error("Socks: failed to write UDP message (", nBytes, " bytes) to ", packet.Destination(), ": ", err)
 		}
-
-		udpPacket := v2net.NewPacket(request.Destination(), request.Data, false)
-		log.Info("Socks: Send packet to ", udpPacket.Destination(), " with ", request.Data.Len(), " bytes")
-		this.udpServer.Dispatch(
-			v2net.UDPDestination(v2net.IPAddress(addr.IP), v2net.Port(addr.Port)), udpPacket,
-			func(packet v2net.Packet) {
-				response := &protocol.Socks5UDPRequest{
-					Fragment: 0,
-					Address:  udpPacket.Destination().Address(),
-					Port:     udpPacket.Destination().Port(),
-					Data:     packet.Chunk(),
-				}
-				log.Info("Socks: Writing back UDP response with ", response.Data.Len(), " bytes to ", packet.Destination())
-
-				udpMessage := alloc.NewSmallBuffer().Clear()
-				response.Write(udpMessage)
-
-				this.udpMutex.RLock()
-				if !this.accepting {
-					this.udpMutex.RUnlock()
-					return
-				}
-				nBytes, err := this.udpConn.WriteToUDP(udpMessage.Value, &net.UDPAddr{
-					IP:   packet.Destination().Address().IP(),
-					Port: int(packet.Destination().Port()),
-				})
-				this.udpMutex.RUnlock()
-				udpMessage.Release()
-				response.Data.Release()
-				if err != nil {
-					log.Error("Socks: failed to write UDP message (", nBytes, " bytes) to ", packet.Destination(), ": ", err)
-				}
-			})
-	}
-	return nil
+	})
 }