astreum 0.3.1__py3-none-any.whl → 0.3.9__py3-none-any.whl

astreum/communication/models/message.py

@@ -1,7 +1,9 @@
+import os
 from enum import IntEnum
 from typing import Optional
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PublicKey
+from astreum.crypto import chacha20poly1305
 
 class MessageTopic(IntEnum):
     PING = 0
@@ -10,92 +12,113 @@ class MessageTopic(IntEnum):
     ROUTE_REQUEST = 3
     ROUTE_RESPONSE = 4
     TRANSACTION = 5
-    STORAGE_REQUEST = 6
-    
 
-class Message:
-    handshake: bool
-    sender: Optional[X25519PublicKey]
-
-    topic: Optional[MessageTopic]
-    content: bytes
 
+class Message:
     def __init__(
         self,
         *,
         handshake: bool = False,
         sender: Optional[X25519PublicKey] = None,
         topic: Optional[MessageTopic] = None,
-        content: bytes = b"",
+        content: Optional[bytes] = None,
         body: Optional[bytes] = None,
+        sender_bytes: Optional[bytes] = None,
+        encrypted: Optional[bytes] = None,
     ) -> None:
         if body is not None:
-            if content and content != b"":
+            if content is not None and content != b"":
                 raise ValueError("specify only one of 'content' or 'body'")
             content = body
 
         self.handshake = handshake
-        self.sender = sender
         self.topic = topic
-        self.content = content or b""
+        self.content = content if content is not None else b""
+        self.encrypted = encrypted
 
         if self.handshake:
-            if self.sender is None:
-                raise ValueError("handshake Message requires a sender public key")
+            if sender_bytes is None and sender is None:
+                raise ValueError("handshake Message requires a sender public key or sender bytes")
             self.topic = None
-            self.content = b""
         else:
-            if self.topic is None:
-                raise ValueError("non-handshake Message requires a topic")
+            if self.topic is None and self.encrypted is None:
+                raise ValueError("non-handshake Message requires a topic or encrypted payload")
+            if sender_bytes is None and sender is None:
+                raise ValueError("non-handshake Message requires a sender public key or sender bytes")
 
