From 9ebfadbe89bde011fd31ed6d6503fac0086b13a7 Mon Sep 17 00:00:00 2001 From: max Date: Fri, 2 Oct 2026 22:43:18 +0200 Subject: [PATCH] Optimize server startup speed, client update caching, and HTTP 304 conditional downloads --- src/Server.py | 53 ++++++++++++++++++++++++++---- src/client_updater.py | 48 ++++++++++++++++++++------- src/storage_engine.py | 33 +++++++++---------- src/update_manager.py | 75 ++++++++++++++++++++++++++++++++++++++----- tests/test_updates.py | 49 ++++++++++++++++++++++++++-- 5 files changed, 213 insertions(+), 45 deletions(-) diff --git a/src/Server.py b/src/Server.py index 797be83..6b055aa 100644 --- a/src/Server.py +++ b/src/Server.py @@ -1497,11 +1497,17 @@ def check_client_update( @app.get("/api/client/update/download") -def download_client_update_package(os: Optional[str] = "windows", os_name: Optional[str] = None): +def download_client_update_package( + request: Request, + os: Optional[str] = "windows", + os_name: Optional[str] = None +): """ Streams the latest client binary or installer package for the specified OS. + Supports HTTP ETag and conditional 304 Not Modified requests to minimize network bandwidth. """ import os as _os + import email.utils target_os = os_name or os or "windows" mgr: ServerUpdateManager = SERVER_STATE.get("update_manager") or UPDATE_MANAGER pkg_tuple = mgr.get_client_package_path(os_type=target_os) @@ -1512,14 +1518,50 @@ def download_client_update_package(os: Optional[str] = "windows", os_name: Optio if not _os.path.exists(file_path): raise HTTPException(status_code=404, detail="Client package file missing on server disk.") + etag = f'"{sha256}"' + if request is not None: + if_none_match = request.headers.get("if-none-match") + if if_none_match and (if_none_match == etag or if_none_match.strip('"') == sha256): + return Response( + status_code=304, + headers={ + "ETag": etag, + "Cache-Control": "public, max-age=300, must-revalidate", + "X-Checksum-SHA256": sha256 + } + ) + + try: + mtime = _os.path.getmtime(file_path) + last_modified = email.utils.formatdate(mtime, usegmt=True) + except Exception: + last_modified = None + + if_modified_since = request.headers.get("if-modified-since") + if if_modified_since and last_modified and if_modified_since == last_modified: + return Response( + status_code=304, + headers={ + "ETag": etag, + "Cache-Control": "public, max-age=300, must-revalidate", + "X-Checksum-SHA256": sha256 + } + ) + + headers = { + "ETag": etag, + "Cache-Control": "public, max-age=300, must-revalidate", + "X-Checksum-SHA256": sha256, + "Content-Disposition": f'attachment; filename="{filename}"' + } + if last_modified: + headers["Last-Modified"] = last_modified + return FileResponse( path=file_path, filename=filename, media_type="application/octet-stream", - headers={ - "X-Checksum-SHA256": sha256, - "Content-Disposition": f'attachment; filename="{filename}"' - } + headers=headers ) @@ -2360,7 +2402,6 @@ def main(): csr_path = os.path.join(cert_dir, "ca.csr") if not os.path.exists(csr_path): enrollment.generate_ca_csr(cert_dir=cert_dir) - export_license_request_lic(config) srv_cert, srv_key, srv_pem, srv_key_pem = enrollment.generate_server_cert_if_needed( ca_cert, ca_key, diff --git a/src/client_updater.py b/src/client_updater.py index 9ad6a08..bfcaff1 100644 --- a/src/client_updater.py +++ b/src/client_updater.py @@ -76,22 +76,48 @@ class ClientUpdater: def download_package(self, download_path_rel: str, expected_sha256: str, target_temp: str, timeout: int = 30) -> bool: """ Downloads the client binary package from the server and validates its SHA-256 checksum. + Uses local binary hash checking and ETag conditional headers to avoid redundant downloads. """ + # 1. If running binary already matches target SHA-256, skip downloading + target_exec = sys.executable + if expected_sha256 and os.path.exists(target_exec): + try: + local_sha = self.calculate_sha256(target_exec) + if local_sha and local_sha == expected_sha256.lower(): + print(f"[*] Local binary already matches target checksum ({local_sha[:12]}...). Skipping redundant download.") + return True + except Exception: + pass + base_url = self.get_server_base_url() download_url = f"{base_url.rstrip('/')}/{download_path_rel.lstrip('/')}" print(f"[*] Downloading client update package from {download_url}...") - req = urllib.request.Request( - download_url, - headers={"User-Agent": f"LOGAR-Client/{self.current_version} ({self.os_type})"} - ) - with urllib.request.urlopen(req, timeout=timeout) as resp: - if resp.status != 200: - print(f"[!] Server returned HTTP {resp.status} during update download.") - return False - with open(target_temp, "wb") as f: - while chunk := resp.read(65536): - f.write(chunk) + headers = {"User-Agent": f"LOGAR-Client/{self.current_version} ({self.os_type})"} + if expected_sha256: + headers["If-None-Match"] = f'"{expected_sha256}"' + + req = urllib.request.Request(download_url, headers=headers) + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + if resp.status == 304: + print("[+] Server returned 304 Not Modified. Local binary package is current.") + return True + if resp.status != 200: + print(f"[!] Server returned HTTP {resp.status} during update download.") + return False + with open(target_temp, "wb") as f: + while chunk := resp.read(65536): + f.write(chunk) + except urllib.error.HTTPError as e: + if e.code == 304: + print("[+] Server returned 304 Not Modified. Local binary package is current.") + return True + print(f"[!] HTTP error during update download: {e}") + return False + except Exception as e: + print(f"[!] Error downloading update package: {e}") + return False # Verify SHA-256 actual_sha256 = self.calculate_sha256(target_temp) diff --git a/src/storage_engine.py b/src/storage_engine.py index 3812d5c..5b23d86 100644 --- a/src/storage_engine.py +++ b/src/storage_engine.py @@ -193,19 +193,24 @@ class StorageEngine: ) """) + def _ensure_columns(tbl: str, cols: List[Tuple[str, str]]): + try: + existing = {row[1] for row in conn.execute(f"PRAGMA table_info({tbl})").fetchall()} + for col, col_def in cols: + if col not in existing: + conn.execute(f"ALTER TABLE {tbl} ADD COLUMN {col} {col_def}") + except Exception: + pass + # Migrations for license_config - for col, col_def in [ + _ensure_columns("license_config", [ ("license_id", "TEXT"), ("license_server_url", "TEXT"), ("issued_at", "TEXT"), ("expires_at", "TEXT"), ("license_status", "TEXT DEFAULT 'unlicensed'"), ("features", "TEXT DEFAULT '[]'") - ]: - try: - conn.execute(f"ALTER TABLE license_config ADD COLUMN {col} {col_def}") - except sqlite3.OperationalError: - pass + ]) # 4. Clients Table conn.execute(""" @@ -221,7 +226,7 @@ class StorageEngine: """) # Migrations for clients table - for col, col_def in [ + _ensure_columns("clients", [ ("metrics_json", "TEXT DEFAULT '{}'"), ("cpu_percent", "REAL DEFAULT 0.0"), ("memory_percent", "REAL DEFAULT 0.0"), @@ -229,11 +234,7 @@ class StorageEngine: ("client_version", "TEXT DEFAULT ''"), ("update_status", "TEXT DEFAULT 'up_to_date'"), ("last_update_check", "TEXT") - ]: - try: - conn.execute(f"ALTER TABLE clients ADD COLUMN {col} {col_def}") - except sqlite3.OperationalError: - pass + ]) # 5. Users Table conn.execute(""" @@ -249,14 +250,10 @@ class StorageEngine: """) # Migrations for users table (2FA) - for col, col_def in [ + _ensure_columns("users", [ ("totp_secret", "TEXT"), ("totp_enabled", "INTEGER DEFAULT 0") - ]: - try: - conn.execute(f"ALTER TABLE users ADD COLUMN {col} {col_def}") - except sqlite3.OperationalError: - pass + ]) # Ensure default admin user exists c = conn.cursor() diff --git a/src/update_manager.py b/src/update_manager.py index ce5ec58..2721591 100644 --- a/src/update_manager.py +++ b/src/update_manager.py @@ -45,6 +45,9 @@ class ServerUpdateManager: # Hosted client packages cache self.hosted_client_version: str = LOGAR_VERSION + self._sha256_cache: Dict[str, Tuple[float, int, str]] = {} + self._package_cache: Dict[str, Dict[str, Any]] = {} + self._package_cache_mtimes: Dict[str, float] = {} self._ensure_package_structure() def _ensure_package_structure(self): @@ -52,21 +55,48 @@ class ServerUpdateManager: for os_name in ["windows", "linux"]: os.makedirs(os.path.join(self.update_dir, os_name), exist_ok=True) - @staticmethod - def calculate_sha256(filepath: str) -> str: - """Computes the SHA-256 digest of a local file.""" + @classmethod + def calculate_sha256(cls, filepath: str) -> str: + """Computes the SHA-256 digest of a local file with stat-based memoization.""" if not os.path.exists(filepath): return "" + try: + stat = os.stat(filepath) + mtime = stat.st_mtime + size = stat.st_size + if hasattr(cls, "_global_sha256_cache"): + cached = cls._global_sha256_cache.get(filepath) + if cached and cached[0] == mtime and cached[1] == size: + return cached[2] + except Exception: + mtime = 0.0 + size = 0 + h = hashlib.sha256() with open(filepath, "rb") as f: while chunk := f.read(65536): h.update(chunk) - return h.hexdigest().lower() + digest = h.hexdigest().lower() + + if not hasattr(cls, "_global_sha256_cache"): + cls._global_sha256_cache = {} + cls._global_sha256_cache[filepath] = (mtime, size, digest) + return digest def set_hosted_client_version(self, version_str: str): """Sets the version tag for currently hosted client packages.""" if version_str: self.hosted_client_version = version_str.strip().lstrip("v") + self._package_cache.clear() + + def invalidate_package_cache(self, os_type: Optional[str] = None): + """Clears in-memory package cache.""" + if os_type: + self._package_cache.pop(os_type.lower(), None) + self._package_cache_mtimes.pop(os_type.lower(), None) + else: + self._package_cache.clear() + self._package_cache_mtimes.clear() def register_client_package( self, @@ -115,16 +145,37 @@ class ServerUpdateManager: with open(manifest_path, "w", encoding="utf-8") as mf: json.dump(manifest, mf, indent=2) + # Update cache directly + cached_entry = dict(manifest) + cached_entry["file_path"] = dest_path + cached_entry["exists"] = True + self._package_cache[norm_os] = cached_entry + self._package_cache_mtimes[norm_os] = os.path.getmtime(dest_path) + return manifest - def get_hosted_packages(self) -> Dict[str, Dict[str, Any]]: + def get_hosted_packages(self, force_refresh: bool = False) -> Dict[str, Dict[str, Any]]: """ Returns info on currently hosted client packages for Windows and Linux. - Scans update_dir and falls back to project dist/ directory if present. + Uses in-memory cache validated against file modification times to minimize disk I/O. """ result: Dict[str, Dict[str, Any]] = {} for os_name in ["windows", "linux"]: manifest_path = os.path.join(self.update_dir, os_name, "manifest.json") + target_file = None + + if not force_refresh and os_name in self._package_cache: + cached = self._package_cache[os_name] + tf = cached.get("file_path") + if tf and os.path.exists(tf): + try: + cur_mtime = os.path.getmtime(tf) + if cur_mtime == self._package_cache_mtimes.get(os_name): + result[os_name] = cached + continue + except Exception: + pass + if os.path.exists(manifest_path): try: with open(manifest_path, "r", encoding="utf-8") as mf: @@ -134,6 +185,8 @@ class ServerUpdateManager: info["exists"] = True info["file_path"] = target_file result[os_name] = info + self._package_cache[os_name] = info + self._package_cache_mtimes[os_name] = os.path.getmtime(target_file) continue except Exception: pass @@ -164,7 +217,7 @@ class ServerUpdateManager: if found_pkg and os.path.exists(found_pkg): sha256 = self.calculate_sha256(found_pkg) - result[os_name] = { + pkg_data = { "os": os_name, "filename": os.path.basename(found_pkg), "version": self.hosted_client_version or LOGAR_VERSION, @@ -174,8 +227,11 @@ class ServerUpdateManager: "exists": True, "updated_at": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC") } + result[os_name] = pkg_data + self._package_cache[os_name] = pkg_data + self._package_cache_mtimes[os_name] = os.path.getmtime(found_pkg) else: - result[os_name] = { + pkg_data = { "os": os_name, "filename": None, "version": self.hosted_client_version or LOGAR_VERSION, @@ -184,6 +240,9 @@ class ServerUpdateManager: "file_path": None, "exists": False } + result[os_name] = pkg_data + self._package_cache[os_name] = pkg_data + self._package_cache_mtimes[os_name] = 0.0 return result diff --git a/tests/test_updates.py b/tests/test_updates.py index 98270aa..81c379d 100644 --- a/tests/test_updates.py +++ b/tests/test_updates.py @@ -137,11 +137,34 @@ class TestUpdateSystem(unittest.TestCase): self.assertEqual(chk_resp["latest_version"], "0.3.6.0") # 4. Download package endpoint - dl_resp = Server.download_client_update_package(os="windows") + from starlette.requests import Request as StarletteRequest + req_scope = { + "type": "http", + "method": "GET", + "path": "/api/client/update/download", + "headers": [], + "query_string": b"os=windows" + } + init_req = StarletteRequest(req_scope) + dl_resp = Server.download_client_update_package(request=init_req, os="windows") self.assertEqual(dl_resp.status_code, 200) self.assertIn("X-Checksum-SHA256", dl_resp.headers) + self.assertIn("ETag", dl_resp.headers) + etag = dl_resp.headers["ETag"] - # 5. WebUI rendering verification + # 5. Conditional GET with ETag (304 Not Modified) + cond_scope = { + "type": "http", + "method": "GET", + "path": "/api/client/update/download", + "headers": [(b"if-none-match", etag.encode("utf-8"))], + "query_string": b"os=windows" + } + cached_req = StarletteRequest(cond_scope) + cached_resp = Server.download_client_update_package(request=cached_req, os="windows") + self.assertEqual(cached_resp.status_code, 304) + + # 6. WebUI rendering verification # Simulate server update available self.update_mgr.server_update_available = True self.update_mgr.latest_upstream_version = "0.3.6.0" @@ -155,6 +178,28 @@ class TestUpdateSystem(unittest.TestCase): self.assertIn("tabPane-updates", html) self.assertIn("Client Package Distribution Repository", html) + def test_update_memoization_and_cache_invalidation(self): + """Tests that package info is served from memory cache and invalidates on update.""" + dummy_pkg = os.path.join(self.test_dir, "LOGAR_Client_Linux") + with open(dummy_pkg, "wb") as f: + f.write(b"LINUX_CLIENT_PAYLOAD_V1") + + self.update_mgr.register_client_package("linux", dummy_pkg, version="0.3.6.0") + pkgs_1 = self.update_mgr.get_hosted_packages() + sha1 = pkgs_1["linux"]["sha256"] + + # Repeated call should hit memory cache with identical hash + pkgs_2 = self.update_mgr.get_hosted_packages() + self.assertEqual(pkgs_2["linux"]["sha256"], sha1) + + # Invalidate cache on new registration + with open(dummy_pkg, "wb") as f: + f.write(b"LINUX_CLIENT_PAYLOAD_V2_UPDATED") + self.update_mgr.register_client_package("linux", dummy_pkg, version="0.3.7.0") + pkgs_3 = self.update_mgr.get_hosted_packages() + self.assertEqual(pkgs_3["linux"]["version"], "0.3.7.0") + self.assertNotEqual(pkgs_3["linux"]["sha256"], sha1) + if __name__ == "__main__": unittest.main()