astreum 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

astreum/node/relay/__init__.py

@@ -0,0 +1,248 @@
+"""
+Relay module for handling network communication in the Astreum node.
+"""
+
+import socket
+import threading
+from queue import Queue
+from typing import Tuple, Callable, Dict, Set, Optional, List
+from .message import Message, Topic
+from .envelope import Envelope
+from .bucket import KBucket
+from .peer import Peer, PeerManager
+from .route import RouteTable
+
+class Relay:
+    def __init__(self, config: dict):
+        self.use_ipv6 = config.get('use_ipv6', False)
+        incoming_port = config.get('incoming_port', 7373)
+        self.max_message_size = config.get('max_message_size', 65536)  # Max UDP datagram size
+        self.num_workers = config.get('num_workers', 4)
+
+        # Routes that this node participates in (0 = peer route, 1 = validation route)
+        self.routes: List[int] = []
+        
+        # Initialize routes from config if provided
+        if config.get('peer_route', False):
+            self.routes.append(0)  # Peer route
+            
+        if config.get('validation_route', False):
+            self.routes.append(1)  # Validation route
+
+        # Choose address family based on IPv4 or IPv6
+        family = socket.AF_INET6 if self.use_ipv6 else socket.AF_INET
+
+        # Create a UDP socket
+        self.incoming_socket = socket.socket(family, socket.SOCK_DGRAM)
+
+        # Allow dual-stack support (IPv4-mapped addresses on IPv6)
+        if self.use_ipv6:
+            self.incoming_socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
+
+        # Bind to an address (IPv6 "::" or IPv4 "0.0.0.0") and port
+        bind_address = "::" if self.use_ipv6 else "0.0.0.0"
+        self.incoming_socket.bind((bind_address, incoming_port or 0))
+
+        # Get the actual port assigned
+        self.incoming_port = self.incoming_socket.getsockname()[1]
+
+        # Create a UDP socket for sending messages
+        self.outgoing_socket = socket.socket(family, socket.SOCK_DGRAM)
+
+        # Message queues
+        self.incoming_queue = Queue()
+        self.outgoing_queue = Queue()
+        
+        # Message handling
+        self.message_handlers: Dict[Topic, Callable] = {}
+
+        # Route buckets (peers for each route)
+        self.peer_route_bucket = KBucket(k=20)  # Bucket for peer route
+        self.validation_route_bucket = KBucket(k=20)  # Bucket for validation route
+
+        # Start worker threads
+        self._start_workers()
+        
+    def is_in_peer_route(self) -> bool:
+        """Check if this node is part of the peer route."""
+        return 0 in self.routes
+        
+    def is_in_validation_route(self) -> bool:
+        """Check if this node is part of the validation route."""
+        return 1 in self.routes
+        
+    def add_to_peer_route(self):
+        """Add this node to the peer route."""
+        if 0 not in self.routes:
+            self.routes.append(0)
+        
+    def add_to_validation_route(self):
+        """Add this node to the validation route."""
+        if 1 not in self.routes:
+            self.routes.append(1)
+        
+    def remove_from_peer_route(self):
+        """Remove this node from the peer route."""
+        if 0 in self.routes:
+            self.routes.remove(0)
+        
+    def remove_from_validation_route(self):
+        """Remove this node from the validation route."""
+        if 1 in self.routes:
+            self.routes.remove(1)
+        
+    def add_peer_to_route(self, peer: Peer, route_types: List[int]):
+        """
+        Add a peer to specified routes.
+        
+        Args:
+            peer (Peer): The peer to add
+            route_types (List[int]): List of route types to add the peer to (0 = peer, 1 = validation)
+        """
+        for route_type in route_types:
+            if route_type == 0:  # Peer route
+                # Add to top of bucket, eject last if at capacity
+                self.peer_route_bucket.add(peer, to_front=True)
+            elif route_type == 1:  # Validation route
+                # Add to top of bucket, eject last if at capacity
+                self.validation_route_bucket.add(peer, to_front=True)
+            
+    def get_route_peers(self, route_type: int) -> List[Peer]:
+        """
+        Get all peers in a specific route.
+        
+        Args:
+            route_type (int): Route type (0 for peer, 1 for validation)
+            
+        Returns:
+            List[Peer]: List of peers in the route
+        """
+        if route_type == 0:  # Peer route
+            return self.peer_route_bucket.get_peers()
+        elif route_type == 1:  # Validation route
+            return self.validation_route_bucket.get_peers()
+        return []
+        
+    def register_message_handler(self, topic: Topic, handler_func):
+        """Register a handler function for a specific message topic."""
+        self.message_handlers[topic] = handler_func
+
+    def _start_workers(self):
+        """Start worker threads for processing incoming and outgoing messages."""
+        self.running = True
+        
+        # Start receiver thread
+        self.receiver_thread = threading.Thread(target=self._receive_messages)
+        self.receiver_thread.daemon = True
+        self.receiver_thread.start()
+
+        # Start sender thread
+        self.sender_thread = threading.Thread(target=self._send_messages)
+        self.sender_thread.daemon = True
+        self.sender_thread.start()
+
+        # Start worker threads for processing incoming messages
+        self.worker_threads = []
+        for _ in range(self.num_workers):
+            thread = threading.Thread(target=self._process_messages)
+            thread.daemon = True
+            thread.start()
+            self.worker_threads.append(thread)
+
+    def _receive_messages(self):
+        """Continuously receive messages and add them to the incoming queue."""
+        while self.running:
+            try:
+                data, addr = self.incoming_socket.recvfrom(self.max_message_size)
+                self.incoming_queue.put((data, addr))
+            except Exception as e:
+                # Log error but continue running
+                print(f"Error receiving message: {e}")
+
+    def _send_messages(self):
+        """Continuously send messages from the outgoing queue."""
+        while self.running:
+            try:
+                data, addr = self.outgoing_queue.get()
+                self.outgoing_socket.sendto(data, addr)
+                self.outgoing_queue.task_done()
+            except Exception as e:
+                # Log error but continue running
+                print(f"Error sending message: {e}")
+
+    def _process_messages(self):
+        """Process messages from the incoming queue."""
+        while self.running:
+            try:
+                data, addr = self.incoming_queue.get()
+                self._handle_message(data, addr)
+                self.incoming_queue.task_done()
+            except Exception as e:
+                # Log error but continue running
+                print(f"Error processing message: {e}")
+
+    def _handle_message(self, data: bytes, addr: Tuple[str, int]):
+        """Handle an incoming message."""
+        envelope = Envelope.from_bytes(data)
+        if envelope and envelope.message.topic in self.message_handlers:
+            # Check if this is a transaction or block message that requires validation route
+            if envelope.message.topic in (Topic.TRANSACTION, Topic.BLOCK):
+                # Only process if we're part of the validation route
+                if self.is_in_validation_route():
+                    self.message_handlers[envelope.message.topic](envelope.message.body, addr, envelope)
+            elif envelope.message.topic == Topic.LATEST_BLOCK:
+                # For latest_block, we only process if we're in the validation route
+                if self.is_in_validation_route():
+                    self.message_handlers[envelope.message.topic](envelope.message.body, addr, envelope)
+            elif envelope.message.topic in (Topic.LATEST_BLOCK_REQUEST, Topic.GET_BLOCKS):
+                # Allow all nodes to request blocks for syncing
+                self.message_handlers[envelope.message.topic](envelope.message.body, addr, envelope)
+            else:
+                # For other message types, always process
+                self.message_handlers[envelope.message.topic](envelope.message.body, addr, envelope)
+
+    def send(self, data: bytes, addr: Tuple[str, int]):
+        """Send raw data to a specific address."""
+        self.outgoing_queue.put((data, addr))
+    
+    def get_address(self) -> Tuple[str, int]:
+        """
+        Get the local address of this relay node.
+        
+        Returns:
+            Tuple[str, int]: The local address (host, port)
+        """
+        # This is a simplification - in a real implementation this would determine the 
+        # actual public-facing IP address, which may be different from the binding address
+        return ("localhost", self.incoming_port)
+        
+    def get_routes(self) -> bytes:
+        """
+        Get the routes this node is part of as a bytes object.
+        
+        Returns:
+            bytes: List of route types (0 for peer, 1 for validation)
+        """
+        return bytes(self.routes)
+        
+    def send_message(self, body: bytes, topic: Topic, addr: Tuple[str, int], encrypted: bool = False, difficulty: int = 1):
+        """
+        Create and send a message to a specific address.
+        
+        Args:
+            body (bytes): The message body
+            topic (Topic): The message topic
+            addr (Tuple[str, int]): The recipient's address (host, port)
+            encrypted (bool): Whether the message is encrypted
+            difficulty (int): Number of leading zero bits required in the nonce hash
+        """
+        envelope = Envelope.create(body, topic, encrypted, difficulty)
+        encoded_data = envelope.to_bytes()
+        self.send(encoded_data, addr)
+
+    def stop(self):
+        """Stop all worker threads."""
+        self.running = False
+        # Wait for queues to be processed
+        self.incoming_queue.join()
+        self.outgoing_queue.join()

