varun-r-mallya/py-libp2p
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/py-libp2p.git synced 2026-03-24 06:01:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add early data support to Noise protocol

Browse Source 
Signed-off-by: varun-r-mallya <varunrmallya@gmail.com>
This commit is contained in:
varun-r-mallya
2025-08-19 04:41:14 +05:30
parent cff0bfc17d
commit e4ab3cb2c5
11 changed files with 238 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

120
libp2p/security/noise/patterns.py
View File

@ -30,6 +30,9 @@ from libp2p.security.secure_session import (
SecureSession,
)
from .early_data import (
EarlyDataHandler,
)
from .exceptions import (
HandshakeHasNotFinished,
InvalidSignature,
@ -45,6 +48,7 @@ from .messages import (
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 @@ class BasePattern(IPattern):
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 @@ class PatternXX(BasePattern):
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 @@ class PatternXX(BasePattern):
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 @@ class PatternXX(BasePattern):
)
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 @@ class PatternXX(BasePattern):
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 @@ class PatternXX(BasePattern):
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,8 +257,15 @@ class PatternXX(BasePattern):
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)
@ -193,6 +273,16 @@ class PatternXX(BasePattern):
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 @@ class PatternXX(BasePattern):
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