]> jfr.im git - yt-dlp.git/blobdiff - test/test_socks.py
[test:download] Raise on network errors (#10283)
[yt-dlp.git] / test / test_socks.py
index d8ac88dad5846751756988d5a1bd8d7f3dc36265..68af19d0ca449d726da43447f8619ce53a82b39a 100644 (file)
@@ -25,7 +25,7 @@
     ThreadingTCPServer,
 )
 
-from test.helper import http_server_port
+from test.helper import http_server_port, verify_address_availability
 from yt_dlp.networking import Request
 from yt_dlp.networking.exceptions import ProxyError, TransportError
 from yt_dlp.socks import (
@@ -95,7 +95,7 @@ def handle(self):
             return
 
         elif Socks5Auth.AUTH_USER_PASS in methods:
-            self.connection.sendall(struct.pack("!BB", SOCKS5_VERSION, Socks5Auth.AUTH_USER_PASS))
+            self.connection.sendall(struct.pack('!BB', SOCKS5_VERSION, Socks5Auth.AUTH_USER_PASS))
 
             _, user_len = struct.unpack('!BB', self.connection.recv(2))
             username = self.connection.recv(user_len).decode()
@@ -174,7 +174,7 @@ def handle(self):
         if 0x0 < dest_ip <= 0xFF:
             use_remote_dns = True
         else:
-            socks_info['ipv4_address'] = socket.inet_ntoa(struct.pack("!I", dest_ip))
+            socks_info['ipv4_address'] = socket.inet_ntoa(struct.pack('!I', dest_ip))
 
         user_id = self._read_until_null().decode()
         if user_id != (self.socks_kwargs.get('user_id') or ''):
@@ -210,6 +210,16 @@ def do_GET(self):
             self.wfile.write(payload.encode())
 
 
+class SocksWebSocketTestRequestHandler(SocksTestRequestHandler):
+    def handle(self):
+        import websockets.sync.server
+        protocol = websockets.ServerProtocol()
+        connection = websockets.sync.server.ServerConnection(socket=self.request, protocol=protocol, close_timeout=0)
+        connection.handshake()
+        connection.send(json.dumps(self.socks_info))
+        connection.close()
+
+
 @contextlib.contextmanager
 def socks_server(socks_server_class, request_handler, bind_ip=None, **socks_server_kwargs):
     server = server_thread = None
@@ -252,8 +262,22 @@ def socks_info_request(self, handler, target_domain=None, target_port=None, **re
         return json.loads(handler.send(request).read().decode())
 
 
+class WebSocketSocksTestProxyContext(SocksProxyTestContext):
+    REQUEST_HANDLER_CLASS = SocksWebSocketTestRequestHandler
+
+    def socks_info_request(self, handler, target_domain=None, target_port=None, **req_kwargs):
+        request = Request(f'ws://{target_domain or "127.0.0.1"}:{target_port or "40000"}', **req_kwargs)
+        handler.validate(request)
+        ws = handler.send(request)
+        ws.send('socks_info')
+        socks_info = ws.recv()
+        ws.close()
+        return json.loads(socks_info)
+
+
 CTX_MAP = {
     'http': HTTPSocksTestProxyContext,
+    'ws': WebSocketSocksTestProxyContext,
 }
 
 
@@ -262,8 +286,14 @@ def ctx(request):
     return CTX_MAP[request.param]()
 
 
+@pytest.mark.parametrize(
+    'handler,ctx', [
+        ('Urllib', 'http'),
+        ('Requests', 'http'),
+        ('Websockets', 'ws'),
+        ('CurlCFFI', 'http'),
+    ], indirect=True)
 class TestSocks4Proxy:
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_socks4_no_auth(self, handler, ctx):
         with handler() as rh:
             with ctx.socks_server(Socks4ProxyHandler) as server_address:
@@ -271,7 +301,6 @@ def test_socks4_no_auth(self, handler, ctx):
                     rh, proxies={'all': f'socks4://{server_address}'})
                 assert response['version'] == 4
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_socks4_auth(self, handler, ctx):
         with handler() as rh:
             with ctx.socks_server(Socks4ProxyHandler, user_id='user') as server_address:
@@ -281,7 +310,6 @@ def test_socks4_auth(self, handler, ctx):
                     rh, proxies={'all': f'socks4://user:@{server_address}'})
                 assert response['version'] == 4
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_socks4a_ipv4_target(self, handler, ctx):
         with ctx.socks_server(Socks4ProxyHandler) as server_address:
             with handler(proxies={'all': f'socks4a://{server_address}'}) as rh:
