mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2025-12-31 20:36:24 +00:00
88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
import logging
|
|
|
|
from multiaddr import Multiaddr
|
|
|
|
from libp2p.abc import IListener, ITransport
|
|
from libp2p.custom_types import THandler
|
|
from libp2p.network.connection.raw_connection import RawConnection
|
|
from libp2p.transport.exceptions import OpenConnectionError
|
|
from libp2p.transport.upgrader import TransportUpgrader
|
|
|
|
from .connection import P2PWebSocketConnection
|
|
from .listener import WebsocketListener
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class WebsocketTransport(ITransport):
|
|
"""
|
|
Libp2p WebSocket transport: dial and listen on /ip4/.../tcp/.../ws
|
|
|
|
Implements production-ready WebSocket transport with:
|
|
- Flow control and buffer management
|
|
- Connection limits and rate limiting
|
|
- Proper error handling and cleanup
|
|
- Support for both WS and WSS protocols
|
|
"""
|
|
|
|
def __init__(
|
|
self, upgrader: TransportUpgrader, max_buffered_amount: int = 4 * 1024 * 1024
|
|
):
|
|
self._upgrader = upgrader
|
|
self._max_buffered_amount = max_buffered_amount
|
|
self._connection_count = 0
|
|
self._max_connections = 1000 # Production limit
|
|
|
|
async def dial(self, maddr: Multiaddr) -> RawConnection:
|
|
"""Dial a WebSocket connection to the given multiaddr."""
|
|
logger.debug(f"WebsocketTransport.dial called with {maddr}")
|
|
|
|
# Extract host and port from multiaddr
|
|
host = (
|
|
maddr.value_for_protocol("ip4")
|
|
or maddr.value_for_protocol("ip6")
|
|
or maddr.value_for_protocol("dns")
|
|
or maddr.value_for_protocol("dns4")
|
|
or maddr.value_for_protocol("dns6")
|
|
)
|
|
port_str = maddr.value_for_protocol("tcp")
|
|
if port_str is None:
|
|
raise ValueError(f"No TCP port found in multiaddr: {maddr}")
|
|
port = int(port_str)
|
|
|
|
# Build WebSocket URL
|
|
ws_url = f"ws://{host}:{port}/"
|
|
logger.debug(f"WebsocketTransport.dial connecting to {ws_url}")
|
|
|
|
try:
|
|
from trio_websocket import open_websocket_url
|
|
|
|
# Check connection limits
|
|
if self._connection_count >= self._max_connections:
|
|
raise OpenConnectionError(
|
|
f"Maximum connections reached: {self._max_connections}"
|
|
)
|
|
|
|
# Use the context manager but don't exit it immediately
|
|
# The connection will be closed when the RawConnection is closed
|
|
ws_context = open_websocket_url(ws_url)
|
|
ws = await ws_context.__aenter__()
|
|
conn = P2PWebSocketConnection(
|
|
ws, ws_context, max_buffered_amount=self._max_buffered_amount
|
|
) # type: ignore[attr-defined]
|
|
|
|
self._connection_count += 1
|
|
logger.debug(f"Total connections: {self._connection_count}")
|
|
|
|
return RawConnection(conn, initiator=True)
|
|
except Exception as e:
|
|
logger.error(f"Failed to dial WebSocket {maddr}: {e}")
|
|
raise OpenConnectionError(f"Failed to dial WebSocket {maddr}: {e}") from e
|
|
|
|
def create_listener(self, handler: THandler) -> IListener: # type: ignore[override]
|
|
"""
|
|
The type checker is incorrectly reporting this as an inconsistent override.
|
|
"""
|
|
logger.debug("WebsocketTransport.create_listener called")
|
|
return WebsocketListener(handler, self._upgrader)
|