4
4
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
5
5
"""Dummy Socks5 server for testing."""
6
6
7
+ import select
7
8
import socket
8
9
import threading
9
10
import queue
10
11
import logging
11
12
13
+ from .netutil import (
14
+ format_addr_port
15
+ )
16
+
12
17
logger = logging .getLogger ("TestFramework.socks5" )
13
18
14
19
# Protocol constants
@@ -32,6 +37,42 @@ def recvall(s, n):
32
37
n -= len (d )
33
38
return rv
34
39
40
+ def sendall (s , data ):
41
+ """Send all data to a socket, or fail."""
42
+ sent = 0
43
+ while sent < len (data ):
44
+ _ , wlist , _ = select .select ([], [s ], [])
45
+ if len (wlist ) > 0 :
46
+ n = s .send (data [sent :])
47
+ if n == 0 :
48
+ raise IOError ('send() on socket returned 0' )
49
+ sent += n
50
+
51
+ def forward_sockets (a , b ):
52
+ """Forward data received on socket a to socket b and vice versa, until EOF is received on one of the sockets."""
53
+ # Mark as non-blocking so that we do not end up in a deadlock-like situation
54
+ # where we block and wait on data from `a` while there is data ready to be
55
+ # received on `b` and forwarded to `a`. And at the same time the application
56
+ # at `a` is not sending anything because it waits for the data from `b` to
57
+ # respond.
58
+ a .setblocking (False )
59
+ b .setblocking (False )
60
+ sockets = [a , b ]
61
+ done = False
62
+ while not done :
63
+ rlist , _ , xlist = select .select (sockets , [], sockets )
64
+ if len (xlist ) > 0 :
65
+ raise IOError ('Exceptional condition on socket' )
66
+ for s in rlist :
67
+ data = s .recv (4096 )
68
+ if data is None or len (data ) == 0 :
69
+ done = True
70
+ break
71
+ if s == a :
72
+ sendall (b , data )
73
+ else :
74
+ sendall (a , data )
75
+
35
76
# Implementation classes
36
77
class Socks5Configuration ():
37
78
"""Proxy configuration."""
@@ -41,6 +82,19 @@ def __init__(self):
41
82
self .unauth = False # Support unauthenticated
42
83
self .auth = False # Support authentication
43
84
self .keep_alive = False # Do not automatically close connections
85
+ # This function is called whenever a new connection arrives to the proxy
86
+ # and it decides where the connection is redirected to. It is passed:
87
+ # - the address the client requested to connect to
88
+ # - the port the client requested to connect to
89
+ # It is supposed to return an object like:
90
+ # {
91
+ # "actual_to_addr": "127.0.0.1"
92
+ # "actual_to_port": 28276
93
+ # }
94
+ # or None.
95
+ # If it returns an object then the connection is redirected to actual_to_addr:actual_to_port.
96
+ # If it returns None, or destinations_factory itself is None then the connection is closed.
97
+ self .destinations_factory = None
44
98
45
99
class Socks5Command ():
46
100
"""Information about an incoming socks5 command."""
@@ -117,6 +171,22 @@ def handle(self):
117
171
cmdin = Socks5Command (cmd , atyp , addr , port , username , password )
118
172
self .serv .queue .put (cmdin )
119
173
logger .debug ('Proxy: %s' , cmdin )
174
+
175
+ requested_to_addr = addr .decode ("utf-8" )
176
+ requested_to = format_addr_port (requested_to_addr , port )
177
+
178
+ if self .serv .conf .destinations_factory is not None :
179
+ dest = self .serv .conf .destinations_factory (requested_to_addr , port )
180
+ if dest is not None :
181
+ logger .debug (f"Serving connection to { requested_to } , will redirect it to "
182
+ f"{ dest ['actual_to_addr' ]} :{ dest ['actual_to_port' ]} instead" )
183
+ with socket .create_connection ((dest ["actual_to_addr" ], dest ["actual_to_port" ])) as conn_to :
184
+ forward_sockets (self .conn , conn_to )
185
+ else :
186
+ logger .debug (f"Closing connection to { requested_to } : the destinations factory returned None" )
187
+ else :
188
+ logger .debug (f"Closing connection to { requested_to } : no destinations factory" )
189
+
120
190
# Fall through to disconnect
121
191
except Exception as e :
122
192
logger .exception ("socks5 request handling failed." )
0 commit comments