networking: fix discv5 and tests

src/lean_spec/subspecs/networking/discovery/codec.py

@@ -20,6 +20,8 @@
 
 from __future__ import annotations
 
+import os
+
 from lean_spec.subspecs.networking.types import SeqNumber
 from lean_spec.types import Uint64, decode_rlp, encode_rlp
 from lean_spec.types.rlp import RLPDecodingError
@@ -298,6 +300,4 @@ def _decode_talkresp(payload: bytes) -> TalkResp:
 
 def generate_request_id() -> RequestId:
     """Generate a random request ID."""
-    import os
-
     return RequestId(data=os.urandom(8))

src/lean_spec/subspecs/networking/discovery/handshake.py

@@ -253,6 +253,8 @@ def create_handshake_response(
         whoareyou: WhoAreYouAuthdata,
         remote_pubkey: bytes,
         challenge_data: bytes,
+        remote_ip: str = "",
+        remote_port: int = 0,
     ) -> tuple[bytes, bytes, bytes]:
         """
         Create a HANDSHAKE packet in response to WHOAREYOU.
@@ -265,6 +267,8 @@ def create_handshake_response(
             remote_pubkey: Remote's 33-byte compressed public key.
             challenge_data: Full WHOAREYOU data for key derivation
                 (masking-iv || static-header || authdata from received packet).
+            remote_ip: Remote peer's IP address for session keying.
+            remote_port: Remote peer's UDP port for session keying.
 
         Returns:
             Tuple of (authdata, send_key, recv_key).
@@ -312,12 +316,14 @@ def create_handshake_response(
             is_initiator=True,
         )
 
-        # Store session.
+        # Store session keyed by (node_id, ip, port).
         self._session_cache.create(
             node_id=remote_node_id,
             send_key=send_key,
             recv_key=recv_key,
             is_initiator=True,
+            ip=remote_ip,
+            port=remote_port,
         )
 
         # Clean up pending handshake.
@@ -330,6 +336,8 @@ def handle_handshake(
         self,
         remote_node_id: bytes,
         handshake: HandshakeAuthdata,
+        remote_ip: str = "",
+        remote_port: int = 0,
     ) -> HandshakeResult:
         """
         Process a received HANDSHAKE packet.
