1 #!/usr/bin/env python3
2 # Copyright (c) 2015-present The Bitcoin Core developers
3 # Distributed under the MIT software license, see the accompanying
4 # file COPYING or http://www.opensource.org/licenses/mit-license.php.
5 """Dummy Socks5 server for testing."""
6 7 import select
8 import socket
9 import threading
10 import queue
11 import logging
12 13 from .netutil import (
14 format_addr_port,
15 format_sock,
16 set_ephemeral_port_range,
17 )
18 19 logger = logging.getLogger("TestFramework.socks5")
20 21 # Protocol constants
22 class Command:
23 CONNECT = 0x01
24 25 class AddressType:
26 IPV4 = 0x01
27 DOMAINNAME = 0x03
28 IPV6 = 0x04
29 30 # Utility functions
31 def recvall(s, n):
32 """Receive n bytes from a socket, or fail."""
33 rv = bytearray()
34 while n > 0:
35 d = s.recv(n)
36 if not d:
37 raise IOError('Unexpected end of stream')
38 rv.extend(d)
39 n -= len(d)
40 return rv
41 42 def sendall(s, data):
43 """Send all data to a socket, or fail."""
44 sent = 0
45 while sent < len(data):
46 _, wlist, _ = select.select([], [s], [])
47 if len(wlist) > 0:
48 n = s.send(data[sent:])
49 if n == 0:
50 raise IOError('send() on socket returned 0')
51 sent += n
52 53 def forward_sockets(a, b, wakeup_socket, serv):
54 """Forwards data between sockets a and b until EOF, error, or shutdown.
55 56 Monitors wakeup_socket for a shutdown signal and checks serv.is_running()
57 to exit gracefully when the server is stopping.
58 """
59 # Prefix messages with e.g.:
60 # forward_sockets(a{remote=127.0.0.1:36935, local=127.0.0.1:9050} <-> b{local=127.0.0.1:33424, remote=127.0.0.1:8333})
61 log_prefix = ("forward_sockets("
62 f"a{{remote={format_sock(a, local=False)}, local={format_sock(a, local=True)}}} <-> "
63 f"b{{local={format_sock(b, local=True)}, remote={format_sock(b, local=False)}}}): ")
64 65 # Mark as non-blocking so that we do not end up in a deadlock-like situation
66 # where we block and wait on data from `a` while there is data ready to be
67 # received on `b` and forwarded to `a`. And at the same time the application
68 # at `a` is not sending anything because it waits for the data from `b` to
69 # respond.
70 a.setblocking(False)
71 b.setblocking(False)
72 sockets = [a, b, wakeup_socket]
73 while True:
74 # Blocking select with timeout
75 rlist, _, xlist = select.select(sockets, [], sockets, 2)
76 if not serv.is_running():
77 logger.debug(f"{log_prefix}Exit due to shutdown")
78 return
79 if len(xlist) > 0:
80 raise IOError(f"{log_prefix}Exceptional condition on socket")
81 for s in rlist:
82 try:
83 data = s.recv(4096)
84 if data is None or len(data) == 0:
85 return
86 if s == a:
87 sendall(b, data)
88 elif s == b:
89 sendall(a, data)
90 except (BrokenPipeError, ConnectionResetError) as e:
91 logger.debug(f"{log_prefix}cannot send or receive data on socket {'a' if s == a else 'b'}: {str(e)}")
92 return
93 94 # Implementation classes
95 class Socks5Configuration():
96 """Proxy configuration."""
97 def __init__(self):
98 self.addr = None # Bind address (must be set)
99 self.af = socket.AF_INET # Bind address family
100 self.unauth = False # Support unauthenticated
101 self.auth = False # Support authentication
102 self.keep_alive = False # Do not automatically close connections
103 # This function is called whenever a new connection arrives to the proxy
104 # and it decides where the connection is redirected to. It is passed:
105 # - the address the client requested to connect to
106 # - the port the client requested to connect to
107 # It is supposed to return an object like:
108 # {
109 # "actual_to_addr": "127.0.0.1"
110 # "actual_to_port": 28276
111 # }
112 # or None.
113 # If it returns an object then the connection is redirected to actual_to_addr:actual_to_port.
114 # If it returns None, or destinations_factory itself is None then the connection is closed.
115 self.destinations_factory = None
116 117 class Socks5Command():
118 """Information about an incoming socks5 command."""
119 def __init__(self, cmd, atyp, addr, port, username, password):
120 self.cmd = cmd # Command (one of Command.*)
121 self.atyp = atyp # Address type (one of AddressType.*)
122 self.addr = addr # Address
123 self.port = port # Port to connect to
124 self.username = username
125 self.password = password
126 def __repr__(self):
127 return 'Socks5Command(%s,%s,%s,%s,%s,%s)' % (self.cmd, self.atyp, self.addr, self.port, self.username, self.password)
128 129 class Socks5Connection():
130 def __init__(self, serv, conn):
131 self.serv = serv
132 self.conn = conn
133 # Socket-pair used to wake up blocking forwarding select
134 # Note: a pipe could be used as well, but that does not work with select() on Windows
135 self.wakeup_socket_pair = socket.socketpair()
136 # Index of this handler (within the server)
137 self.handler_index = None
138 139 def handle(self):
140 """Handle socks5 request according to RFC1928."""
141 log_exception_prefix = "Socks5Connection.handle(): "
142 try:
143 log_exception_prefix = ("Socks5Connection.handle("
144 f"client={format_sock(self.conn, local=False)}, "
145 f"proxy={format_sock(self.conn, local=True)}): ")
146 147 # Verify socks version
148 ver = recvall(self.conn, 1)[0]
149 if ver != 0x05:
150 raise IOError('Invalid socks version %i' % ver)
151 # Choose authentication method
152 nmethods = recvall(self.conn, 1)[0]
153 methods = bytearray(recvall(self.conn, nmethods))
154 method = None
155 if 0x02 in methods and self.serv.conf.auth:
156 method = 0x02 # username/password
157 elif 0x00 in methods and self.serv.conf.unauth:
158 method = 0x00 # unauthenticated
159 if method is None:
160 raise IOError('No supported authentication method was offered')
161 # Send response
162 self.conn.sendall(bytearray([0x05, method]))
163 # Read authentication (optional)
164 username = None
165 password = None
166 if method == 0x02:
167 ver = recvall(self.conn, 1)[0]
168 if ver != 0x01:
169 raise IOError('Invalid auth packet version %i' % ver)
170 ulen = recvall(self.conn, 1)[0]
171 username = str(recvall(self.conn, ulen))
172 plen = recvall(self.conn, 1)[0]
173 password = str(recvall(self.conn, plen))
174 # Send authentication response
175 self.conn.sendall(bytearray([0x01, 0x00]))
176 177 # Read connect request
178 ver, cmd, _, atyp = recvall(self.conn, 4)
179 if ver != 0x05:
180 raise IOError('Invalid socks version %i in connect request' % ver)
181 if cmd != Command.CONNECT:
182 raise IOError('Unhandled command %i in connect request' % cmd)
183 184 if atyp == AddressType.IPV4:
185 addr = recvall(self.conn, 4)
186 elif atyp == AddressType.DOMAINNAME:
187 n = recvall(self.conn, 1)[0]
188 addr = recvall(self.conn, n)
189 elif atyp == AddressType.IPV6:
190 addr = recvall(self.conn, 16)
191 else:
192 raise IOError('Unknown address type %i' % atyp)
193 port_hi,port_lo = recvall(self.conn, 2)
194 port = (port_hi << 8) | port_lo
195 196 # Send dummy response
197 self.conn.sendall(bytearray([0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]))
198 199 cmdin = Socks5Command(cmd, atyp, addr, port, username, password)
200 self.serv.queue.put(cmdin)
201 logger.debug('Proxy: %s', cmdin)
202 203 requested_to_addr = addr.decode("utf-8")
204 requested_to = format_addr_port(requested_to_addr, port)
205 206 if self.serv.is_running():
207 if self.serv.conf.destinations_factory is not None:
208 dest = self.serv.conf.destinations_factory(requested_to_addr, port)
209 if dest is not None:
210 logger.debug(f"Serving connection to {requested_to}, will redirect it to "
211 f"{dest['actual_to_addr']}:{dest['actual_to_port']} instead")
212 with socket.create_connection((dest["actual_to_addr"], dest["actual_to_port"])) as conn_to:
213 forward_sockets(self.conn, conn_to, self.wakeup_socket_pair[1], self.serv)
214 conn_to.close()
215 else:
216 logger.debug(f"Can't serve the connection to {requested_to}: the destinations factory returned None")
217 else:
218 logger.debug(f"Can't serve the connection to {requested_to}: no destinations factory")
219 220 # Disconnect happens in the "finally" block below.
221 222 except (BrokenPipeError, ConnectionResetError) as e:
223 logger.debug(f"{log_exception_prefix}abnormal connection close: {str(e)}")
224 except Exception as e:
225 logger.exception(f"{log_exception_prefix}exception: {str(e)} (running {self.serv.is_running()})")
226 if self.serv.is_running():
227 self.serv.queue.put(e)
228 finally:
229 if not self.serv.keep_alive:
230 self.conn.close()
231 else:
232 logger.debug("Keeping client connection alive")
233 s0 = self.wakeup_socket_pair[0]
234 s1 = self.wakeup_socket_pair[1]
235 self.wakeup_socket_pair = None
236 try:
237 s0.close()
238 s1.close()
239 except OSError:
240 pass
241 self.serv.remove_handler(self.handler_index)
242 self.handler_index = None
243 244 def wakeup(self):
245 # Wake up the blocking forwarding select by writing to the wake-up socket
246 try:
247 socket_pair = self.wakeup_socket_pair
248 if socket_pair is not None:
249 socket_pair[0].send("CloseWakeup".encode())
250 logger.debug("Waking up forwarding thread")
251 except OSError as e:
252 logger.warning(f"Error waking up forwarding thread: {e}")
253 pass
254 255 256 # Wrapper for thread.join(), which may throw for daemon threads (in late stages of finalization).
257 # Return True if the thread is no longer active (join succeeded), False otherwise
258 # See PR #34863 for more details on using daemon threads.
259 def try_join_daemon_thread(thread, timeout=0) -> bool:
260 try:
261 thread.join(timeout=timeout)
262 return not thread.is_alive()
263 except Exception as e:
264 logger.debug(f"Exception in thread.join, {e}")
265 return True
266 267 class Socks5Server():
268 def __init__(self, conf):
269 self.conf = conf
270 self.s = socket.socket(conf.af)
271 self.s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
272 # When using dynamic port allocation (port=0), ensure we don't get a
273 # port that conflicts with the test framework's static port range.
274 if conf.addr[1] == 0:
275 set_ephemeral_port_range(self.s)
276 self.s.bind(conf.addr)
277 # When port=0, the OS assigns an available port. Update conf.addr
278 # to reflect the actual bound address so callers can use it.
279 self.conf.addr = self.s.getsockname()
280 self.s.listen(5)
281 # Set to False when stop is initiated
282 self._running = False
283 self._running_lock = threading.Lock()
284 self.thread = None
285 self.queue = queue.Queue() # report connections and exceptions to client
286 self.keep_alive = conf.keep_alive
287 # Store the background handlers, needed for clean shutdown
288 # Append-only array, completed handlers are set to None
289 self._handlers = []
290 self._handlers_lock = threading.Lock()
291 292 def is_running(self) -> bool:
293 with self._running_lock:
294 return self._running
295 296 def set_running(self, new_value: bool):
297 with self._running_lock:
298 self._running = new_value
299 300 def run(self):
301 while self.is_running():
302 (sockconn, _) = self.s.accept()
303 if self.is_running():
304 conn = Socks5Connection(self, sockconn)
305 # Use "daemon" threads, see PR #34863 for more discussion.
306 thread = threading.Thread(None, conn.handle, daemon=True)
307 with self._handlers_lock:
308 conn.handler_index = len(self._handlers)
309 self._handlers.append((thread, conn))
310 assert(conn.handler_index < len(self._handlers))
311 thread.start()
312 313 def remove_handler(self, handler_index):
314 with self._handlers_lock:
315 if handler_index < len(self._handlers):
316 if self._handlers[handler_index] is not None:
317 self._handlers[handler_index] = None
318 logger.debug(f"Handler {handler_index} removed")
319 320 def start(self):
321 assert not self.is_running()
322 self.set_running(True)
323 self.thread = threading.Thread(None, self.run, daemon=True)
324 self.thread.start()
325 326 def stop(self):
327 self.set_running(False)
328 # connect to self to end run loop
329 s = socket.socket(self.conf.af)
330 s.connect(self.conf.addr)
331 s.close()
332 self.thread.join()
333 # if there are active handlers, close them
334 with self._handlers_lock:
335 items = list(self._handlers)
336 for i, item in enumerate(items):
337 if item is None:
338 continue
339 thread, conn = item
340 # check if thread is still active
341 if not try_join_daemon_thread(thread, timeout=0):
342 conn.wakeup()
343 if try_join_daemon_thread(thread, timeout=2):
344 logger.debug(f"Stop(): Handler {i} thread joined")
345 else:
346 logger.warning(f"Stop(): Handler thread {i} didn't finish after force close")
347 348 def start_socks5_server(destinations_factory):
349 config = Socks5Configuration()
350 config.addr = ("127.0.0.1", 0) # Use port=0 to let the OS pick one. The actual port is later in server.conf.addr[1].
351 config.unauth = True
352 config.auth = True
353 config.destinations_factory = destinations_factory
354 355 server = Socks5Server(config)
356 server.start()
357 358 return server
359