Skip to content

Commit a38db82

Browse files
committed
new UDP over TCP protocol, merge master
1 parent d475076 commit a38db82

File tree

5 files changed

+1199
-215
lines changed

5 files changed

+1199
-215
lines changed

shadowsocks/asyncdns.py

Lines changed: 29 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
from __future__ import absolute_import, division, print_function, \
1919
with_statement
2020

21-
import time
2221
import os
2322
import socket
2423
import struct
@@ -256,7 +255,6 @@ def __init__(self):
256255
self._hostname_to_cb = {}
257256
self._cb_to_hostname = {}
258257
self._cache = lru_cache.LRUCache(timeout=300)
259-
self._last_time = time.time()
260258
self._sock = None
261259
self._servers = None
262260
self._parse_resolv()
@@ -304,16 +302,16 @@ def _parse_hosts(self):
304302
except IOError:
305303
self._hosts['localhost'] = '127.0.0.1'
306304

307-
def add_to_loop(self, loop, ref=False):
305+
def add_to_loop(self, loop):
308306
if self._loop:
309307
raise Exception('already add to loop')
310308
self._loop = loop
311309
# TODO when dns server is IPv6
312310
self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
313311
socket.SOL_UDP)
314312
self._sock.setblocking(False)
315-
loop.add(self._sock, eventloop.POLL_IN)
316-
loop.add_handler(self.handle_events, ref=ref)
313+
loop.add(self._sock, eventloop.POLL_IN, self)
314+
loop.add_periodic(self.handle_periodic)
317315

318316
def _call_callback(self, hostname, ip, error=None):
319317
callbacks = self._hostname_to_cb.get(hostname, [])
@@ -354,30 +352,27 @@ def _handle_data(self, data):
354352
self._call_callback(hostname, None)
355353
break
356354

357-
def handle_events(self, events):
358-
for sock, fd, event in events:
359-
if sock != self._sock:
360-
continue
361-
if event & eventloop.POLL_ERR:
362-
logging.error('dns socket err')
363-
self._loop.remove(self._sock)
364-
self._sock.close()
365-
# TODO when dns server is IPv6
366-
self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
367-
socket.SOL_UDP)
368-
self._sock.setblocking(False)
369-
self._loop.add(self._sock, eventloop.POLL_IN)
370-
else:
371-
data, addr = sock.recvfrom(1024)
372-
if addr[0] not in self._servers:
373-
logging.warn('received a packet other than our dns')
374-
break
375-
self._handle_data(data)
376-
break
377-
now = time.time()
378-
if now - self._last_time > CACHE_SWEEP_INTERVAL:
379-
self._cache.sweep()
380-
self._last_time = now
355+
def handle_event(self, sock, fd, event):
356+
if sock != self._sock:
357+
return
358+
if event & eventloop.POLL_ERR:
359+
logging.error('dns socket err')
360+
self._loop.remove(self._sock)
361+
self._sock.close()
362+
# TODO when dns server is IPv6
363+
self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
364+
socket.SOL_UDP)
365+
self._sock.setblocking(False)
366+
self._loop.add(self._sock, eventloop.POLL_IN, self)
367+
else:
368+
data, addr = sock.recvfrom(1024)
369+
if addr[0] not in self._servers:
370+
logging.warn('received a packet other than our dns')
371+
return
372+
self._handle_data(data)
373+
374+
def handle_periodic(self):
375+
self._cache.sweep()
381376

382377
def remove_callback(self, callback):
383378
hostname = self._cb_to_hostname.get(callback)
@@ -430,14 +425,17 @@ def resolve(self, hostname, callback):
430425

431426
def close(self):
432427
if self._sock:
428+
if self._loop:
429+
self._loop.remove_periodic(self.handle_periodic)
430+
self._loop.remove(self._sock)
433431
self._sock.close()
434432
self._sock = None
435433

436434

437435
def test():
438436
dns_resolver = DNSResolver()
439437
loop = eventloop.EventLoop()
440-
dns_resolver.add_to_loop(loop, ref=True)
438+
dns_resolver.add_to_loop(loop)
441439

442440
global counter
443441
counter = 0
@@ -451,8 +449,8 @@ def callback(result, error):
451449
print(result, error)
452450
counter += 1
453451
if counter == 9:
454-
loop.remove_handler(dns_resolver.handle_events)
455452
dns_resolver.close()
453+
loop.stop()
456454
a_callback = callback
457455
return a_callback
458456

