mirror of https://github.com/go-gost/gost.git
136 lines
3.5 KiB
Python
136 lines
3.5 KiB
Python
import socket
|
|
import struct
|
|
import sys
|
|
|
|
|
|
UDP_ECHO_HOST = "udp-echo"
|
|
UDP_ECHO_PORT = 5679
|
|
|
|
|
|
def encode_socks5_addr(host, port):
|
|
"""Encode (host, port) as SOCKS5 ADDRESS + PORT.
|
|
|
|
Uses ATYP=3 (domain name) so Docker DNS resolves the host.
|
|
Returns (atyp, address_bytes, port_bytes).
|
|
"""
|
|
host_bytes = host.encode()
|
|
addr_bytes = struct.pack("!B", len(host_bytes)) + host_bytes
|
|
port_bytes = struct.pack("!H", port)
|
|
return 0x03, addr_bytes, port_bytes
|
|
|
|
|
|
def build_udp_frame(payload, host, port):
|
|
"""Build a SOCKS5 UDP relay frame over TCP.
|
|
|
|
gost HTTP UDP relay uses RSV=data-length and FRAG=0xff.
|
|
Frame: [RSV:2][FRAG:1][ATYP:1][DST.ADDR][DST.PORT:2][DATA]
|
|
"""
|
|
atyp, addr_bytes, port_bytes = encode_socks5_addr(host, port)
|
|
rsv = struct.pack("!H", len(payload))
|
|
frag = b"\xff"
|
|
atyp_byte = struct.pack("!B", atyp)
|
|
return rsv + frag + atyp_byte + addr_bytes + port_bytes + payload
|
|
|
|
|
|
def recvn(sock, n):
|
|
"""Receive exactly n bytes from socket."""
|
|
buf = b""
|
|
while len(buf) < n:
|
|
chunk = sock.recv(n - len(buf))
|
|
if not chunk:
|
|
raise ConnectionError("connection closed while reading")
|
|
buf += chunk
|
|
return buf
|
|
|
|
|
|
def read_socks5_frame(sock):
|
|
"""Read one SOCKS5 UDP relay frame from the TCP connection.
|
|
|
|
Returns (addr, port, data).
|
|
"""
|
|
# Read header: RSV(2) + FRAG(1) + ATYP(1)
|
|
header = recvn(sock, 4)
|
|
rsv = struct.unpack("!H", header[:2])[0]
|
|
frag = header[2]
|
|
atyp = header[3]
|
|
|
|
# Read address based on ATYP
|
|
if atyp == 1: # IPv4
|
|
addr_bytes = recvn(sock, 4)
|
|
addr = socket.inet_ntoa(addr_bytes)
|
|
elif atyp == 3: # Domain name
|
|
domain_len = recvn(sock, 1)
|
|
addr = recvn(sock, domain_len[0]).decode()
|
|
elif atyp == 4: # IPv6
|
|
addr_bytes = recvn(sock, 16)
|
|
addr = socket.inet_ntop(socket.AF_INET6, addr_bytes)
|
|
else:
|
|
raise ValueError(f"unknown ATYP: {atyp}")
|
|
|
|
port_bytes = recvn(sock, 2)
|
|
port = struct.unpack("!H", port_bytes)[0]
|
|
|
|
# Read data
|
|
if rsv > 0:
|
|
data = recvn(sock, rsv)
|
|
else:
|
|
# Standard SOCKS5 UDP: read remaining
|
|
data = b""
|
|
while True:
|
|
chunk = sock.recv(4096)
|
|
if not chunk:
|
|
break
|
|
data += chunk
|
|
|
|
return addr, port, data
|
|
|
|
|
|
def main():
|
|
host = sys.argv[1] if len(sys.argv) > 1 else "127.0.0.1"
|
|
port = int(sys.argv[2]) if len(sys.argv) > 2 else 8080
|
|
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
s.settimeout(10)
|
|
s.connect((host, port))
|
|
|
|
# Send CONNECT with X-Gost-Protocol: udp
|
|
req = (
|
|
b"CONNECT 0.0.0.0:0 HTTP/1.1\r\n"
|
|
b"Host: 0.0.0.0:0\r\n"
|
|
b"X-Gost-Protocol: udp\r\n"
|
|
b"Proxy-Connection: keep-alive\r\n"
|
|
b"\r\n"
|
|
)
|
|
s.sendall(req)
|
|
|
|
# Read 200 OK
|
|
resp = b""
|
|
while b"\r\n\r\n" not in resp:
|
|
chunk = s.recv(4096)
|
|
if not chunk:
|
|
break
|
|
resp += chunk
|
|
|
|
if b"200" not in resp:
|
|
print(f"FAIL: expected 200, got {resp.decode(errors='replace')}")
|
|
sys.exit(1)
|
|
|
|
# Build and send UDP relay frame
|
|
payload = b"hello-gost"
|
|
frame = build_udp_frame(payload, UDP_ECHO_HOST, UDP_ECHO_PORT)
|
|
s.sendall(frame)
|
|
|
|
# Read response frame
|
|
addr, rport, data = read_socks5_frame(s)
|
|
|
|
if b"hello-gost" in data:
|
|
print(f"PASS: received expected data from {addr}:{rport}")
|
|
sys.exit(0)
|
|
else:
|
|
print(f"FAIL: expected hello-gost in response, got {data!r}")
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|