tftp: negotiate blocksize and timeout

Signed-off-by: Benny Zlotnik <[email protected]>

Author: bennyz

contrib/drivers/tftp/src/jumpstarter_driver_tftp/server.py

@@ -14,6 +14,7 @@ class Opcode(IntEnum):
     DATA = 3
     ACK = 4
     ERROR = 5
+    OACK = 6
 
 
 class TftpErrorCode(IntEnum):
@@ -136,7 +137,6 @@ def datagram_received(self, data: bytes, addr: Tuple[str, int]):
 
     async def _handle_read_request(self, data: bytes, addr: Tuple[str, int]):
         try:
-            # Parse filename and mode from request
             parts = data[2:].split(b'\x00')
             if len(parts) < 2:
                 self.logger.error(f"Invalid RRQ format from {addr}")
@@ -145,7 +145,18 @@ async def _handle_read_request(self, data: bytes, addr: Tuple[str, int]):
             filename = parts[0].decode('utf-8')
             mode = parts[1].decode('utf-8').lower()
 
-            self.logger.info(f"RRQ from {addr}: '{filename}' in mode '{mode}'")
+            options = {}
+            i = 2
+            while i < len(parts) - 1:
+                try:
+                    opt_name = parts[i].decode('utf-8').lower()
+                    opt_value = parts[i + 1].decode('utf-8')
+                    options[opt_name] = opt_value
+                    i += 2
+                except IndexError:
+                    break
+
+            self.logger.info(f"RRQ from {addr}: '{filename}' in mode '{mode}' with options {options}")
 
             if mode not in ('netascii', 'octet'):
                 self.logger.warning(f"Unsupported transfer mode '{mode}' from {addr}")
@@ -162,17 +173,40 @@ async def _handle_read_request(self, data: bytes, addr: Tuple[str, int]):
                 return
 
             if not is_subpath(resolved_path, self.server.root_dir):
-                self.logger.error(f"Access violation: {resolved_path} is outside the root directory")
+                self.logger.error(f"Access violation: {resolved_path} is outside root directory")
                 self._send_error(addr, TftpErrorCode.ACCESS_VIOLATION, "Access denied")
                 return
 
+            negotiated_options = {}
+
+            if 'blksize' in options:
+                try:
+                    requested_blksize = int(options['blksize'])
+                    if 512 <= requested_blksize <= 8192:
+                        negotiated_options['blksize'] = requested_blksize
+                    else:
+                        negotiated_options['blksize'] = 512
+                except ValueError:
+                    negotiated_options['blksize'] = 512
+            else:
+                negotiated_options['blksize'] = self.server.block_size
+
+            if 'timeout' in options:
+                try:
+                    requested_timeout = int(options['timeout'])
+                    if 1 <= requested_timeout <= 255:
+                        negotiated_options['timeout'] = requested_timeout
+                except ValueError:
+                    pass
+
             transfer = TftpReadTransfer(
                 server=self.server,
                 filepath=resolved_path,
                 client_addr=addr,
-                block_size=self.server.block_size,
-                timeout=self.server.timeout,
-                retries=self.server.retries
+                block_size=negotiated_options['blksize'],
+                timeout=negotiated_options.get('timeout', self.server.timeout),
+                retries=self.server.retries,
+                negotiated_options=negotiated_options if options else None
             )
             self.server.register_transfer(transfer)
             asyncio.create_task(transfer.start())
@@ -181,6 +215,16 @@ async def _handle_read_request(self, data: bytes, addr: Tuple[str, int]):
             self.logger.error(f"Error handling RRQ from {addr}: {e}")
             self._send_error(addr, TftpErrorCode.NOT_DEFINED, str(e))
 
