Add early data support to Noise protocol

examples/doc-examples/example_encryption_noise.py

@@ -24,13 +24,8 @@ async def main():
     noise_transport = NoiseTransport(
         # local_key_pair: The key pair used for libp2p identity and authentication
         libp2p_keypair=key_pair,
-        # noise_privkey: The private key used for Noise protocol encryption
         noise_privkey=key_pair.private_key,
-        # early_data: Optional data to send during the handshake
-        # (None means no early data)
-        early_data=None,
-        # with_noise_pipes: Whether to use Noise pipes for additional security features
-        with_noise_pipes=False,
+        # TODO: add early data
     )
 
     # Create a security options dictionary mapping protocol ID to transport

examples/doc-examples/example_multiplexer.py

@@ -28,9 +28,7 @@ async def main():
         noise_privkey=key_pair.private_key,
         # early_data: Optional data to send during the handshake
         # (None means no early data)
-        early_data=None,
-        # with_noise_pipes: Whether to use Noise pipes for additional security features
-        with_noise_pipes=False,
+        # TODO: add early data
     )
 
     # Create a security options dictionary mapping protocol ID to transport

examples/doc-examples/example_peer_discovery.py

@@ -31,9 +31,7 @@ async def main():
         noise_privkey=key_pair.private_key,
         # early_data: Optional data to send during the handshake
         # (None means no early data)
-        early_data=None,
-        # with_noise_pipes: Whether to use Noise pipes for additional security features
-        with_noise_pipes=False,
+        # TODO: add early data
     )
 
     # Create a security options dictionary mapping protocol ID to transport

examples/doc-examples/example_running.py

@@ -28,9 +28,7 @@ async def main():
         noise_privkey=key_pair.private_key,
         # early_data: Optional data to send during the handshake
         # (None means no early data)
-        early_data=None,
-        # with_noise_pipes: Whether to use Noise pipes for additional security features
-        with_noise_pipes=False,
+        # TODO: add early data
     )
 
     # Create a security options dictionary mapping protocol ID to transport

libp2p/security/noise/early_data.py

@@ -0,0 +1,68 @@
+from abc import ABC, abstractmethod
+
+from libp2p.abc import IRawConnection
+from libp2p.custom_types import TProtocol
+from libp2p.peer.id import ID
+
+from .pb import noise_pb2 as noise_pb
+
+
+class EarlyDataHandler(ABC):
+    """Interface for handling early data during Noise handshake"""
+
+    @abstractmethod
+    async def send(
+        self, conn: IRawConnection, peer_id: ID
+    ) -> noise_pb.NoiseExtensions | None:
+        """Called to generate early data to send during handshake"""
+        pass
+
+    @abstractmethod
+    async def received(
+        self, conn: IRawConnection, extensions: noise_pb.NoiseExtensions | None
+    ) -> None:
+        """Called when early data is received during handshake"""
+        pass
+
+
+class TransportEarlyDataHandler(EarlyDataHandler):
+    """Default early data handler for muxer negotiation"""
+
+    def __init__(self, supported_muxers: list[TProtocol]):
+        self.supported_muxers = supported_muxers
+        self.received_muxers: list[TProtocol] = []
+
+    async def send(
+        self, conn: IRawConnection, peer_id: ID
+    ) -> noise_pb.NoiseExtensions | None:
+        """Send our supported muxers list"""
+        if not self.supported_muxers:
+            return None
+
+        extensions = noise_pb.NoiseExtensions()
+        # Convert TProtocol to string for serialization
+        extensions.stream_muxers[:] = [str(muxer) for muxer in self.supported_muxers]
+        return extensions
+
+    async def received(
+        self, conn: IRawConnection, extensions: noise_pb.NoiseExtensions | None
+    ) -> None:
+        """Store received muxers list"""
+        if extensions and extensions.stream_muxers:
+            self.received_muxers = [
+                TProtocol(muxer) for muxer in extensions.stream_muxers
+            ]
+
+    def match_muxers(self, is_initiator: bool) -> TProtocol | None:
+        """Find first common muxer between local and remote"""
+        if is_initiator:
+            # Initiator: find first local muxer that remote supports
+            for local_muxer in self.supported_muxers:
+                if local_muxer in self.received_muxers:
+                    return local_muxer
+        else:
+            # Responder: find first remote muxer that we support
+            for remote_muxer in self.received_muxers:
+                if remote_muxer in self.supported_muxers:
+                    return remote_muxer
+        return None

libp2p/security/noise/patterns.py

@@ -30,6 +30,9 @@
     SecureSession,
 )
 
+from .early_data import (
+    EarlyDataHandler,
+)
 from .exceptions import (
     HandshakeHasNotFinished,
     InvalidSignature,
@@ -45,6 +48,7 @@
     make_handshake_payload_sig,
     verify_handshake_payload_sig,
 )
+from .pb import noise_pb2 as noise_pb
 
 
 class IPattern(ABC):
