plato-sdk-v2 2.7.1__py3-none-any.whl → 2.7.2__py3-none-any.whl

plato/v2/utils/db_cleanup.py

@@ -12,6 +12,7 @@ from sqlalchemy.ext.asyncio import create_async_engine
 
 from plato._generated.api.v1.simulator import get_db_config
 from plato._generated.models import DbConfigResponse
+from plato.v2.utils.gateway_tunnel import GatewayTunnel, find_free_port
 from plato.v2.utils.models import (
     ApiCleanupResult,
     DatabaseCleanupResult,
@@ -19,12 +20,10 @@ from plato.v2.utils.models import (
     EnvironmentInfo,
     SessionCleanupResult,
 )
-from plato.v2.utils.proxy_tunnel import ProxyTunnel, find_free_port, make_db_url
+from plato.v2.utils.proxy_tunnel import make_db_url
 
 logger = logging.getLogger(__name__)
 
-TEMP_PASSWORD = "newpass"
-
 
 class DatabaseCleaner:
     """Handles database audit_log cleanup operations."""
@@ -145,14 +144,13 @@ class DatabaseCleaner:
         config: DbConfigResponse,
         local_port: int,
     ) -> DatabaseCleanupResult:
-        """Connect to a single DB via tunnel and truncate audit_log tables."""
+        """Connect to a single DB via gateway tunnel and truncate audit_log tables."""
         db_port = config.db_port
 
