SameerMesiah97 commented on code in PR #72165:
URL: https://github.com/apache/airflow/pull/72165#discussion_r3874976233


##########
providers/http/tests/unit/http/hooks/test_http.py:
##########
@@ -898,3 +903,61 @@ async def test_build_request_url_from_endpoint_param(self):
             async with aiohttp.ClientSession() as session:
                 await hook.run(session=session, 
endpoint="test.com:8080/v1/test")
                 assert mocked_function.call_args.args[0] == 
"http://test.com:8080/v1/test";
+
+
+class TestConnectionHeadersAcrossRedirects:
+    """Connection-supplied headers must not follow a redirect to a different 
host.
+
+    ``requests`` strips only ``Authorization`` (``Session.rebuild_auth``), so 
a credential
+    carried in a differently-named header — the documented ``extra``-field 
pattern — would
+    otherwise be replayed to whatever host the redirect points at.
+    """
+
+    @staticmethod
+    def _redirect(session, old_url, new_url):
+        original = requests.Request("GET", old_url).prepare()
+        response = Response()
+        response.request = original
+        prepared = requests.Request("GET", new_url).prepare()
+        prepared.headers.update(session.headers)
+        session.rebuild_auth(prepared, response)
+        return prepared.headers
+
+    def test_connection_headers_dropped_on_cross_host_redirect(self):
+        session = _ConnectionHeaderSession()
+        session.headers.update({"X-API-Key": "secret", "Accept": 
"application/json"})
+        session.connection_header_keys.update({"X-API-Key", "Accept"})
+
+        headers = self._redirect(session, "https://original.example.com/a";, 
"https://evil.example.com/b";)
+
+        assert "X-API-Key" not in headers
+        assert "Accept" not in headers
+
+    def test_connection_headers_kept_on_same_host_redirect(self):
+        session = _ConnectionHeaderSession()
+        session.headers.update({"X-API-Key": "secret"})
+        session.connection_header_keys.update({"X-API-Key"})
+
+        headers = self._redirect(session, "https://example.com/a";, 
"https://example.com/b";)
+
+        assert headers["X-API-Key"] == "secret"
+
+    def test_caller_supplied_headers_are_not_dropped(self):
+        """Only headers that came from the connection are stripped."""
+        session = _ConnectionHeaderSession()
+        session.headers.update({"X-Caller": "kept"})
+
+        headers = self._redirect(session, "https://original.example.com/a";, 
"https://evil.example.com/b";)
+
+        assert headers["X-Caller"] == "kept"
+
+    @mock.patch(
+        "airflow.hooks.base.BaseHook.get_connection",

Review Comment:
   This patch could be causing one of the failing tests. I would try the below:
   
   `"airflow.providers.http.hooks.http.HttpHook.get_connection"`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to