Refactor header handling in AuthenticatedOTLPExporter to prevent critical headers from being overridden by user-supplied values. Update tests to verify protection of critical headers and ensure proper JWT token usage.

Author: Dwij1704

agentops/client/api/base.py

@@ -99,7 +99,7 @@ async def async_request(
         url = self._get_full_url(path)
 
         try:
-            response_data = await HttpClient.async_request(
+            response_data = await self.http_client.async_request(
                 method=method, url=url, data=data, headers=headers, timeout=timeout
             )
             return response_data

agentops/sdk/exporters.py

@@ -54,10 +54,13 @@ def __init__(
         # Store any additional kwargs for potential future use
         self._custom_kwargs = kwargs
 
+        # Filter headers to prevent override of critical headers
+        filtered_headers = self._filter_user_headers(headers) if headers else None
+
         # Initialize parent with only known parameters
         parent_kwargs = {}
-        if headers is not None:
-            parent_kwargs["headers"] = headers
+        if filtered_headers is not None:
+            parent_kwargs["headers"] = filtered_headers
         if timeout is not None:
             parent_kwargs["timeout"] = timeout
         if compression is not None:
@@ -66,24 +69,49 @@ def __init__(
         super().__init__(endpoint=endpoint, **parent_kwargs)
 
     def _get_current_jwt(self) -> Optional[str]:
-        """Get the current JWT token from the provider."""
+        """Get the current JWT token from the provider or stored JWT."""
         if self._jwt_provider:
             try:
                 return self._jwt_provider()
             except Exception as e:
                 logger.warning(f"Failed to get JWT token: {e}")
-        return None
+        return self._jwt
+
+    def _filter_user_headers(self, headers: Optional[Dict[str, str]]) -> Optional[Dict[str, str]]:
+        """Filter user-supplied headers to prevent override of critical headers."""
+        if not headers:
+            return None
+
+        # Define critical headers that cannot be overridden by user-supplied headers
+        PROTECTED_HEADERS = {
+            "authorization",
+            "content-type",
+            "user-agent",
+            "x-api-key",
+            "api-key",
+            "bearer",
+            "x-auth-token",
+            "x-session-token",
+        }
+
+        filtered_headers = {}
+        for key, value in headers.items():
+            if key.lower() not in PROTECTED_HEADERS:
+                filtered_headers[key] = value
+
+        return filtered_headers if filtered_headers else None
 
     def _prepare_headers(self, headers: Optional[Dict[str, str]] = None) -> Dict[str, str]:
         """Prepare headers with current JWT token."""
         # Start with base headers
         prepared_headers = dict(self._headers)
 
-        # Add any additional headers
-        if headers:
-            prepared_headers.update(headers)
+        # Add any additional headers, but only allow non-critical headers
+        filtered_headers = self._filter_user_headers(headers)
+        if filtered_headers:
+            prepared_headers.update(filtered_headers)
 
-        # Add current JWT token if available
+        # Add current JWT token if available (this ensures Authorization cannot be overridden)
         jwt_token = self._get_current_jwt()
         if jwt_token:
             prepared_headers["Authorization"] = f"Bearer {jwt_token}"

tests/unit/sdk/test_exporters.py

@@ -172,11 +172,30 @@ def test_headers_merging(self):
         # Verify the exporter was created successfully
         self.assertIsInstance(exporter, AuthenticatedOTLPExporter)
 
-    def test_headers_override_authorization(self):
-        """Test that custom Authorization header overrides the default one."""
-        custom_headers = {"Authorization": "Custom-Auth custom-token", "X-Custom-Header": "test-value"}
-
-        exporter = AuthenticatedOTLPExporter(endpoint=self.endpoint, jwt=self.jwt, headers=custom_headers)
+    def test_headers_protected_from_override(self):
+        """Test that critical headers cannot be overridden by user-supplied headers."""
+        # Attempt to override critical headers
+        malicious_headers = {
+            "Authorization": "Malicious-Auth malicious-token",
+            "Content-Type": "text/plain",
+            "User-Agent": "malicious-agent",
+            "X-API-Key": "malicious-key",
+            "X-Custom-Header": "test-value",  # This should be allowed
+        }
+
+        exporter = AuthenticatedOTLPExporter(endpoint=self.endpoint, jwt=self.jwt, headers=malicious_headers)
+
+        # Test the _prepare_headers method directly to verify protection
+        prepared_headers = exporter._prepare_headers(malicious_headers)
+
+        # Critical headers should not be overridden
+        self.assertEqual(prepared_headers["Authorization"], f"Bearer {self.jwt}")
+        self.assertNotEqual(prepared_headers.get("Content-Type"), "text/plain")
+        self.assertNotEqual(prepared_headers.get("User-Agent"), "malicious-agent")
+        self.assertNotIn("X-API-Key", prepared_headers)  # Should be filtered out
+
+        # Non-critical headers should be allowed
+        self.assertEqual(prepared_headers.get("X-Custom-Header"), "test-value")
 
         # Verify the exporter was created successfully
         self.assertIsInstance(exporter, AuthenticatedOTLPExporter)