Only access port forwarding via PortScope

This is in preparation of adding custom port-manager subclasses for
LinuxSSH and ADB platforms.

Change-Id: If7ebad712732920ca4edf59237b25da0e6f5b229
Reviewed-on: https://chromium-review.googlesource.com/c/crossbench/+/6652280
Commit-Queue: Camillo Bruni <[email protected]>
Reviewed-by: Patrick Thier <[email protected]>

Author: camillobruni

crossbench/plt/base.py

@@ -35,7 +35,7 @@
 from crossbench.plt import proc_helper
 from crossbench.plt.arch import MachineArch
 from crossbench.plt.bin import Binary
-from crossbench.plt.port_manager import PortManager
+from crossbench.plt.port_manager import PortManager, PortScope
 from crossbench.plt.remote import RemotePopen
 
 if TYPE_CHECKING:
@@ -571,8 +571,8 @@ def default_tmp_dir(self) -> pth.AnyPath:
     return self.path(tempfile.gettempdir())
 
   @property
-  def ports(self) -> PortManager:
-    return self._default_port_manager
+  def ports(self) -> PortScope:
+    return self._default_port_manager.scope
 
   def port_forward(self, local_port: int, remote_port: int) -> int:
     """ Forwards a device remote_port to a local port."""

crossbench/plt/port_manager.py

@@ -27,14 +27,18 @@ class PortScope:
   """
 
   def __init__(self,
-               platform: Platform,
+               manager: PortManager,
                parent_scope: Self | None = None) -> None:
-    self._platform = platform
+    self._manager: PortManager = manager
     self._parent_scope: Self | None = parent_scope
     assert parent_scope is not self
     self.forwarded_ports: dict[int, int] = {}
     self.reverse_forwarded_ports: dict[int, int] = {}
 
+  @property
+  def platform(self) -> Platform:
+    return self._manager.platform
+
   @property
   def is_nested(self) -> bool:
     return self._parent_scope is not None
@@ -43,6 +47,11 @@ def is_nested(self) -> bool:
   def is_empty(self) -> bool:
     return not self.forwarded_ports and not self.reverse_forwarded_ports
 
+  @contextlib.contextmanager
+  def nested(self) -> Iterator[PortScope]:
+    with self._manager.nested() as scope:
+      yield scope
+
   def is_forwarded_port_used(self, local_port: int) -> bool:
     return bool(self.lookup_forwarded_port(local_port))
 
@@ -72,7 +81,7 @@ def forward(self, local_port: int, remote_port: int) -> int:
       raise PortForwardException(
           f"Cannot forward local port {local_port} twice, "
           "it is already forwarded.")
-    local_port = self._platform.port_forward(local_port, remote_port)
+    local_port = self.platform.port_forward(local_port, remote_port)
     self.forwarded_ports[local_port] = remote_port
     return local_port
 
@@ -82,14 +91,14 @@ def stop_forward(self, local_port: int) -> None:
           f"Cannot stop forwarding local port {local_port}, "
           f"it was never forwarded.")
     del self.forwarded_ports[local_port]
-    self._platform.stop_port_forward(local_port)
+    self.platform.stop_port_forward(local_port)
 
   def reverse_forward(self, remote_port: int, local_port: int) -> int:
     if self.is_reverse_forwarded_port_used(remote_port):
       raise PortForwardException(
           f"Cannot reverse forward remote port {remote_port} twice, "
           "it is already forwarded.")
-    remote_port = self._platform.reverse_port_forward(remote_port, local_port)
+    remote_port = self.platform.reverse_port_forward(remote_port, local_port)
     self.reverse_forwarded_ports[remote_port] = local_port
     return remote_port
 
@@ -99,7 +108,7 @@ def stop_reverse_forward(self, remote_port: int) -> None:
           f"Cannot stop reverse forwarding remote port {remote_port}, "
           f"it was never forwarded.")
     del self.reverse_forwarded_ports[remote_port]
-    self._platform.stop_reverse_port_forward(remote_port)
+    self.platform.stop_reverse_port_forward(remote_port)
 
 
 class PortManager:
@@ -120,11 +129,11 @@ class PortManager:
   """
 
   def __init__(self, platform: Platform, throw: bool = False) -> None:
-    self._platform = platform
-    self._throw = throw
+    self._platform: Platform = platform
+    self._throw: bool = throw
     self._is_active: bool = False
     # Keeps track of scoped ports.
-    self._port_scope: PortScope = PortScope(self._platform, None)
+    self._port_scope: PortScope = PortScope(self, None)
     self._start()
 
   def _start(self):
@@ -137,14 +146,18 @@ def _start(self):
   def scope(self) -> PortScope:
     return self._port_scope
 
+  @property
+  def platform(self) -> Platform:
+    return self._platform
+
   @contextlib.contextmanager
