tftp: remove checksum caching

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

Author: bennyz

packages/jumpstarter-driver-tftp/jumpstarter_driver_tftp/driver.py

@@ -17,36 +17,28 @@
 
 class TftpError(Exception):
     """Base exception for TFTP server errors"""
-
     pass
 
-
 class ServerNotRunning(TftpError):
     """Server is not running"""
-
     pass
 
-
 class FileNotFound(TftpError):
     """File not found"""
-
     pass
 
-
 @dataclass(kw_only=True)
 class Tftp(Driver):
     """TFTP Server driver for Jumpstarter"""
 
     root_dir: str = "/var/lib/tftpboot"
     host: str = field(default='')
     port: int = 69
-    checksum_suffix: str = ".sha256"
     server: Optional["TftpServer"] = field(init=False, default=None)
     server_thread: Optional[threading.Thread] = field(init=False, default=None)
     _shutdown_event: threading.Event = field(init=False, default_factory=threading.Event)
     _loop_ready: threading.Event = field(init=False, default_factory=threading.Event)
     _loop: Optional[asyncio.AbstractEventLoop] = field(init=False, default=None)
-    _checksums: dict[str, str] = field(default_factory=dict)
 
     def __post_init__(self):
         super().__post_init__()
@@ -73,10 +65,7 @@ def _start_server(self):
         asyncio.set_event_loop(self._loop)
         self.server = TftpServer(host=self.host, port=self.port, root_dir=self.root_dir)
         try:
-            # Signal that the loop is ready
             self._loop_ready.set()
-
-            # Run the server until shutdown is requested
             self._loop.run_until_complete(self._run_server())
         except Exception as e:
             self.logger.error(f"Error running TFTP server: {e}")
@@ -86,7 +75,6 @@ def _start_server(self):
                 self._loop.close()
             except Exception as e:
                 self.logger.error(f"Error during event loop cleanup: {e}")
-
             self._loop = None
             self.logger.info("TFTP server thread completed")
 
@@ -111,11 +99,9 @@ def start(self):
             self.logger.warning("TFTP server is already running")
             return
 
-        # Clear any previous shutdown state
         self._shutdown_event.clear()
         self._loop_ready.clear()
 
-        # Start the server thread
         self.server_thread = threading.Thread(target=self._start_server, daemon=True)
         self.server_thread.start()
 
@@ -133,7 +119,6 @@ def stop(self):
             return
 
         self.logger.info("Initiating TFTP server shutdown")
-
         self._shutdown_event.set()
         self.server_thread.join(timeout=10)
         if self.server_thread.is_alive():
@@ -148,7 +133,6 @@ def list_files(self) -> list[str]:
 
     @export
     async def put_file(self, filename: str, src_stream, client_checksum: str):
-        """Compute and store checksum at write time"""
         file_path = os.path.join(self.root_dir, filename)
 
         try:
@@ -160,47 +144,37 @@ async def put_file(self, filename: str, src_stream, client_checksum: str):
                     async for chunk in src:
                         await dst.send(chunk)
 
-            current_checksum = self._compute_checksum(file_path)
-            self._checksums[filename] = current_checksum
-            self._write_checksum_file(filename, current_checksum)
             return filename
         except Exception as e:
             raise TftpError(f"Failed to upload file: {str(e)}") from e
 
-
     @export
     def delete_file(self, filename: str):
-        """Delete file and its checksum file"""
         file_path = os.path.join(self.root_dir, filename)
-        checksum_path = self._get_checksum_path(filename)
 
         if not os.path.exists(file_path):
             raise FileNotFound(f"File {filename} not found")
 
         try:
             os.remove(file_path)
-            if os.path.exists(checksum_path):
-                os.remove(checksum_path)
-            self._checksums.pop(filename, None)
+            return filename
         except Exception as e:
             raise TftpError(f"Failed to delete {filename}") from e
 
     @export
     def check_file_checksum(self, filename: str, client_checksum: str) -> bool:
-        """
-        check if the checksum of the file matches the client checksum
-        """
-
         file_path = os.path.join(self.root_dir, filename)
+        self.logger.debug(f"checking checksum for file: {filename}")
+        self.logger.debug(f"file path: {file_path}")
+
         if not os.path.exists(file_path):
+            self.logger.debug(f"File {filename} does not exist")
             return False
 
         current_checksum = self._compute_checksum(file_path)
+        self.logger.debug(f"Computed checksum: {current_checksum}")
+        self.logger.debug(f"Client checksum: {client_checksum}")
 
-        self._checksums[filename] = current_checksum
-        self._write_checksum_file(filename, current_checksum)
-
-        self.logger.debug(f"Client checksum: {client_checksum}, server checksum: {current_checksum}")
         return current_checksum == client_checksum
 
     @export
@@ -216,28 +190,6 @@ def close(self):
             self.stop()
         super().close()
 
-    def _get_checksum_path(self, filename: str) -> str:
-        return os.path.join(self.root_dir, f"{filename}{self.checksum_suffix}")
-
-    def _read_checksum_file(self, filename: str) -> Optional[str]:
-        try:
-            checksum_path = self._get_checksum_path(filename)
-            if os.path.exists(checksum_path):
-                with open(checksum_path, 'r') as f:
-                    return f.read().strip()
-        except Exception as e:
-            self.logger.warning(f"Failed to read checksum file for {filename}: {e}")
-        return None
-
-    def _write_checksum_file(self, filename: str, checksum: str):
-        """Write checksum to the checksum file"""
-        try:
-            checksum_path = self._get_checksum_path(filename)
-            with open(checksum_path, 'w') as f:
-                f.write(f"{checksum}\n")
-        except Exception as e:
-            self.logger.error(f"Failed to write checksum file for {filename}: {e}")
-
     def _compute_checksum(self, path: str) -> str:
         hasher = hashlib.sha256()
         with open(path, "rb") as f:

