From 7e70fe290ac6ef3161eb53641942c1b865738f62 Mon Sep 17 00:00:00 2001
From: clowwindy <clowwindy42@gmail.com>
Date: Thu, 24 Apr 2014 15:07:18 +0800
Subject: [PATCH] fix local

---
 shadowsocks/local.py    | 67 +++++++++++++++++++++++++++++++++++++----
 shadowsocks/udprelay.py | 16 ++++++----
 2 files changed, 71 insertions(+), 12 deletions(-)

diff --git a/shadowsocks/local.py b/shadowsocks/local.py
index 39d725d..88e0749 100755
--- a/shadowsocks/local.py
+++ b/shadowsocks/local.py
@@ -47,6 +47,7 @@ import logging
 import getopt
 import encrypt
 import utils
+import udprelay
 
 
 def send_all(sock, data):
@@ -120,12 +121,63 @@ class Socks5Server(SocketServer.StreamRequestHandler):
         try:
             self.encryptor = encrypt.Encryptor(KEY, METHOD)
             sock = self.connection
-            sock.recv(262)
-            sock.send("\x05\x00")
+            data = sock.recv(262)
+            print '1.', len(data), data.encode('hex')
+            if not data:
+                sock.close()
+                return
+            if len(data) < 3:
+                return
+            method = ord(data[2])
+            if method == 2:
+                logging.warn('client tries to use username/password auth, prete'
+                             'nding the password is OK')
+                sock.send('\x05\x02')
+                try:
+                    ver_ulen = sock.recv(2)
+                    print ver_ulen.encode('hex')
+                    ulen = ord(ver_ulen[1])
+                    if ulen:
+                        username = sock.recv(ulen)
+                        print username
+                        assert(ulen == len(username))
+                    plen = ord(sock.recv(1))
+                    if plen:
+                        _password = sock.recv(plen)
+                        print _password
+                        assert(plen == len(_password))
+                    sock.send('\x01\x00')
+                except Exception as e:
+                    logging.error(e)
+                    return
+            elif method == 0:
+                sock.send("\x05\x00")
+            else:
+                logging.error('unsupported method %d' % method)
+                return
             data = self.rfile.read(4) or '\x00' * 4
+            print '2.', len(data), data.encode('hex')
             mode = ord(data[1])
-            if mode != 1:
-                logging.warn('mode != 1')
+            if mode == 1:
+                pass
+            elif mode == 3:
+                # UDP
+                logging.debug('UDP assc request')
+                if sock.family == socket.AF_INET6:
+                    header = '\x05\x00\x00\x04'
+                else:
+                    header = '\x05\x00\x00\x01'
+                addr, port = sock.getsockname()
+                addr_to_send = socket.inet_pton(sock.family, addr)
+                port_to_send = struct.pack('>H', port)
+                sock.send(header + addr_to_send + port_to_send)
+                while True:
+                    data = sock.recv(4096)
+                    if not data:
+                        break
+                return
+            else:
+                logging.warn('unknown mode %d' % mode)
                 return
             addrtype = ord(data[3])
             addr_to_send = data[3]
@@ -222,7 +274,7 @@ def main():
             elif key == '-m':
                 config['method'] = value
             elif key == '-b':
-                config['local'] = value
+                config['local_address'] = value
             elif key == '-6':
                 IPv6 = True
     except getopt.GetoptError:
@@ -234,7 +286,8 @@ def main():
     PORT = config['local_port']
     KEY = config['password']
     METHOD = config.get('method', None)
-    LOCAL = config.get('local', '127.0.0.1')
+    LOCAL = config.get('local_address', '127.0.0.1')
+    TIMEOUT = config.get('timeout', 600)
 
     if not KEY and not config_path:
         sys.exit('config not specified, please read '
@@ -247,6 +300,8 @@ def main():
     if IPv6:
         ThreadingTCPServer.address_family = socket.AF_INET6
     try:
+        udprelay.UDPRelay(LOCAL, int(PORT), SERVER, REMOTE_PORT, KEY, METHOD,
+                          int(TIMEOUT), True).start()
         server = ThreadingTCPServer((LOCAL, PORT), Socks5Server)
         logging.info("starting local at %s:%d" %
                      tuple(server.server_address[:2]))
diff --git a/shadowsocks/udprelay.py b/shadowsocks/udprelay.py
index 1e72742..92e89cf 100644
--- a/shadowsocks/udprelay.py
+++ b/shadowsocks/udprelay.py
@@ -188,10 +188,6 @@ class UDPRelay(object):
                 return
             self._eventloop.add(client, eventloop.POLL_IN)
 
-            # prevent from recv other sources
-            if self._is_local:
-                client.connect((server_addr, server_port))
-
         data = data[header_length:]
         if not data:
             return
@@ -199,7 +195,14 @@ class UDPRelay(object):
             data = encrypt.encrypt_all(self._password, self._method, 1, data)
             if not data:
                 return
-        client.sendto(data, (server_addr, server_port))
+        try:
+            client.sendto(data, (server_addr, server_port))
+        except IOError as e:
+            err = eventloop.errno_from_exception(e)
+            if err in (errno.EINPROGRESS, errno.EAGAIN):
+                pass
+            else:
+                logging.error(e)
 
     def _handle_client(self, sock):
         data, r_addr = sock.recvfrom(BUF_SIZE)
@@ -223,7 +226,7 @@ class UDPRelay(object):
             if header_result is None:
                 return
             # addrtype, dest_addr, dest_port, header_length = header_result
-            response = '\x00\x00\0x00' + data
+            response = '\x00\x00\x00' + data
         client_addr = self._client_fd_to_server_addr.get(sock.fileno(), None)
         if client_addr:
             self._server_socket.sendto(response, client_addr)
@@ -271,6 +274,7 @@ class UDPRelay(object):
         self._server_socket = server_socket
 
         t = threading.Thread(target=self._run)
+        t.setName('UDPThread')
         t.setDaemon(True)
         t.start()
         self._thread = t