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

astreum/communication/handlers/handshake.py

@@ -1,83 +1,89 @@
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Sequence
-
-from cryptography.hazmat.primitives import serialization
-
-from ..models.peer import Peer
-from ..models.message import Message
-
-if TYPE_CHECKING:
-    from .... import Node
-
-
-def handle_handshake(node: "Node", addr: Sequence[object], message: Message) -> bool:
-    """Handle incoming handshake messages.
-
-    Returns True if the outer loop should `continue`, False otherwise.
-    """
-    logger = node.logger
-
-    sender_key = message.sender
-    try:
-        sender_public_key_bytes = sender_key.public_bytes(
-            encoding=serialization.Encoding.Raw,
-            format=serialization.PublicFormat.Raw,
-        )
-    except Exception as exc:
-        logger.warning("Error extracting sender key bytes: %s", exc)
-        return True
-
-    try:
-        host, port = addr[0], int(addr[1])
-    except Exception:
-        return True
-    address_key = (host, port)
-
-    old_key_bytes = node.addresses.get(address_key)
-    node.addresses[address_key] = sender_public_key_bytes
-
-    if old_key_bytes is None:
-        try:
-            peer = Peer(node.relay_secret_key, sender_key)
-        except Exception:
-            return True
-        peer.address = address_key
-
-        node.peers[sender_public_key_bytes] = peer
-        node.peer_route.add_peer(sender_public_key_bytes, peer)
-
-        logger.info(
-            "Handshake accepted from %s:%s; peer added",
-            address_key[0],
-            address_key[1],
-        )
-        response = Message(handshake=True, sender=node.relay_public_key)
-        node.outgoing_queue.put((response.to_bytes(), address_key))
-        return True
-
-    if old_key_bytes == sender_public_key_bytes:
-        peer = node.peers.get(sender_public_key_bytes)
-        if peer is not None:
-            peer.address = address_key
-        return False
-
-    node.peers.pop(old_key_bytes, None)
-    try:
-        node.peer_route.remove_peer(old_key_bytes)
-    except Exception:
-        pass
-    try:
-        peer = Peer(node.relay_secret_key, sender_key)
-    except Exception:
-        return True
-    peer.address = address_key
-
-    node.peers[sender_public_key_bytes] = peer
-    node.peer_route.add_peer(sender_public_key_bytes, peer)
-    logger.info(
-        "Peer at %s:%s replaced due to key change",
-        address_key[0],
-        address_key[1],
-    )
-    return False
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PublicKey
+
+from ..models.peer import Peer
+from ..models.message import Message
+
+if TYPE_CHECKING:
+    from .... import Node
+
+
+def handle_handshake(node: "Node", addr: Sequence[object], message: Message) -> bool:
+    """Handle incoming handshake messages.
+
+    Returns True if the outer loop should `continue`, False otherwise.
+    """
+    sender_public_key_bytes = message.sender_bytes
+    try:
+        sender_key = X25519PublicKey.from_public_bytes(sender_public_key_bytes)
+    except Exception as exc:
+        node.logger.warning("Error extracting sender key bytes: %s", exc)
+        return True
+
+    try:
+        host = addr[0]
+        port = int.from_bytes(message.content[:2], "big", signed=False)
+    except Exception:
+        return True
+    peer_address = (host, port)
+
+    old_key_bytes = node.addresses.get(peer_address)
+    node.addresses[peer_address] = sender_public_key_bytes
+
+    if old_key_bytes is None:
+        try:
+            peer = Peer(
+                node_secret_key=node.relay_secret_key,
+                peer_public_key=sender_key,
+                address=peer_address,
+            )
+        except Exception:
+            return True
+
+        node.add_peer(sender_public_key_bytes, peer)
+        node.peer_route.add_peer(sender_public_key_bytes, peer)
+
+        node.logger.info(
+            "Handshake accepted from %s:%s; peer added",
+            peer_address[0],
+            peer_address[1],
+        )
+        response = Message(
+            handshake=True,
+            sender=node.relay_public_key,
+            content=int(node.config["incoming_port"]).to_bytes(2, "big", signed=False),
+        )
+        node.outgoing_queue.put((response.to_bytes(), peer_address))
+        return True
+
+    if old_key_bytes == sender_public_key_bytes:
+        peer = node.get_peer(sender_public_key_bytes)
+        if peer is not None:
+            peer.address = peer_address
+        return False
+
+    try:
+        node.peer_route.remove_peer(old_key_bytes)
+    except Exception:
+        pass
+    try:
+        peer = Peer(
+            node_secret_key=node.relay_secret_key,
+            peer_public_key=sender_key,
+            address=peer_address,
+        )
+    except Exception:
+        return True
+
+    node.replace_peer(old_key_bytes, sender_public_key_bytes, peer)
+    node.peer_route.add_peer(sender_public_key_bytes, peer)
+    node.logger.info(
+        "Peer at %s:%s replaced due to key change",
+        peer_address[0],
+        peer_address[1],
+    )
+    return False

astreum/communication/handlers/object_request.py

