astreum 0.3.16__py3-none-any.whl → 0.3.46__py3-none-any.whl

astreum/communication/handlers/object_response.py

@@ -2,6 +2,7 @@ import socket
 from enum import IntEnum
 from typing import Tuple, TYPE_CHECKING
 
+from ..outgoing_queue import enqueue_outgoing
 from ..models.message import Message, MessageTopic
 from ...storage.models.atom import Atom
 
@@ -109,7 +110,12 @@ def handle_object_response(node: "Node", peer: "Peer", message: Message) -> None
                 sender=node.relay_public_key,
             )
             obj_req_msg.encrypt(peer.shared_key_bytes)
-            node.outgoing_queue.put((obj_req_msg.to_bytes(), (provider_address, provider_port)))
+            enqueue_outgoing(
+                node,
+                (provider_address, provider_port),
+                message=obj_req_msg,
+                difficulty=1,
+            )
 
         case ObjectResponseType.OBJECT_NEAREST_PEER:
             node.logger.debug("Ignoring OBJECT_NEAREST_PEER response from %s", peer.address)

astreum/communication/handlers/ping.py

@@ -20,6 +20,15 @@ def handle_ping(node: "Node", peer: Peer, payload: bytes) -> None:
 
     peer.timestamp = datetime.now(timezone.utc)
     peer.latest_block = ping.latest_block
+    peer.difficulty = ping.difficulty
+    if peer.is_default_seed and ping.latest_block:
+        if getattr(node, "latest_block_hash", None) != ping.latest_block:
+            node.latest_block_hash = ping.latest_block
+            node.latest_block = None
+            node.logger.info(
+                "Updated latest block hash from default seed %s",
+                peer.address[0] if peer.address else "unknown",
+            )
 
     validation_route = node.validation_route
     if validation_route is None:

astreum/communication/handlers/route_request.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import socket
 
+from ..outgoing_queue import enqueue_outgoing
 from ..models.message import Message, MessageTopic
 from ..util import xor_distance
 
@@ -73,4 +74,9 @@ def handle_route_request(node: "Node", peer: "Peer", message: Message) -> None:
         sender=node.relay_public_key,
     )
     response.encrypt(peer.shared_key_bytes)
-    node.outgoing_queue.put((response.to_bytes(), peer.address))
+    enqueue_outgoing(
+        node,
+        peer.address,
+        message=response,
+        difficulty=peer.difficulty,
+    )

astreum/communication/handlers/route_response.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import socket
 
+from ..outgoing_queue import enqueue_outgoing
 from ..models.message import Message
 
 from typing import TYPE_CHECKING
@@ -50,4 +51,9 @@ def handle_route_response(node: "Node", peer: "Peer", message: Message) -> None:
         content=int(node.config["incoming_port"]).to_bytes(2, "big", signed=False),
     )
     for host, port in decoded_addresses:
-        node.outgoing_queue.put((handshake_message.to_bytes(), (host, port)))
+        enqueue_outgoing(
+            node,
+            (host, port),
+            message=handshake_message,
+            difficulty=1,
+        )