+    def _send_oack(self, addr: Tuple[str, int], options: dict):
+        """Send Option Acknowledgment (OACK) packet."""
+        oack_data = Opcode.OACK.to_bytes(2, 'big')
+        for opt_name, opt_value in options.items():
+            oack_data += f"{opt_name}\0{str(opt_value)}\0".encode('utf-8')
+
+        if self.transport:
+            self.transport.sendto(oack_data, addr)
+            self.logger.debug(f"Sent OACK to {addr} with options {options}")
+
     def _send_error(self, addr: Tuple[str, int], error_code: TftpErrorCode, message: str):
         error_packet = (
             Opcode.ERROR.to_bytes(2, 'big') +
@@ -232,16 +276,14 @@ async def cleanup(self):
 
 
 class TftpReadTransfer(TftpTransfer):
-    """
-    Handles a TFTP Read (RRQ) transfer.
-    """
-
     def __init__(self, server: TftpServer, filepath: pathlib.Path, client_addr: Tuple[str, int],
-                 block_size: int, timeout: float, retries: int):
+                 block_size: int, timeout: float, retries: int, negotiated_options: Optional[dict] = None):
         super().__init__(server, filepath, client_addr, block_size, timeout, retries)
-        self.block_num = 1
+        self.block_num = 0
         self.ack_received = asyncio.Event()
         self.last_ack = 0
+        self.negotiated_options = negotiated_options
+        self.oack_confirmed = False
 
     async def start(self):
         self.logger.info(f"Starting read transfer of '{self.filepath.name}' to {self.client_addr}")
@@ -256,89 +298,117 @@ async def start(self):
         self.logger.debug(f"Transfer bound to local {local_addr}")
 
         try:
+            if self.negotiated_options:
+                oack_packet = self._create_oack_packet()
+                if not await self._send_with_retries(oack_packet, is_oack=True):
+                    self.logger.error("Failed to get acknowledgment for OACK")
+                    return
+                self.block_num = 1
+
             async with aiofiles.open(self.filepath, 'rb') as f:
                 while True:
                     if self.server.shutdown_event.is_set():
                         self.logger.info(f"Server shutdown detected, stopping transfer to {self.client_addr}")
                         break
+
                     data = await f.read(self.block_size)
-                    if data:
-                        packet = (
-                            Opcode.DATA.to_bytes(2, 'big') +
-                            self.block_num.to_bytes(2, 'big') +
-                            data
-                        )
+                    if not data and self.block_num == 1:
+                        # Empty file case
+                        packet = self._create_data_packet(b'')
+                        await self._send_with_retries(packet)
+                        break
+                    elif data:
+                        packet = self._create_data_packet(data)
                         success = await self._send_with_retries(packet)
                         if not success:
                             self.logger.error(f"Failed to send block {self.block_num} to {self.client_addr}")
                             break
+
                         self.logger.debug(f"Block {self.block_num} sent successfully")
                         self.block_num += 1
 
-                        # If the data read is less than block_size, this is the last packet
                         if len(data) < self.block_size:
-                            self.logger.info(f"Final block {self.block_num - 1} reached for {self.client_addr}")
+                            self.logger.info(f"Final block {self.block_num - 1} sent")
                             break
                     else:
-                        # If no data is returned, it means the file size is an exact multiple of block_size
-                        # Send an extra empty DATA packet to signal end of transfer
-                        packet = (
-                            Opcode.DATA.to_bytes(2, 'big') +
-                            self.block_num.to_bytes(2, 'big') +
-                            b''
-                        )
+                        # End of file reached
+                        packet = self._create_data_packet(b'')
                         success = await self._send_with_retries(packet)
                         if not success:
-                            self.logger.error(
-                                f"Failed to send final empty block {self.block_num} "
-                                f"to {self.client_addr}"
-                            )
+                            self.logger.error(f"Failed to send final block {self.block_num}")
                             break
-                        self.logger.info(f"Transfer complete to {self.client_addr}, final block {self.block_num}")
+                        self.logger.info(f"Transfer complete, final block {self.block_num}")
                         break
 
         except Exception as e:
             self.logger.error(f"Error during read transfer: {e}")
         finally:
             await self.cleanup()
 