@@ -62,7 +66,8 @@ class BasePattern(IPattern):
     noise_static_key: PrivateKey
     local_peer: ID
     libp2p_privkey: PrivateKey
-    early_data: bytes | None
+    initiator_early_data_handler: EarlyDataHandler | None
+    responder_early_data_handler: EarlyDataHandler | None
 
     def create_noise_state(self) -> NoiseState:
         noise_state = NoiseState.from_name(self.protocol_name)
@@ -73,11 +78,50 @@ def create_noise_state(self) -> NoiseState:
             raise NoiseStateError("noise_protocol is not initialized")
         return noise_state
 
-    def make_handshake_payload(self) -> NoiseHandshakePayload:
+    async def make_handshake_payload(
+        self, conn: IRawConnection, peer_id: ID, is_initiator: bool
+    ) -> NoiseHandshakePayload:
         signature = make_handshake_payload_sig(
             self.libp2p_privkey, self.noise_static_key.get_public_key()
         )
-        return NoiseHandshakePayload(self.libp2p_privkey.get_public_key(), signature)
+
+        # NEW: Get early data from appropriate handler
+        extensions = None
+        if is_initiator and self.initiator_early_data_handler:
+            extensions = await self.initiator_early_data_handler.send(conn, peer_id)
+        elif not is_initiator and self.responder_early_data_handler:
+            extensions = await self.responder_early_data_handler.send(conn, peer_id)
+
+        # NEW: Serialize extensions into early_data field
+        early_data = None
+        if extensions:
+            early_data = extensions.SerializeToString()
+
+        return NoiseHandshakePayload(
+            self.libp2p_privkey.get_public_key(),
+            signature,
+            early_data,  # ← This is the key addition
+        )
+
+    async def handle_received_payload(
+        self, conn: IRawConnection, payload: NoiseHandshakePayload, is_initiator: bool
+    ) -> None:
+        """Process early data from received payload"""
+        if not payload.early_data:
+            return
+
+        # Deserialize the NoiseExtensions from early_data field
+        try:
+            extensions = noise_pb.NoiseExtensions.FromString(payload.early_data)
+        except Exception:
+            # Invalid extensions, ignore silently
+            return
+
+        # Pass to appropriate handler
+        if is_initiator and self.initiator_early_data_handler:
+            await self.initiator_early_data_handler.received(conn, extensions)
+        elif not is_initiator and self.responder_early_data_handler:
+            await self.responder_early_data_handler.received(conn, extensions)
 
 
 class PatternXX(BasePattern):
@@ -86,13 +130,15 @@ def __init__(
         local_peer: ID,
         libp2p_privkey: PrivateKey,
         noise_static_key: PrivateKey,
-        early_data: bytes | None = None,
+        initiator_early_data_handler: EarlyDataHandler | None,
+        responder_early_data_handler: EarlyDataHandler | None,
     ) -> None:
         self.protocol_name = b"Noise_XX_25519_ChaChaPoly_SHA256"
         self.local_peer = local_peer
         self.libp2p_privkey = libp2p_privkey
         self.noise_static_key = noise_static_key
-        self.early_data = early_data
+        self.initiator_early_data_handler = initiator_early_data_handler
+        self.responder_early_data_handler = responder_early_data_handler
 
     async def handshake_inbound(self, conn: IRawConnection) -> ISecureConn:
         noise_state = self.create_noise_state()
@@ -106,18 +152,23 @@ async def handshake_inbound(self, conn: IRawConnection) -> ISecureConn:
 
         read_writer = NoiseHandshakeReadWriter(conn, noise_state)
 
-        # Consume msg#1.
+        # 1. Consume msg#1 (just empty bytes)
         await read_writer.read_msg()
 
-        # Send msg#2, which should include our handshake payload.
-        our_payload = self.make_handshake_payload()
+        # 2. Send msg#2 with our payload INCLUDING EARLY DATA
+        our_payload = await self.make_handshake_payload(
+            conn,
+            self.local_peer,  # We send our own peer ID in responder role
+            is_initiator=False,
+        )
         msg_2 = our_payload.serialize()
         await read_writer.write_msg(msg_2)
 
-        # Receive and consume msg#3.
+        # 3. Receive msg#3
         msg_3 = await read_writer.read_msg()
         peer_handshake_payload = NoiseHandshakePayload.deserialize(msg_3)
 
+        # Extract remote pubkey from noise handshake state
         if handshake_state.rs is None:
             raise NoiseStateError(
                 "something is wrong in the underlying noise `handshake_state`: "
@@ -126,14 +177,31 @@ async def handshake_inbound(self, conn: IRawConnection) -> ISecureConn:
             )
         remote_pubkey = self._get_pubkey_from_noise_keypair(handshake_state.rs)
 
+        # 4. Verify signature (unchanged)
         if not verify_handshake_payload_sig(peer_handshake_payload, remote_pubkey):
             raise InvalidSignature
