Working websockets HTTP/S proxy

test/test_http_proxy.py

@@ -116,6 +116,9 @@ if urllib3:
         @_io_refs.setter
         def _io_refs(self, value):
             self.socket._io_refs = value
+
+        def shutdown(self, *args, **kwargs):
+            self.socket.shutdown(*args, **kwargs)
 else:
     SSLTransport = None
 
@@ -142,13 +145,14 @@ class WebSocketProxyHandler(BaseRequestHandler):
         protocol = websockets.ServerProtocol()
         connection = websockets.sync.server.ServerConnection(socket=self.request, protocol=protocol, close_timeout=0)
         connection.handshake()
-        connection.send(json.dumps(self.proxy_info))
+        for message in connection:
+            if message == 'proxy_info':
+                connection.send(json.dumps(self.proxy_info))
         connection.close()
 
 
 class WebSocketSecureProxyHandler(WebSocketProxyHandler):
-    def __init__(self, request, *args, proxy_info=None, **kwargs):
+    def __init__(self, request, *args, **kwargs):
-        self.proxy_info = proxy_info
         certfn = os.path.join(TEST_DIR, 'testcert.pem')
         sslctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
         sslctx.load_cert_chain(certfn, None)
@@ -218,7 +222,7 @@ def proxy_server(proxy_server_class, request_handler, bind_ip=None, **proxy_serv
     finally:
         server.shutdown()
         server.server_close()
-        server_thread.join(2.0)
+        server_thread.join()
 
 
 class HTTPProxyTestContext(abc.ABC):
@@ -297,6 +301,7 @@ class TestHTTPProxy:
                 proxy_info = ctx.proxy_info_request(rh)
                 assert proxy_info['connect'] is False
                 assert 'Proxy-Authorization' not in proxy_info['headers']
+                assert proxy_info['proxy'] == server_address
 
     def test_http_auth(self, handler, ctx):
         with ctx.http_server(HTTPProxyHandler, username='test', password='test') as server_address:
@@ -318,8 +323,9 @@ class TestHTTPProxy:
             verify_address_availability(source_address)
             with handler(proxies={ctx.REQUEST_PROTO: f'http://{server_address}'},
                          source_address=source_address) as rh:
-                response = ctx.proxy_info_request(rh)
+                proxy_info = ctx.proxy_info_request(rh)
-                assert response['client_address'][0] == source_address
+                assert proxy_info['client_address'][0] == source_address
+                assert proxy_info['proxy'] == server_address
 
     @pytest.mark.skip_handler('Urllib', 'urllib does not support https proxies')
     def test_https(self, handler, ctx):
@@ -328,6 +334,7 @@ class TestHTTPProxy:
                 proxy_info = ctx.proxy_info_request(rh)
                 assert proxy_info['connect'] is False
                 assert 'Proxy-Authorization' not in proxy_info['headers']
+                assert proxy_info['proxy'] == server_address
 
     @pytest.mark.skip_handler('Urllib', 'urllib does not support https proxies')
     def test_https_verify_failed(self, handler, ctx):
@@ -345,6 +352,7 @@ class TestHTTPProxy:
                 proxy_info = ctx.proxy_info_request(rh, target_domain='中文.tw')
                 assert proxy_info['path'].startswith('http://xn--fiq228c.tw')
                 assert proxy_info['headers']['Host'].split(':', 1)[0] == 'xn--fiq228c.tw'
+                assert proxy_info['proxy'] == server_address
 
 
 @pytest.mark.parametrize(
@@ -361,12 +369,14 @@ class TestHTTPConnectProxy:
                 proxy_info = ctx.proxy_info_request(rh)
                 assert proxy_info['connect'] is True
                 assert 'Proxy-Authorization' not in proxy_info['headers']
+                assert proxy_info['proxy'] == server_address
 
     def test_http_connect_auth(self, handler, ctx):
         with ctx.http_server(HTTPConnectProxyHandler, username='test', password='test') as server_address:
             with handler(verify=False, proxies={ctx.REQUEST_PROTO: f'http://test:test@{server_address}'}) as rh:
                 proxy_info = ctx.proxy_info_request(rh)
                 assert 'Proxy-Authorization' in proxy_info['headers']
+                assert proxy_info['proxy'] == server_address
 
     def test_http_connect_bad_auth(self, handler, ctx):
         with ctx.http_server(HTTPConnectProxyHandler, username='test', password='test') as server_address:
@@ -381,8 +391,9 @@ class TestHTTPConnectProxy:
             with handler(proxies={ctx.REQUEST_PROTO: f'http://{server_address}'},
                          source_address=source_address,
                          verify=False) as rh:
-                response = ctx.proxy_info_request(rh)
+                proxy_info = ctx.proxy_info_request(rh)
-                assert response['client_address'][0] == source_address
+                assert proxy_info['client_address'][0] == source_address
+                assert proxy_info['proxy'] == server_address
 
     @pytest.mark.skipif(urllib3 is None, reason='requires urllib3 to test')
     def test_https_connect_proxy(self, handler, ctx):
@@ -391,6 +402,7 @@ class TestHTTPConnectProxy:
                 proxy_info = ctx.proxy_info_request(rh)
                 assert proxy_info['connect'] is True
                 assert 'Proxy-Authorization' not in proxy_info['headers']
+                assert proxy_info['proxy'] == server_address
 
     @pytest.mark.skipif(urllib3 is None, reason='requires urllib3 to test')
     def test_https_connect_verify_failed(self, handler, ctx):
@@ -408,3 +420,4 @@ class TestHTTPConnectProxy:
             with handler(verify=False, proxies={ctx.REQUEST_PROTO: f'https://test:test@{server_address}'}) as rh:
                 proxy_info = ctx.proxy_info_request(rh)
                 assert 'Proxy-Authorization' in proxy_info['headers']
+                assert proxy_info['proxy'] == server_address

test/test_socks.py

@@ -216,7 +216,9 @@ class SocksWebSocketTestRequestHandler(SocksTestRequestHandler):
         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))