-    async def _send_with_retries(self, packet: bytes) -> bool:
+    def _create_oack_packet(self) -> bytes:
+        """Create OACK packet with negotiated options."""
+        packet = Opcode.OACK.to_bytes(2, 'big')
+        for opt_name, opt_value in self.negotiated_options.items():
+            packet += f"{opt_name}\0{str(opt_value)}\0".encode('utf-8')
+        return packet
+
+    def _create_data_packet(self, data: bytes) -> bytes:
+        """Create DATA packet with block number and data."""
+        return (
+            Opcode.DATA.to_bytes(2, 'big') +
+            self.block_num.to_bytes(2, 'big') +
+            data
+        )
+
+    def _send_packet(self, packet: bytes):
+        """
+        Sends a packet to the client.
+        """
+        self.transport.sendto(packet)
+        if packet[0:2] == Opcode.DATA.to_bytes(2, 'big'):
+            block = int.from_bytes(packet[2:4], 'big')
+            data_length = len(packet) - 4
+            self.logger.debug(f"Sent DATA block {block} ({data_length} bytes) to {self.client_addr}")
+        elif packet[0:2] == Opcode.OACK.to_bytes(2, 'big'):
+            self.logger.debug(f"Sent OACK to {self.client_addr}")
+
+    async def _send_with_retries(self, packet: bytes, is_oack: bool = False) -> bool:
         self.current_packet = packet
+        expected_block = 0 if is_oack else self.block_num
+
         for attempt in range(1, self.retries + 1):
             try:
                 self._send_packet(packet)
-                self.logger.debug(f"Sent DATA block {self.block_num}, waiting for ACK (Attempt {attempt})")
+                self.logger.debug(f"Sent {'OACK' if is_oack else 'DATA'} block {expected_block}, waiting for ACK (Attempt {attempt})")
                 self.ack_received.clear()
                 await asyncio.wait_for(self.ack_received.wait(), timeout=self.timeout)
 
-                if self.last_ack == self.block_num:
-                    self.logger.debug(f"ACK received for block {self.block_num}")
+                if self.last_ack == expected_block:
+                    self.logger.debug(f"ACK received for block {expected_block}")
                     return True
                 else:
-                    self.logger.warning(f"Received wrong ACK: expected {self.block_num}, got {self.last_ack}")
+                    self.logger.warning(f"Received wrong ACK: expected {expected_block}, got {self.last_ack}")
 
             except asyncio.TimeoutError:
-                self.logger.warning(f"Timeout waiting for ACK of block {self.block_num} (Attempt {attempt})")
+                self.logger.warning(f"Timeout waiting for ACK of block {expected_block} (Attempt {attempt})")
 
         return False
 
-    def _send_packet(self, packet: bytes):
-        """
-        Sends a DATA packet to the client.
-        """
-        self.transport.sendto(packet)
-        block = int.from_bytes(packet[2:4], 'big')
-        data_length = len(packet) - 4
-        self.logger.debug(f"Sent DATA block {block} ({data_length} bytes) to {self.client_addr}")
-
     def handle_ack(self, block_num: int):
         self.logger.debug(f"Received ACK for block {block_num} from {self.client_addr}")
+
+        # special handling for OACK acknowledgment
+        if not self.oack_confirmed and self.negotiated_options and block_num == 0:
+            self.oack_confirmed = True
+            self.last_ack = block_num
+            self.ack_received.set()
+            return
+
         if block_num == self.block_num:
             self.last_ack = block_num
             self.ack_received.set()
         elif block_num == self.block_num - 1:
-            # Duplicate ACK for previous block, resend current packet
-            self.logger.warning(f"Duplicate ACK for block {block_num} received, resending DATA block {self.block_num}")
+            self.logger.warning(f"Duplicate ACK for block {block_num} received, resending block {self.block_num}")
             self.transport.sendto(self.current_packet)
         else:
             self.logger.warning(f"Out of sequence ACK: expected {self.block_num}, got {block_num}")