-        tunnel = ProxyTunnel(
-            env_id=job_id,
-            db_port=db_port,
-            temp_password=TEMP_PASSWORD,
-            host_port=local_port,
+        tunnel = GatewayTunnel(
+            job_id=job_id,
+            remote_port=db_port,
+            local_port=local_port,
         )
 
         try:

plato/v2/utils/gateway_tunnel.py

@@ -0,0 +1,221 @@
+"""TLS + SNI gateway tunnel for database connections.
+
+Routes traffic through gateway.plato.so:443 using SNI-based routing,
+replacing the deprecated HTTP CONNECT proxy approach.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import socket
+import ssl
+
+logger = logging.getLogger(__name__)
+
+# Default gateway configuration
+DEFAULT_GATEWAY_HOST = "gateway.plato.so"
+DEFAULT_GATEWAY_PORT = 443
+
+
+def get_gateway_config() -> tuple[str, int]:
+    """Get gateway host and port from environment or defaults.
+
+    Returns:
+        Tuple of (host, port) for the gateway.
+    """
+    host = os.environ.get("PLATO_GATEWAY_HOST", DEFAULT_GATEWAY_HOST)
+    port = int(os.environ.get("PLATO_GATEWAY_PORT", str(DEFAULT_GATEWAY_PORT)))
+    return host, port
+
+
+def find_free_port(start_port: int = 55432) -> int:
+    """Find the first available TCP port starting from start_port."""
+    port = start_port
+    while port < 65535:
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                s.bind(("127.0.0.1", port))
+                return port
+        except OSError:
+            port += 1
+    raise RuntimeError(f"No free port found starting from {start_port}")
+
+
+class GatewayTunnel:
+    """Async TLS + SNI gateway tunnel for database connections.
+
+    Routes local connections through gateway.plato.so using SNI-based routing.
+    This replaces the deprecated HTTP CONNECT proxy tunnel.
+
+    Usage:
+        tunnel = GatewayTunnel(job_id="abc123", remote_port=5432, local_port=55432)
+        await tunnel.start()
+        # Connect to localhost:55432 to reach the VM's port 5432
+        await tunnel.stop()
+    """
+
+    def __init__(
+        self,
+        job_id: str,
+        remote_port: int,
+        local_port: int,
+        gateway_host: str | None = None,
+        gateway_port: int | None = None,
+        verify_ssl: bool = True,
+    ):
+        """Initialize the gateway tunnel.
+
+        Args:
+            job_id: The job/environment ID to connect to.
+            remote_port: Port on the VM to forward to.
+            local_port: Local port to listen on.
+            gateway_host: Gateway hostname (default: from env or gateway.plato.so).
+            gateway_port: Gateway port (default: from env or 443).
+            verify_ssl: Whether to verify SSL certificates.
+        """
+        self.job_id = job_id
+        self.remote_port = remote_port
+        self.local_port = local_port
+        self.verify_ssl = verify_ssl
+
+        # Get gateway config
+        default_host, default_port = get_gateway_config()
+        self.gateway_host = gateway_host or default_host
+        self.gateway_port = gateway_port or default_port
+
+        # SNI for routing: {job_id}--{port}.gateway.plato.so
+        self.sni = f"{job_id}--{remote_port}.{self.gateway_host}"
+
+        self._server: asyncio.AbstractServer | None = None
+        self._client_tasks: set[asyncio.Task] = set()
+
+    async def _open_gateway_connection(
+        self,
+        timeout: float = 30.0,
+    ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+        """Open a TLS connection to the gateway with SNI routing.
+
+        Returns:
+            Tuple of (reader, writer) for the gateway connection.
+        """
+        # Create SSL context
+        ssl_ctx = ssl.create_default_context()
+        if not self.verify_ssl:
+            ssl_ctx.check_hostname = False
+            ssl_ctx.verify_mode = ssl.CERT_NONE
+
+        # Connect with TLS, using SNI for routing
+        reader, writer = await asyncio.wait_for(
+            asyncio.open_connection(
+                self.gateway_host,
+                self.gateway_port,
+                ssl=ssl_ctx,
+                server_hostname=self.sni,  # SNI determines which VM/port to route to
+            ),
+            timeout=timeout,
+        )
+
+        # Enable TCP keepalive
+        sock = writer.get_extra_info("socket")
+        if isinstance(sock, socket.socket):
+            try:
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+                # macOS/BSD keepalive idle time
+                TCP_KEEPALIVE = getattr(socket, "TCP_KEEPALIVE", 0x10)
+                sock.setsockopt(socket.IPPROTO_TCP, TCP_KEEPALIVE, 30)
+            except OSError:
+                pass  # Best effort
+
+        return reader, writer
+
+    async def _pipe(
+        self,
+        src: asyncio.StreamReader,
+        dst: asyncio.StreamWriter,
+    ) -> None:
+        """Forward data from src to dst until EOF."""
+        try:
+            while True:
+                data = await src.read(65536)
+                if not data:
+                    break
+                dst.write(data)
+                await dst.drain()
+        except (ConnectionResetError, BrokenPipeError, OSError):
+            pass
+        finally:
+            try:
+                dst.close()
+                await dst.wait_closed()
+            except Exception:
+                pass
+
+    async def _handle_client(
+        self,
+        client_reader: asyncio.StreamReader,
+        client_writer: asyncio.StreamWriter,
+    ) -> None:
+        """Handle a single client connection by forwarding through gateway."""
+        task = asyncio.current_task()
+        if task:
+            self._client_tasks.add(task)
+
+        try:
+            # Connect to gateway via TLS with SNI
+            gateway_reader, gateway_writer = await self._open_gateway_connection()
+
+            # Bidirectional forwarding
+            await asyncio.gather(
+                self._pipe(client_reader, gateway_writer),
+                self._pipe(gateway_reader, client_writer),
+            )
+        except Exception as e:
+            logger.warning(f"Gateway tunnel error: {e}")
+            try:
+                client_writer.close()
+                await client_writer.wait_closed()
+            except Exception:
+                pass
+        finally:
+            if task:
+                self._client_tasks.discard(task)
+
+    async def start(self) -> None:
+        """Start the gateway tunnel server."""
+        logger.info(
+            f"Starting gateway tunnel: localhost:{self.local_port} -> "
+            f"{self.job_id}:{self.remote_port} via {self.gateway_host}"
+        )
+
+        self._server = await asyncio.start_server(
+            self._handle_client,
+            host="127.0.0.1",
+            port=self.local_port,
+        )
+
+        # Small delay to ensure binding is settled
+        await asyncio.sleep(0.1)
+
+        if not self._server.sockets:
+            raise RuntimeError("Gateway tunnel failed to start: no listening sockets")
+
+        logger.info(f"Gateway tunnel established on port {self.local_port}")
+
+    async def stop(self) -> None:
+        """Stop the gateway tunnel server."""
+        if self._server is not None:
+            logger.info("Stopping gateway tunnel")
+
+            # Stop accepting new connections
+            self._server.close()
+            await self._server.wait_closed()
+            self._server = None
+
+            # Cancel active client tasks
+            for t in list(self._client_tasks):
+                t.cancel()
+            self._client_tasks.clear()
+
+            logger.info("Gateway tunnel stopped")

{plato_sdk_v2-2.7.1.dist-info → plato_sdk_v2-2.7.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plato-sdk-v2
-Version: 2.7.1
+Version: 2.7.2
 Summary: Python SDK for the Plato API
 Author-email: Plato <support@plato.so>
 License-Expression: MIT

{plato_sdk_v2-2.7.1.dist-info → plato_sdk_v2-2.7.2.dist-info}/RECORD

@@ -494,7 +494,8 @@ plato/v2/sync/environment.py,sha256=WnDzbyEHpwCSEP8XnfNSjIYS7rt7lYR4HGJjzprZmTQ,
 plato/v2/sync/flow_executor.py,sha256=N41-WCWIJVcCR2UmPUEiK7roNacYoeONkRXpR7lUgT8,13941
 plato/v2/sync/session.py,sha256=okXqF-CjMmA82WRy2zPXaGidbovgjAENSqiuvE4_jKE,30420
 plato/v2/utils/__init__.py,sha256=XLeFFsjXkm9g2raMmo7Wt4QN4hhCrNZDJKnpffJ4LtM,38
-plato/v2/utils/db_cleanup.py,sha256=lnI5lsMHNHpG85Y99MaE4Rzc3618piuzhvH-uXO1zIc,8702
+plato/v2/utils/db_cleanup.py,sha256=JMzAAJz0ZnoUXtd8F4jpQmBpJpos2__RkgN_cuEearg,8692
+plato/v2/utils/gateway_tunnel.py,sha256=eWgwf4VV8-jx6iCuHFgCISsAOVmNOOjCB56EuZLsnOA,7171
 plato/v2/utils/models.py,sha256=PwehSSnIRG-tM3tWL1PzZEH77ZHhIAZ9R0UPs6YknbM,1441
 plato/v2/utils/proxy_tunnel.py,sha256=8ZTd0jCGSfIHMvSv1fgEyacuISWnGPHLPbDglWroTzY,10463
 plato/worlds/README.md,sha256=XFOkEA3cNNcrWkk-Cxnsl-zn-y0kvUENKQRSqFKpdqw,5479
@@ -503,7 +504,7 @@ plato/worlds/base.py,sha256=-RR71bSxEFI5yydtrtq-AAbuw98CIjvmrbztqzB9oIc,31041
 plato/worlds/build_hook.py,sha256=KSoW0kqa5b7NyZ7MYOw2qsZ_2FkWuz0M3Ru7AKOP7Qw,3486
 plato/worlds/config.py,sha256=O1lUXzxp-Z_M7izslT8naXgE6XujjzwYFFrDDzUOueI,12736
 plato/worlds/runner.py,sha256=r9B2BxBae8_dM7y5cJf9xhThp_I1Qvf_tlPq2rs8qC8,4013
-plato_sdk_v2-2.7.1.dist-info/METADATA,sha256=W64dXq4E_YTbyTp5SBJJBm3sxSryOHQSB6oXU8x5_mI,8652
-plato_sdk_v2-2.7.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-plato_sdk_v2-2.7.1.dist-info/entry_points.txt,sha256=upGMbJCx6YWUTKrPoYvYUYfFCqYr75nHDwhA-45m6p8,136
-plato_sdk_v2-2.7.1.dist-info/RECORD,,
+plato_sdk_v2-2.7.2.dist-info/METADATA,sha256=fv3apoCuAT5uplSgW9dbuMmduW2z7NqhVSOVErRIb5k,8652
+plato_sdk_v2-2.7.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+plato_sdk_v2-2.7.2.dist-info/entry_points.txt,sha256=upGMbJCx6YWUTKrPoYvYUYfFCqYr75nHDwhA-45m6p8,136
+plato_sdk_v2-2.7.2.dist-info/RECORD,,