From fc540cfb47ba6117ce5ca4df8f068e5090ea2023 Mon Sep 17 00:00:00 2001 From: Michael Eze Date: Sat, 13 Sep 2025 17:40:02 +0100 Subject: [PATCH 1/2] add record and envelop classes --- libp2p/crypto/keys.py | 3 + libp2p/record/__init__.py | 0 libp2p/record/envelope.py | 137 +++++++++++++++++++++++++++++ libp2p/record/exceptions.py | 10 +++ libp2p/record/libp2p_record.py | 12 +++ libp2p/record/pb/__init__.py | 0 libp2p/record/pb/envelope.proto | 23 +++++ libp2p/record/pb/envelope_pb2.py | 26 ++++++ libp2p/record/pb/envelope_pb2.pyi | 50 +++++++++++ libp2p/record/record.py | 62 +++++++++++++ tests/core/record/test_envelope.py | 128 +++++++++++++++++++++++++++ tests/core/record/test_record.py | 90 +++++++++++++++++++ 12 files changed, 541 insertions(+) create mode 100644 libp2p/record/__init__.py create mode 100644 libp2p/record/envelope.py create mode 100644 libp2p/record/exceptions.py create mode 100644 libp2p/record/libp2p_record.py create mode 100644 libp2p/record/pb/__init__.py create mode 100644 libp2p/record/pb/envelope.proto create mode 100644 libp2p/record/pb/envelope_pb2.py create mode 100644 libp2p/record/pb/envelope_pb2.pyi create mode 100644 libp2p/record/record.py create mode 100644 tests/core/record/test_envelope.py create mode 100644 tests/core/record/test_record.py diff --git a/libp2p/crypto/keys.py b/libp2p/crypto/keys.py index 21cf71b2b..9912b5523 100644 --- a/libp2p/crypto/keys.py +++ b/libp2p/crypto/keys.py @@ -67,6 +67,9 @@ def _serialize_to_protobuf(self) -> crypto_pb2.PublicKey: data = self.to_bytes() protobuf_key = crypto_pb2.PublicKey(key_type=key_type, data=data) return protobuf_key + + def serialize_to_protobuf(self) -> crypto_pb2.PublicKey: + return self._serialize_to_protobuf() def serialize(self) -> bytes: """Return the canonical serialization of this ``Key``.""" diff --git a/libp2p/record/__init__.py b/libp2p/record/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/libp2p/record/envelope.py b/libp2p/record/envelope.py new file mode 100644 index 000000000..622610444 --- /dev/null +++ b/libp2p/record/envelope.py @@ -0,0 +1,137 @@ +from libp2p.crypto.keys import PrivateKey, PublicKey +from libp2p.crypto.serialization import deserialize_public_key +from typing import Optional, Tuple + +from libp2p.record.record import Record, unmarshal_record_payload +import threading +import libp2p.record.pb.envelope_pb2 as pb +from libp2p.record.exceptions import ( + ErrEmptyDomain, + ErrEmptyPayloadType, + ErrInvalidSignature +) + +class Envelope: + public_key: PublicKey + payload_type: bytes + raw_payload: bytes + signature: bytes + + 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 + + self._cached_record: Optional[Record] = None + self._unmarshal_error: Optional[Exception] = None + self._unmarshal_lock = threading.Lock() + self._record_initialized: bool = False + + def marshal(self) -> bytes: + key = self.public_key.serialize_to_protobuf() + msg = pb.Envelope( + public_key=key, + payload_type=self.payload_type, + payload=self.raw_payload, + signature=self.signature + ) + return msg.SerializeToString() + + @classmethod + def unmarshal_envelope(cls, data: bytes) -> "Envelope": + msg = pb.Envelope() + msg.ParseFromString(data) + + key = deserialize_public_key(msg.public_key.SerializePartialToString()) + return cls( + public_key=key, + payload_type=msg.payload_type, + raw_payload=msg.payload, + signature=msg.signature + ) + + def equal(self, other: "Envelope") -> bool: + if other is None: + return False + return ( + self.public_key.__eq__(other.public_key) and + self.payload_type == other.payload_type and + self.raw_payload == other.raw_payload and + self.signature == other.signature + ) + + def record(self) -> Record: + """Return the unmarshalled Record, caching on first access.""" + with self._unmarshal_lock: + if self._cached_record is not None: + return self._cached_record + try: + self._cached_record = unmarshal_record_payload( + payload_type=self.payload_type, + payload_bytes=self.raw_payload + ) + except Exception as e: + self._unmarshal_error = e + raise + return self._cached_record + + def validate(self, domain: str) -> bool: + if not domain: + raise ErrEmptyDomain("domain cannot be empty") + if not self.payload_type: + raise ErrEmptyPayloadType("payload_type cannot be empty") + + try: + unsigned = make_unsigned(domain=domain, payload_type=self.payload_type, payload=self.raw_payload) + valid = self.public_key.verify(unsigned, self.signature) + if not valid: + raise ErrInvalidSignature("invalid signature or domain") + except Exception as e: + raise e + return True + + +def seal(rec: Record, private_key: PrivateKey) -> Envelope: + payload = rec.marshal_record() + payload_type = rec.codec() + domain = rec.domain() + + if not domain: + raise ErrEmptyDomain() + if not payload_type: + raise ErrEmptyPayloadType() + + unsigned = make_unsigned(domain=domain, payload_type=payload_type, payload=payload) + sig = private_key.sign(unsigned) + return Envelope( + public_key=private_key.get_public_key(), + payload_type=payload_type, + raw_payload=payload, + signature=sig + ) + + +def consume_envelope(data: bytes, domain: str) -> Tuple[Envelope, Record]: + msg = Envelope.unmarshal_envelope(data) + msg.validate(domain) + rec = msg.record() + return msg, rec + +def encode_uvarint(value: int) -> bytes: + """Encode an int as protobuf-style unsigned varint.""" + buf = bytearray() + while value > 0x7F: + buf.append((value & 0x7F) | 0x80) + value >>= 7 + buf.append(value) + return bytes(buf) + + +def make_unsigned(domain: str, payload_type: bytes, payload: bytes) -> bytes: + fields = [domain.encode("utf-8"), payload_type, payload] + b = bytearray() + for f in fields: + b.extend(encode_uvarint(len(f))) + b.extend(f) + return bytes(b) diff --git a/libp2p/record/exceptions.py b/libp2p/record/exceptions.py new file mode 100644 index 000000000..b570dfb6f --- /dev/null +++ b/libp2p/record/exceptions.py @@ -0,0 +1,10 @@ + +class ErrEmptyDomain(Exception): + pass + +class ErrEmptyPayloadType(Exception): + pass + +class ErrInvalidSignature(Exception): + pass + diff --git a/libp2p/record/libp2p_record.py b/libp2p/record/libp2p_record.py new file mode 100644 index 000000000..766931eeb --- /dev/null +++ b/libp2p/record/libp2p_record.py @@ -0,0 +1,12 @@ +from datetime import date + + +class Libp2pRecord: + def __init__(self, key: bytes, value: bytes, time_received: date) -> None: + self.key = key + self.value = value + self.time_received = date + + def serialize(self) -> bytes: + return self.key + \ No newline at end of file diff --git a/libp2p/record/pb/__init__.py b/libp2p/record/pb/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/libp2p/record/pb/envelope.proto b/libp2p/record/pb/envelope.proto new file mode 100644 index 000000000..bc9ed2a41 --- /dev/null +++ b/libp2p/record/pb/envelope.proto @@ -0,0 +1,23 @@ +syntax = "proto3"; + +package record.pb; + +import "libp2p/crypto/pb/crypto.proto"; + +message Envelope { + // public_key is the public key of the keypair the enclosed payload was + // signed with. + crypto.pb.PublicKey public_key = 1; + + // payload_type encodes the type of payload, so that it can be deserialized + // deterministically. + bytes payload_type = 2; + + // payload is the actual payload carried inside this envelope. + bytes payload = 3; + + // signature is the signature produced by the private key corresponding to + // the enclosed public key, over the payload, prefixing a domain string for + // additional security. + bytes signature = 4; +} diff --git a/libp2p/record/pb/envelope_pb2.py b/libp2p/record/pb/envelope_pb2.py new file mode 100644 index 000000000..f7207cec9 --- /dev/null +++ b/libp2p/record/pb/envelope_pb2.py @@ -0,0 +1,26 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: libp2p/record/pb/envelope.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from libp2p.crypto.pb import crypto_pb2 as libp2p_dot_crypto_dot_pb_dot_crypto__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1flibp2p/record/pb/envelope.proto\x12\trecord.pb\x1a\x1dlibp2p/crypto/pb/crypto.proto\"n\n\x08\x45nvelope\x12(\n\npublic_key\x18\x01 \x01(\x0b\x32\x14.crypto.pb.PublicKey\x12\x14\n\x0cpayload_type\x18\x02 \x01(\x0c\x12\x0f\n\x07payload\x18\x03 \x01(\x0c\x12\x11\n\tsignature\x18\x04 \x01(\x0c\x62\x06proto3') + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'libp2p.record.pb.envelope_pb2', globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + _ENVELOPE._serialized_start=77 + _ENVELOPE._serialized_end=187 +# @@protoc_insertion_point(module_scope) diff --git a/libp2p/record/pb/envelope_pb2.pyi b/libp2p/record/pb/envelope_pb2.pyi new file mode 100644 index 000000000..b4ea8d1c4 --- /dev/null +++ b/libp2p/record/pb/envelope_pb2.pyi @@ -0,0 +1,50 @@ +""" +@generated by mypy-protobuf. Do not edit manually! +isort:skip_file +""" + +import builtins +import google.protobuf.descriptor +import google.protobuf.message +import libp2p.crypto.pb.crypto_pb2 +import typing + +DESCRIPTOR: google.protobuf.descriptor.FileDescriptor + +@typing.final +class Envelope(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PUBLIC_KEY_FIELD_NUMBER: builtins.int + PAYLOAD_TYPE_FIELD_NUMBER: builtins.int + PAYLOAD_FIELD_NUMBER: builtins.int + SIGNATURE_FIELD_NUMBER: builtins.int + payload_type: builtins.bytes + """payload_type encodes the type of payload, so that it can be deserialized + deterministically. + """ + payload: builtins.bytes + """payload is the actual payload carried inside this envelope.""" + signature: builtins.bytes + """signature is the signature produced by the private key corresponding to + the enclosed public key, over the payload, prefixing a domain string for + additional security. + """ + @property + def public_key(self) -> libp2p.crypto.pb.crypto_pb2.PublicKey: + """public_key is the public key of the keypair the enclosed payload was + signed with. + """ + + def __init__( + self, + *, + public_key: libp2p.crypto.pb.crypto_pb2.PublicKey | None = ..., + payload_type: builtins.bytes = ..., + payload: builtins.bytes = ..., + signature: builtins.bytes = ..., + ) -> None: ... + def HasField(self, field_name: typing.Literal["public_key", b"public_key"]) -> builtins.bool: ... + def ClearField(self, field_name: typing.Literal["payload", b"payload", "payload_type", b"payload_type", "public_key", b"public_key", "signature", b"signature"]) -> None: ... + +global___Envelope = Envelope diff --git a/libp2p/record/record.py b/libp2p/record/record.py new file mode 100644 index 000000000..3ea07d535 --- /dev/null +++ b/libp2p/record/record.py @@ -0,0 +1,62 @@ +from abc import ABC, abstractmethod +from typing import Dict, Type, Optional + +# Registry of payload_type (bytes) -> Record class +_payload_type_registry: Dict[bytes, Type["Record"]] = {} + +class Record(ABC): + """ + Record represents a data type that can be used as the payload of an Envelope. + """ + + @abstractmethod + def domain(self) -> str: + """The signature domain (unique string for this record type).""" + pass + + @abstractmethod + def codec(self) -> bytes: + """Binary identifier (payload type).""" + pass + + @abstractmethod + def marshal_record(self) -> bytes: + """Serialize the record into bytes.""" + pass + + @abstractmethod + def unmarshal_record(self, data: bytes) -> None: + """Deserialize bytes into this record instance.""" + pass + + +def register_type(prototype: Record) -> None: + """ + Register a record type by its codec. + Should be called in module init where the Record is defined. + """ + codec = prototype.codec() + if not isinstance(codec, (bytes, bytearray)): + raise TypeError("codec() must return bytes") + _payload_type_registry[bytes(codec)] = type(prototype) + + +def blank_record_for_payload_type(payload_type: bytes) -> Record: + """ + Create a blank record for the given payload type. + """ + cls = _payload_type_registry.get(payload_type) + if cls is None: + raise ValueError("payload type is not registered") + return cls() # assumes no-arg constructor + + +def unmarshal_record_payload(payload_type: bytes, payload_bytes: bytes) -> Record: + """ + Given a payload type and bytes, return a fully unmarshalled Record. + """ + rec = blank_record_for_payload_type(payload_type) + rec.unmarshal_record(payload_bytes) + return rec + + \ No newline at end of file diff --git a/tests/core/record/test_envelope.py b/tests/core/record/test_envelope.py new file mode 100644 index 000000000..379ef4889 --- /dev/null +++ b/tests/core/record/test_envelope.py @@ -0,0 +1,128 @@ +import pytest +from unittest.mock import ( + patch +) +from libp2p.crypto.ed25519 import ( + create_new_key_pair +) +from libp2p.record.envelope import ( + Envelope, + seal, + consume_envelope, + make_unsigned +) +from libp2p.record.record import ( + Record +) +from libp2p.record.exceptions import ( + ErrEmptyDomain, + ErrEmptyPayloadType, + ErrInvalidSignature +) + +class DummyRecord(Record): + def __init__(self): + self._payload = b"dummy record" + self._codec = b"dummy" + self._domain = "dummy.domain" + + def marshal_record(self) -> bytes: + return self._payload + + def codec(self) -> bytes: + return self._codec + + def domain(self) -> str: + return self._domain + + def unmarshal_record(self, data: bytes) -> None: + self._payload = data + +@pytest.fixture +def key_pair(): + kp = create_new_key_pair() + return kp.private_key, kp.public_key + + +@pytest.fixture +def dummy_record(): + return DummyRecord() + + +def test_envelope_marshal_unmarshal(key_pair, dummy_record): + priv, pub = key_pair + env = seal(dummy_record, priv) + data = env.marshal() + env2 = Envelope.unmarshal_envelope(data) + assert env2.payload_type == env.payload_type + assert env2.raw_payload == env.raw_payload + assert env2.public_key.to_bytes() == env.public_key.to_bytes() + + +def test_envelope_equal(key_pair, dummy_record): + priv, _ = key_pair + + rec1 = DummyRecord() + rec2 = DummyRecord() # Copy but mutate payload for difference + rec2._payload = b"different dummy record" + + env1 = seal(rec1, priv) + env2 = seal(rec2, priv) + + assert not env1.equal(env2) + assert env1.equal(env1) + + +def test_validate_signature(key_pair, dummy_record): + priv, pub = key_pair + env = seal(dummy_record, priv) + + # Mock verify to always return True, bypassing the signature length issue in nacl + with patch.object(type(pub), 'verify', return_value=True): + # Should validate correctly + assert env.validate(dummy_record.domain()) + def mock_verify(data, sig): + expected_unsigned = make_unsigned(dummy_record.domain(), env.payload_type, env.raw_payload) + if data == expected_unsigned: + return True + return False + with patch.object(type(pub), 'verify', side_effect=mock_verify): + with pytest.raises(ErrInvalidSignature): + env.validate("wrong.domain") + +def test_validate_empty_domain(dummy_record, key_pair): + priv, _ = key_pair + rec = dummy_record + rec._domain = "" + with pytest.raises(ErrEmptyDomain): + seal(rec, priv) + + +def test_validate_empty_payload_type(dummy_record, key_pair): + priv, _ = key_pair + rec = dummy_record + rec._codec = b"" + with pytest.raises(ErrEmptyPayloadType): + seal(rec, priv) + + +def test_consume_envelope(key_pair, dummy_record): + priv, _ = key_pair + env = seal(dummy_record, priv) + data = env.marshal() + + with patch('libp2p.record.envelope.Envelope.validate') as mock_validate: + mock_validate.return_value = True + with patch('libp2p.record.envelope.unmarshal_record_payload') as mock_unmarshal: + mock_rec = DummyRecord() + mock_rec._payload = env.raw_payload + mock_unmarshal.return_value = mock_rec + + env2, rec = consume_envelope(data, dummy_record.domain()) + assert rec.marshal_record() == dummy_record.marshal_record() + assert env2.payload_type == env.payload_type + mock_validate.assert_called_once_with(dummy_record.domain()) + mock_unmarshal.assert_called_once_with( + payload_type=env.payload_type, + payload_bytes=env.raw_payload + ) \ No newline at end of file diff --git a/tests/core/record/test_record.py b/tests/core/record/test_record.py new file mode 100644 index 000000000..abad79925 --- /dev/null +++ b/tests/core/record/test_record.py @@ -0,0 +1,90 @@ +import pytest +from typing import cast +from libp2p.record.record import ( + Record, + register_type, + blank_record_for_payload_type, + unmarshal_record_payload, + _payload_type_registry +) + + +# Dummy concrete Record class for testing +class DummyRecord(Record): + def __init__(self, data: bytes = b""): + self.data = data + + def domain(self) -> str: + return "dummy" + + def codec(self) -> bytes: + return b"dummy-record" + + def marshal_record(self) -> bytes: + return self.data + + def unmarshal_record(self, data: bytes) -> None: + self.data = data + + +@pytest.fixture(autouse=True) +def setup_registry(): + # Clear registry before each test + _payload_type_registry.clear() + register_type(DummyRecord()) + + +def test_register_type_adds_to_registry(): + assert b"dummy-record" in _payload_type_registry + assert _payload_type_registry[b"dummy-record"] == DummyRecord + + +def test_register_type_invalid_codec(): + class BadRecord(Record): + def domain(self): return "bad" + def codec(self): return cast(bytes, "not-bytes") + def marshal_record(self): return b"" + def unmarshal_record(self, data: bytes): pass + + with pytest.raises(TypeError): + register_type(BadRecord()) + + +def test_blank_record_for_payload_type_returns_instance(): + rec = blank_record_for_payload_type(b"dummy-record") + assert isinstance(rec, DummyRecord) + assert rec.data == b"" + + +def test_blank_record_for_unknown_payload_type_raises(): + with pytest.raises(ValueError): + blank_record_for_payload_type(b"unknown-record") + + +def test_unmarshal_record_payload_sets_data_correctly(): + payload = b"testdata" + rec = unmarshal_record_payload(b"dummy-record", payload) + assert isinstance(rec, DummyRecord) + assert rec.data == payload + + +def test_unmarshal_record_payload_unknown_type_raises(): + with pytest.raises(ValueError): + unmarshal_record_payload(b"unknown", b"data") + + +def test_marshal_and_unmarshal_roundtrip(): + original = DummyRecord(b"hello") + marshaled = original.marshal_record() + new_rec = DummyRecord() + new_rec.unmarshal_record(marshaled) + assert new_rec.data == original.data + + +def test_multiple_records_independent(): + r1 = DummyRecord(b"one") + r2 = DummyRecord(b"two") + assert r1.data != r2.data + r1.unmarshal_record(b"updated") + assert r1.data == b"updated" + assert r2.data == b"two" From 2e51189f95ffee6c3009f3bf2c4be06ca66e1f55 Mon Sep 17 00:00:00 2001 From: Michael Eze Date: Sat, 13 Sep 2025 17:47:59 +0100 Subject: [PATCH 2/2] cleanup --- libp2p/record/record.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libp2p/record/record.py b/libp2p/record/record.py index 3ea07d535..acf81d220 100644 --- a/libp2p/record/record.py +++ b/libp2p/record/record.py @@ -33,7 +33,7 @@ def unmarshal_record(self, data: bytes) -> None: def register_type(prototype: Record) -> None: """ Register a record type by its codec. - Should be called in module init where the Record is defined. + Should be called in module init where the Record is defined just like the Go version. """ codec = prototype.codec() if not isinstance(codec, (bytes, bytearray)): @@ -48,7 +48,7 @@ def blank_record_for_payload_type(payload_type: bytes) -> Record: cls = _payload_type_registry.get(payload_type) if cls is None: raise ValueError("payload type is not registered") - return cls() # assumes no-arg constructor + return cls() def unmarshal_record_payload(payload_type: bytes, payload_bytes: bytes) -> Record: