dispatch udp response

common/net/packet.go

@@ -12,11 +12,11 @@ func NewTCPPacket(dest Destination) *TCPPacket {
 	}
 }
 
-func NewUDPPacket(dest Destination, data []byte, id uint16) *UDPPacket {
+func NewUDPPacket(dest Destination, data []byte, token uint16) *UDPPacket {
 	return &UDPPacket{
 		basePacket: basePacket{destination: dest},
 		data:       data,
-		id:         id,
+		token:      token,
 	}
 }
 
@@ -42,12 +42,12 @@ func (packet *TCPPacket) MoreChunks() bool {
 
 type UDPPacket struct {
 	basePacket
-	data []byte
-	id   uint16
+	data  []byte
+	token uint16
 }
 
-func (packet *UDPPacket) ID() uint16 {
-	return packet.id
+func (packet *UDPPacket) Token() uint16 {
+	return packet.token
 }
 
 func (packet *UDPPacket) Chunk() []byte {

proxy/socks/udp.go

@@ -56,8 +56,21 @@ func (m *portMap) removePorts(removedPorts <-chan interface{}) {
 	}
 }
 
+func (m *portMap) popPort(token uint16) *net.UDPAddr {
+	m.access.Lock()
+	defer m.access.Unlock()
+	addr, exists := m.data[token]
+	if !exists {
+		return nil
+	}
+	delete(m.data, token)
+	return addr
+}
+
 var (
 	ports = newPortMap()
+
+	udpConn *net.UDPConn
 )
 
 func (server *SocksServer) ListenUDP(port uint16) error {
@@ -73,6 +86,7 @@ func (server *SocksServer) ListenUDP(port uint16) error {
 	}
 
 	go server.AcceptPackets(conn)
+  udpConn = conn
 	return nil
 }
 
@@ -100,3 +114,14 @@ func (server *SocksServer) AcceptPackets(conn *net.UDPConn) error {
 		server.vPoint.DispatchToOutbound(udpPacket)
 	}
 }
+
+func (server *SocksServer) Dispatch(packet v2net.Packet) {
+	if udpPacket, ok := packet.(*v2net.UDPPacket); ok {
+		token := udpPacket.Token()
+		addr := ports.popPort(token)
+		if udpConn != nil {
+			udpConn.WriteToUDP(udpPacket.Chunk(), addr)
+		}
+	}
+	// We don't expect TCP Packets here
+}