+        for message in connection:
+            if message == 'socks_info':
+                connection.send(json.dumps(self.socks_info))
         connection.close()
 
 

yt_dlp/networking/_websockets.py

@@ -118,7 +118,7 @@ class WebsocketsRH(WebSocketRequestHandler):
         for name in ('websockets.client', 'websockets.server'):
             logger = logging.getLogger(name)
             handler = logging.StreamHandler(stream=sys.stdout)
-            handler.setFormatter(logging.Formatter(f'{self.RH_NAME}: %(message)s'))
+            handler.setFormatter(logging.Formatter(f'{self.RH_NAME}: [{name}] %(message)s'))
             self.__logging_handlers[name] = handler
             logger.addHandler(handler)
             if self.verbose:
@@ -152,7 +152,7 @@ class WebsocketsRH(WebSocketRequestHandler):
                     **create_conn_kwargs
                 )
 
-            elif parsed_proxy_url.scheme.startswith('http'):
+            elif parsed_proxy_url.scheme in ('http', 'https'):
                 return create_http_connect_conn(
                     proxy_url=proxy,
                     url=url,
@@ -177,6 +177,7 @@ class WebsocketsRH(WebSocketRequestHandler):
                 headers['cookie'] = cookie_header
 
         proxy = select_proxy(request.url, self._get_proxies(request))
+
         try:
             conn = websockets.sync.client.connect(
                 sock=self._make_sock(proxy, request.url, timeout),
@@ -184,7 +185,10 @@ class WebsocketsRH(WebSocketRequestHandler):
                 additional_headers=headers,
                 open_timeout=timeout,
                 user_agent_header=None,
-                ssl_context=self._make_sslcontext() if parse_uri(request.url).secure else None,
+                ssl_context=(
+                    WebsocketsSSLContext(self._make_sslcontext())
+                    if parse_uri(request.url).secure else None
+                ),
                 close_timeout=0,  # not ideal, but prevents yt-dlp hanging
             )
             return WebsocketsResponseAdapter(conn, url=request.url)
@@ -218,12 +222,34 @@ class NoCloseHTTPResponse(HTTPResponse):
         if not self._check_close() and not self.chunked and self.length is None:
             self.will_close = False
 
-class CustomSSLTransport(SSLTransport):
+
+# todo: only define if urllib3 is available
+class WebsocketsSSLTransport(SSLTransport):
+    """
+    Modified version of urllib3 SSLTransport to support additional operations used by websockets
+    """
     def setsockopt(self, *args, **kwargs):
         self.socket.setsockopt(*args, **kwargs)
+
     def shutdown(self, *args, **kwargs):
+        self.unwrap()
         self.socket.shutdown(*args, **kwargs)
 
+
+class WebsocketsSSLContext:
+    """
+    Dummy SSL Context for websockets which returns a WebsocketsSSLTransport instance
+    for wrap socket when using TLS-in-TLS.
+    """
+    def __init__(self, ssl_context: ssl.SSLContext):
+        self.ssl_context = ssl_context
+
+    def wrap_socket(self, sock, server_hostname=None):
+        if isinstance(sock, ssl.SSLSocket):
+            return WebsocketsSSLTransport(sock, self.ssl_context, server_hostname=server_hostname)
+        return self.ssl_context.wrap_socket(sock, server_hostname=server_hostname)
+
+
 def create_http_connect_conn(
     proxy_url,
     url,
@@ -256,17 +282,18 @@ def create_http_connect_conn(
     if source_address is not None:
         conn.source_address = (source_address, 0)
 
-    conn.debuglevel=2
     try:
         conn.connect()
         if ssl_context:
-            conn.sock = CustomSSLTransport(conn.sock, ssl_context, server_hostname=proxy_url_parsed.hostname)
+            conn.sock = ssl_context.wrap_socket(conn.sock, server_hostname=proxy_url_parsed.hostname)
-
+        conn.request(
-        conn.request(method='CONNECT', url=f'{request_url_parsed.host}:{request_url_parsed.port}', headers=proxy_headers)
+            method='CONNECT',
+            url=f'{request_url_parsed.host}:{request_url_parsed.port}',
+            headers=proxy_headers)
         response = conn.getresponse()
     except OSError as e:
         conn.close()
-        raise TransportError('Unable to connect to proxy', cause=e) from e
+        raise ProxyError('Unable to connect to proxy', cause=e) from e
 
     if response.status == 200:
         return conn.sock