astreum/communication/incoming_queue.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Tuple
+
+from blake3 import blake3
+
+from .difficulty import message_difficulty
+from .message_pow import NONCE_SIZE, _leading_zero_bits
+
+if TYPE_CHECKING:
+    from .. import Node
+
+
+INCOMING_QUEUE_ITEM_OVERHEAD_BYTES = 6
+
+
+def enqueue_incoming(
+    node: "Node",
+    address: Tuple[str, int],
+    payload: bytes,
+) -> bool:
+    """Enqueue an incoming UDP payload while tracking queued bytes.
+    Increments `node.incoming_queue_size` by `len(payload) + 6` and enforces
+    `node.incoming_queue_size_limit` (bytes) as a soft cap by dropping enqueues that
+    would exceed the limit. If `node.incoming_queue_timeout` is > 0, it waits up to
+    that many seconds (using `communication_stop_event.wait`) for space before dropping.
+    """
+    required_difficulty = message_difficulty(node)
+    if len(payload) <= NONCE_SIZE:
+        node.logger.warning(
+            "Incoming payload too short for difficulty check (len=%s, required=%s)",
+            len(payload),
+            required_difficulty,
+        )
+        return False
+
+    nonce_bytes = payload[:NONCE_SIZE]
+    message_bytes = payload[NONCE_SIZE:]
+    digest = blake3(message_bytes + nonce_bytes).digest()
+    zeros = _leading_zero_bits(digest)
+    if zeros < required_difficulty:
+        node.logger.warning(
+            "Incoming payload failed difficulty check (zeros=%s required=%s bytes=%s)",
+            zeros,
+            required_difficulty,
+            len(payload),
+        )
+        return False
+
+    accounted_size = len(payload) + INCOMING_QUEUE_ITEM_OVERHEAD_BYTES
+    timeout = float(node.incoming_queue_timeout or 0)
+
+    with node.incoming_queue_size_lock:
+        current_size = int(node.incoming_queue_size)
+        limit = int(node.incoming_queue_size_limit)
+        projected_size = current_size + accounted_size
+        if projected_size > limit:
+            if timeout <= 0:
+                node.logger.warning(
+                    "Incoming queue size limit reached (%s > %s); dropping inbound payload (bytes=%s)",
+                    projected_size,
+                    limit,
+                    len(payload),
+                )
+                return False
+            wait_for_space = True
+        else:
+            node.incoming_queue_size = projected_size
+            wait_for_space = False
+
+    if wait_for_space:
+        if node.communication_stop_event.wait(timeout):
+            return False
+        with node.incoming_queue_size_lock:
+            current_size = int(node.incoming_queue_size)
+            limit = int(node.incoming_queue_size_limit)
+            projected_size = current_size + accounted_size
+            if projected_size > limit:
+                node.logger.warning(
+                    "Incoming queue still full after waiting %ss (%s > %s); dropping inbound payload (bytes=%s)",
+                    timeout,
+                    projected_size,
+                    limit,
+                    len(payload),
+                )
+                return False
+            node.incoming_queue_size = projected_size
+
+    try:
+        node.incoming_queue.put((message_bytes, address, accounted_size))
+    except Exception:
+        with node.incoming_queue_size_lock:
+            node.incoming_queue_size = max(0, int(node.incoming_queue_size) - accounted_size)
+        raise
+
+    return True

astreum/communication/message_pow.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from blake3 import blake3
+
+NONCE_SIZE = 8
+MAX_MESSAGE_NONCE = (1 << (NONCE_SIZE * 8)) - 1
+
+
+def _leading_zero_bits(buf: bytes) -> int:
+    """Return the number of leading zero bits in the provided buffer."""
+    zeros = 0
+    for byte in buf:
+        if byte == 0:
+            zeros += 8
+            continue
+        zeros += 8 - int(byte).bit_length()
+        break
+    return zeros
+
+
+def calculate_message_nonce(message_bytes: bytes, difficulty: int) -> int:
+    """Find a nonce such that blake3(message_bytes + nonce_bytes) meets difficulty.
+
+    message_bytes should exclude any nonce prefix that will be added on the wire.
+    """
+    target = max(1, int(difficulty))
+    nonce = 0
+    message_bytes = bytes(message_bytes)
+    while True:
+        if nonce > MAX_MESSAGE_NONCE:
+            raise ValueError("nonce search exhausted")
+        nonce_bytes = int(nonce).to_bytes(NONCE_SIZE, "big", signed=False)
+        digest = blake3(message_bytes + nonce_bytes).digest()
+        if _leading_zero_bits(digest) >= target:
+            return nonce
+        nonce += 1

astreum/communication/models/peer.py

@@ -13,11 +13,15 @@ class Peer:
         peer_public_key: X25519PublicKey,
         latest_block: Optional[bytes] = None,
         address: Optional[Tuple[str, int]] = None,
+        is_default_seed: bool = False,
+        difficulty: int = 1,
     ):
         self.shared_key_bytes = node_secret_key.exchange(peer_public_key)
         self.timestamp = datetime.now(timezone.utc)
         self.latest_block = latest_block
+        self.difficulty = max(1, int(difficulty or 1))
         self.address = address
