Implemented: Envelope wrapper class + linter hacks for protobuf checks

libp2p/peer/envelope.py

@@ -0,0 +1,252 @@
+from typing import Any
+
+from libp2p.crypto.ed25519 import Ed25519PublicKey
+from libp2p.crypto.keys import PrivateKey, PublicKey
+from libp2p.crypto.rsa import RSAPublicKey
+from libp2p.crypto.secp256k1 import Secp256k1PublicKey
+import libp2p.peer.pb.crypto_pb2 as cryto_pb
+import libp2p.peer.pb.envelope_pb2 as pb
+import libp2p.peer.pb.peer_record_pb2 as record_pb
+from libp2p.peer.peer_record import PeerRecord, peer_record_from_protobuf
+from libp2p.utils.varint import encode_uvarint
+
+ENVELOPE_DOMAIN = "libp2p-peer-record"
+PEER_RECORD_CODEC = b"\x03\x01"
+
+
+class Envelope:
+    """
+    A signed wrapper around a serialized libp2p record.
+
+    Envelopes are cryptographically signed by the author's private key
+    and are scoped to a specific 'domain' to prevent cross-protocol replay.
+
+    Attributes:
+        public_key: The public key that can verify the envelope's signature.
+        payload_type: A multicodec code identifying the type of payload inside.
+        raw_payload: The raw serialized record data.
+        signature: Signature over the domain-scoped payload content.
+
+    """
+
+    public_key: PublicKey
+    payload_type: bytes
+    raw_payload: bytes
+    signature: bytes
+
+    _cached_record: PeerRecord | None = None
+    _unmarshal_error: Exception | None = None
+
+    def __init__(
+        self,
+        public_key: PublicKey,
+        payload_type: bytes,
+        raw_payload: bytes,
+        signature: bytes,
+    ):
+        self.public_key = public_key
+        self.payload_type = payload_type
+        self.raw_payload = raw_payload
+        self.signature = signature
+
+    def marshal_envelope(self) -> bytes:
+        """
+        Serialize this Envelope into its protobuf wire format.
+
+        Converts all envelope fields into a `pb.Envelope` protobuf message
+        and returns the serialized bytes.
+
+        :return: Serialized envelope as bytes.
+        """
+        pb_env = pb.Envelope(  # type: ignore[attr-defined]
+            public_key=pub_key_to_protobuf(self.public_key),
+            payload_type=self.payload_type,
+            payload=self.raw_payload,
+            signature=self.signature,
+        )
+        return pb_env.SerializeToString()
+
+    def validate(self, domain: str) -> None:
+        """
+        Verify the envelope's signature within the given domain scope.
+
+        This ensures that the envelope has not been tampered with
+        and was signed under the correct usage context.
+
+        :param domain: Domain string that contextualizes the signature.
+        :raises ValueError: If the signature is invalid.
+        """
+        unsigned = make_unsigned(domain, self.payload_type, self.raw_payload)
+        if not self.public_key.verify(unsigned, self.signature):
+            raise ValueError("Invalid envelope signature")
+
+    def record(self) -> PeerRecord:
+        """
+        Lazily decode and return the embedded PeerRecord.
+
+        This method unmarshals the payload bytes into a `PeerRecord` instance,
+        using the registered codec to identify the type. The decoded result
+        is cached for future use.
+
+        :return: Decoded PeerRecord object.
+        :raises Exception: If decoding fails or payload type is unsupported.
+        """
+        if self._cached_record is not None:
+            return self._cached_record
+
+        try:
+            if self.payload_type != PEER_RECORD_CODEC:
+                raise ValueError("Unsuported payload type in envelope")
+            msg = record_pb.PeerRecord()  # type: ignore[attr-defined]
+            msg.ParseFromString(self.raw_payload)
+
+            self._cached_record = peer_record_from_protobuf(msg)
+            return self._cached_record
+        except Exception as e:
+            self._unmarshal_error = e
+            raise
+
+    def equal(self, other: Any) -> bool:
+        """
+        Compare this Envelope with another for structural equality.
+
+        Two envelopes are considered equal if:
+        - They have the same public key
+        - The payload type and payload bytes match
+        - Their signatures are identical
+
+        :param other: Another object to compare.
+        :return: True if equal, False otherwise.
+        """
+        if isinstance(other, Envelope):
+            return (
+                self.public_key.__eq__(other.public_key)
+                and self.payload_type == other.payload_type
+                and self.signature == other.signature
+                and self.raw_payload == other.raw_payload
+            )
+        return False
+
+
+def pub_key_to_protobuf(pub_key: PublicKey) -> cryto_pb.PublicKey:  # type: ignore[name-defined]
+    """
+    Convert a Python PublicKey object to its protobuf equivalent.
+
+    :param pub_key: A libp2p-compatible PublicKey instance.
+    :return: Serialized protobuf PublicKey message.
+    """
+    key_type = pub_key.get_type().value
+    data = pub_key.to_bytes()
+    protobuf_key = cryto_pb.PublicKey(Type=key_type, Data=data)  # type: ignore[attr-defined]
+    return protobuf_key
+
+
+def pub_key_from_protobuf(pb_key: cryto_pb.PublicKey) -> PublicKey:  # type: ignore[name-defined]
+    """
+    Parse a protobuf PublicKey message into a native libp2p PublicKey.
+
+    Supports Ed25519, RSA, and Secp256k1 key types.
+
+    :param pb_key: Protobuf representation of a public key.
+    :return: Parsed PublicKey object.
+    :raises ValueError: If the key type is unrecognized.
+    """
+    if pb_key.Type == cryto_pb.KeyType.Ed25519:  # type: ignore[attr-defined]
+        return Ed25519PublicKey.from_bytes(pb_key.Data)
+    elif pb_key.Type == cryto_pb.KeyType.RSA:  # type: ignore[attr-defined]
+        return RSAPublicKey.from_bytes(pb_key.Data)
+    elif pb_key.Type == cryto_pb.KeyType.Secp256k1:  # type: ignore[attr-defined]
+        return Secp256k1PublicKey.from_bytes(pb_key.Data)
+    # TODO: Add suport fot ECDSA parsing also
+    else:
+        raise ValueError(f"Unknown key type: {pb_key.Type}")
+
+
+def seal_record(record: PeerRecord, private_key: PrivateKey) -> Envelope:
+    """
+    Create and sign a new Envelope from a PeerRecord.
+
+    The record is serialized and signed in the scope of its domain and codec.
+    The result is a self-contained, verifiable Envelope.
+
+    :param record: A PeerRecord to encapsulate.
+    :param private_key: The signer's private key.
+    :return: A signed Envelope instance.
+    """
+    payload = record.marshal_record()
+
+    unsigned = make_unsigned(record.domain(), record.codec(), payload)
+    signature = private_key.sign(unsigned)
+
+    return Envelope(
+        public_key=private_key.get_public_key(),
+        payload_type=record.codec(),
+        raw_payload=payload,
+        signature=signature,
+    )
+
+
+def consume_envelope(data: bytes, domain: str) -> tuple[Envelope, PeerRecord]:
+    """
+    Parse, validate, and decode an Envelope from bytes.
+
+    Validates the envelope's signature using the given domain and decodes
+    the inner payload into a PeerRecord.
+
+    :param data: Serialized envelope bytes.
+    :param domain: Domain string to verify signature against.
+    :return: Tuple of (Envelope, PeerRecord).
+    :raises ValueError: If signature validation or decoding fails.
+    """
+    env = unmarshal_envelope(data)
+    env.validate(domain)
+    record = env.record()
+    return env, record
+
+
+def unmarshal_envelope(data: bytes) -> Envelope:
+    """
+    Deserialize an Envelope from its wire format.
+
+    This parses the protobuf fields without verifying the signature.
+
+    :param data: Serialized envelope bytes.
+    :return: Parsed Envelope object.
+    :raises DecodeError: If protobuf parsing fails.
+    """
+    pb_env = pb.Envelope()  # type: ignore[attr-defined]
+    pb_env.ParseFromString(data)
+    pk = pub_key_from_protobuf(pb_env.public_key)
+
+    return Envelope(
+        public_key=pk,
+        payload_type=pb_env.payload_type,
+        raw_payload=pb_env.payload,
+        signature=pb_env.signature,
+    )
+
+
+def make_unsigned(domain: str, payload_type: bytes, payload: bytes) -> bytes:
+    """
+    Build a byte buffer to be signed for an Envelope.
+
+    The unsigned byte structure is:
+        varint(len(domain)) || domain ||
+        varint(len(payload_type)) || payload_type ||
+        varint(len(payload)) || payload
+
+    This is the exact input used during signing and verification.
+
+    :param domain: Domain string for signature scoping.
+    :param payload_type: Identifier for the type of payload.
+    :param payload: Raw serialized payload bytes.
+    :return: Byte buffer to be signed or verified.
+    """
+    fields = [domain.encode(), payload_type, payload]
+    buf = bytearray()
+
+    for field in fields:
+        buf.extend(encode_uvarint(len(field)))
+        buf.extend(field)
+
+    return bytes(buf)