astreum/node/relay/bucket.py

@@ -0,0 +1,80 @@
+"""
+K-bucket implementation for Kademlia-style routing in Astreum node.
+"""
+
+from typing import List, Tuple
+
+class KBucket:
+    """
+    A Kademlia k-bucket that stores peers.
+    
+    K-buckets are used to store contact information for nodes in the DHT.
+    When a new node is added, it's placed at the tail of the list.
+    If a node is already in the list, it is moved to the tail.
+    This creates a least-recently seen eviction policy.
+    """
+    
+    def __init__(self, size: int):
+        """
+        Initialize a k-bucket with a fixed size.
+        
+        Args:
+            size (int): Maximum number of peers in the bucket
+        """
+        self.size = size
+        self.peers: List[Tuple[str, int]] = []
+
+    def add(self, peer: Tuple[str, int]) -> bool:
+        """
+        Add peer to bucket if not full or if peer exists.
+        
+        Args:
+            peer (Tuple[str, int]): Peer address (host, port)
+            
+        Returns:
+            bool: True if added/exists, False if bucket full and peer not in bucket
+        """
+        if peer in self.peers:
+            # Move to end (most recently seen)
+            self.peers.remove(peer)
+            self.peers.append(peer)
+            return True
+        
+        if len(self.peers) < self.size:
+            self.peers.append(peer)
+            return True
+            
+        return False
+
+    def remove(self, peer: Tuple[str, int]) -> bool:
+        """
+        Remove peer from bucket.
+        
+        Args:
+            peer (Tuple[str, int]): Peer address to remove
+            
+        Returns:
+            bool: True if peer was removed, False if peer not in bucket
+        """
+        if peer in self.peers:
+            self.peers.remove(peer)
+            return True
+        return False
+        
+    def get_peers(self) -> List[Tuple[str, int]]:
+        """
+        Get all peers in the bucket.
+        
+        Returns:
+            List[Tuple[str, int]]: List of peer addresses
+        """
+        return self.peers.copy()
+        
+    def __len__(self) -> int:
+        """
+        Get the number of peers in the bucket.
+        
+        Returns:
+            int: Number of peers
+        """
+        return len(self.peers)

