fix: socks5 udp domain handling bug & add test case for it

This commit is contained in:
tobyxdd 2023-06-11 16:34:45 -07:00
parent fefabaf739
commit b27628607a
2 changed files with 31 additions and 8 deletions

View file

@ -254,7 +254,12 @@ func (s *Server) udpServer(udpConn *net.UDPConn, hyUDP client.HyUDPConn) error {
if err != nil {
continue
}
addr = addr[1:] // Remove the leading length byte
if atyp == socks5.ATYPDomain {
// socks5.ParseAddress adds a leading byte for domains,
// but socks5.NewDatagram will add it again as it expects a raw domain.
// So we must remove it here.
addr = addr[1:]
}
d := socks5.NewDatagram(atyp, addr, port, bs)
_, _ = udpConn.WriteToUDP(d.Bytes(), clientAddr)
}

View file

@ -6,33 +6,51 @@ ADDR = "127.0.0.1"
PORT = 11080
def test_tcp(size, count, it):
def test_tcp(size, count, it, domain=False):
for i in range(it):
s = socks.socksocket(socket.AF_INET, socket.SOCK_STREAM)
s.set_proxy(socks.SOCKS5, ADDR, PORT)
s.connect(("test_tcp", 12345))
if domain:
s.connect(("test.tcp.com", 12345))
else:
s.connect(("1.2.3.4", 12345))
for j in range(count):
payload = os.urandom(size)
s.send(payload)
rsp = s.recv(size)
assert rsp == payload
s.close()
def test_udp(size, count, it):
def test_udp(size, count, it, domain=False):
for i in range(it):
s = socks.socksocket(socket.AF_INET, socket.SOCK_DGRAM)
s.set_proxy(socks.SOCKS5, ADDR, PORT)
for j in range(count):
payload = os.urandom(size)
s.sendto(payload, ("test_udp", 12345))
if domain:
s.sendto(payload, ("test.udp.com", 12345))
else:
s.sendto(payload, ("1.2.3.4", 12345))
rsp, addr = s.recvfrom(size)
assert rsp == payload and addr == (b"test_udp", 12345)
assert rsp == payload
if domain:
assert addr == (b"test.udp.com", 12345)
else:
assert addr == ("1.2.3.4", 12345)
s.close()
if __name__ == "__main__":
test_tcp(1024, 1024, 10)
test_udp(1024, 1024, 10)
test_tcp(1024, 1024, 10, domain=False)
test_tcp(1024, 1024, 10, domain=True)
test_udp(1024, 1024, 10, domain=False)
test_udp(1024, 1024, 10, domain=True)