@@ -0,0 +1,176 @@
+import logging
+import socket
+from enum import IntEnum
+from typing import TYPE_CHECKING, Tuple
+
+from .object_response import ObjectResponse, ObjectResponseType
+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
+
+
+class ObjectRequest:
+    type: ObjectRequestType
+    data: bytes
+    atom_id: bytes
+
+    def __init__(self, type: ObjectRequestType, data: bytes, atom_id: bytes = None):
+        self.type = type
+        self.data = data
+        self.atom_id = atom_id
+
+    def to_bytes(self):
+        return bytes([self.type.value]) + self.atom_id + self.data
+
+    @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}")
+
+        atom_id_bytes = data[1:33]
+        payload    = data[33:]
+        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:
+        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)
+
+            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
+                )
+                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,
+                )
+                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)

astreum/communication/handlers/object_response.py

@@ -0,0 +1,115 @@
+import socket
+from enum import IntEnum
+from typing import Tuple, TYPE_CHECKING
+
+from ..models.message import Message, MessageTopic
+from ...storage.models.atom import Atom
+
+if TYPE_CHECKING:
+    from .. import Node
+    from ..models.peer import Peer
+
+
+class ObjectResponseType(IntEnum):
+    OBJECT_FOUND = 0
+    OBJECT_PROVIDER = 1
+    OBJECT_NEAREST_PEER = 2
+
+
+class ObjectResponse:
+    type: ObjectResponseType
+    data: bytes
+    atom_id: bytes
+
+    def __init__(self, type: ObjectResponseType, data: bytes, atom_id: bytes = None):
+        self.type = type
+        self.data = data
+        self.atom_id = atom_id
+
+    def to_bytes(self):
+        return bytes([self.type.value]) + self.atom_id + self.data
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> "ObjectResponse":
+        # need at least 1 byte for type + 32 bytes for atom id
+        if len(data) < 1 + 32:
+            raise ValueError(f"Too short to be a valid ObjectResponse ({len(data)} bytes)")
+
+        type_val = data[0]
+        try:
+            resp_type = ObjectResponseType(type_val)
+        except ValueError:
+            raise ValueError(f"Unknown ObjectResponseType: {type_val}")
+
+        atom_id = data[1:33]
+        payload   = data[33:]
+        return cls(resp_type, payload, atom_id)
+
+
+def decode_object_provider(payload: bytes) -> Tuple[bytes, str, int]:
+    expected_len = 32 + 4 + 2
+    if len(payload) < expected_len:
+        raise ValueError("provider payload too short")
+
+    provider_public_key = payload[:32]
+    provider_ip_bytes = payload[32:36]
+    provider_port_bytes = payload[36:38]
+
+    provider_address = socket.inet_ntoa(provider_ip_bytes)
+    provider_port = int.from_bytes(provider_port_bytes, byteorder="big", signed=False)
+    return provider_public_key, provider_address, provider_port
+
+
+def handle_object_response(node: "Node", peer: "Peer", message: Message) -> None:
+    if message.content is None:
+        node.logger.warning("OBJECT_RESPONSE from %s missing content", peer.address)
+        return
+
+    try:
+        object_response = ObjectResponse.from_bytes(message.content)
+    except Exception as exc:
+        node.logger.warning("Error decoding OBJECT_RESPONSE from %s: %s", peer.address, exc)
+        return
+
+    if not node.has_atom_req(object_response.atom_id):
+        return
+
+    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:
+                node.logger.warning(
+                    "OBJECT_FOUND atom ID mismatch (expected=%s got=%s)",
+                    object_response.atom_id.hex(),
+                    atom_id.hex(),
+                )
+
+        case ObjectResponseType.OBJECT_PROVIDER:
+            try:
+                _, provider_address, provider_port = decode_object_provider(object_response.data)
+            except Exception as exc:
+                node.logger.warning("Invalid OBJECT_PROVIDER payload from %s: %s", peer.address, exc)
+                return
+
+            from .object_request import ObjectRequest, ObjectRequestType
+
+            obj_req = ObjectRequest(
+                type=ObjectRequestType.OBJECT_GET,
+                data=b"",
+                atom_id=object_response.atom_id,
+            )
+            obj_req_bytes = obj_req.to_bytes()
+            obj_req_msg = Message(
+                topic=MessageTopic.OBJECT_REQUEST,
+                body=obj_req_bytes,
+                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)))
+
+        case ObjectResponseType.OBJECT_NEAREST_PEER:
+            node.logger.debug("Ignoring OBJECT_NEAREST_PEER response from %s", peer.address)

astreum/communication/handlers/ping.py

@@ -1,35 +1,21 @@
 from __future__ import annotations
 
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from ..models.ping import Ping
+from ..models.peer import Peer
 
 if TYPE_CHECKING:
     from .... import Node
 
 
-def handle_ping(node: "Node", addr: Sequence[object], payload: bytes) -> None:
+def handle_ping(node: "Node", peer: Peer, payload: bytes) -> None:
     """Update peer and validation state based on an incoming ping message."""
