add socks5 version check

shadowsocks/tcprelay.py

@@ -37,6 +37,9 @@
 
 MSG_FASTOPEN = 0x20000000
 
+# SOCKS METHOD definition
+METHOD_NOAUTH = 0
+
 # SOCKS command definition
 CMD_CONNECT = 1
 CMD_BIND = 2
@@ -123,6 +126,12 @@ def isExceed(self):
             return self.sum_len >= self.max_speed
         return False
 
+class BadSocksHeader(Exception):
+    pass
+
+
+class NoAcceptableMethods(Exception):
+    pass
 class TCPRelayHandler(object):
     def __init__(self, server, fd_to_handlers, loop, local_sock, config,
                  dns_resolver, is_local):
@@ -828,7 +837,42 @@ def _get_read_size(self, sock, recv_buffer_size, up):
         if buffer_size > frame_size:
             buffer_size = int(buffer_size / frame_size) * frame_size
         return buffer_size
+    def _check_auth_method(self, data):
+        # VER, NMETHODS, and at least 1 METHODS
+        if len(data) < 3:
+            logging.warning('method selection header too short')
+            raise BadSocksHeader
+        socks_version = common.ord(data[0])
+        nmethods = common.ord(data[1])
+        if socks_version != 5:
+            logging.warning('unsupported SOCKS protocol version ' +
+                            str(socks_version))
+            raise BadSocksHeader
+        if nmethods < 1 or len(data) != nmethods + 2:
+            logging.warning('NMETHODS and number of METHODS mismatch')
+            raise BadSocksHeader
+        noauth_exist = False
+        for method in data[2:]:
+            if common.ord(method) == METHOD_NOAUTH:
+                noauth_exist = True
+                break
+        if not noauth_exist:
+            logging.warning('none of SOCKS METHOD\'s '
+                            'requested by client is supported')
+            raise NoAcceptableMethods
+    def _handle_stage_init(self, data):
+        try:
+            self._check_auth_method(data)
+        except BadSocksHeader:
+            self.destroy()
+            return
+        except NoAcceptableMethods:
+            self._write_to_sock(b'\x05\xff', self._local_sock)
+            self.destroy()
+            return
 
+        self._write_to_sock(b'\x05\00', self._local_sock)
+        self._stage = STAGE_ADDR
     def _on_local_read(self):
         # handle all local read events and dispatch them to methods for
         # each stage
@@ -917,9 +961,7 @@ def _on_local_read(self):
                     data = self._obfs.client_encode(data)
             self._write_to_sock(data, self._remote_sock)
         elif is_local and self._stage == STAGE_INIT:
-            # TODO check auth method
-            self._write_to_sock(b'\x05\00', self._local_sock)
-            self._stage = STAGE_ADDR
+            self._handle_stage_init(data)
         elif self._stage == STAGE_CONNECTING:
             self._handle_stage_connecting(data)
         elif (is_local and self._stage == STAGE_ADDR) or \