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,63 +9,42 @@ 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()
-		this.udpMutex.RLock()
-		if !this.accepting {
-			this.udpMutex.RUnlock()
-			return nil
-		}
-		nBytes, addr, err := this.udpConn.ReadFromUDP(buffer.Value)
-		this.udpMutex.RUnlock()
-		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()
+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)
-			continue
+		log.Error("Socks: Failed to parse UDP request: ", err)
+		return
 	}
-		if request.Data == nil || request.Data.Len() == 0 {
-			continue
+	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()
-			continue
+		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(
-			v2net.UDPDestination(v2net.IPAddress(addr.IP), v2net.Port(addr.Port)), udpPacket,
-			func(packet v2net.Packet) {
+	this.udpServer.Dispatch(source, udpPacket, func(packet v2net.Packet) {
 		response := &protocol.Socks5UDPRequest{
 			Fragment: 0,
 			Address:  udpPacket.Destination().Address(),
@@ -84,10 +61,7 @@ func (this *SocksServer) AcceptPackets() error {
 			this.udpMutex.RUnlock()
 			return
 		}
-				nBytes, err := this.udpConn.WriteToUDP(udpMessage.Value, &net.UDPAddr{
-					IP:   packet.Destination().Address().IP(),
-					Port: int(packet.Destination().Port()),
-				})
+		nBytes, err := this.udpHub.WriteTo(udpMessage.Value, packet.Destination())
 		this.udpMutex.RUnlock()
 		udpMessage.Release()
 		response.Data.Release()
@@ -95,6 +69,4 @@ func (this *SocksServer) AcceptPackets() error {
 			log.Error("Socks: failed to write UDP message (", nBytes, " bytes) to ", packet.Destination(), ": ", err)
 		}
 	})
-	}
-	return nil
 }