-    def to_bytes(self):
-        if self.handshake:
-            # handshake byte (1) + raw public key bytes
-            return bytes([1]) + self.sender.public_bytes(
+        if sender_bytes is not None:
+            self.sender_bytes = sender_bytes
+        else:
+            if sender is None:
+                raise ValueError("sender public key required to derive sender bytes")
+            self.sender_bytes = sender.public_bytes(
                 encoding=serialization.Encoding.Raw,
-                format=serialization.PublicFormat.Raw
+                format=serialization.PublicFormat.Raw,
             )
+
+    def to_bytes(self):
+        if self.handshake:
+            # handshake byte (1) + raw public key bytes + payload
+            return bytes([1]) + self.sender_bytes + self.content
         else:
-            # normal message: 0 + topic + content
-            return bytes([0, self.topic.value]) + self.content
+            # normal message: 0 + sender + encrypted payload (nonce + ciphertext)
+            if not self.encrypted:
+                raise ValueError("non-handshake Message missing encrypted payload; call encrypt() first")
+            return bytes([0]) + self.sender_bytes + self.encrypted
 
     @classmethod
     def from_bytes(cls, data: bytes) -> "Message":
         if len(data) < 1:
             raise ValueError("Cannot parse Message: no data")
-        flag = data[0]
-        # create empty instance
-        msg = cls.__new__(cls)
-
-        if flag == 1:
-            # handshake message: the rest is the peer’s public key
-            key_bytes = data[1:]
-            if not key_bytes:
-                raise ValueError("Handshake message missing sender public key bytes")
-            try:
-                sender = X25519PublicKey.from_public_bytes(key_bytes)
-            except ValueError:
-                raise ValueError("Invalid public key bytes")
-            if sender is None:
-                raise ValueError("Handshake message missing sender public key")
-            msg.handshake = True
-            msg.sender = sender
-            msg.topic = None
-            msg.content = b''
-        elif flag == 0:
-            # normal message: next byte is topic, rest is content
-            if len(data) < 2:
-                raise ValueError("Cannot parse Message: missing topic byte")
-            topic_val = data[1]
-            try:
-                topic = MessageTopic(topic_val)
-            except ValueError:
-                raise ValueError(f"Unknown MessageTopic: {topic_val}")
-            msg.handshake = False
-            msg.sender = None
-            msg.topic = topic
-            msg.content = data[2:]
+
+        if len(data) < 33:
+            raise ValueError("Cannot parse Message: missing sender bytes")
+
+        if data[0] == 1:
+            return Message(
+                handshake=True,
+                sender_bytes=data[1:33],
+                content=data[33:],
+            )
+
         else:
-            raise ValueError(f"Invalid handshake flag: {flag}")
+            if len(data) <= 33:
+                raise ValueError("Cannot parse Message: missing encrypted payload")
+
+            return Message(
+                handshake=False,
+                sender_bytes=data[1:33],
+                encrypted=data[33:],
+            )
+
+    def encrypt(self, shared_key_bytes: bytes) -> None:
+        if self.handshake:
+            return
+        
+        if len(shared_key_bytes) != 32:
+            raise ValueError("Shared key must be 32 bytes for ChaCha20-Poly1305")
+        
+        if self.topic is None:
+            raise ValueError("Cannot encrypt message without a topic")
+
+        nonce = os.urandom(12)
+        data_to_encrypt = bytes([self.topic.value]) + self.content
+        ciphertext = chacha20poly1305.encrypt(shared_key_bytes, nonce, data_to_encrypt)
+        self.encrypted = nonce + ciphertext
+
+    def decrypt(self, shared_key_bytes: bytes) -> None:
+        if self.handshake:
+            return
+        
+        if len(shared_key_bytes) != 32:
+            raise ValueError("Shared key must be 32 bytes for ChaCha20-Poly1305")
+        
+        if not self.encrypted or len(self.encrypted) < 13:
+            raise ValueError("Encrypted content missing or too short")
 
-        return msg
+        nonce = self.encrypted[:12]
+        ciphertext = self.encrypted[12:]
+        decrypted = chacha20poly1305.decrypt(shared_key_bytes, nonce, ciphertext)
+        topic_value = decrypted[0]
+        self.topic = MessageTopic(topic_value)
+        self.content = decrypted[1:]

astreum/communication/models/peer.py

@@ -1,23 +1,51 @@
 from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
 from cryptography.hazmat.primitives import serialization
 from datetime import datetime, timezone
-from typing import Optional, Tuple
+from typing import Optional, Tuple, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ...node import Node
 
 class Peer:
-    shared_key: bytes
-    timestamp: datetime
-    latest_block: bytes
-    address: Optional[Tuple[str, int]]
-    public_key: X25519PublicKey
-    public_key_bytes: bytes
-
-    def __init__(self, my_sec_key: X25519PrivateKey, peer_pub_key: X25519PublicKey):
-        self.shared_key = my_sec_key.exchange(peer_pub_key)
+    def __init__(
+        self,
+        node_secret_key: X25519PrivateKey,
+        peer_public_key: X25519PublicKey,
+        latest_block: Optional[bytes] = None,
+        address: Optional[Tuple[str, int]] = None,
+    ):
+        self.shared_key_bytes = node_secret_key.exchange(peer_public_key)
         self.timestamp = datetime.now(timezone.utc)
-        self.latest_block = b""
-        self.address = None
-        self.public_key = peer_pub_key
-        self.public_key_bytes = peer_pub_key.public_bytes(
+        self.latest_block = latest_block
+        self.address = address
+        self.public_key_bytes = peer_public_key.public_bytes(
             encoding=serialization.Encoding.Raw,
             format=serialization.PublicFormat.Raw,
         )
+
+
+def add_peer(node: "Node", peer_key, peer: "Peer") -> "Peer":
+    """Register a peer entry on the node under lock."""
+    with node.peers_lock:
+        node.peers[peer_key] = peer
+    return peer
+
+
+def replace_peer(node: "Node", old_key, peer_key, peer: "Peer") -> "Peer":
+    """Replace an existing peer entry (if any) with a new one under lock."""
+    with node.peers_lock:
+        node.peers.pop(old_key, None)
+        node.peers[peer_key] = peer
+    return peer
+
+
+def get_peer(node: "Node", peer_key):
+    """Retrieve a peer entry under lock."""
+    with node.peers_lock:
+        return node.peers.get(peer_key)
+
+
+def remove_peer(node: "Node", peer_key):
+    """Remove a peer entry under lock."""
+    with node.peers_lock:
+        return node.peers.pop(peer_key, None)

astreum/communication/models/route.py

@@ -2,6 +2,7 @@ from typing import Dict, List, Optional, Union
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PublicKey
 from .peer import Peer
+from ..util import xor_distance
 
 PeerKey = Union[X25519PublicKey, bytes, bytearray]
 
@@ -39,12 +40,6 @@ class Route:
             return key_bytes
         raise TypeError("peer_public_key must be raw bytes or X25519PublicKey")
 
-    @staticmethod
-    def _xor_distance(a: bytes, b: bytes) -> int:
-        if len(a) != len(b):
-            raise ValueError("xor distance requires equal-length operands")
-        return int.from_bytes(bytes(x ^ y for x, y in zip(a, b)), "big", signed=False)
-
     def add_peer(self, peer_public_key: PeerKey, peer: Optional[Peer] = None):
         peer_public_key_bytes = self._normalize_peer_key(peer_public_key)
         bucket_idx = self._matching_leading_bits(self.relay_public_key_bytes, peer_public_key_bytes)
@@ -82,7 +77,7 @@ class Route:
         for bucket in self.buckets.values():
             for peer_key in bucket:
                 try:
-                    distance = self._xor_distance(target, peer_key)
+                    distance = xor_distance(target, peer_key)
                 except ValueError:
                     continue
                 if closest_distance is None or distance < closest_distance:

astreum/communication/processors/incoming.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ..handlers.handshake import handle_handshake
+from ..handlers.object_request import handle_object_request
+from ..handlers.object_response import handle_object_response
+from ..handlers.ping import handle_ping
+from ..handlers.route_request import handle_route_request
+from ..handlers.route_response import handle_route_response
+from ..models.message import Message, MessageTopic
+from ..models.peer import Peer
+from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PublicKey
+
+if TYPE_CHECKING:
+    from .. import Node
+
+
+def process_incoming_messages(node: "Node") -> None:
+    """Process incoming messages (placeholder)."""
+    while True:
+        try:
+            data, addr = node.incoming_queue.get()
+        except Exception as exc:
+            node.logger.exception("Error taking from incoming queue")
+            continue
+
+        try:
+            message = Message.from_bytes(data)
+        except Exception as exc:
+            node.logger.warning("Error decoding message: %s", exc)
+            continue
+
+        if message.handshake:
+            if handle_handshake(node, addr, message):
+                continue
+        
+        peer = None
+        try:
+            peer = node.get_peer(message.sender_bytes)
+        except Exception:
+            peer = None
+        if peer is None:
+            try:
+                peer_key = X25519PublicKey.from_public_bytes(message.sender_bytes)
+                host, port = addr[0], int(addr[1])
+                peer = Peer(
+                    node_secret_key=node.relay_secret_key,
+                    peer_public_key=peer_key,
+                    address=(host, port),
+                )
+            except Exception:
+                peer = None
+
+        if peer is None:
+            node.logger.debug("Unable to resolve peer for message from %s", addr)
+            continue
+
+        # decrypt message payload before dispatch
+        try:
+            message.decrypt(peer.shared_key_bytes)
+        except Exception as exc:
+            node.logger.warning("Error decrypting message from %s: %s", peer.address, exc)
+            continue
+
+        match message.topic:
+            case MessageTopic.PING:
+                handle_ping(node, peer, message.content)
+
+            case MessageTopic.OBJECT_REQUEST:
+                handle_object_request(node, peer, message)
+
+            case MessageTopic.OBJECT_RESPONSE:
+                handle_object_response(node, peer, message)
+
+            case MessageTopic.ROUTE_REQUEST:
+                handle_route_request(node, peer, message)
+
+            case MessageTopic.ROUTE_RESPONSE:
+                handle_route_response(node, peer, message)
+
+            case MessageTopic.TRANSACTION:
+                if node.validation_secret_key is None:
+                    continue
+                node._validation_transaction_queue.put(message.content)
+
+            case _:
+                continue
+
+
+def populate_incoming_messages(node: "Node") -> None:
+    """Receive UDP packets and feed the incoming queue."""
+    while True:
+        try:
+            data, addr = node.incoming_socket.recvfrom(4096)
+            node.incoming_queue.put((data, addr))
+        except Exception as exc:
+            node.logger.warning("Error populating incoming queue: %s", exc)

astreum/communication/processors/outgoing.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Tuple
+
+if TYPE_CHECKING:
+    from .. import Node
+
+def process_outgoing_messages(node: "Node") -> None:
+    """Send queued outbound packets."""
+    while True:
+        try:
+            payload, addr = node.outgoing_queue.get()
+        except Exception:
+            node.logger.exception("Error taking from outgoing queue")
+            continue
+
+        try:
+            node.outgoing_socket.sendto(payload, addr)
+        except Exception as exc:
+            node.logger.warning("Error sending message to %s: %s", addr, exc)

astreum/communication/setup.py

@@ -13,10 +13,11 @@ if TYPE_CHECKING:
     from .. import Node
 
 from . import Route, Message
-from .handlers.handshake import handle_handshake
-from .handlers.ping import handle_ping
-from .handlers.storage_request import handle_storage_request
-from .models.message import MessageTopic
+from .processors.incoming import (
+    process_incoming_messages,
+    populate_incoming_messages,
+)
+from .processors.outgoing import process_outgoing_messages
 from .util import address_str_to_host_and_port
 from ..utils.bytes import hex_to_bytes
 
@@ -40,77 +41,15 @@ def make_routes(
     val_rt  = Route(val_sk.public_key()) if val_sk else None
     return peer_rt, val_rt
 
-def setup_outgoing(
-    use_ipv6: bool
-) -> Tuple[socket.socket, Queue, threading.Thread]:
-    fam  = socket.AF_INET6 if use_ipv6 else socket.AF_INET
-    sock = socket.socket(fam, socket.SOCK_DGRAM)
-    q    = Queue()
-    thr  = threading.Thread(target=lambda: None, daemon=True)
-    thr.start()
-    return sock, q, thr
-
 def make_maps():
     """Empty lookup maps: peers and addresses."""
     return
 
 
-def process_incoming_messages(node: "Node") -> None:
-    """Process incoming messages (placeholder)."""
-    node_logger = node.logger
-    while True:
-        try:
-            data, addr = node.incoming_queue.get()
-        except Exception as exc:
-            node_logger.exception("Error taking from incoming queue")
-            continue
-
-        try:
-            message = Message.from_bytes(data)
-        except Exception as exc:
-            node_logger.warning("Error decoding message: %s", exc)
-            continue
-
-        if message.handshake:
-            if handle_handshake(node, addr, message):
-                continue
-
-        match message.topic:
-            case MessageTopic.PING:
-                handle_ping(node, addr, message.content)
-            case MessageTopic.OBJECT_REQUEST:
-                pass
-            case MessageTopic.OBJECT_RESPONSE:
-                pass
-            case MessageTopic.ROUTE_REQUEST:
-                pass
-            case MessageTopic.ROUTE_RESPONSE:
-                pass
-            case MessageTopic.TRANSACTION:
-                if node.validation_secret_key is None:
-                    continue
-                node._validation_transaction_queue.put(message.content)
-
-            case MessageTopic.STORAGE_REQUEST:
-                handle_storage_request(node, addr, message)
-            
-            case _:
-                continue
-
-
-def populate_incoming_messages(node: "Node") -> None:
-    """Receive UDP packets and feed the incoming queue (placeholder)."""
-    node_logger = node.logger
-    while True:
-        try:
-            data, addr = node.incoming_socket.recvfrom(4096)
-            node.incoming_queue.put((data, addr))
-        except Exception as exc:
-            node_logger.warning("Error populating incoming queue: %s", exc)
-
 def communication_setup(node: "Node", config: dict):
     node.logger.info("Setting up node communication")
     node.use_ipv6              = config.get('use_ipv6', False)
+    node.peers_lock = threading.RLock()
 
     # key loading
     node.relay_secret_key      = load_x25519(config.get('relay_secret_key'))
@@ -118,6 +57,10 @@ def communication_setup(node: "Node", config: dict):
 
     # derive pubs + routes
     node.relay_public_key      = node.relay_secret_key.public_key()
+    node.relay_public_key_bytes = node.relay_public_key.public_bytes(
+        encoding=serialization.Encoding.Raw,
+        format=serialization.PublicFormat.Raw,
+    )
     node.validation_public_key = (
         node.validation_secret_key.public_key().public_bytes(
             encoding=serialization.Encoding.Raw,
@@ -131,6 +74,11 @@ def communication_setup(node: "Node", config: dict):
         node.validation_secret_key
     )
 
+    # connection state & atom request tracking
+    node.is_connected = False
+    node.atom_requests = set()
+    node.atom_requests_lock = threading.RLock()
+
     # sockets + queues + threads
     incoming_port = config.get('incoming_port', 7373)
     fam = socket.AF_INET6 if node.use_ipv6 else socket.AF_INET
@@ -158,20 +106,29 @@ def communication_setup(node: "Node", config: dict):
     node.incoming_populate_thread.start()
     node.incoming_process_thread.start()
 
-    (node.outgoing_socket,
-        node.outgoing_queue,
-        node.outgoing_thread
-    ) = setup_outgoing(node.use_ipv6)
+    node.outgoing_socket = socket.socket(
+        socket.AF_INET6 if node.use_ipv6 else socket.AF_INET,
+        socket.SOCK_DGRAM,
+    )
+    node.outgoing_queue = Queue()
+
+    node.outgoing_thread = threading.Thread(
+        target=process_outgoing_messages,
+        args=(node,),
+        daemon=True,
+    )
+    node.outgoing_thread.start()
 
     # other workers & maps
-    node.object_request_queue = Queue()
+    # track atom requests we initiated; guarded by atom_requests_lock on the node
     node.peer_manager_thread  = threading.Thread(
         target=node._relay_peer_manager,
         daemon=True
     )
     node.peer_manager_thread.start()
 
-    node.peers, node.addresses = {}, {} # peers: Dict[bytes,Peer], addresses: Dict[(str,int),bytes]
+    with node.peers_lock:
+        node.peers, node.addresses = {}, {} # peers: Dict[bytes,Peer], addresses: Dict[(str,int),bytes]
 
     latest_block_hex = config.get("latest_block_hash")
     if latest_block_hex:
@@ -192,8 +149,11 @@ def communication_setup(node: "Node", config: dict):
             node.logger.warning("Invalid bootstrap address %s: %s", addr, exc)
             continue
 
-        handshake_message = Message(handshake=True, sender=node.relay_public_key)
-        
+        handshake_message = Message(
+            handshake=True,
+            sender=node.relay_public_key,
+            content=int(node.config["incoming_port"]).to_bytes(2, "big", signed=False),
+        )
         node.outgoing_queue.put((handshake_message.to_bytes(), (host, port)))
         node.logger.info("Sent bootstrap handshake to %s:%s", host, port)
 
@@ -203,3 +163,4 @@ def communication_setup(node: "Node", config: dict):
         node.outgoing_socket is not None,
         len(bootstrap_peers),
     )
+    node.is_connected = True

astreum/communication/start.py

@@ -1,20 +1,19 @@
 def connect_to_network_and_verify(self):
     """Initialize communication and consensus components, then load latest block state."""
-    node_logger = self.logger
-    node_logger.info("Starting communication and consensus setup")
+    self.logger.info("Starting communication and consensus setup")
     try:
         from astreum.communication import communication_setup  # type: ignore
         communication_setup(node=self, config=self.config)
-        node_logger.info("Communication setup completed")
+        self.logger.info("Communication setup completed")
     except Exception:
-        node_logger.exception("Communication setup failed")
+        self.logger.exception("Communication setup failed")
 
     try:
         from astreum.consensus import consensus_setup  # type: ignore
         consensus_setup(node=self, config=self.config)
-        node_logger.info("Consensus setup completed")
+        self.logger.info("Consensus setup completed")
     except Exception:
-        node_logger.exception("Consensus setup failed")
+        self.logger.exception("Consensus setup failed")
 
     # Load latest_block_hash from config
     self.latest_block_hash = getattr(self, "latest_block_hash", None)
@@ -25,14 +24,14 @@ def connect_to_network_and_verify(self):
         try:
             from astreum.utils.bytes import hex_to_bytes
             self.latest_block_hash = hex_to_bytes(latest_block_hex, expected_length=32)
-            node_logger.debug("Loaded latest_block_hash override from config")
+            self.logger.debug("Loaded latest_block_hash override from config")
         except Exception as exc:
-            node_logger.error("Invalid latest_block_hash in config: %s", exc)
+            self.logger.error("Invalid latest_block_hash in config: %s", exc)
 
     if self.latest_block_hash and self.latest_block is None:
         try:
             from astreum.consensus.models.block import Block
             self.latest_block = Block.from_atom(self, self.latest_block_hash)
-            node_logger.info("Loaded latest block %s from storage", self.latest_block_hash.hex())
+            self.logger.info("Loaded latest block %s from storage", self.latest_block_hash.hex())
         except Exception as exc:
-            node_logger.warning("Could not load latest block from storage: %s", exc)
+            self.logger.warning("Could not load latest block from storage: %s", exc)

astreum/communication/util.py

@@ -40,3 +40,10 @@ def address_str_to_host_and_port(address: str) -> Tuple[str, int]:
         raise ValueError(f"port out of range: {port}")
 
     return host, port
+
+
+def xor_distance(a: bytes, b: bytes) -> int:
+    """Return the unsigned integer XOR distance between two equal-length identifiers."""
+    if len(a) != len(b):
+        raise ValueError("xor distance requires operands of equal length")
+    return int.from_bytes(bytes(x ^ y for x, y in zip(a, b)), "big", signed=False)