+
+        # NEW: Process early data from msg#3 AFTER signature verification
+        await self.handle_received_payload(
+            conn, peer_handshake_payload, is_initiator=False
+        )
+
         remote_peer_id_from_pubkey = ID.from_pubkey(peer_handshake_payload.id_pubkey)
 
         if not noise_state.handshake_finished:
             raise HandshakeHasNotFinished(
                 "handshake is done but it is not marked as finished in `noise_state`"
             )
+
+        # NEW: Get negotiated muxer for connection state
+        # negotiated_muxer = None
+        if self.responder_early_data_handler and hasattr(
+            self.responder_early_data_handler, "match_muxers"
+        ):
+            # negotiated_muxer =
+            # self.responder_early_data_handler.match_muxers(is_initiator=False)
+            pass
+
         transport_read_writer = NoiseTransportReadWriter(conn, noise_state)
         return SecureSession(
             local_peer=self.local_peer,
@@ -142,6 +210,8 @@ async def handshake_inbound(self, conn: IRawConnection) -> ISecureConn:
             remote_permanent_pubkey=remote_pubkey,
             is_initiator=False,
             conn=transport_read_writer,
+            # NOTE: negotiated_muxer would need to be added to SecureSession constructor
+            # For now, store it in connection metadata or similar
         )
 
     async def handshake_outbound(
@@ -158,24 +228,27 @@ async def handshake_outbound(
         if handshake_state is None:
             raise NoiseStateError("Handshake state is not initialized")
 
-        # Send msg#1, which is *not* encrypted.
+        # 1. Send msg#1 (empty) - no early data possible in XX pattern
         msg_1 = b""
         await read_writer.write_msg(msg_1)
 
-        # Read msg#2 from the remote, which contains the public key of the peer.
+        # 2. Read msg#2 from responder
         msg_2 = await read_writer.read_msg()
         peer_handshake_payload = NoiseHandshakePayload.deserialize(msg_2)
 
+        # Extract remote pubkey from noise handshake state
         if handshake_state.rs is None:
             raise NoiseStateError(
                 "something is wrong in the underlying noise `handshake_state`: "
-                "we received and consumed msg#3, which should have included the "
+                "we received and consumed msg#2, which should have included the "
                 "remote static public key, but it is not present in the handshake_state"
             )
         remote_pubkey = self._get_pubkey_from_noise_keypair(handshake_state.rs)
 
+        # Verify signature BEFORE processing early data (security)
         if not verify_handshake_payload_sig(peer_handshake_payload, remote_pubkey):
             raise InvalidSignature
+
         remote_peer_id_from_pubkey = ID.from_pubkey(peer_handshake_payload.id_pubkey)
         if remote_peer_id_from_pubkey != remote_peer:
             raise PeerIDMismatchesPubkey(
@@ -184,15 +257,32 @@ async def handshake_outbound(
                 f"remote_peer_id_from_pubkey={remote_peer_id_from_pubkey}"
             )
 
-        # Send msg#3, which includes our encrypted payload and our noise static key.
-        our_payload = self.make_handshake_payload()
+        # NEW: Process early data from msg#2 AFTER verification
+        await self.handle_received_payload(
+            conn, peer_handshake_payload, is_initiator=True
+        )
+
+        # 3. Send msg#3 with our payload INCLUDING EARLY DATA
+        our_payload = await self.make_handshake_payload(
+            conn, remote_peer, is_initiator=True
+        )
         msg_3 = our_payload.serialize()
         await read_writer.write_msg(msg_3)
 
         if not noise_state.handshake_finished:
             raise HandshakeHasNotFinished(
                 "handshake is done but it is not marked as finished in `noise_state`"
             )
+
+        # NEW: Get negotiated muxer
+        # negotiated_muxer = None
+        if self.initiator_early_data_handler and hasattr(
+            self.initiator_early_data_handler, "match_muxers"
+        ):
+            pass
+            # negotiated_muxer =
+            # self.initiator_early_data_handler.match_muxers(is_initiator=True)
+
         transport_read_writer = NoiseTransportReadWriter(conn, noise_state)
         return SecureSession(
             local_peer=self.local_peer,
@@ -201,6 +291,8 @@ async def handshake_outbound(
             remote_permanent_pubkey=remote_pubkey,
             is_initiator=True,
             conn=transport_read_writer,
+            # NOTE: negotiated_muxer would need to be added to SecureSession constructor
+            # For now, store it in connection metadata or similar
         )
 
     @staticmethod

libp2p/security/noise/pb/noise.proto

@@ -1,8 +1,13 @@
-syntax = "proto3";
+syntax = "proto2";
 package pb;
 
+message NoiseExtensions {
+    repeated bytes webtransport_certhashes = 1;
+    repeated string stream_muxers = 2;
+}
+
 message NoiseHandshakePayload {
-	bytes identity_key = 1;
-	bytes identity_sig = 2;
-	bytes data = 3;
+	optional bytes identity_key = 1;
+	optional bytes identity_sig = 2;
+	optional bytes data = 3;
 }