@@ -289,7 +317,6 @@ def test_socks4a_ipv4_target(self, handler, ctx):
                 assert response['version'] == 4
                 assert (response['ipv4_address'] == '127.0.0.1') != (response['domain_address'] == '127.0.0.1')
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_socks4a_domain_target(self, handler, ctx):
         with ctx.socks_server(Socks4ProxyHandler) as server_address:
             with handler(proxies={'all': f'socks4a://{server_address}'}) as rh:
@@ -298,17 +325,16 @@ def test_socks4a_domain_target(self, handler, ctx):
                 assert response['ipv4_address'] is None
                 assert response['domain_address'] == 'localhost'
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_ipv4_client_source_address(self, handler, ctx):
         with ctx.socks_server(Socks4ProxyHandler) as server_address:
             source_address = f'127.0.0.{random.randint(5, 255)}'
+            verify_address_availability(source_address)
             with handler(proxies={'all': f'socks4://{server_address}'},
                          source_address=source_address) as rh:
                 response = ctx.socks_info_request(rh)
                 assert response['client_address'][0] == source_address
                 assert response['version'] == 4
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     @pytest.mark.parametrize('reply_code', [
         Socks4CD.REQUEST_REJECTED_OR_FAILED,
         Socks4CD.REQUEST_REJECTED_CANNOT_CONNECT_TO_IDENTD,
@@ -320,7 +346,6 @@ def test_socks4_errors(self, handler, ctx, reply_code):
                 with pytest.raises(ProxyError):
                     ctx.socks_info_request(rh)
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_ipv6_socks4_proxy(self, handler, ctx):
         with ctx.socks_server(Socks4ProxyHandler, bind_ip='::1') as server_address:
             with handler(proxies={'all': f'socks4://{server_address}'}) as rh:
@@ -329,7 +354,6 @@ def test_ipv6_socks4_proxy(self, handler, ctx):
                 assert response['ipv4_address'] == '127.0.0.1'
                 assert response['version'] == 4
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_timeout(self, handler, ctx):
         with ctx.socks_server(Socks4ProxyHandler, sleep=2) as server_address:
             with handler(proxies={'all': f'socks4://{server_address}'}, timeout=0.5) as rh:
@@ -337,9 +361,15 @@ def test_timeout(self, handler, ctx):
                     ctx.socks_info_request(rh)
 
 
+@pytest.mark.parametrize(
+    'handler,ctx', [
+        ('Urllib', 'http'),
+        ('Requests', 'http'),
+        ('Websockets', 'ws'),
+        ('CurlCFFI', 'http'),
+    ], indirect=True)
 class TestSocks5Proxy:
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_socks5_no_auth(self, handler, ctx):
         with ctx.socks_server(Socks5ProxyHandler) as server_address:
             with handler(proxies={'all': f'socks5://{server_address}'}) as rh:
@@ -347,7 +377,6 @@ def test_socks5_no_auth(self, handler, ctx):
                 assert response['auth_methods'] == [0x0]
                 assert response['version'] == 5
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_socks5_user_pass(self, handler, ctx):
         with ctx.socks_server(Socks5ProxyHandler, auth=('test', 'testpass')) as server_address:
             with handler() as rh:
@@ -360,7 +389,6 @@ def test_socks5_user_pass(self, handler, ctx):
                 assert response['auth_methods'] == [Socks5Auth.AUTH_NONE, Socks5Auth.AUTH_USER_PASS]
                 assert response['version'] == 5
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_socks5_ipv4_target(self, handler, ctx):
         with ctx.socks_server(Socks5ProxyHandler) as server_address:
             with handler(proxies={'all': f'socks5://{server_address}'}) as rh:
@@ -368,7 +396,6 @@ def test_socks5_ipv4_target(self, handler, ctx):
                 assert response['ipv4_address'] == '127.0.0.1'
                 assert response['version'] == 5
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_socks5_domain_target(self, handler, ctx):
         with ctx.socks_server(Socks5ProxyHandler) as server_address:
             with handler(proxies={'all': f'socks5://{server_address}'}) as rh:
@@ -376,7 +403,6 @@ def test_socks5_domain_target(self, handler, ctx):
                 assert (response['ipv4_address'] == '127.0.0.1') != (response['ipv6_address'] == '::1')
                 assert response['version'] == 5
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_socks5h_domain_target(self, handler, ctx):
         with ctx.socks_server(Socks5ProxyHandler) as server_address:
             with handler(proxies={'all': f'socks5h://{server_address}'}) as rh:
@@ -385,7 +411,6 @@ def test_socks5h_domain_target(self, handler, ctx):
                 assert response['domain_address'] == 'localhost'
                 assert response['version'] == 5
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_socks5h_ip_target(self, handler, ctx):
         with ctx.socks_server(Socks5ProxyHandler) as server_address:
             with handler(proxies={'all': f'socks5h://{server_address}'}) as rh:
@@ -394,7 +419,6 @@ def test_socks5h_ip_target(self, handler, ctx):
                 assert response['domain_address'] is None
                 assert response['version'] == 5
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_socks5_ipv6_destination(self, handler, ctx):
         with ctx.socks_server(Socks5ProxyHandler) as server_address:
             with handler(proxies={'all': f'socks5://{server_address}'}) as rh:
@@ -402,7 +426,6 @@ def test_socks5_ipv6_destination(self, handler, ctx):
                 assert response['ipv6_address'] == '::1'
                 assert response['version'] == 5
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_ipv6_socks5_proxy(self, handler, ctx):
         with ctx.socks_server(Socks5ProxyHandler, bind_ip='::1') as server_address:
             with handler(proxies={'all': f'socks5://{server_address}'}) as rh:
@@ -413,16 +436,15 @@ def test_ipv6_socks5_proxy(self, handler, ctx):
 
     # XXX: is there any feasible way of testing IPv6 source addresses?
     # Same would go for non-proxy source_address test...
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     def test_ipv4_client_source_address(self, handler, ctx):
         with ctx.socks_server(Socks5ProxyHandler) as server_address:
             source_address = f'127.0.0.{random.randint(5, 255)}'
+            verify_address_availability(source_address)
             with handler(proxies={'all': f'socks5://{server_address}'}, source_address=source_address) as rh:
                 response = ctx.socks_info_request(rh)
                 assert response['client_address'][0] == source_address
                 assert response['version'] == 5
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http')], indirect=True)
     @pytest.mark.parametrize('reply_code', [
         Socks5Reply.GENERAL_FAILURE,
         Socks5Reply.CONNECTION_NOT_ALLOWED,
@@ -439,7 +461,6 @@ def test_socks5_errors(self, handler, ctx, reply_code):
                 with pytest.raises(ProxyError):
                     ctx.socks_info_request(rh)
 
-    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http')], indirect=True)
     def test_timeout(self, handler, ctx):
         with ctx.socks_server(Socks5ProxyHandler, sleep=2) as server_address:
             with handler(proxies={'all': f'socks5://{server_address}'}, timeout=1) as rh: