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

astreum/communication/handlers/object_request.py

@@ -1,176 +1,264 @@
-import logging
-import socket
+import logging
+import socket
 from enum import IntEnum
-from typing import TYPE_CHECKING, Tuple
+from typing import Optional, TYPE_CHECKING, Tuple
 
-from .object_response import ObjectResponse, ObjectResponseType
+from .object_response import (
+    ObjectResponse,
+    ObjectResponseType,
+    OBJECT_FOUND_ATOM_PAYLOAD,
+    OBJECT_FOUND_LIST_PAYLOAD,
+    encode_object_found_atom_payload,
+    encode_object_found_list_payload,
+)
+from ..outgoing_queue import enqueue_outgoing
 from ..models.message import Message, MessageTopic
 from ..util import xor_distance
-
-if TYPE_CHECKING:
-    from .. import Node
-    from ..models.peer import Peer
-
-
-class ObjectRequestType(IntEnum):
-    OBJECT_GET = 0
-    OBJECT_PUT = 1
-
-
+from ...storage.providers import provider_id_for_payload, provider_payload_for_id
+
+if TYPE_CHECKING:
+    from .. import Node
+    from ..models.peer import Peer
+
+
+class ObjectRequestType(IntEnum):
+    OBJECT_GET = 0
+    OBJECT_PUT = 1
+
+
 class ObjectRequest:
     type: ObjectRequestType
     data: bytes
     atom_id: bytes
+    payload_type: Optional[int]
 
-    def __init__(self, type: ObjectRequestType, data: bytes, atom_id: bytes = None):
+    def __init__(
+        self,
+        type: ObjectRequestType,
+        data: bytes = b"",
+        atom_id: bytes = None,
+        payload_type: Optional[int] = None,
+    ):
         self.type = type
         self.data = data
         self.atom_id = atom_id
+        self.payload_type = payload_type
 
     def to_bytes(self):
-        return bytes([self.type.value]) + self.atom_id + self.data
+        if self.type == ObjectRequestType.OBJECT_GET and self.payload_type is not None:
+            payload = bytes([self.payload_type]) + self.data
+        else:
+            payload = self.data
+        return bytes([self.type.value]) + self.atom_id + payload
 
     @classmethod
     def from_bytes(cls, data: bytes) -> "ObjectRequest":
         # need at least 1 byte for type + 32 bytes for hash
-        if len(data) < 1 + 32:
-            raise ValueError(f"Too short for ObjectRequest ({len(data)} bytes)")
-
-        type_val = data[0]
-        try:
-            req_type = ObjectRequestType(type_val)
-        except ValueError:
-            raise ValueError(f"Unknown ObjectRequestType: {type_val!r}")
-
+        if len(data) < 1 + 32:
+            raise ValueError(f"Too short for ObjectRequest ({len(data)} bytes)")
+
+        type_val = data[0]
+        try:
+            req_type = ObjectRequestType(type_val)
+        except ValueError:
+            raise ValueError(f"Unknown ObjectRequestType: {type_val!r}")
+
         atom_id_bytes = data[1:33]
-        payload    = data[33:]
+        payload = data[33:]
+        if req_type == ObjectRequestType.OBJECT_GET:
+            if payload:
+                payload_type = payload[0]
+                payload = payload[1:]
+            else:
+                payload_type = None
+            return cls(req_type, payload, atom_id_bytes, payload_type=payload_type)
         return cls(req_type, payload, atom_id_bytes)
-
-
-def encode_peer_contact_bytes(peer: "Peer") -> bytes:
-    """Return a fixed-width peer contact payload (32-byte key + IPv4 + port)."""
-    host, port = peer.address
-    key_bytes = peer.public_key_bytes
-    try:
-        ip_bytes = socket.inet_aton(host)
-    except OSError as exc:  # pragma: no cover - inet_aton raises for invalid hosts
-        raise ValueError(f"invalid IPv4 address: {host}") from exc
-    if not (0 <= port <= 0xFFFF):
-        raise ValueError(f"port out of range (0-65535): {port}")
-    port_bytes = int(port).to_bytes(2, "big", signed=False)
-    return key_bytes + ip_bytes + port_bytes
-
-
-def handle_object_request(node: "Node", peer: "Peer", message: Message) -> None:
-    if message.content is None:
-        node.logger.warning("OBJECT_REQUEST from %s missing content", peer.address)
-        return
-
-    try:
-        object_request = ObjectRequest.from_bytes(message.content)
-    except Exception as exc:
-        node.logger.warning("Error decoding OBJECT_REQUEST from %s: %s", peer.address, exc)
-        return
-
-    match object_request.type:
+
+
+def encode_peer_contact_bytes(peer: "Peer") -> bytes:
+    """Return a fixed-width peer contact payload (32-byte key + IPv4 + port)."""
+    host, port = peer.address
+    key_bytes = peer.public_key_bytes
+    try:
+        ip_bytes = socket.inet_aton(host)
+    except OSError as exc:  # pragma: no cover - inet_aton raises for invalid hosts
+        raise ValueError(f"invalid IPv4 address: {host}") from exc
+    if not (0 <= port <= 0xFFFF):
+        raise ValueError(f"port out of range (0-65535): {port}")
+    port_bytes = int(port).to_bytes(2, "big", signed=False)
+    return key_bytes + ip_bytes + port_bytes
+
+
+def handle_object_request(node: "Node", peer: "Peer", message: Message) -> None:
+    if message.content is None:
+        node.logger.warning("OBJECT_REQUEST from %s missing content", peer.address)
+        return
+
+    try:
+        object_request = ObjectRequest.from_bytes(message.content)
+    except Exception as exc:
+        node.logger.warning("Error decoding OBJECT_REQUEST from %s: %s", peer.address, exc)
+        return
+
+    match object_request.type:
         case ObjectRequestType.OBJECT_GET:
             atom_id = object_request.atom_id
             node.logger.debug("Handling OBJECT_GET for %s from %s", atom_id.hex(), peer.address)
+            payload_type = object_request.payload_type
+            if payload_type is None:
+                payload_type = OBJECT_FOUND_ATOM_PAYLOAD
 
