From 5e2583ec8d9e3230a72a52ed9248a0561a592392 Mon Sep 17 00:00:00 2001
From: v2ray <admin@v2ray.com>
Date: Mon, 1 Feb 2016 21:34:07 +0100
Subject: [PATCH] fix udp for dokodemo and shadowsocks

---
 proxy/dokodemo/dokodemo.go       | 14 ++++++--------
 proxy/shadowsocks/shadowsocks.go | 16 ++++++++--------
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go
index d811a5ad..6ba7ca73 100644
--- a/proxy/dokodemo/dokodemo.go
+++ b/proxy/dokodemo/dokodemo.go
@@ -23,6 +23,7 @@ type DokodemoDoor struct {
 	packetDispatcher dispatcher.PacketDispatcher
 	tcpListener      *hub.TCPHub
 	udpHub           *hub.UDPHub
+	udpServer        *hub.UDPServer
 	listeningPort    v2net.Port
 }
 
@@ -89,26 +90,23 @@ func (this *DokodemoDoor) ListenUDP(port v2net.Port) error {
 	}
 	this.udpMutex.Lock()
 	this.udpHub = udpHub
+	this.udpServer = hub.NewUDPServer(this.packetDispatcher)
 	this.udpMutex.Unlock()
 	return nil
 }
 
 func (this *DokodemoDoor) handleUDPPackets(payload *alloc.Buffer, dest v2net.Destination) {
 	packet := v2net.NewPacket(v2net.UDPDestination(this.address, this.port), payload, false)
-	ray := this.packetDispatcher.DispatchToOutbound(packet)
-	close(ray.InboundInput())
-
-	for resp := range ray.InboundOutput() {
+	this.udpServer.Dispatch(dest, packet, func(packet v2net.Packet) {
+		defer packet.Chunk().Release()
 		this.udpMutex.RLock()
 		if !this.accepting {
 			this.udpMutex.RUnlock()
-			resp.Release()
 			return
 		}
-		this.udpHub.WriteTo(resp.Value, dest)
+		this.udpHub.WriteTo(packet.Chunk().Value, packet.Destination())
 		this.udpMutex.RUnlock()
-		resp.Release()
-	}
+	})
 }
 
 func (this *DokodemoDoor) ListenTCP(port v2net.Port) error {
diff --git a/proxy/shadowsocks/shadowsocks.go b/proxy/shadowsocks/shadowsocks.go
index de491097..c8a1ade2 100644
--- a/proxy/shadowsocks/shadowsocks.go
+++ b/proxy/shadowsocks/shadowsocks.go
@@ -26,6 +26,7 @@ type Shadowsocks struct {
 	accepting        bool
 	tcpHub           *hub.TCPHub
 	udpHub           *hub.UDPHub
+	udpServer        *hub.UDPServer
 }
 
 func NewShadowsocks(config *Config, packetDispatcher dispatcher.PacketDispatcher) *Shadowsocks {
@@ -77,6 +78,7 @@ func (this *Shadowsocks) Listen(port v2net.Port) error {
 			log.Error("Shadowsocks: Failed to listen UDP on port ", port, ": ", err)
 		}
 		this.udpHub = udpHub
+		this.udpServer = hub.NewUDPServer(this.packetDispatcher)
 	}
 
 	return nil
@@ -107,12 +109,12 @@ func (this *Shadowsocks) handlerUDPPayload(payload *alloc.Buffer, source v2net.D
 	log.Info("Shadowsocks: Tunnelling request to ", dest)
 
 	packet := v2net.NewPacket(dest, request.UDPPayload, false)
-	ray := this.packetDispatcher.DispatchToOutbound(packet)
-	close(ray.InboundInput())
-
-	for respChunk := range ray.InboundOutput() {
+	this.udpServer.Dispatch(source, packet, func(packet v2net.Packet) {
+		defer packet.Chunk().Release()
 
 		response := alloc.NewBuffer().Slice(0, this.config.Cipher.IVSize())
+		defer response.Release()
+
 		rand.Read(response.Value)
 		respIv := response.Value
 
@@ -135,8 +137,7 @@ func (this *Shadowsocks) handlerUDPPayload(payload *alloc.Buffer, source v2net.D
 		}
 
 		writer.Write(request.Port.Bytes())
-		writer.Write(respChunk.Value)
-		respChunk.Release()
+		writer.Write(packet.Chunk().Value)
 
 		if request.OTA {
 			respAuth := NewAuthenticator(HeaderKeyGenerator(key, respIv))
@@ -144,8 +145,7 @@ func (this *Shadowsocks) handlerUDPPayload(payload *alloc.Buffer, source v2net.D
 		}
 
 		this.udpHub.WriteTo(response.Value, source)
-		response.Release()
-	}
+	})
 }
 
 func (this *Shadowsocks) handleConnection(conn *hub.TCPConn) {