更新

backend/tests/test_egress_channels.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import os
 import sys
 import time
@@ -25,6 +26,7 @@ from rpa_engine.egress_channels import (
     reset_egress_cache_for_tests,
     resolve_send_channels,
 )
+from rpa_engine.source_bound_proxy import SourceBoundProxy
 from models.db_migrate import migrate_accounts_table
 from models.models import Account
 
@@ -99,6 +101,49 @@ class EgressChannelTests(unittest.IsolatedAsyncioTestCase):
             with self.assertRaises(EgressChannelUnavailable):
                 await resolve_send_channels("198.51.100.99", 2)
 
+    async def test_browser_proxy_binds_selected_source_address(self):
+        observed_peer = asyncio.get_running_loop().create_future()
+
+        async def target_handler(reader, writer):
+            if not observed_peer.done():
+                observed_peer.set_result(writer.get_extra_info("peername")[0])
+            payload = await reader.readexactly(4)
+            writer.write(payload)
+            await writer.drain()
+            writer.close()
+            await writer.wait_closed()
+
+        target = await asyncio.start_server(target_handler, "127.0.0.1", 0)
+        target_port = target.sockets[0].getsockname()[1]
+        proxy = await SourceBoundProxy("127.0.0.2").start()
+        writer = None
+        try:
+            reader, writer = await asyncio.open_connection(
+                "127.0.0.1",
+                int(proxy.server_url.rpartition(":")[2]),
+            )
+            writer.write(
+                (
+                    f"CONNECT 127.0.0.1:{target_port} HTTP/1.1\r\n"
+                    f"Host: 127.0.0.1:{target_port}\r\n\r\n"
+                ).encode("ascii")
+            )
+            await writer.drain()
+            response = await reader.readuntil(b"\r\n\r\n")
+            self.assertIn(b"200 Connection Established", response)
+
+            writer.write(b"ping")
+            await writer.drain()
+            self.assertEqual(await reader.readexactly(4), b"ping")
+            self.assertEqual(await asyncio.wait_for(observed_peer, 1), "127.0.0.2")
+        finally:
+            if writer is not None:
+                writer.close()
+                await writer.wait_closed()
+            await proxy.close()
+            target.close()
+            await target.wait_closed()
+
 
 class EgressMigrationTests(unittest.TestCase):
     def test_mysql_accounts_uses_longtext_for_browser_payloads(self):