-    logger = node.logger
-    try:
-        host, port = addr[0], int(addr[1])
-    except Exception:
-        return
-
-    address_key = (host, port)
-    sender_public_key_bytes = node.addresses.get(address_key)
-    if sender_public_key_bytes is None:
-        return
-
-    peer = node.peers.get(sender_public_key_bytes)
-    if peer is None:
-        return
-
     try:
         ping = Ping.from_bytes(payload)
     except Exception as exc:
-        logger.warning("Error decoding ping: %s", exc)
+        node.logger.warning("Error decoding ping: %s", exc)
         return
 
     peer.timestamp = datetime.now(timezone.utc)
@@ -41,8 +27,8 @@ def handle_ping(node: "Node", addr: Sequence[object], payload: bytes) -> None:
 
     try:
         if ping.is_validator:
-            validation_route.add_peer(sender_public_key_bytes)
+            validation_route.add_peer(peer.public_key_bytes)
         else:
-            validation_route.remove_peer(sender_public_key_bytes)
+            validation_route.remove_peer(peer.public_key_bytes)
     except Exception:
         pass

astreum/communication/handlers/route_request.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import socket
+
+from ..models.message import Message, MessageTopic
+from ..util import xor_distance
+
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from .... import Node
+    from ..models.peer import Peer
+
+
+def handle_route_request(node: "Node", peer: "Peer", message: Message) -> None:
+    sender_public_key = getattr(peer, "public_key_bytes", None)
+    if not sender_public_key:
+        node.logger.warning("Unknown sender for ROUTE_REQUEST from %s", peer.address)
+        return
+
+    if not message.content:
+        node.logger.warning("ROUTE_REQUEST missing route id from %s", peer.address)
+        return
+    route_id = message.content[0]
+    if route_id == 0:
+        route = node.peer_route
+    elif route_id == 1:
+        route = node.validation_route
+        if route is None:
+            node.logger.warning("Validation route not initialized for %s", peer.address)
+            return
+    else:
+        node.logger.warning("Unknown route id %s in ROUTE_REQUEST from %s", route_id, peer.address)
+        return
+
+    payload_parts = []
+    for bucket in route.buckets.values():
+        closest_key = None
+        closest_distance = None
+
+        for peer_key in bucket:
+            try:
+                distance = xor_distance(sender_public_key, peer_key)
+            except ValueError:
+                continue
+
+            if closest_distance is None or distance < closest_distance:
+                closest_distance = distance
+                closest_key = peer_key
+
+        if closest_key is None:
+            continue
+
+        bucket_peer = node.get_peer(closest_key)
+        if bucket_peer is None or bucket_peer.address is None:
+            continue
+
+        host, port = bucket_peer.address
+        try:
+            address_bytes = socket.inet_pton(socket.AF_INET, host)
+        except OSError:
+            try:
+                address_bytes = socket.inet_pton(socket.AF_INET6, host)
+            except OSError as exc:
+                node.logger.warning("Invalid peer address %s: %s", bucket_peer.address, exc)
+                continue
+
+        port_bytes = int(port).to_bytes(2, "big", signed=False)
+        payload_parts.append(address_bytes + port_bytes)
+
+    response = Message(
+        topic=MessageTopic.ROUTE_RESPONSE,
+        content=b"".join(payload_parts),
+        sender=node.relay_public_key,
+    )
+    response.encrypt(peer.shared_key_bytes)
+    node.outgoing_queue.put((response.to_bytes(), peer.address))

astreum/communication/handlers/route_response.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import socket
+
+from ..models.message import Message
+
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from .... import Node
+    from ..models.peer import Peer
+
+
+def handle_route_response(node: "Node", peer: "Peer", message: Message) -> None:
+    payload = message.content
+    if not payload:
+        return
+    host_len = 16 if node.use_ipv6 else 4
+    chunk_size = host_len + 2
+    if len(payload) % chunk_size != 0:
+        node.logger.warning(
+            "ROUTE_RESPONSE payload size mismatch (%s bytes) from %s",
+            len(payload),
+            peer.address,
+        )
+        return
+
+    decoded_addresses = []
+    family = socket.AF_INET6 if node.use_ipv6 else socket.AF_INET
+    for index in range(0, len(payload), chunk_size):
+        host_bytes = payload[index : index + host_len]
+        port_bytes = payload[index + host_len : index + chunk_size]
+        try:
+            host = socket.inet_ntop(family, host_bytes)
+        except OSError as exc:
+            node.logger.warning(
+                "Invalid host bytes in ROUTE_RESPONSE from %s: %s",
+                peer.address,
+                exc,
+            )
+            continue
+        port = int.from_bytes(port_bytes, "big", signed=False)
+        decoded_addresses.append((host, port))
+    if not decoded_addresses:
+        return
+    node.logger.debug("Decoded %s addresses from ROUTE_RESPONSE", len(decoded_addresses))
+
+    handshake_message = Message(
+        handshake=True,
+        sender=node.relay_public_key,
+        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)))