shadowsocks/common.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -151,15 +151,24 @@ def pre_parse_header(data):
151151
data = data[rand_data_size + 2:]
152152
elif datatype == 0x81:
153153
data = data[1:]
154+
elif datatype == 0x82 :
155+
if len(data) <= 3:
156+
return None
157+
rand_data_size = struct.unpack('>H', data[1:3])[0]
158+
if rand_data_size + 3 >= len(data):
159+
logging.warn('header too short, maybe wrong password or '
160+
'encryption method')
161+
return None
162+
data = data[rand_data_size + 3:]
154163
return data
155164

156165
def parse_header(data):
157166
addrtype = ord(data[0])
158167
dest_addr = None
159168
dest_port = None
160169
header_length = 0
161-
connecttype = (addrtype & 8) and 1 or 0
162-
addrtype &= ~8
170+
connecttype = (addrtype & 0x10) and 1 or 0
171+
addrtype &= ~0x10
163172
if addrtype == ADDRTYPE_IPV4:
164173
if len(data) >= 7:
165174
dest_addr = socket.inet_ntoa(data[1:5])
@@ -173,7 +182,7 @@ def parse_header(data):
173182
if len(data) >= 2 + addrlen:
174183
dest_addr = data[2:2 + addrlen]
175184
dest_port = struct.unpack('>H', data[2 + addrlen:4 +
176-
addrlen])[0]
185+
addrlen])[0]
177186
header_length = 4 + addrlen
178187
else:
179188
logging.warn('header is too short')

shadowsocks/eventloop.py

Lines changed: 52 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
with_statement
2323

2424
import os
25+
import time
2526
import socket
2627
import select
2728
import errno
@@ -51,23 +52,8 @@
5152
POLL_NVAL: 'POLL_NVAL',
5253
}
5354

54-
55-
class EpollLoop(object):
56-
57-
def __init__(self):
58-
self._epoll = select.epoll()
59-
60-
def poll(self, timeout):
61-
return self._epoll.poll(timeout)
62-
63-
def add_fd(self, fd, mode):
64-
self._epoll.register(fd, mode)
65-
66-
def remove_fd(self, fd):
67-
self._epoll.unregister(fd)
68-
69-
def modify_fd(self, fd, mode):
70-
self._epoll.modify(fd, mode)
55+
# we check timeouts every TIMEOUT_PRECISION seconds
56+
TIMEOUT_PRECISION = 10
7157

7258

7359
class KqueueLoop(object):
@@ -100,17 +86,17 @@ def poll(self, timeout):
10086
results[fd] |= POLL_OUT
10187
return results.items()
10288

103-
def add_fd(self, fd, mode):
89+
def register(self, fd, mode):
10490
self._fds[fd] = mode
10591
self._control(fd, mode, select.KQ_EV_ADD)
10692

107-
def remove_fd(self, fd):
93+
def unregister(self, fd):
10894
self._control(fd, self._fds[fd], select.KQ_EV_DELETE)
10995
del self._fds[fd]
11096

111-
def modify_fd(self, fd, mode):
112-
self.remove_fd(fd)
113-
self.add_fd(fd, mode)
97+
def modify(self, fd, mode):
98+
self.unregister(fd)
99+
self.register(fd, mode)
114100

115101

116102
class SelectLoop(object):
@@ -129,32 +115,31 @@ def poll(self, timeout):
129115
results[fd] |= p[1]
130116
return results.items()
131117

132-
def add_fd(self, fd, mode):
118+
def register(self, fd, mode):
133119
if mode & POLL_IN:
134120
self._r_list.add(fd)
135121
if mode & POLL_OUT:
136122
self._w_list.add(fd)
137123
if mode & POLL_ERR:
138124
self._x_list.add(fd)
139125

140-
def remove_fd(self, fd):
126+
def unregister(self, fd):
141127
if fd in self._r_list:
142128
self._r_list.remove(fd)
143129
if fd in self._w_list:
144130
self._w_list.remove(fd)
145131
if fd in self._x_list:
146132
self._x_list.remove(fd)
147133

148-
def modify_fd(self, fd, mode):
149-
self.remove_fd(fd)
150-
self.add_fd(fd, mode)
134+
def modify(self, fd, mode):
135+
self.unregister(fd)
136+
self.register(fd, mode)
151137