-  def nested(self) -> Iterator[PortManager]:
+  def nested(self) -> Iterator[PortScope]:
     """Open a nested port scope, all forwarded ports that were opened
     during this scope will be closed when leaving the scope. """
     old_scope = self._port_scope
-    self._port_scope = PortScope(self._platform, self._port_scope)
+    self._port_scope = PortScope(self, self._port_scope)
     try:
-      yield self
+      yield self._port_scope
     finally:
       try:
         self._stop_current_scoped_ports()

crossbench/probes/web_page_replay/recorder.py

@@ -213,7 +213,7 @@ def __init__(self, probe: WebPageReplayProbe, run: Run) -> None:
     })
     self._recorder = WprRecorder(**kwargs)
     self._browser_platform = run.browser_platform
-    self._ports: PortScope = run.host_platform.ports.scope
+    self._ports: PortScope = run.host_platform.ports
 
   @override
   def setup(self) -> None:
@@ -245,7 +245,7 @@ def _setup_port_forwarding(self) -> None:
     if self._browser_platform.is_remote:
       # TODO: Fix run.setup and teardown layering so they they're called with
       # the same active port scope.
-      self._ports = self._browser_platform.ports.scope
+      self._ports = self._browser_platform.ports
       self._ports.reverse_forward(self._recorder.http_port,
                                   self._recorder.http_port)
       self._ports.reverse_forward(self._recorder.https_port,

tests/crossbench/mock_helper.py

@@ -27,6 +27,7 @@
 from crossbench.plt.linux import LinuxPlatform, RemoteLinuxPlatform
 from crossbench.plt.linux_ssh import LinuxSshPlatform
 from crossbench.plt.macos import MacOSPlatform
+from crossbench.plt.port_manager import PortManager
 from crossbench.plt.win import WinPlatform
 from crossbench.runner.run import Run
 from crossbench.stories.story import Story
@@ -94,6 +95,10 @@ def __init__(self, *args, is_battery_powered=False, **kwargs):
   def has_display(self) -> bool:
     return True
 
+  @property
+  def port_manager(self) -> PortManager:
+    return self._default_port_manager
+
   def os_details(self):
     return {
         "system": "mock os system",

tests/crossbench/plt/test_port_manager.py

@@ -6,7 +6,7 @@
 
 import unittest
 
-from crossbench.plt.port_manager import PortForwardException
+from crossbench.plt.port_manager import PortForwardException, PortManager
 from tests import test_helper
 from tests.crossbench.mock_helper import LinuxMockPlatform
 
@@ -49,33 +49,35 @@ class PortManagerTestCase(unittest.TestCase):
   def setUp(self):
     super().setUp()
     self.platform = FakePortLinuxMockPlatform()
-    self.port_manager = self.platform.ports
-    self.port_manager.assert_is_active()
+    self.port_scope = self.platform.ports
+    self.port_manager: PortManager = self.platform.port_manager
 
   def tearDown(self):
     self.assertFalse(self.platform.forwarded_ports)
     self.assertFalse(self.platform.reverse_forwarded_ports)
-    self.assertTrue(self.port_manager.is_empty)
+    self.assertTrue(self.port_scope.is_empty)
     super().tearDown()
 
   def test_default(self):
-    self.assertTrue(self.port_manager.is_empty)
+    self.assertTrue(self.port_scope.is_empty)
     self.assertFalse(self.port_manager.has_nested_scopes)
-    self.port_manager.assert_is_active()
 
   def test_nested(self):
-    with self.port_manager.nested():
+    self.assertTrue(self.port_scope.is_empty)
+    with self.port_scope.nested() as scope:
       self.assertFalse(self.port_manager.is_empty)
+      self.assertTrue(self.port_scope.is_empty)
+      self.assertTrue(scope.is_empty)
       self.assertTrue(self.port_manager.has_nested_scopes)
-      self.port_manager.assert_is_active()
+    self.assertTrue(self.port_scope.is_empty)
 
   def test_stop(self):
     self.port_manager.stop()
-    self.assertTrue(self.port_manager.is_empty)
+    self.assertTrue(self.port_scope.is_empty)
     self.assertFalse(self.port_manager.has_nested_scopes)
 
   def test_forward_port(self):
-    with self.port_manager.nested() as port_scope:
+    with self.port_scope.nested() as port_scope:
       returned_local_port = port_scope.forward(12345, 8080)
       self.assertEqual(returned_local_port, 12345)
       self.assertIn(12345, self.platform.forwarded_ports)
@@ -85,90 +87,90 @@ def test_forward_port(self):
     self.assertTrue(port_scope.is_empty)
 
   def test_forward_port_auto_assign(self):
-    with self.port_manager.nested() as port_scope:
+    with self.port_scope.nested() as port_scope:
       returned_local_port = port_scope.forward(0, 8080)
       self.assertEqual(returned_local_port, 60001)
       self.assertIn(60001, self.platform.forwarded_ports)
 
   def test_stop_forward_port(self):
-    with self.port_manager.nested() as port_scope:
+    with self.port_scope.nested() as port_scope:
       port_scope.forward(12345, 8080)
       port_scope.stop_forward(12345)
       self.assertNotIn(12345, self.platform.forwarded_ports)
 
   def test_forward_port_conflict(self):
-    with self.port_manager.nested() as port_scope:
+    with self.port_scope.nested() as port_scope:
       port_scope.forward(12345, 8080)
       with self.assertRaises(PortForwardException):
         # Try to forward same local port
         port_scope.forward(12345, 8081)
 
   def test_stop_forward_port_not_forwarded(self):
-    with self.port_manager.nested() as port_scope:
+    with self.port_scope.nested() as port_scope:
       with self.assertRaises(PortForwardException):
         port_scope.stop_forward(12345)
 
   def test_reverse_forward_port(self):
-    with self.port_manager.nested() as port_scope:
+    with self.port_scope.nested() as port_scope:
       returned_remote_port = port_scope.reverse_forward(54321, 8081)
       self.assertEqual(returned_remote_port, 54321)
       self.assertIn(54321, self.platform.reverse_forwarded_ports)
       self.assertEqual(self.platform.reverse_forwarded_ports[54321], 8081)
       self.assertFalse(port_scope.is_empty)
 
   def test_reverse_forward_port_auto_assign(self):
-    with self.port_manager.nested() as port_scope:
+    with self.port_scope.nested() as port_scope:
       returned_remote_port = port_scope.reverse_forward(0, 8081)
       self.assertEqual(returned_remote_port, 60001)
       self.assertIn(60001, self.platform.reverse_forwarded_ports)
 
   def test_stop_reverse_forward_port(self):
-    with self.port_manager.nested() as port_scope:
+    with self.port_scope.nested() as port_scope:
       port_scope.reverse_forward(54321, 8081)
       port_scope.stop_reverse_forward(54321)
       self.assertNotIn(54321, self.platform.reverse_forwarded_ports)
 
   def test_reverse_forward_port_conflict(self):
-    with self.port_manager.nested() as port_scope:
+    with self.port_scope.nested() as port_scope:
       port_scope.reverse_forward(54321, 8081)
       with self.assertRaises(PortForwardException):
         # Try to reverse forward same remote port
         port_scope.reverse_forward(54321, 8082)
 
   def test_stop_reverse_forward_port_not_forwarded(self):
-    with self.port_manager.nested() as port_scope:
+    with self.port_scope.nested() as port_scope:
       with self.assertRaises(PortForwardException):
         port_scope.stop_reverse_forward(54321)
 
   def test_nested_cleanup(self):
-    self.port_manager.forward(1111, 2222)
-    with self.port_manager.nested() as port_scope:
+    self.port_scope.forward(1111, 2222)
+    with self.port_scope.nested() as port_scope:
       port_scope.forward(3333, 4444)
     self.assertIn(1111, self.platform.forwarded_ports)
     self.assertNotIn(3333, self.platform.forwarded_ports)
     self.assertFalse(self.port_manager.has_nested_scopes)
     self.port_manager.stop()
-    self.assertTrue(self.port_manager.is_empty)
+    self.assertTrue(self.port_scope.is_empty)
     self.assertNotIn(1111, self.platform.forwarded_ports)
     self.assertNotIn(3333, self.platform.forwarded_ports)
 
   def test_forward_nested_cleanup_stop_outer(self):
-    self.port_manager.forward(1111, 2222)
-    with self.port_manager.nested() as port_scope:
+    self.port_scope.forward(1111, 2222)
+    with self.port_scope.nested() as port_scope:
       port_scope.forward(3333, 4444)
       port_scope.stop_forward(3333)
       with self.assertRaisesRegex(PortForwardException, "1111"):
         port_scope.stop_forward(1111)
-    self.port_manager.stop_forward(1111)
+    self.port_scope.stop_forward(1111)
 
   def test_reverse_forward_nested_cleanup_stop_outer(self):
-    self.port_manager.reverse_forward(1111, 2222)
-    with self.port_manager.nested() as port_scope:
+    self.port_scope.reverse_forward(1111, 2222)
+    with self.port_scope.nested() as port_scope:
       port_scope.reverse_forward(3333, 4444)
       port_scope.stop_reverse_forward(3333)
       with self.assertRaisesRegex(PortForwardException, "1111"):
         port_scope.stop_reverse_forward(1111)
-    self.port_manager.stop_reverse_forward(1111)
+    self.port_scope.stop_reverse_forward(1111)
 
 
 if __name__ == "__main__":