diff --git a/common/net/packet.go b/common/net/packet.go index 8f235bf2e..c8b4410da 100644 --- a/common/net/packet.go +++ b/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 { diff --git a/proxy/socks/udp.go b/proxy/socks/udp.go index 28f153ef6..4c088172e 100644 --- a/proxy/socks/udp.go +++ b/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 +}