+        self.is_default_seed = bool(is_default_seed)
         self.public_key_bytes = peer_public_key.public_bytes(
             encoding=serialization.Encoding.Raw,
             format=serialization.PublicFormat.Raw,

astreum/communication/models/ping.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Optional
 
 
 class PingFormatError(ValueError):
@@ -10,24 +11,44 @@ class PingFormatError(ValueError):
 @dataclass
 class Ping:
     is_validator: bool
-    latest_block: bytes
+    difficulty: int
+    latest_block: Optional[bytes]
 
-    PAYLOAD_SIZE = 33
+    PAYLOAD_SIZE = 34
+    ZERO_BLOCK = b"\x00" * 32
 
     def __post_init__(self) -> None:
-        lb = bytes(self.latest_block or b"")
+        self.difficulty = int(self.difficulty)
+        if self.difficulty < 1 or self.difficulty > 255:
+            raise ValueError("difficulty must be between 1 and 255")
+        if self.latest_block is None:
+            return
+        lb = bytes(self.latest_block)
         if len(lb) != 32:
             raise ValueError("latest_block must be exactly 32 bytes")
         self.latest_block = lb
 
     def to_bytes(self) -> bytes:
-        return (b"\x01" if self.is_validator else b"\x00") + self.latest_block
+        flag = b"\x01" if self.is_validator else b"\x00"
+        difficulty = bytes([self.difficulty])
+        latest_block = self.latest_block if self.latest_block is not None else self.ZERO_BLOCK
+        return flag + difficulty + latest_block
 
     @classmethod
     def from_bytes(cls, data: bytes) -> "Ping":
         if len(data) != cls.PAYLOAD_SIZE:
-            raise PingFormatError("ping payload must be exactly 33 bytes")
+            raise PingFormatError("ping payload must be 34 bytes")
         flag = data[0]
         if flag not in (0, 1):
             raise PingFormatError("ping validator flag must be 0 or 1")
-        return cls(is_validator=bool(flag), latest_block=data[1:])
+        difficulty = data[1]
+        if difficulty < 1:
+            raise PingFormatError("ping difficulty must be >= 1")
+        latest_block = data[2:]
+        if latest_block == cls.ZERO_BLOCK:
+            latest_block = None
+        return cls(
+            is_validator=bool(flag),
+            difficulty=difficulty,
+            latest_block=latest_block,
+        )

astreum/communication/models/route.py

@@ -42,6 +42,8 @@ class Route:
 
     def add_peer(self, peer_public_key: PeerKey, peer: Optional[Peer] = None):
         peer_public_key_bytes = self._normalize_peer_key(peer_public_key)
+        if peer_public_key_bytes == self.relay_public_key_bytes:
+            return
         bucket_idx = self._matching_leading_bits(self.relay_public_key_bytes, peer_public_key_bytes)
         if len(self.buckets[bucket_idx]) < self.bucket_size:
             bucket = self.buckets[bucket_idx]
@@ -52,6 +54,8 @@ class Route:
 
     def remove_peer(self, peer_public_key: PeerKey):
         peer_public_key_bytes = self._normalize_peer_key(peer_public_key)
+        if peer_public_key_bytes == self.relay_public_key_bytes:
+            return
         bucket_idx = self._matching_leading_bits(self.relay_public_key_bytes, peer_public_key_bytes)
         bucket = self.buckets.get(bucket_idx)
         if not bucket:

astreum/communication/{start.py → node.py}

@@ -1,19 +1,18 @@
-def connect_to_network_and_verify(self):
+
+def connect_node(self):
     """Initialize communication and consensus components, then load latest block state."""
+    if self.is_connected:
+        self.logger.debug("Node already connected; skipping communication setup")
+        return
+
     self.logger.info("Starting communication and consensus setup")
     try:
         from astreum.communication import communication_setup  # type: ignore
         communication_setup(node=self, config=self.config)
         self.logger.info("Communication setup completed")
-    except Exception:
-        self.logger.exception("Communication setup failed")
-
-    try:
-        from astreum.consensus import consensus_setup  # type: ignore
-        consensus_setup(node=self, config=self.config)
-        self.logger.info("Consensus setup completed")
-    except Exception:
-        self.logger.exception("Consensus setup failed")
+    except Exception as exc:
+        self.logger.exception("Communication setup failed: %s", exc)
+        return exc
 
     # Load latest_block_hash from config
     self.latest_block_hash = getattr(self, "latest_block_hash", None)
@@ -30,7 +29,7 @@ def connect_to_network_and_verify(self):
 
     if self.latest_block_hash and self.latest_block is None:
         try:
-            from astreum.consensus.models.block import Block
+            from astreum.validation.models.block import Block
             self.latest_block = Block.from_atom(self, self.latest_block_hash)
             self.logger.info("Loaded latest block %s from storage", self.latest_block_hash.hex())
         except Exception as exc:

astreum/communication/outgoing_queue.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, Tuple
+
+from .message_pow import NONCE_SIZE, calculate_message_nonce
+
+if TYPE_CHECKING:
+    from .models.message import Message
+    from .. import Node
+
+
+OUTGOING_QUEUE_ITEM_OVERHEAD_BYTES = 6
+def enqueue_outgoing(
+    node: "Node",
+    address: Tuple[str, int],
+    message: Optional["Message"] = None,
+    message_bytes: Optional[bytes] = None,
+    difficulty: int = 1,
+) -> bool:
+    """Enqueue an outgoing UDP payload while tracking queued bytes.
+    When used, it increments `node.outgoing_queue_size` by `len(payload) + 6` and enforces
+    `node.outgoing_queue_size_limit` (bytes) as a soft cap by dropping enqueues that
+    would exceed the limit. If `node.outgoing_queue_timeout` is > 0, it waits up to
+    that many seconds (using `communication_stop_event.wait`) for space before dropping.
+    """
+    if not node.is_connected:
+        raise RuntimeError("node is not connected; call node.connect() (communication_setup) first")
+
+    if message is not None and message_bytes is not None:
+        raise ValueError("Specify only one of message or message_bytes")
+
+    if message_bytes is not None:
+        payload = message_bytes
+    elif message is not None:
+        payload = message.to_bytes()
+    else:
+        raise ValueError("Either message or message_bytes must be provided")
+
+    try:
+        difficulty_value = int(difficulty)
+    except Exception:
+        difficulty_value = 1
+    if difficulty_value < 1:
+        difficulty_value = 1
+
+    try:
+        nonce = calculate_message_nonce(payload, difficulty_value)
+    except Exception as exc:
+        node.logger.warning(
+            "Failed generating message nonce (difficulty=%s bytes=%s): %s",
+            difficulty_value,
+            len(payload),
+            exc,
+        )
+        return False
+
+    payload = int(nonce).to_bytes(NONCE_SIZE, "big", signed=False) + payload
+
+    accounted_size = len(payload) + OUTGOING_QUEUE_ITEM_OVERHEAD_BYTES
+
+    timeout = float(node.outgoing_queue_timeout or 0)
+
+    with node.outgoing_queue_size_lock:
+        current_size = int(node.outgoing_queue_size)
+        limit = int(node.outgoing_queue_size_limit)
+        projected_size = current_size + accounted_size
+        if projected_size > limit:
+            if timeout <= 0:
+                node.logger.warning(
+                    "Outgoing queue size limit reached (%s > %s); dropping outbound payload (bytes=%s)",
+                    projected_size,
+                    limit,
+                    len(payload),
+                )
+                return False
+            wait_for_space = True
+        else:
+            node.outgoing_queue_size = projected_size
+            wait_for_space = False
+
+    if wait_for_space:
+        if node.communication_stop_event.wait(timeout):
+            return False
+        if not node.is_connected:
+            return False
+        with node.outgoing_queue_size_lock:
+            current_size = int(node.outgoing_queue_size)
+            limit = int(node.outgoing_queue_size_limit)
+            projected_size = current_size + accounted_size
+            if limit and projected_size > limit:
+                node.logger.warning(
+                    "Outgoing queue still full after waiting %ss (%s > %s); dropping outbound payload (bytes=%s)",
+                    timeout,
+                    projected_size,
+                    limit,
+                    len(payload),
+                )
+                return False
+            node.outgoing_queue_size = projected_size
+
+    try:
+        node.outgoing_queue.put((payload, address, accounted_size))
+    except Exception:
+        with node.outgoing_queue_size_lock:
+            node.outgoing_queue_size = max(0, int(node.outgoing_queue_size) - accounted_size)
+        raise
+
+    return True

astreum/communication/processors/incoming.py

@@ -1,30 +1,56 @@
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-from ..handlers.handshake import handle_handshake
+from __future__ import annotations
+
+import socket
+from queue import Empty
+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 ..incoming_queue import enqueue_incoming
 from ..models.message import Message, MessageTopic
 from ..models.peer import Peer
+from ..outgoing_queue import enqueue_outgoing
 from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PublicKey
 
 if TYPE_CHECKING:
-    from .. import Node
-
-
+    from .. import Node
+
+
 def process_incoming_messages(node: "Node") -> None:
     """Process incoming messages (placeholder)."""
-    while True:
+    stop = getattr(node, "communication_stop_event", None)
+    while stop is None or not stop.is_set():
         try:
-            data, addr = node.incoming_queue.get()
-        except Exception as exc:
+            item = node.incoming_queue.get(timeout=0.5)
+        except Empty:
+            continue
+        except Exception:
             node.logger.exception("Error taking from incoming queue")
             continue
 
+        data = None
+        addr = None
+        accounted_size = None
+
+        if isinstance(item, tuple) and len(item) == 3:
+            data, addr, accounted_size = item
+        else:
+            node.logger.warning("Incoming queue item has unexpected shape: %r", item)
+            continue
+
+        if stop is not None and stop.is_set():
+            if accounted_size is not None:
+                try:
+                    with node.incoming_queue_size_lock:
+                        node.incoming_queue_size = max(0, node.incoming_queue_size - int(accounted_size))
+                except Exception:
+                    node.logger.exception("Failed updating incoming_queue_size on shutdown")
+            break
+
         try:
             message = Message.from_bytes(data)
         except Exception as exc:
@@ -34,7 +60,7 @@ def process_incoming_messages(node: "Node") -> None:
         if message.handshake:
             if handle_handshake(node, addr, message):
                 continue
-        
+
         peer = None
         try:
             peer = node.get_peer(message.sender_bytes)
@@ -44,10 +70,13 @@ def process_incoming_messages(node: "Node") -> None:
             try:
                 peer_key = X25519PublicKey.from_public_bytes(message.sender_bytes)
                 host, port = addr[0], int(addr[1])
+                default_seed_ips = getattr(node, "default_seed_ips", None)
+                is_default_seed = bool(default_seed_ips) and host in default_seed_ips
                 peer = Peer(
                     node_secret_key=node.relay_secret_key,
                     peer_public_key=peer_key,
                     address=(host, port),
+                    is_default_seed=is_default_seed,
                 )
             except Exception:
                 peer = None
@@ -60,39 +89,83 @@ def process_incoming_messages(node: "Node") -> None:
         try:
             message.decrypt(peer.shared_key_bytes)
         except Exception as exc:
-            node.logger.warning("Error decrypting message from %s: %s", peer.address, exc)
+            node.logger.warning(
+                "Error decrypting message from %s (len=%s, enc_len=%s, exc=%s)",
+                peer.address,
+                len(data),
+                len(message.encrypted) if message.encrypted is not None else None,
+                exc,
+            )
+            try:
+                host, port = addr[0], int(addr[1])
+                handshake_message = Message(
+                    handshake=True,
+                    sender=node.relay_public_key,
+                    content=int(node.config["incoming_port"]).to_bytes(2, "big", signed=False),
+                )
+                enqueue_outgoing(
+                    node,
+                    (host, port),
+                    message=handshake_message,
+                    difficulty=1,
+                )
+            except Exception as handshake_exc:
+                node.logger.debug(
+                    "Failed queueing rekey handshake to %s: %s",
+                    addr,
+                    handshake_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)
+        try:
+            match message.topic:
+                case MessageTopic.PING:
+                    handle_ping(node, peer, message.content)
 
-            case MessageTopic.ROUTE_REQUEST:
-                handle_route_request(node, peer, message)
+                case MessageTopic.OBJECT_REQUEST:
+                    handle_object_request(node, peer, message)
 
-            case MessageTopic.ROUTE_RESPONSE:
-                handle_route_response(node, peer, message)
+                case MessageTopic.OBJECT_RESPONSE:
+                    handle_object_response(node, peer, message)
 
-            case MessageTopic.TRANSACTION:
-                if node.validation_secret_key is None:
-                    continue
-                node._validation_transaction_queue.put(message.content)
+                case MessageTopic.ROUTE_REQUEST:
+                    handle_route_request(node, peer, message)
 
-            case _:
-                continue
+                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)
 
-def populate_incoming_messages(node: "Node") -> None:
+                case _:
+                    continue
+        finally:
+            if accounted_size is not None:
+                try:
+                    with node.incoming_queue_size_lock:
+                        node.incoming_queue_size = max(0, node.incoming_queue_size - int(accounted_size))
+                except Exception:
+                    node.logger.exception("Failed updating incoming_queue_size")
+
+    node.logger.info("Incoming message processor stopped")
+
+
+def populate_incoming_messages(node: "Node") -> None:
     """Receive UDP packets and feed the incoming queue."""
-    while True:
+    stop = getattr(node, "communication_stop_event", None)
+    while stop is None or not stop.is_set():
         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)
+            enqueue_incoming(node, addr, payload=data)
+        except socket.timeout:
+            continue
+        except OSError:
+            if stop is not None and stop.is_set():
+                break
+            node.logger.warning("Error populating incoming queue: socket closed")
+        except Exception as exc:
+            node.logger.warning("Error populating incoming queue: %s", exc)
+
+    node.logger.info("Incoming message populator stopped")