152138

153139
class EventLoop(object):
154140
def __init__(self):
155-
self._iterating = False
156141
if hasattr(select, 'epoll'):
157-
self._impl = EpollLoop()
142+
self._impl = select.epoll()
158143
model = 'epoll'
159144
elif hasattr(select, 'kqueue'):
160145
self._impl = KqueueLoop()
@@ -165,72 +150,71 @@ def __init__(self):
165150
else:
166151
raise Exception('can not find any available functions in select '
167152
'package')
168-
self._fd_to_f = {}
169-
self._handlers = []
170-
self._ref_handlers = []
171-
self._handlers_to_remove = []
153+
self._fdmap = {} # (f, handler)
154+
self._last_time = time.time()
155+
self._periodic_callbacks = []
156+
self._stopping = False
172157
logging.debug('using event model: %s', model)
173158

174159
def poll(self, timeout=None):
175160
events = self._impl.poll(timeout)
176-
return [(self._fd_to_f[fd], fd, event) for fd, event in events]
161+
return [(self._fdmap[fd][0], fd, event) for fd, event in events]
177162

178-
def add(self, f, mode):
163+
def add(self, f, mode, handler):
179164
fd = f.fileno()
180-
self._fd_to_f[fd] = f
181-
self._impl.add_fd(fd, mode)
165+
self._fdmap[fd] = (f, handler)
166+
self._impl.register(fd, mode)
182167

183168
def remove(self, f):
184169
fd = f.fileno()
185-
del self._fd_to_f[fd]
186-
self._impl.remove_fd(fd)
170+
del self._fdmap[fd]
171+
self._impl.unregister(fd)
172+
173+
def add_periodic(self, callback):
174+
self._periodic_callbacks.append(callback)
175+
176+
def remove_periodic(self, callback):
177+
self._periodic_callbacks.remove(callback)
187178

188179
def modify(self, f, mode):
189180
fd = f.fileno()
190-
self._impl.modify_fd(fd, mode)
191-
192-
def add_handler(self, handler, ref=True):
193-
self._handlers.append(handler)
194-
if ref:
195-
# when all ref handlers are removed, loop stops
196-
self._ref_handlers.append(handler)
197-
198-
def remove_handler(self, handler):
199-
if handler in self._ref_handlers:
200-
self._ref_handlers.remove(handler)
201-
if self._iterating:
202-
self._handlers_to_remove.append(handler)
203-
else:
204-
self._handlers.remove(handler)
181+
self._impl.modify(fd, mode)
182+
183+
def stop(self):
184+
self._stopping = True
205185

206186
def run(self):
207187
events = []
208-
while self._ref_handlers:
188+
while not self._stopping:
189+
asap = False
209190
try:
210-
events = self.poll(1)
191+
events = self.poll(TIMEOUT_PRECISION)
211192
except (OSError, IOError) as e:
212193
if errno_from_exception(e) in (errno.EPIPE, errno.EINTR):
213194
# EPIPE: Happens when the client closes the connection
214195
# EINTR: Happens when received a signal
215196
# handles them as soon as possible
197+
asap = True
216198
logging.debug('poll:%s', e)
217199
else:
218200
logging.error('poll:%s', e)
219201
import traceback
220202
traceback.print_exc()
221203
continue
222-
self._iterating = True
223-
for handler in self._handlers:
224-
# TODO when there are a lot of handlers
225-
try:
226-
handler(events)
227-
except (OSError, IOError) as e:
228-
shell.print_exception(e)
229-
if self._handlers_to_remove:
230-
for handler in self._handlers_to_remove:
231-
self._handlers.remove(handler)
232-
self._handlers_to_remove = []
233-
self._iterating = False
204+
205+
for sock, fd, event in events:
206+
handler = self._fdmap.get(fd, None)
207+
if handler is not None:
208+
handler = handler[1]
209+
try:
210+
handler.handle_event(sock, fd, event)
211+
except (OSError, IOError) as e:
212+
shell.print_exception(e)
213+
now = time.time()
214+
if asap or now - self._last_time >= TIMEOUT_PRECISION:
215+
for callback in self._periodic_callbacks:
216+
callback()
217+
self._last_time = now
234218

235219

236220
# from tornado

0 commit comments

Comments
 (0)