# core/sfp_transport.py """ Provides a reusable transport layer for the Simple Fragmentation Protocol (SFP). This module handles UDP socket communication, SFP header parsing, fragment reassembly, and ACK generation. It is application-agnostic and uses a callback/handler system to pass fully reassembled payloads to the application layer based on the SFP_FLOW identifier. """ import socket import logging import threading import time from typing import Dict, Callable, Optional from controlpanel import config from controlpanel.utils.network import create_udp_socket, close_udp_socket from controlpanel.core.sfp_structures import SFPHeader # Define a type hint for payload handlers PayloadHandler = Callable[[bytearray], None] class SfpTransport: """Manages SFP communication and payload reassembly.""" def __init__(self, host: str, port: int, payload_handlers: Dict[int, PayloadHandler]): """ Initializes the SFP Transport layer. Args: host (str): The local IP address to bind the UDP socket to. port (int): The local port to listen on. payload_handlers (Dict[int, PayloadHandler]): A dictionary mapping SFP_FLOW IDs (as integers) to their corresponding handler functions. Each handler will be called with the complete payload (bytearray). """ self._log_prefix = "[SfpTransport]" logging.info(f"{self._log_prefix} Initializing for {host}:{port}...") self._host = host self._port = port self._payload_handlers = payload_handlers self._socket: Optional[socket.socket] = None self._receiver_thread: Optional[threading.Thread] = None self._stop_event = threading.Event() # Reassembly state dictionaries, managed by this transport layer self._fragments: Dict[tuple, Dict[int, int]] = {} self._buffers: Dict[tuple, bytearray] = {} logging.debug( f"{self._log_prefix} Registered handlers for flows: " f"{[chr(k) if 32 <= k <= 126 else k for k in self._payload_handlers.keys()]}" ) def start(self) -> bool: """ Starts the transport layer by creating the socket and launching the receiver thread. Returns: bool: True if started successfully, False otherwise. """ if self._receiver_thread is not None and self._receiver_thread.is_alive(): logging.warning(f"{self._log_prefix} Start called, but receiver is already running.") return True self._socket = create_udp_socket(self._host, self._port) if not self._socket: logging.critical(f"{self._log_prefix} Failed to create and bind socket. Cannot start.") return False self._stop_event.clear() self._receiver_thread = threading.Thread( target=self._receive_loop, name="SfpTransportThread", daemon=True ) self._receiver_thread.start() logging.info(f"{self._log_prefix} Receiver thread started.") return True def shutdown(self): """Stops the receiver thread and closes the socket.""" logging.info(f"{self._log_prefix} Shutdown initiated.") self._stop_event.set() # The socket is closed here to interrupt the blocking recvfrom call if self._socket: close_udp_socket(self._socket) self._socket = None if self._receiver_thread and self._receiver_thread.is_alive(): logging.debug(f"{self._log_prefix} Waiting for receiver thread to join...") self._receiver_thread.join(timeout=2.0) if self._receiver_thread.is_alive(): logging.warning(f"{self._log_prefix} Receiver thread did not join cleanly.") logging.info(f"{self._log_prefix} Shutdown complete.") def _receive_loop(self): """The main loop that listens for UDP packets and processes them.""" log_prefix = f"{self._log_prefix} Loop" logging.info(f"{log_prefix} Starting receive loop.") while not self._stop_event.is_set(): if not self._socket: logging.error(f"{log_prefix} Socket is not available. Stopping loop.") break try: data, addr = self._socket.recvfrom(65535) if not data: continue except socket.timeout: continue except OSError: # This is expected when the socket is closed during shutdown if not self._stop_event.is_set(): logging.error(f"{log_prefix} Socket error.", exc_info=True) break except Exception: logging.exception(f"{log_prefix} Unexpected error in recvfrom.") time.sleep(0.01) continue self._process_packet(data, addr) logging.info(f"{log_prefix} Receive loop terminated.") def _process_packet(self, data: bytes, addr: tuple): """Parses an SFP packet and handles fragment reassembly.""" header_size = SFPHeader.size() if len(data) < header_size: logging.warning(f"Packet from {addr} is too small for SFP header. Ignoring.") return try: header = SFPHeader.from_buffer_copy(data) except (ValueError, TypeError): logging.error(f"Failed to parse SFP header from {addr}. Ignoring.") return flow, tid = header.SFP_FLOW, header.SFP_TID frag, total_frags = header.SFP_FRAG, header.SFP_TOTFRGAS pl_size, pl_offset = header.SFP_PLSIZE, header.SFP_PLOFFSET total_size = header.SFP_TOTSIZE key = (flow, tid) # Handle ACK Request if header.SFP_FLAGS & 0x01: self._send_ack(addr, data[:header_size]) # Validate packet metadata if total_frags == 0 or total_frags > 60000 or total_size <= 0: logging.warning(f"Invalid metadata for {key}: total_frags={total_frags}, total_size={total_size}. Ignoring.") return # Start of a new transaction if frag == 0: self._cleanup_lingering_transactions(flow, tid) logging.debug(f"New transaction started for key={key}. Total size: {total_size} bytes.") self._fragments[key] = {} try: self._buffers[key] = bytearray(total_size) except (MemoryError, ValueError): logging.error(f"Failed to allocate {total_size} bytes for key={key}. Ignoring transaction.") self._fragments.pop(key, None) return # Check if we are tracking this transaction if key not in self._buffers or key not in self._fragments: logging.debug(f"Ignoring fragment {frag} for untracked transaction key={key}.") return # Store fragment info and copy payload self._fragments[key][frag] = total_frags payload = data[header_size:] bytes_to_copy = min(pl_size, len(payload)) if (pl_offset + bytes_to_copy) > len(self._buffers[key]): logging.error(f"Payload for key={key}, frag={frag} would overflow buffer. Ignoring.") return self._buffers[key][pl_offset : pl_offset + bytes_to_copy] = payload[:bytes_to_copy] # Check for completion if len(self._fragments[key]) == total_frags: #logging.info(f"Transaction complete for key={key}. Handing off to application layer.") # Retrieve completed buffer and clean up state for this key completed_payload = self._buffers.pop(key) self._fragments.pop(key) # Find and call the appropriate handler handler = self._payload_handlers.get(flow) if handler: try: handler(completed_payload) except Exception: logging.exception(f"Error executing payload handler for flow {flow}.") else: logging.warning(f"No payload handler registered for flow ID {flow}.") def _send_ack(self, dest_addr: tuple, original_header_bytes: bytes): """Sends an SFP ACK packet back to the sender.""" log_prefix = f"{self._log_prefix} ACK" if not self._socket: return try: ack_header = bytearray(original_header_bytes) flow = ack_header[SFPHeader.get_field_offset("SFP_FLOW")] # Determine window size based on flow window_size = 0 if flow == ord('M'): # MFD window_size = config.ACK_WINDOW_SIZE_MFD elif flow == ord('S'): # SAR window_size = config.ACK_WINDOW_SIZE_SAR # Modify header for ACK response ack_header[SFPHeader.get_field_offset("SFP_DIRECTION")] = 0x3C # '<' ack_header[SFPHeader.get_field_offset("SFP_WIN")] = window_size original_flags = ack_header[SFPHeader.get_field_offset("SFP_FLAGS")] ack_header[SFPHeader.get_field_offset("SFP_FLAGS")] = (original_flags | 0x01) & ~0x02 self._socket.sendto(ack_header, dest_addr) logging.debug(f"{log_prefix} Sent to {dest_addr} for flow {chr(flow) if 32<=flow<=126 else flow}.") except Exception: logging.exception(f"{log_prefix} Failed to send to {dest_addr}.") def _cleanup_lingering_transactions(self, current_flow: int, current_tid: int): """Removes old, incomplete transactions for the same flow.""" # This is a simplified cleanup. The original was more complex for stats. keys_to_remove = [ key for key in self._fragments if key[0] == current_flow and key[1] != current_tid ] for key in keys_to_remove: logging.warning(f"Cleaning up lingering/incomplete transaction for key={key}.") self._fragments.pop(key, None) self._buffers.pop(key, None)