flashers: fail early on header parsing validation

Signed-off-by: Benny Zlotnik <bzlotnik@redhat.com>
(cherry picked from commit f381274)

Author: bennyz

packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/client.py

@@ -70,7 +70,7 @@ def bootloader_shell(self):
                 pass
             yield self.serial
 
-    def flash(
+    def flash(  # noqa: C901
         self,
         path: PathBuf,
         *,
@@ -87,6 +87,9 @@ def flash(
         if bearer_token:
             bearer_token = self._validate_bearer_token(bearer_token)
 
+        if headers:
+            headers = self._validate_header_dict(headers)
+
         """Flash image to DUT"""
         should_download_to_httpd = True
         image_url = ""
@@ -689,15 +692,15 @@ def _validate_header_dict(self, header_map: dict[str, str]) -> dict[str, str]:
             key = key.strip()
             value = value.strip()
             if not key:
-                raise click.ClickException(f"Invalid header key: '{key}'")
+                raise ArgumentError(f"Invalid header key: '{key}'")
 
             if not token_re.match(key):
-                raise click.ClickException(f"Invalid header name '{key}': must be an HTTP token (RFC7230)")
+                raise ArgumentError(f"Invalid header name '{key}': must be an HTTP token (RFC7230)")
             if any(c in ("\r", "\n") for c in key) or any(c in ("\r", "\n") for c in value):
-                raise click.ClickException("Header names/values must not contain CR/LF")
+                raise ArgumentError("Header names/values must not contain CR/LF")
             kl = key.lower()
             if kl in seen:
-                raise click.ClickException(f"Duplicate header '{key}'")
+                raise ArgumentError(f"Duplicate header '{key}'")
             seen.add(kl)
         return header_map
 
@@ -710,7 +713,10 @@ def _parse_headers(self, headers: list[str]) -> dict[str, str]:
             key, value = h.split(":", 1)
             header_map[key.strip()] = value.strip()
 
-        return self._validate_header_dict(header_map)
+        try:
+            return self._validate_header_dict(header_map)
+        except ArgumentError as e:
+            raise click.ClickException(str(e)) from e
 
     def _prepare_headers(self, headers: dict[str, str] | None, bearer_token: str | None) -> str:
         all_headers = headers.copy() if headers else {}
@@ -720,8 +726,11 @@ def _prepare_headers(self, headers: dict[str, str] | None, bearer_token: str | N
             else:
                 all_headers["Authorization"] = f"Bearer {bearer_token}"
 
-        validated_headers = self._validate_header_dict(all_headers)
-        return self._curl_header_args(validated_headers)
+        if bearer_token and "Authorization" not in (headers or {}):
+            auth_header = {"Authorization": all_headers["Authorization"]}
+            self._validate_header_dict(auth_header)
+
+        return self._curl_header_args(all_headers)
 
     def _validate_bearer_token(self, token: str | None) -> str | None:
         if token is None:

packages/jumpstarter-driver-flashers/jumpstarter_driver_flashers/client_test.py

@@ -2,6 +2,7 @@
 import pytest
 
 from .client import BaseFlasherClient
+from jumpstarter.common.exceptions import ArgumentError
 
 
 class MockFlasherClient(BaseFlasherClient):
@@ -14,6 +15,9 @@ def __init__(self):
             "MockLogger", (), {"warning": lambda msg: None, "info": lambda msg: None, "error": lambda msg: None}
         )()
 
+    def close(self):
+        pass
+
 
 def test_validate_bearer_token_fails_invalid():
     """Test bearer token validation fails with invalid tokens"""
@@ -37,3 +41,11 @@ def test_curl_header_args_handles_quotes():
     assert "'\"'\"'" in result
     assert result.startswith("-H '")
     assert result.endswith("'")
+
+
+def test_flash_fails_with_invalid_headers():
+    """Test flash method fails early with invalid headers"""
+    client = MockFlasherClient()
+
+    with pytest.raises(ArgumentError, match="Invalid header name 'Invalid Header': must be an HTTP token"):
+        client.flash("test.raw", headers={"Invalid Header": "value"})