packages/jumpstarter-driver-tftp/jumpstarter_driver_tftp/driver_test.py

@@ -2,7 +2,6 @@
 import os
 import tempfile
 from pathlib import Path
-from typing import Optional
 from uuid import uuid4
 
 import anyio
@@ -12,7 +11,6 @@
 from jumpstarter_driver_tftp.driver import (
     FileNotFound,
     Tftp,
-    TftpError,
 )
 
 from jumpstarter.common.resources import ClientStreamResource
@@ -23,14 +21,12 @@ def temp_dir():
     with tempfile.TemporaryDirectory() as tmpdir:
         yield tmpdir
 
-
 @pytest.fixture
 def server(temp_dir):
     server = Tftp(root_dir=temp_dir, host="127.0.0.1")
     yield server
     server.close()
 
-
 @pytest.mark.anyio
 async def test_tftp_file_operations(server):
     filename = "test.txt"
@@ -64,103 +60,48 @@ async def send_data():
     with pytest.raises(FileNotFound):
         server.delete_file("nonexistent.txt")
 
-
 def test_tftp_host_config(temp_dir):
     custom_host = "192.168.1.1"
     server = Tftp(root_dir=temp_dir, host=custom_host)
     assert server.get_host() == custom_host
 
-
 def test_tftp_root_directory_creation(temp_dir):
     new_dir = os.path.join(temp_dir, "new_tftp_root")
     server = Tftp(root_dir=new_dir)
     assert os.path.exists(new_dir)
     server.close()
 
-
-@pytest.mark.anyio
-async def test_tftp_checksum_validation(server):
-    filename = "test_checksum.txt"
-    test_data = b"Hello world"
-    modified_data = b"Modified content"
-
-    def compute_checksum(data: bytes) -> str:
-        return hashlib.sha256(data).hexdigest()
-
-    initial_checksum = await _upload_file(server, filename, test_data)
-    assert filename in server.list_files()
-    assert compute_checksum(test_data) == initial_checksum
-
-    # Second upload with same data should be skipped
-    same_data_checksum = await _upload_file(server, filename, test_data)
-    assert same_data_checksum == initial_checksum
-
-    modified_checksum = await _upload_file(server, filename, modified_data)
-    assert modified_checksum != initial_checksum
-    assert Path(server.root_dir).joinpath(filename).read_bytes() == modified_data
-
-    empty_checksum = await _upload_file(server, "empty.txt", b"")
-    assert empty_checksum == hashlib.sha256(b"").hexdigest()
-
 @pytest.mark.anyio
 async def test_tftp_detect_corrupted_file(server):
     filename = "corrupted.txt"
     original_data = b"Original Data"
     client_checksum = hashlib.sha256(original_data).hexdigest()
 
     await _upload_file(server, filename, original_data)
+
     assert server.check_file_checksum(filename, client_checksum)
 
     file_path = Path(server.root_dir, filename)
-    with open(file_path, "wb") as f:
-        f.write(b"Corrupted Data")
+    file_path.write_bytes(b"corrupted Data")
 
     assert not server.check_file_checksum(filename, client_checksum)
 
-@pytest.mark.anyio
-async def test_tftp_reupload_different_checksum(server):
-    filename = "reupload.txt"
-    initial_data = b"Initial Data"
-    updated_data = b"Updated Data"
-    initial_checksum = hashlib.sha256(initial_data).hexdigest()
-    updated_checksum = hashlib.sha256(updated_data).hexdigest()
-
-    await _upload_file(server, filename, initial_data)
-    assert server.check_file_checksum(filename, initial_checksum)
-    assert Path(server.root_dir, filename).read_bytes() == initial_data
-
-    await _upload_file(server, filename, updated_data, client_checksum=updated_checksum)
-    assert server.check_file_checksum(filename, updated_checksum)
-    assert Path(server.root_dir, filename).read_bytes() == updated_data
-
 @pytest.fixture
 def anyio_backend():
     return "asyncio"
 
-async def _upload_file(server, filename: str, data: bytes, client_checksum: Optional[str] = None) -> str:
+async def _upload_file(server, filename: str, data: bytes) -> str:
     send_stream, receive_stream = create_memory_object_stream()
     resource_uuid = uuid4()
     server.resources[resource_uuid] = receive_stream
     resource_handle = ClientStreamResource(uuid=resource_uuid).model_dump(mode="json")
-    client_checksum = client_checksum or hashlib.sha256(data).hexdigest()
 
     async def send_data():
         await send_stream.send(data)
         await send_stream.aclose()
 
     async with anyio.create_task_group() as tg:
         tg.start_soon(send_data)
-        await server.put_file(filename, resource_handle, client_checksum)
-
-    return server._compute_checksum(os.path.join(server.root_dir, filename))
-
-@pytest.mark.anyio
-async def test_tftp_path_traversal_attempt(server):
-    malicious_filename = "../../evil.txt"
-
-    resource_uuid = uuid4()
-    server.resources[resource_uuid] = None
-    resource_handle = ClientStreamResource(uuid=resource_uuid).model_dump(mode="json")
+        await server.put_file(filename, resource_handle, hashlib.sha256(data).hexdigest())
 
-    with pytest.raises(TftpError, match="Invalid target path"):
-        await server.put_file(malicious_filename, resource_handle, "checksum")
+    return hashlib.sha256(data).hexdigest()