test: testing router

leoslf · Jun 9, 2024 · 8cbf208 · 8cbf208
1 parent b48388f
commit 8cbf208
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 10 deletions.
diff --git a/socks_router/models.py b/socks_router/models.py
@@ -106,7 +106,7 @@ class Socks5Method(IntEnum):
     USERNAME_PASSWORD = 0x02
     # IANA_ASSIGNED = 0x03
     # RESERVED_FOR_PRIVATE_METHODS = frozenset(range(0x80, 0xFF)) # 0x80..0xFE
-    NO_ACCEPTABLE_METHDOS = 0xFF
+    NO_ACCEPTABLE_METHODS = 0xFF
 
 
 class Socks5Command(IntEnum):
@@ -135,7 +135,7 @@ class Socks5MethodSelectionRequest:
     | 1 byte  | 1 byte       | [method_count] bytes |
     """
 
-    version: Literal[5]
+    version: int
     methods: list[int]
 
     def __bytes__(self) -> bytes:
@@ -152,7 +152,7 @@ class Socks5MethodSelectionResponse:
     | 1 byte  | 1 byte |
     """
 
-    version: Literal[5]
+    version: int
     method: Socks5Method
 
     def __bytes__(self) -> bytes:
@@ -163,7 +163,7 @@ def __bytes__(self) -> bytes:
 class Socks5Request:
     """SEE: https://datatracker.ietf.org/doc/html/rfc1928#section-4"""
 
-    version: Literal[5]
+    version: int
     command: Socks5Command
     reserved: Literal[0x00]
     address_type: Socks5AddressType
@@ -199,11 +199,13 @@ class Socks5Reply:
     | 1 byte  | 1 byte | 0x00 | 1 byte | 4-255 bytes | 2 bytes  |
     """
 
+    version: Literal[5]
     reply: Socks5ReplyType
+    reserved: Literal[0] = 0x00
     address: Address = IPv4("0.0.0.0", 0)
 
     def __bytes__(self) -> bytes:
-        return struct.pack("!BBB", SOCKS_VERSION, self.reply, 0x00) + bytes(self.address)
+        return struct.pack("!BBB", self.version, self.reply, self.reserved) + bytes(self.address)
 
 
 class Socks5State(StrEnum):

diff --git a/socks_router/router.py b/socks_router/router.py
@@ -316,7 +316,7 @@ def handshake(self):
 
         # none of the methods listed by the client are acceptable
         # notify the client
-        self.connection.sendall(bytes(Socks5MethodSelectionResponse(SOCKS_VERSION, Socks5Method.NO_ACCEPTABLE_METHDOS)))
+        self.connection.sendall(bytes(Socks5MethodSelectionResponse(SOCKS_VERSION, Socks5Method.NO_ACCEPTABLE_METHODS)))
         # the client MUST close the connection
         self.state = Socks5State.CLOSED
 
@@ -337,7 +337,7 @@ def handle_request(self):
                     self.logger.info(
                         f"Connected to destination {request.destination}, binding client socket: {self.remote.getsockname()}"
                     )
-                    self.connection.sendall(bytes(Socks5Reply(Socks5ReplyType.SUCCEEDED)))
+                    self.connection.sendall(bytes(Socks5Reply(SOCKS_VERSION, Socks5ReplyType.SUCCEEDED)))
                     self.state = Socks5State.ESTABLISHED
                     return
                 case _:

diff --git a/tests/test_router.py b/tests/test_router.py
@@ -15,9 +15,12 @@
 import socks
 
 from socks_router.models import (
+    SOCKS_VERSION,
+    Socks5Method,
+    Socks5MethodSelectionRequest,
+    Socks5MethodSelectionResponse,
     Socks5Command,
     Socks5AddressType,
-    Socks5MethodSelectionRequest,
     Socks5Request,
     IPv4,
     Host,
@@ -28,7 +31,6 @@
     ApplicationContext,
 )
 from socks_router.router import (
-    SOCKS_VERSION,
     free_port,
     read_request,
     read_address,
@@ -233,7 +235,17 @@ def it_should_relay_through_upstream(destination: Address):
     def when_client_attempt_to_use_socks4():
         def it_should_close_socket(destination):
             with daemonize() as proxy:
-                with pytest.raises(requests.exceptions.ConnectionError):
+                with pytest.raises(requests.exceptions.ConnectionError, match=r".*Connection reset by peer.*"):
                     requests.get(
                         f"http://{destination}/", proxies={type: f"socks4://{proxy.address}" for type in ["http", "https"]}
                     ).json()
+
+    def when_client_attempt_to_use_unacceptable_methods():
+        def it_should_reply_no_acceptable_methods():
+            with daemonize() as proxy:
+                # handshake
+                with connect_remote(proxy.address) as client:
+                    client.sendall(bytes(Socks5MethodSelectionRequest(SOCKS_VERSION, methods=[Socks5Method.USERNAME_PASSWORD])))
+                    assert bytes(
+                        Socks5MethodSelectionResponse(SOCKS_VERSION, method=Socks5Method.NO_ACCEPTABLE_METHODS)
+                    ) == client.recv(2)