@@ -339,6 +347,8 @@ def handle_handshake(
         Args:
             remote_node_id: 32-byte node ID from packet source.
             handshake: Decoded HANDSHAKE authdata.
+            remote_ip: Remote peer's IP address for session keying.
+            remote_port: Remote peer's UDP port for session keying.
 
         Returns:
             HandshakeResult with established session.
@@ -409,12 +419,14 @@ def handle_handshake(
             is_initiator=False,
         )
 
-        # Create session.
+        # Create session keyed by (node_id, ip, port).
         session = self._session_cache.create(
             node_id=remote_node_id,
             send_key=send_key,
             recv_key=recv_key,
             is_initiator=False,
+            ip=remote_ip,
+            port=remote_port,
         )
 
         # Clean up pending handshake.

src/lean_spec/subspecs/networking/discovery/messages.py

@@ -291,20 +291,3 @@ class StaticHeader(StrictBaseModel):
 
     authdata_size: Uint16
     """Byte length of the authdata section following this header."""
-
-
-class WhoAreYouAuthdata(StrictBaseModel):
-    """
-    Authdata for WHOAREYOU packets (flag=1).
-
-    Sent when the recipient cannot decrypt an incoming message packet.
-    The nonce in the packet header is set to the nonce of the failed message.
-
-    Total size: 24 bytes (16 + 8).
-    """
-
-    id_nonce: IdNonce
-    """128-bit random value for identity verification."""
-
-    enr_seq: SeqNumber
-    """Recipient's known ENR sequence for the sender. 0 if unknown."""

src/lean_spec/subspecs/networking/discovery/packet.py

@@ -32,7 +32,6 @@
 import os
 import struct
 from dataclasses import dataclass
-from enum import IntEnum
 
 from lean_spec.types import Bytes12, Bytes16, Uint64
 
@@ -63,14 +62,6 @@
 """Fixed portion of handshake authdata: src-id (32) + sig-size (1) + eph-key-size (1)."""
 
 
-class PacketType(IntEnum):
-    """Packet type aliases matching PacketFlag for clarity."""
-
-    MESSAGE = 0
-    WHOAREYOU = 1
-    HANDSHAKE = 2
-
-
 @dataclass(frozen=True, slots=True)
 class PacketHeader:
     """Decoded packet header."""
@@ -135,6 +126,7 @@ def encode_packet(
     authdata: bytes,
     message: bytes,
     encryption_key: bytes | None = None,
+    masking_iv: Bytes16 | None = None,
 ) -> bytes:
     """
     Encode a Discovery v5 packet.
@@ -147,6 +139,8 @@ def encode_packet(
         authdata: Authentication data (varies by packet type).
         message: Message payload (plaintext for WHOAREYOU, encrypted otherwise).
         encryption_key: 16-byte key for message encryption (None for WHOAREYOU).
+        masking_iv: Optional 16-byte IV for header masking. Random if not provided.
+            Must be provided for WHOAREYOU to match the IV used in challenge_data.
 
     Returns:
         Complete encoded packet ready for UDP transmission.
@@ -156,13 +150,14 @@ def encode_packet(
     if len(nonce) != GCM_NONCE_SIZE:
         raise ValueError(f"Nonce must be {GCM_NONCE_SIZE} bytes, got {len(nonce)}")
 
-    # Fresh random IV for header masking.
-    #
-    # Using dest_node_id as the masking key is deterministic,
-    # so the IV MUST be random to prevent ciphertext patterns.
-    # Without randomness, identical packets would produce
-    # identical masked headers, enabling traffic analysis.
-    masking_iv = Bytes16(os.urandom(CTR_IV_SIZE))
+    if masking_iv is None:
+        # Fresh random IV for header masking.
+        #
+        # Using dest_node_id as the masking key is deterministic,
+        # so the IV MUST be random to prevent ciphertext patterns.
+        # Without randomness, identical packets would produce
+        # identical masked headers, enabling traffic analysis.
+        masking_iv = Bytes16(os.urandom(CTR_IV_SIZE))
 
     static_header = _encode_static_header(flag, nonce, len(authdata))
     header = static_header + authdata
@@ -182,24 +177,25 @@ def encode_packet(
         if encryption_key is None:
             raise ValueError("Encryption key required for non-WHOAREYOU packets")
 
-        # Masked header as AAD prevents header tampering.
+        # Per spec: message-ad = masking-iv || header (plaintext).
         #
-        # The recipient verifies the header wasn't modified
-        # without having to decrypt the payload first.
+        # The AAD binds the plaintext header to the encrypted message.
+        # The recipient reconstructs this from the decoded header.
+        message_ad = bytes(masking_iv) + header
         encrypted_message = aes_gcm_encrypt(
-            Bytes16(encryption_key), Bytes12(nonce), message, masked_header
+            Bytes16(encryption_key), Bytes12(nonce), message, message_ad
         )
 
     # Assemble packet.
-    packet = masking_iv + masked_header + encrypted_message
+    packet = bytes(masking_iv) + masked_header + encrypted_message
 
     if len(packet) > MAX_PACKET_SIZE:
         raise ValueError(f"Packet exceeds max size: {len(packet)} > {MAX_PACKET_SIZE}")
 
     return packet
 
 
-def decode_packet_header(local_node_id: bytes, data: bytes) -> tuple[PacketHeader, bytes]:
+def decode_packet_header(local_node_id: bytes, data: bytes) -> tuple[PacketHeader, bytes, bytes]:
     """
     Decode and unmask a Discovery v5 packet header.
 
@@ -208,7 +204,8 @@ def decode_packet_header(local_node_id: bytes, data: bytes) -> tuple[PacketHeade
         data: Raw packet bytes.
 
     Returns:
-        Tuple of (header, message_bytes).
+        Tuple of (header, message_bytes, message_ad).
+        message_ad is masking-iv || plaintext header, used as AAD for decryption.
 
     Raises:
         ValueError: If packet is malformed.
@@ -224,8 +221,7 @@ def decode_packet_header(local_node_id: bytes, data: bytes) -> tuple[PacketHeade
     masked_data = data[CTR_IV_SIZE:]
 
     # Decrypt static header first to get authdata size.
-    static_header_masked = masked_data[:STATIC_HEADER_SIZE]
-    static_header = aes_ctr_decrypt(masking_key, masking_iv, static_header_masked)
+    static_header = aes_ctr_decrypt(masking_key, masking_iv, masked_data[:STATIC_HEADER_SIZE])
 
     # Parse static header.
     protocol_id = static_header[:6]
@@ -245,15 +241,19 @@ def decode_packet_header(local_node_id: bytes, data: bytes) -> tuple[PacketHeade
     if len(data) < header_end:
         raise ValueError(f"Packet truncated: need {header_end}, have {len(data)}")
 
-    # Decrypt the full header including authdata.
-    full_masked_header = masked_data[: STATIC_HEADER_SIZE + authdata_size]
-    full_header = aes_ctr_decrypt(masking_key, masking_iv, full_masked_header)
+    # Decrypt the full header (static header + authdata) in one pass.
+    full_header = aes_ctr_decrypt(
+        masking_key, masking_iv, masked_data[: STATIC_HEADER_SIZE + authdata_size]
+    )
     authdata = full_header[STATIC_HEADER_SIZE:]
 
     # Message bytes are everything after the header.
     message_bytes = data[header_end:]
 
-    return PacketHeader(flag=flag, nonce=nonce, authdata=authdata), message_bytes
+    # Per spec: message-ad = masking-iv || header (plaintext).
+    message_ad = bytes(masking_iv) + full_header
+
+    return PacketHeader(flag=flag, nonce=nonce, authdata=authdata), message_bytes, message_ad
 
 
 def decode_message_authdata(authdata: bytes) -> MessageAuthdata:
@@ -310,7 +310,7 @@ def decrypt_message(
     encryption_key: bytes,
     nonce: bytes,
     ciphertext: bytes,
-    masked_header: bytes,
+    message_ad: bytes,
 ) -> bytes:
     """
     Decrypt an encrypted message payload.
@@ -319,12 +319,12 @@ def decrypt_message(
         encryption_key: 16-byte session key.
         nonce: 12-byte nonce from packet header.
         ciphertext: Encrypted message with GCM tag.
-        masked_header: Masked header bytes (used as AAD).
+        message_ad: Additional authenticated data (masking-iv || plaintext header).
 
     Returns:
         Decrypted message plaintext.
     """
-    return aes_gcm_decrypt(Bytes16(encryption_key), Bytes12(nonce), ciphertext, masked_header)
+    return aes_gcm_decrypt(Bytes16(encryption_key), Bytes12(nonce), ciphertext, message_ad)
 
 
 def encode_message_authdata(src_id: bytes) -> bytes:

src/lean_spec/subspecs/networking/discovery/routing.py

@@ -4,7 +4,7 @@
 Kademlia-style routing table for Node Discovery Protocol v5.1.
 
 Node Table Structure
---------------------
+
 Nodes keep information about other nodes in their neighborhood. Neighbor nodes
 are stored in a routing table consisting of 'k-buckets'. For each 0 <= i < 256,
 every node keeps a k-bucket for nodes of logdistance(self, n) == i.
@@ -14,7 +14,7 @@
 seen at tail.
 
 Distance Metric
----------------
+
 The 'distance' between two node IDs is the bitwise XOR of the IDs, interpreted
 as a big-endian number:
 
@@ -26,7 +26,7 @@
     logdistance(n1, n2) = log2(distance(n1, n2))
 
 Bucket Eviction Policy
-----------------------
+
 When a new node N1 is encountered, it can be inserted into the corresponding
 bucket.
 
@@ -37,7 +37,7 @@
 removed, and N1 added to the front of the bucket.
 
 Liveness Verification
----------------------
+
 Implementations should perform liveness checks asynchronously and occasionally
 verify that a random node in a random bucket is live by sending PING. When
 responding to FINDNODE, implementations must avoid relaying any nodes whose
@@ -148,13 +148,13 @@ class KBucket:
     - New nodes added to tail, eviction candidates at head
 
     Eviction Policy
-    ---------------
+
     When full, ping the head node (least-recently seen).
     - If it responds, keep it and discard the new node.
     - If it fails, evict it and add the new node.
 
     Replacement Cache
-    -----------------
+
     Implementations should maintain a 'replacement cache' alongside each bucket.
     This cache holds recently-seen nodes which would fall into the corresponding
     bucket but cannot become a member because it is at capacity. Once a bucket
@@ -257,7 +257,6 @@ class RoutingTable:
     Bucket i contains nodes with log2(distance) == i + 1.
 
     Fork Filtering
-    --------------
 
     When local_fork_digest is set:
 
@@ -266,7 +265,6 @@ class RoutingTable:
     - Requires eth2 ENR data to be present
 
     Lookup Algorithm
-    ----------------
 
     Locates the k closest nodes to a target ID:
 
@@ -277,7 +275,6 @@ class RoutingTable:
     5. Stop when k closest have been queried
 
     Table Maintenance
-    -----------------
 
     - Track close neighbors
     - Regularly refresh stale buckets

src/lean_spec/subspecs/networking/discovery/service.py

@@ -24,7 +24,9 @@
 from __future__ import annotations
 
 import asyncio
+import ipaddress
 import logging
+import os
 import random
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Callable
@@ -558,8 +560,6 @@ async def _refresh_loop(self) -> None:
             await asyncio.sleep(REFRESH_INTERVAL_SECS)
             try:
                 # Perform lookup for random target.
-                import os
-
                 target = os.urandom(32)
                 await self.find_node(target)
             except Exception as e:
@@ -602,18 +602,7 @@ def _encode_ip_address(self, ip_str: str) -> bytes:
         Returns:
             Raw bytes representation of the IP address.
         """
-        import ipaddress
-
-        try:
-            # Try IPv4 first.
-            addr = ipaddress.ip_address(ip_str)
-            return addr.packed
-        except ValueError:
-            # Fall back to returning as-is if somehow already bytes.
-            if isinstance(ip_str, bytes):
-                return ip_str
-            # Last resort: encode as UTF-8 (shouldn't happen with valid IPs).
-            return ip_str.encode()
+        return ipaddress.ip_address(ip_str).packed
 
     def _enr_to_entry(self, enr: ENR) -> NodeEntry:
         """Convert an ENR to a NodeEntry."""