-            local_atom = node.local_get(atom_id)
-            if local_atom is not None:
-                node.logger.debug("Object %s found locally; returning to %s", atom_id.hex(), peer.address)
-                resp = ObjectResponse(
-                    type=ObjectResponseType.OBJECT_FOUND,
-                    data=local_atom.to_bytes(),
-                    atom_id=atom_id
-                )
-                obj_res_msg  = Message(
-                    topic=MessageTopic.OBJECT_RESPONSE,
-                    body=resp.to_bytes(),
-                    sender=node.relay_public_key,
-                )
-                obj_res_msg.encrypt(peer.shared_key_bytes)
-                node.outgoing_queue.put((obj_res_msg.to_bytes(), peer.address))
-                return
-
-            if atom_id in node.storage_index:
-                node.logger.debug("Known provider for %s; informing %s", atom_id.hex(), peer.address)
-                provider_bytes = node.storage_index[atom_id]
-                resp = ObjectResponse(
-                    type=ObjectResponseType.OBJECT_PROVIDER,
-                    data=provider_bytes,
-                    atom_id=atom_id
-                )
-                obj_res_msg = Message(
-                    topic=MessageTopic.OBJECT_RESPONSE,
-                    body=resp.to_bytes(),
-                    sender=node.relay_public_key,
-                )
-                obj_res_msg.encrypt(peer.shared_key_bytes)
-                node.outgoing_queue.put((obj_res_msg.to_bytes(), peer.address))
-                return
-
-            nearest_peer = node.peer_route.closest_peer_for_hash(atom_id)
-            if nearest_peer:
-                node.logger.debug("Forwarding requester %s to nearest peer for %s", peer.address, atom_id.hex())
-                peer_info = encode_peer_contact_bytes(nearest_peer)
-                resp = ObjectResponse(
-                    type=ObjectResponseType.OBJECT_PROVIDER,
-                    # type=ObjectResponseType.OBJECT_NEAREST_PEER,
-                    data=peer_info,
-                    atom_id=atom_id
+            if payload_type == OBJECT_FOUND_ATOM_PAYLOAD:
+                local_atom = node.get_atom_from_local_storage(atom_id=atom_id)
+                if local_atom is not None:
+                    node.logger.debug("Object %s found locally; returning to %s", atom_id.hex(), peer.address)
+                    resp = ObjectResponse(
+                        type=ObjectResponseType.OBJECT_FOUND,
+                        data=encode_object_found_atom_payload(local_atom),
+                        atom_id=atom_id
+                    )
+                    obj_res_msg  = Message(
+                        topic=MessageTopic.OBJECT_RESPONSE,
+                        body=resp.to_bytes(),
+                        sender=node.relay_public_key,
+                    )
+                    obj_res_msg.encrypt(peer.shared_key_bytes)
+                    enqueue_outgoing(
+                        node,
+                        peer.address,
+                        message=obj_res_msg,
+                        difficulty=peer.difficulty,
+                    )
+                    return
+            elif payload_type == OBJECT_FOUND_LIST_PAYLOAD:
+                local_atoms = node.get_atom_list_from_local_storage(root_hash=atom_id)
+                if local_atoms is not None:
+                    node.logger.debug("Object list %s found locally; returning to %s", atom_id.hex(), peer.address)
+                    resp = ObjectResponse(
+                        type=ObjectResponseType.OBJECT_FOUND,
+                        data=encode_object_found_list_payload(local_atoms),
+                        atom_id=atom_id
+                    )
+                    obj_res_msg  = Message(
+                        topic=MessageTopic.OBJECT_RESPONSE,
+                        body=resp.to_bytes(),
+                        sender=node.relay_public_key,
+                    )
+                    obj_res_msg.encrypt(peer.shared_key_bytes)
+                    enqueue_outgoing(
+                        node,
+                        peer.address,
+                        message=obj_res_msg,
+                        difficulty=peer.difficulty,
+                    )
+                    return
+            else:
+                node.logger.warning(
+                    "Unknown OBJECT_GET payload type %s for %s",
+                    payload_type,
+                    atom_id.hex(),
                 )
+
+            if atom_id in node.storage_index:
+                provider_id = node.storage_index[atom_id]
+                provider_bytes = provider_payload_for_id(node, provider_id)
+                if provider_bytes is not None:
+                    node.logger.debug("Known provider for %s; informing %s", atom_id.hex(), peer.address)
+                    resp = ObjectResponse(
+                        type=ObjectResponseType.OBJECT_PROVIDER,
+                        data=provider_bytes,
+                        atom_id=atom_id
+                    )
+                    obj_res_msg = Message(
+                        topic=MessageTopic.OBJECT_RESPONSE,
+                        body=resp.to_bytes(),
+                        sender=node.relay_public_key,
+                    )
+                    obj_res_msg.encrypt(peer.shared_key_bytes)
+                    enqueue_outgoing(
+                        node,
+                        peer.address,
+                        message=obj_res_msg,
+                        difficulty=peer.difficulty,
+                    )
+                    return
+                node.logger.warning(
+                    "Unknown provider id %s for %s",
+                    provider_id,
+                    atom_id.hex(),
+                )
+
+            nearest_peer = node.peer_route.closest_peer_for_hash(atom_id)
+            if nearest_peer:
+                node.logger.debug("Forwarding requester %s to nearest peer for %s", peer.address, atom_id.hex())
+                peer_info = encode_peer_contact_bytes(nearest_peer)
+                resp = ObjectResponse(
+                    type=ObjectResponseType.OBJECT_PROVIDER,
+                    # type=ObjectResponseType.OBJECT_NEAREST_PEER,
+                    data=peer_info,
+                    atom_id=atom_id
+                )
                 obj_res_msg = Message(
                     topic=MessageTopic.OBJECT_RESPONSE,
                     body=resp.to_bytes(),
                     sender=node.relay_public_key,
                 )
                 obj_res_msg.encrypt(nearest_peer.shared_key_bytes)
-                node.outgoing_queue.put((obj_res_msg.to_bytes(), peer.address))
-
-        case ObjectRequestType.OBJECT_PUT:
-            node.logger.debug("Handling OBJECT_PUT for %s from %s", object_request.atom_id.hex(), peer.address)
-
-            nearest_peer = node.peer_route.closest_peer_for_hash(object_request.atom_id)
-            is_self_closest = False
-            if nearest_peer is None or nearest_peer.address is None:
-                is_self_closest = True
-            else:
-                try:
-                    self_distance = xor_distance(object_request.atom_id, node.relay_public_key_bytes)
-                    peer_distance = xor_distance(object_request.atom_id, nearest_peer.public_key_bytes)
-                except Exception as exc:
-                    node.logger.warning(
-                        "Failed distance comparison for OBJECT_PUT %s: %s",
-                        object_request.atom_id.hex(),
-                        exc,
-                    )
-                    is_self_closest = True
-                else:
-                    is_self_closest = self_distance <= peer_distance
-
-            if is_self_closest:
-                node.logger.debug("Storing provider info for %s locally", object_request.atom_id.hex())
-                node.storage_index[object_request.atom_id] = object_request.data
-            else:
-                node.logger.debug(
-                    "Forwarding OBJECT_PUT for %s to nearer peer %s",
-                    object_request.atom_id.hex(),
-                    nearest_peer.address,
-                )
-                fwd_req = ObjectRequest(
-                    type=ObjectRequestType.OBJECT_PUT,
-                    data=object_request.data,
-                    atom_id=object_request.atom_id,
+                enqueue_outgoing(
+                    node,
+                    peer.address,
+                    message=obj_res_msg,
+                    difficulty=peer.difficulty,
                 )
+
+        case ObjectRequestType.OBJECT_PUT:
+            node.logger.debug("Handling OBJECT_PUT for %s from %s", object_request.atom_id.hex(), peer.address)
+
+            nearest_peer = node.peer_route.closest_peer_for_hash(object_request.atom_id)
+            is_self_closest = False
+            if nearest_peer is None or nearest_peer.address is None:
+                is_self_closest = True
+            else:
+                try:
+                    self_distance = xor_distance(object_request.atom_id, node.relay_public_key_bytes)
+                    peer_distance = xor_distance(object_request.atom_id, nearest_peer.public_key_bytes)
+                except Exception as exc:
+                    node.logger.warning(
+                        "Failed distance comparison for OBJECT_PUT %s: %s",
+                        object_request.atom_id.hex(),
+                        exc,
+                    )
+                    is_self_closest = True
+                else:
+                    is_self_closest = self_distance <= peer_distance
+
+            if is_self_closest:
+                node.logger.debug("Storing provider info for %s locally", object_request.atom_id.hex())
+                provider_id = provider_id_for_payload(node, object_request.data)
+                node.storage_index[object_request.atom_id] = provider_id
+            else:
+                node.logger.debug(
+                    "Forwarding OBJECT_PUT for %s to nearer peer %s",
+                    object_request.atom_id.hex(),
+                    nearest_peer.address,
+                )
+                fwd_req = ObjectRequest(
+                    type=ObjectRequestType.OBJECT_PUT,
+                    data=object_request.data,
+                    atom_id=object_request.atom_id,
+                )
                 obj_req_msg = Message(
                     topic=MessageTopic.OBJECT_REQUEST,
                     body=fwd_req.to_bytes(),
                     sender=node.relay_public_key,
                 )
                 obj_req_msg.encrypt(nearest_peer.shared_key_bytes)
-                node.outgoing_queue.put((obj_req_msg.to_bytes(), nearest_peer.address))
-
-        case _:
-            node.logger.warning("Unknown ObjectRequestType %s from %s", object_request.type, peer.address)
+                enqueue_outgoing(
+                    node,
+                    nearest_peer.address,
+                    message=obj_req_msg,
+                    difficulty=nearest_peer.difficulty,
+                )
+
+        case _:
+            node.logger.warning("Unknown ObjectRequestType %s from %s", object_request.type, peer.address)

astreum/communication/handlers/object_response.py

@@ -1,9 +1,11 @@
 import socket
 from enum import IntEnum
-from typing import Tuple, TYPE_CHECKING
+from typing import List, Tuple, TYPE_CHECKING
 
+from ..outgoing_queue import enqueue_outgoing
 from ..models.message import Message, MessageTopic
 from ...storage.models.atom import Atom
+from ...storage.requests import get_atom_req_payload
 
 if TYPE_CHECKING:
     from .. import Node
@@ -16,6 +18,10 @@ class ObjectResponseType(IntEnum):
     OBJECT_NEAREST_PEER = 2
 
 
+OBJECT_FOUND_ATOM_PAYLOAD = 1
+OBJECT_FOUND_LIST_PAYLOAD = 2
+
+
 class ObjectResponse:
     type: ObjectResponseType
     data: bytes
@@ -46,6 +52,37 @@ class ObjectResponse:
         return cls(resp_type, payload, atom_id)
 
 
+def encode_object_found_atom_payload(atom: Atom) -> bytes:
+    return bytes([OBJECT_FOUND_ATOM_PAYLOAD]) + atom.to_bytes()
+
+
+def encode_object_found_list_payload(atoms: List[Atom]) -> bytes:
+    parts = [bytes([OBJECT_FOUND_LIST_PAYLOAD])]
+    for atom in atoms:
+        atom_bytes = atom.to_bytes()
+        parts.append(len(atom_bytes).to_bytes(4, "big", signed=False))
+        parts.append(atom_bytes)
+    return b"".join(parts)
+
+
+def decode_object_found_list_payload(payload: bytes) -> List[Atom]:
+    atoms: List[Atom] = []
+    offset = 0
+    while offset < len(payload):
+        if len(payload) - offset < 4:
+            raise ValueError("truncated atom length")
+        atom_len = int.from_bytes(payload[offset : offset + 4], "big", signed=False)
+        offset += 4
+        if atom_len <= 0:
+            raise ValueError("invalid atom length")
+        end = offset + atom_len
+        if end > len(payload):
+            raise ValueError("truncated atom payload")
+        atoms.append(Atom.from_bytes(payload[offset:end]))
+        offset = end
+    return atoms
+
+
 def decode_object_provider(payload: bytes) -> Tuple[bytes, str, int]:
     expected_len = 32 + 4 + 2
     if len(payload) < expected_len:
@@ -76,17 +113,78 @@ def handle_object_response(node: "Node", peer: "Peer", message: Message) -> None
 
     match object_response.type:
         case ObjectResponseType.OBJECT_FOUND:
-            atom = Atom.from_bytes(object_response.data)
-            atom_id = atom.object_id()
-            if object_response.atom_id == atom_id:
-                node.pop_atom_req(atom_id)
-                node._hot_storage_set(atom_id, atom)
-            else:
+            payload = object_response.data
+            if not payload:
                 node.logger.warning(
-                    "OBJECT_FOUND atom ID mismatch (expected=%s got=%s)",
+                    "OBJECT_FOUND payload for %s missing content",
                     object_response.atom_id.hex(),
-                    atom_id.hex(),
                 )
+                return
+
+            payload_type = payload[0]
+            body = payload[1:]
+
+            if payload_type == OBJECT_FOUND_ATOM_PAYLOAD:
+                try:
+                    atom = Atom.from_bytes(body)
+                except Exception as exc:
+                    node.logger.warning(
+                        "Invalid OBJECT_FOUND atom payload for %s: %s",
+                        object_response.atom_id.hex(),
+                        exc,
+                    )
+                    return
+
+                atom_id = atom.object_id()
+                if object_response.atom_id != atom_id:
+                    node.logger.warning(
+                        "OBJECT_FOUND atom ID mismatch (expected=%s got=%s)",
+                        object_response.atom_id.hex(),
+                        atom_id.hex(),
+                    )
+                    return
+
+                node.pop_atom_req(atom_id)
+                node._hot_storage_set(atom_id, atom)
+                return
+
+            if payload_type == OBJECT_FOUND_LIST_PAYLOAD:
+                try:
+                    atoms = decode_object_found_list_payload(body)
+                except Exception as exc:
+                    node.logger.warning(
+                        "Invalid OBJECT_FOUND list payload for %s: %s",
+                        object_response.atom_id.hex(),
+                        exc,
+                    )
+                    return
+
+                if not atoms:
+                    node.logger.warning(
+                        "OBJECT_FOUND list payload for %s contained no atoms",
+                        object_response.atom_id.hex(),
+                    )
+                    return
+
+                root_id = atoms[0].object_id()
+                if object_response.atom_id != root_id:
+                    node.logger.warning(
+                        "OBJECT_FOUND list root ID mismatch (expected=%s got=%s)",
+                        object_response.atom_id.hex(),
+                        root_id.hex(),
+                    )
+                    return
+
+                node.pop_atom_req(root_id)
+                for atom in atoms:
+                    node._hot_storage_set(atom.object_id(), atom)
+                return
+
+            node.logger.warning(
+                "Unknown OBJECT_FOUND payload type %s for %s",
+                payload_type,
+                object_response.atom_id.hex(),
+            )
 
         case ObjectResponseType.OBJECT_PROVIDER:
             try:
@@ -97,10 +195,15 @@ def handle_object_response(node: "Node", peer: "Peer", message: Message) -> None
 
             from .object_request import ObjectRequest, ObjectRequestType
 
+            payload_type = get_atom_req_payload(node, object_response.atom_id)
+            if payload_type is None:
+                payload_type = OBJECT_FOUND_ATOM_PAYLOAD
+
             obj_req = ObjectRequest(
                 type=ObjectRequestType.OBJECT_GET,
                 data=b"",
                 atom_id=object_response.atom_id,
+                payload_type=payload_type,
             )
             obj_req_bytes = obj_req.to_bytes()
             obj_req_msg = Message(
@@ -109,7 +212,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,
+        )