astreum/node/relay/envelope.py

@@ -0,0 +1,280 @@
+"""
+Envelope related classes and utilities for Astreum node network.
+
+Message Structure:
+                    + - - - - - - - +
+                    |   Envelope    |
+                    + - - - - - - - +
+                            ^
+                    . - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - .
+                    ^                       ^                       ^                       ^
+            + - - - - - - - +       + - - - - - - - +       + - - - - - - - +       + - - - - - - - +
+            |     Time      |       |   Encrypted   |       |     Nonce     |       |    Message    |
+            + - - - - - - - +       + - - - - - - - +       + - - - - - - - +       + - - - - - - - +
+                                                                                            ^
+                                                                                . - - - - - - - - - - - .
+                                                                                ^                       ^
+                                                                         + - - - - - - - +       + - - - - - - - +
+                                                                         |     Topic     |       |     Body      |
+                                                                         + - - - - - - - +       + - - - - - - - +
+
+The Envelope uses a Merkle tree structure with the following leaves:
+- Timestamp
+- Encrypted flag
+- Nonce
+- Message bytes
+
+The root hash of this Merkle tree must have a specified number of leading zero bits,
+determined by the difficulty parameter. The nonce is adjusted until this requirement is met.
+"""
+
+import struct
+import time
+import os
+import hashlib
+from dataclasses import dataclass
+from typing import Optional, Tuple, List
+from .message import Message, Topic
+from astreum.utils.bytes_format import encode, decode
+
+@dataclass
+class Envelope:
+    """
+    Represents an envelope that wraps a message with additional metadata.
+    
+    Attributes:
+        encrypted (bool): True if the message is encrypted, False otherwise
+        message (Message): The message being sent
+        nonce (bytes): Nonce for encryption and proof of work
+        timestamp (int): Time when the envelope was created
+    """
+    encrypted: bool
+    message: Message
+    nonce: bytes
+    timestamp: int
+    
+    @classmethod
+    def create(cls, body: bytes, topic: Topic, encrypted: bool = False, difficulty: int = 1) -> 'Envelope':
+        """
+        Create a new envelope with the current timestamp and a nonce that satisfies
+        the given difficulty level using a Merkle tree structure.
+        
+        Args:
+            body (bytes): The message body
+            topic (Topic): The message topic
+            encrypted (bool): Whether the message is encrypted
+            difficulty (int): Number of leading zero bits required in the Merkle root hash
+            
+        Returns:
+            Envelope: A new envelope with a valid nonce
+        """
+        timestamp = int(time.time())
+        message = Message(body=body, topic=topic)
+        
+        # Generate a valid nonce for the Merkle tree
+        nonce = cls._generate_nonce(message, timestamp, encrypted, difficulty)
+        
+        return cls(
+            encrypted=encrypted,
+            message=message,
+            nonce=nonce,
+            timestamp=timestamp
+        )
+    
+    @staticmethod
+    def _generate_nonce(message: Message, timestamp: int, encrypted: bool, difficulty: int) -> bytes:
+        """
+        Generate a nonce that results in a Merkle tree root hash with the specified
+        number of leading zero bits.
+        
+        Args:
+            message (Message): The message to include in the Merkle tree
+            timestamp (int): The timestamp to include in the Merkle tree
+            encrypted (bool): Whether the message is encrypted
+            difficulty (int): Number of leading zero bits required
+            
+        Returns:
+            bytes: A valid nonce
+        """
+        # Prepare the message data
+        message_data = message.to_bytes()
+        timestamp_data = struct.pack('!Q', timestamp)
+        encrypted_flag = b'\x01' if encrypted else b'\x00'
+        
+        # Calculate how many bytes need to be zero
+        zero_bytes = difficulty // 8
+        # Calculate how many bits in the last byte need to be zero
+        remaining_bits = difficulty % 8
+        
+        # Create a mask for the remaining bits
+        mask = 0
+        if remaining_bits > 0:
+            mask = 0xFF >> remaining_bits
+        
+        while True:
+            # Generate a random nonce
+            nonce = os.urandom(32)
+            
+            # Calculate the Merkle root using the leaves
+            merkle_root = Envelope._calculate_merkle_root([
+                timestamp_data,
+                encrypted_flag,
+                nonce,
+                message_data
+            ])
+            
+            # Check if it meets the difficulty requirement
+            valid = True
+            
+            # Check full zero bytes
+            for i in range(zero_bytes):
+                if merkle_root[i] != 0:
+                    valid = False
+                    break
+            
+            # If we need to check partial bits in a byte
+            if valid and remaining_bits > 0:
+                # The next byte should have required number of leading zeros
+                if (merkle_root[zero_bytes] & (0xFF ^ mask)) != 0:
+                    valid = False
+            
+            if valid:
+                return nonce
+    
+    @staticmethod
+    def _calculate_merkle_root(leaves: List[bytes]) -> bytes:
+        """
+        Calculate the Merkle root hash from a list of leaf node data.
+        
+        Args:
+            leaves (List[bytes]): List of leaf node data
+            
+        Returns:
+            bytes: The Merkle root hash
+        """
+        if not leaves:
+            return hashlib.sha256(b'').digest()
+        
+        if len(leaves) == 1:
+            return hashlib.sha256(leaves[0]).digest()
+            
+        # Hash all leaf nodes
+        hashed_leaves = [hashlib.sha256(leaf).digest() for leaf in leaves]
+        
+        # Build the Merkle tree
+        while len(hashed_leaves) > 1:
+            if len(hashed_leaves) % 2 != 0:
+                # Duplicate the last element if there's an odd number
+                hashed_leaves.append(hashed_leaves[-1])
+                
+            # Combine adjacent pairs and hash them
+            next_level = []
+            for i in range(0, len(hashed_leaves), 2):
+                combined = hashed_leaves[i] + hashed_leaves[i+1]
+                next_level.append(hashlib.sha256(combined).digest())
+                
+            hashed_leaves = next_level
+            
+        # Return the root hash
+        return hashed_leaves[0]
+
+    def verify_nonce(self, difficulty: int = 1) -> bool:
+        """
+        Verify that the nonce produces a valid Merkle tree root hash
+        with the specified number of leading zero bits.
+        
+        Args:
+            difficulty (int): Number of leading zero bits required in the root hash
+        
+        Returns:
+            bool: True if the nonce is valid, False otherwise
+        """
+        # Prepare the message data
+        message_data = self.message.to_bytes()
+        timestamp_data = struct.pack('!Q', self.timestamp)
+        encrypted_flag = b'\x01' if self.encrypted else b'\x00'
+        
+        # Calculate the Merkle root
+        merkle_root = self._calculate_merkle_root([
+            timestamp_data,
+            encrypted_flag,
+            self.nonce,
+            message_data
+        ])
+        
+        # Calculate how many bytes need to be zero
+        zero_bytes = difficulty // 8
+        # Calculate how many bits in the last byte need to be zero
+        remaining_bits = difficulty % 8
+        
+        # Create a mask for the remaining bits
+        mask = 0
+        if remaining_bits > 0:
+            mask = 0xFF >> remaining_bits
+        
+        # Check if it meets the difficulty requirement
+        valid = True
+        
+        # Check full zero bytes
+        for i in range(zero_bytes):
+            if merkle_root[i] != 0:
+                valid = False
+                break
+        
+        # If we need to check partial bits in a byte
+        if valid and remaining_bits > 0:
+            # The next byte should have required number of leading zeros
+            if (merkle_root[zero_bytes] & (0xFF ^ mask)) != 0:
+                valid = False
+        
+        return valid
+    
+    def to_bytes(self) -> bytes:
+        """
+        Convert this Envelope to bytes.
+        
+        Returns:
+            bytes: Serialized envelope
+        """
+        return encode([
+            struct.pack('!Q', self.timestamp),
+            b'\x01' if self.encrypted else b'\x00',
+            self.nonce,
+            self.message.to_bytes()
+        ])
+    
+    @classmethod
+    def from_bytes(cls, data: bytes) -> Optional['Envelope']:
+        """
+        Create an Envelope from its serialized form.
+        
+        Args:
+            data (bytes): Serialized envelope
+            
+        Returns:
+            Optional[Envelope]: The deserialized envelope, or None if the data is invalid
+        """
+        try:
+            parts = decode(data)
+            if len(parts) != 4:
+                return None
+                
+            timestamp_data, encrypted_flag, nonce, message_data = parts
+            
+            timestamp = struct.unpack('!Q', timestamp_data)[0]
+            encrypted = encrypted_flag == b'\x01'
+            nonce = nonce
+            message = Message.from_bytes(message_data)
+            
+            if not message:
+                return None
+                
+            return cls(
+                encrypted=encrypted,
+                message=message,
+                nonce=nonce,
+                timestamp=timestamp
+            )
+        except (ValueError, struct.error) as e:
+            print(f"Error deserializing envelope: {e}")
+            return None