Add dialer and listener to webrtc-direct

Author: sukhman-sukh

libp2p/transport/webrtc/__init__.py

@@ -8,7 +8,7 @@
 import sys
 from .private_to_private.transport import WebRTCTransport
 from .private_to_public.transport import WebRTCDirectTransport
-from ..constants import (
+from .constants import (
     DEFAULT_ICE_SERVERS,
     SIGNALING_PROTOCOL,
     MUXER_PROTOCOL,

libp2p/transport/constants.py renamed to libp2p/transport/webrtc/constants.py

@@ -41,7 +41,6 @@
 logging.basicConfig(level=logging.INFO)
 SIGNAL_PROTOCOL: TProtocol = TProtocol("/libp2p/webrtc/signal/1.0.0")
 
-
 class WebRTCListener(IListener):
     """
     WebRTC Listener Implementation.

libp2p/transport/webrtc/listener.py

@@ -33,10 +33,9 @@ def webrtc_direct_decode(b: ByteString) -> str:
     return ""
 
 
-def certhash_encode(s: str) -> ByteString:
-    """Encode certificate hash component."""
+def certhash_decode(s: str) -> Tuple[int, bytes]:
     if not s:
-        return b""
+        raise ValueError("Empty certhash string.")
 
     # Remove multibase prefix if present
     if s.startswith("uEi"):
@@ -46,16 +45,29 @@ def certhash_encode(s: str) -> ByteString:
 
     # Decode base64url encoded hash
     try:
-        # Ensure s is bytes for base64 decoding
-        s_bytes = s.encode("ascii") if isinstance(s, str) else s
+        s_bytes = s.encode("ascii")
         # Add padding if needed
         padding = 4 - (len(s_bytes) % 4)
         if padding != 4:
             s_bytes += b"=" * padding
-        return base64.urlsafe_b64decode(s_bytes)
-    except Exception:
-        # Fallback to raw bytes
-        return s.encode("utf-8")
+        raw_bytes = base64.urlsafe_b64decode(s_bytes)
+    except Exception as e:
+        raise ValueError("Invalid base64url certhash") from e
+
+    if len(raw_bytes) < 2:
+        raise ValueError("Decoded certhash is too short to contain multihash header")
+
+    # Multihash format: <code><length><digest>
+    code = raw_bytes[0]
+    length = raw_bytes[1]
+    digest = raw_bytes[2:]
+
+    if len(digest) != length:
+        raise ValueError(
+            f"Digest length mismatch: expected {length}, got {len(digest)}"
+        )
+
+    return code, digest
 
 
 def certhash_decode(b: ByteString) -> str:
@@ -73,6 +85,6 @@ def certhash_decode(b: ByteString) -> str:
     "webrtc_decode",
     "webrtc_direct_encode",
     "webrtc_direct_decode",
-    "certhash_encode",
-    "certhash_decode",
+    # "certhash_encode",
+    # "certhash_decode",
 ]

libp2p/transport/webrtc/multiaddr_codecs.py

@@ -17,7 +17,7 @@
 
 from ..async_bridge import TrioSafeWebRTCOperations
 from ..connection import WebRTCRawConnection
-from ...constants import (
+from ..constants import (
     DEFAULT_DIAL_TIMEOUT,
     SIGNALING_PROTOCOL,
     SDPHandshakeError,

libp2p/transport/webrtc/private_to_private/initiate_connection.py

@@ -14,7 +14,7 @@
 )
 from libp2p.relay.circuit_v2.config import RelayConfig
 
-from ...constants import (
+from ..constants import (
     DEFAULT_DIAL_TIMEOUT,
     DEFAULT_ICE_SERVERS,
     SIGNALING_PROTOCOL,

libp2p/transport/webrtc/private_to_private/listener.py

@@ -15,7 +15,7 @@
 from libp2p.peer.id import ID
 
 from ..connection import WebRTCRawConnection
-from ...constants import WebRTCError
+from ..constants import WebRTCError
 from .pb import Message
 
 logger = logging.getLogger("webrtc.private.signaling_stream_handler")

libp2p/transport/webrtc/private_to_private/signaling_stream_handler.py

@@ -17,13 +17,13 @@
 from libp2p.host.basic_host import IHost
 from libp2p.transport.exceptions import OpenConnectionError
 
-from ...constants import (
+from ..constants import (
     DEFAULT_DIAL_TIMEOUT,
     DEFAULT_ICE_SERVERS,
     SIGNALING_PROTOCOL,
     WebRTCError,
 )
-from ..util import (
+from ..private_to_public.util import (
     pick_random_ice_servers,
 )
 from .initiate_connection import initiate_connection

libp2p/transport/webrtc/private_to_private/transport.py

@@ -0,0 +1,97 @@
+import trio
+from aiortc import RTCDataChannel, RTCSessionDescription
+from .direct_rtc_connection import DirectPeerConnection
+from libp2p.transport.webrtc.private_to_public.util import (
+    SDP,
+    generate_noise_prologue,
+    fingerprint_to_multiaddr,
+)
+from trio_asyncio import aio_as_trio
+from libp2p.transport.webrtc.noise_handshake import (
+    generate_noise_prologue,
+    NoiseEncrypter,
+)
+from libp2p.transport.webrtc.connection import WebRTCMultiaddrConnection
+from libp2p.transport.webrtc.muxer import DataChannelMuxerFactory
+from libp2p.transport.webrtc.constants import WEBRTC_CONNECTION_STATES
+import logging
+
+logger = logging.getLogger("libp2p.transport.webrtc.private_to_public")
+
+async def connect(
+    peer_connection: DirectPeerConnection,
+    ufrag: str,
+    role: str
+):
+    """
+    Establish a WebRTC-Direct connection, perform the noise handshake, and return the upgraded connection.
+    """
+    
+    # Create data channel for noise handshake (negotiated, id=0)
+    handshake_channel: RTCDataChannel = peer_connection.peer_connection.createDataChannel(
+        "", negotiated=True, id=0
+    )
+
+    try:
+        if role == "client":
+            logger.debug("client creating local offer")
+            offer = await peer_connection.createOffer()
+            logger.debug("client created local offer %s", offer.sdp)
+            munged_offer = SDP.munge_offer(offer, ufrag)
+            logger.debug("client setting local offer %s", munged_offer.sdp)
+            await aio_as_trio(peer_connection.setLocalDescription(munged_offer))
+
+            answer_sdp = SDP.server_answer_from_multiaddr(remote_addr, ufrag)
+            logger.debug("client setting server description %s", answer_sdp.sdp)
+            await aio_as_trio(peer_connection.setRemoteDescription(answer_sdp))
+        else:
+            offer_sdp = SDP.client_offer_from_multiaddr(remote_addr, ufrag)
+            logger.debug("server setting client %s %s", offer_sdp.type, offer_sdp.sdp)
+            await aio_as_trio(peer_connection.setRemoteDescription(offer_sdp))
+
+            logger.debug("server creating local answer")
+            answer = await peer_connection.createAnswer()
+            logger.debug("server created local answer")
+            munged_answer = SDP.munge_offer(answer, ufrag)
+            logger.debug("server setting local description %s", munged_answer.sdp)
+            await aio_as_trio(peer_connection.setLocalDescription(munged_answer))
+
+    # TODO: Fix this
+        # Wait for handshake channel to open
+        if handshake_channel.readyState != "open":
+            logger.debug(
+                "%s wait for handshake channel to open, starting status %s",
+                role,
+                handshake_channel.readyState,
+            )
+            # Wait for the 'open' event or signal cancellation
+            open_event = trio.Event()
+
+            def on_open():
+                open_event.set()
+
+            handshake_channel.on("open", on_open)
+            with trio.move_on_after(30):  # 30s timeout
+                await open_event.wait()
+            if handshake_channel.readyState != "open":
+                raise Exception("Handshake data channel did not open in time")
+
+        logger.debug("%s handshake channel opened", role)
+
+        if role == "server":
+            remote_fingerprint = peer_connection.remoteFingerprint().value
+            remote_addr = fingerprint_to_multiaddr(remote_fingerprint)
+
+        # Get local fingerprint
+        local_desc = peer_connection.localDescription
+        local_fingerprint = SDP.get_fingerprint_from_sdp(local_desc.sdp)
+        if local_fingerprint is None:
+            raise Exception("Could not get fingerprint from local description sdp")
+
+        logger.debug("%s performing noise handshake", role)
+        #TODO: Complete the noise handshake and connection authentication
+        noiseProlouge = generate_noise_prologue(local_fingerprint, remote_addr, role)
+        
+    except Exception as e:
+        logger.error("%s noise handshake failed: %s", role, e)
+        raise

libp2p/transport/webrtc/private_to_public/connect.py

@@ -0,0 +1,107 @@
+from aiortc import (RTCConfiguration, RTCPeerConnection, RTCSessionDescription, RTCDtlsFingerprint)
+from trio_asyncio import aio_as_trio
+from dataclasses import dataclass
+from .gen_certificate import WebRTCCertificate
+import datetime
+from ..constants import MAX_MESSAGE_SIZE
+
+@dataclass
+class DirectRTCConfiguration:
+    ufrag: str
+    peer_connection: RTCPeerConnection
+    rtc_config: RTCConfiguration
+    
+class DirectPeerConnection(RTCPeerConnection):
+    def __init__(self, direct_config: DirectRTCConfiguration):
+        self.ufrag = direct_config.ufrag
+        self.peer_connection = direct_config.peer_connection
+        super().__init__(direct_config.rtc_config)
+    
+    async def createOffer(self) -> RTCSessionDescription:
+        """
+        Create SDP offer, patching ICE ufrag and pwd to self.ufrag and self.upwd,
+        set as local description, and return the patched RTCSessionDescription.
+        """
+        offer = await aio_as_trio(super().createOffer())
+
+        sdp_lines = offer.sdp.splitlines()
+        new_lines = []
+        for line in sdp_lines:
+            if line.startswith("a=ice-ufrag:"):
+                new_lines.append(f"a=ice-ufrag:{getattr(self, 'ufrag', self.ufrag)}")
+            elif line.startswith("a=ice-pwd:"):
+                new_lines.append(f"a=ice-pwd:{getattr(self, 'ufrag', self.ufrag)}")
+            else:
+                new_lines.append(line)
+        patched_sdp = "\r\n".join(new_lines) + "\r\n"
+
+        patched_offer = RTCSessionDescription(sdp=patched_sdp, type=offer.type)
+        await aio_as_trio(self.setLocalDescription(patched_offer))
+        return patched_offer
+
+    async def createAnswer(self) -> RTCSessionDescription:
+        """
+        Create SDP answer, patching ICE ufrag and pwd to self.ufrag and self.upwd,
+        set as local description, and return the patched RTCSessionDescription.
+        """
+        answer = await aio_as_trio(super().createAnswer())
+
+        sdp_lines = answer.sdp.splitlines()
+        new_lines = []
+        for line in sdp_lines:
+            if line.startswith("a=ice-ufrag:"):
+                new_lines.append(f"a=ice-ufrag:{getattr(self, 'ufrag', self.ufrag)}")
+            elif line.startswith("a=ice-pwd:"):
+                new_lines.append(f"a=ice-pwd:{getattr(self, 'ufrag', self.ufrag)}")
+            else:
+                new_lines.append(line)
+        patched_sdp = "\r\n".join(new_lines) + "\r\n"
+
+        patched_answer = RTCSessionDescription(sdp=patched_sdp, type=answer.type)
+        await aio_as_trio(self.setLocalDescription(patched_answer))
+        return patched_answer
+
+    
+    def remoteFingerprint(self) -> RTCDtlsFingerprint:
+        pass
+        # return self.peer_connection.
+    
+    @staticmethod
+    async def create_dialer_rtc_peer_connection(
+        role: str,
+        ufrag: str,
+        rtc_configuration: RTCConfiguration,
+        certificate: WebRTCCertificate | None = None,
+    ):
+        """
+        Create a DirectRTCPeerConnection for dialing, similar to the JS createDialerRTCPeerConnection.
+        """
+       
+        if certificate is None:
+            certificate = WebRTCCertificate.generate()
+
+        # TODO: ICE servers. Should we use the ones from the rtc_configuration?
+        
+        # # ICE servers
+        # ice_servers = rtc_config.get("iceServers") if isinstance(rtc_config, dict) else getattr(rtc_config, "iceServers", None)
+        # if ice_servers is None and default_ice_servers is not None:
+        #     ice_servers = default_ice_servers
+
+        # if map_ice_servers is not None:
+        #     mapped_ice_servers = map_ice_servers(ice_servers)
+        # else:
+        #     mapped_ice_servers = ice_servers
+
+        peer_connection = RTCPeerConnection(
+            RTCConfiguration(
+                f"{role}-{(datetime.datetime.now(datetime.timezone.utc).timestamp() * 1000)}",
+                disable_fingerprint_verification=True,
+                disable_auto_negotiation=True,
+                certificate_pem_file=certificate.to_pem()[0],
+                key_pem_file=certificate.to_pem()[1],
+                enable_ice_udp_mux=(role == "server"),
+                max_message_size=MAX_MESSAGE_SIZE,
+                # ice_servers=mapped_ice_servers,
+            )
+        )
+        return DirectPeerConnection(DirectRTCConfiguration(ufrag, peer_connection, rtc_configuration))