Skip to content

Commit b64ec6e

Browse files
adrian05-msadrian05-ms
committed
updating func to only use 2 instead of 3 params
1 parent 648d31b commit b64ec6e

3 files changed

Lines changed: 97 additions & 86 deletions

File tree

‎packages/http/httpx/kiota_http/middleware/options/redirect_handler_option.py‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,21 +4,24 @@
44
import httpx
55

66
# Type alias for the scrub sensitive headers callback
7-
ScrubSensitiveHeadersCallback = Callable[[httpx.Headers, httpx.URL, httpx.URL], None]
7+
ScrubSensitiveHeadersCallback = Callable[[httpx.Request, httpx.URL], None]
88

99

1010
def default_scrub_sensitive_headers(
11-
headers: httpx.Headers, original_url: httpx.URL, new_url: httpx.URL
11+
new_request: httpx.Request, original_url: httpx.URL
1212
) -> None:
1313
"""
1414
The default implementation for scrubbing sensitive headers during redirects.
1515
This method removes Authorization and Cookie headers when the host, scheme, or port changes.
1616
Args:
17-
headers: The headers object to modify
17+
new_request: The new redirect request to modify
1818
original_url: The original request URL
19-
new_url: The new redirect URL
2019
"""
21-
if not headers or not original_url or not new_url:
20+
if not new_request or not original_url:
21+
return
22+
23+
new_url = new_request.url
24+
if not new_url:
2225
return
2326

2427
# Remove Authorization and Cookie headers if the request's scheme, host, or port changes
@@ -29,8 +32,8 @@ def default_scrub_sensitive_headers(
2932
)
3033

3134
if is_different_origin:
32-
headers.pop("Authorization", None)
33-
headers.pop("Cookie", None)
35+
new_request.headers.pop("Authorization", None)
36+
new_request.headers.pop("Cookie", None)
3437

3538
# Note: Proxy-Authorization is not handled here as proxy configuration in httpx
3639
# is managed at the transport level and not accessible to middleware.

‎packages/http/httpx/kiota_http/middleware/redirect_handler.py‎

Lines changed: 18 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -124,15 +124,31 @@ def _build_redirect_request(
124124
"""
125125
method = self._redirect_method(request, response)
126126
url = self._redirect_url(request, response, options)
127-
headers = self._redirect_headers(request, url, method, options)
128127
stream = self._redirect_stream(request, method)
128+
129+
# Create the new request with the redirect URL and original headers
129130
new_request = httpx.Request(
130131
method=method,
131132
url=url,
132-
headers=headers,
133+
headers=request.headers.copy(),
133134
stream=stream,
134135
extensions=request.extensions,
135136
)
137+
138+
# Scrub sensitive headers before following the redirect
139+
options.scrub_sensitive_headers(new_request, request.url)
140+
141+
# Update the Host header if not same origin
142+
if not self._same_origin(url, request.url):
143+
new_request.headers["Host"] = url.netloc.decode("ascii")
144+
145+
# Handle 303 See Other and other method changes
146+
if method != request.method and method == "GET":
147+
# If we've switched to a 'GET' request, strip any headers which
148+
# are only relevant to the request body.
149+
new_request.headers.pop("Content-Length", None)
150+
new_request.headers.pop("Transfer-Encoding", None)
151+
136152
if hasattr(request, "context"):
137153
new_request.context = request.context #type: ignore
138154
new_request.options = {} #type: ignore
@@ -200,28 +216,6 @@ def _redirect_url(
200216

201217
return url
202218

203-
def _redirect_headers(
204-
self, request: httpx.Request, url: httpx.URL, method: str, options: RedirectHandlerOption
205-
) -> httpx.Headers:
206-
"""
207-
Return the headers that should be used for the redirect request.
208-
"""
209-
headers = httpx.Headers(request.headers)
210-
211-
# Scrub sensitive headers before following the redirect
212-
options.scrub_sensitive_headers(headers, request.url, url)
213-
214-
# Update the Host header if not same origin
215-
if not self._same_origin(url, request.url):
216-
headers["Host"] = url.netloc.decode("ascii")
217-
218-
if method != request.method and method == "GET":
219-
# If we've switch to a 'GET' request, then strip any headers which
220-
# are only relevant to the request body.
221-
headers.pop("Content-Length", None)
222-
headers.pop("Transfer-Encoding", None)
223-
224-
return headers
225219

226220
def _redirect_stream(
227221
self, request: httpx.Request, method: str

‎packages/http/httpx/tests/middleware_tests/test_redirect_handler.py‎

Lines changed: 69 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -423,7 +423,7 @@ def request_handler(request: httpx.Request):
423423
async def test_redirect_with_custom_scrubber():
424424
"""Test that custom scrubber can be provided and is used"""
425425

426-
def custom_scrubber(headers, original_url, new_url):
426+
def custom_scrubber(new_request, original_url):
427427
# Custom logic: never remove headers
428428
pass
429429

@@ -457,110 +457,124 @@ def test_default_scrub_sensitive_headers_removes_on_host_change():
457457
"""Test that default scrubber removes Authorization and Cookie when host changes"""
458458
from kiota_http.middleware.options.redirect_handler_option import default_scrub_sensitive_headers
459459

460-
headers = httpx.Headers({
461-
AUTHORIZATION_HEADER: "Bearer token",
462-
"Cookie": "session=SECRET",
463-
"Content-Type": "application/json"
464-
})
465460
original_url = httpx.URL("https://example.com/v1/api")
466-
new_url = httpx.URL("https://other.com/api")
461+
new_request = httpx.Request(
462+
"GET",
463+
"https://other.com/api",
464+
headers={
465+
AUTHORIZATION_HEADER: "Bearer token",
466+
"Cookie": "session=SECRET",
467+
"Content-Type": "application/json"
468+
}
469+
)
467470

468-
default_scrub_sensitive_headers(headers, original_url, new_url)
471+
default_scrub_sensitive_headers(new_request, original_url)
469472

470-
assert AUTHORIZATION_HEADER not in headers
471-
assert "Cookie" not in headers
472-
assert "Content-Type" in headers # Other headers should remain
473+
assert AUTHORIZATION_HEADER not in new_request.headers
474+
assert "Cookie" not in new_request.headers
475+
assert "Content-Type" in new_request.headers # Other headers should remain
473476

474477

475478
def test_default_scrub_sensitive_headers_removes_on_scheme_change():
476479
"""Test that default scrubber removes Authorization and Cookie when scheme changes"""
477480
from kiota_http.middleware.options.redirect_handler_option import default_scrub_sensitive_headers
478481

479-
headers = httpx.Headers({
480-
AUTHORIZATION_HEADER: "Bearer token",
481-
"Cookie": "session=SECRET",
482-
"Content-Type": "application/json"
483-
})
484482
original_url = httpx.URL("https://example.com/v1/api")
485-
new_url = httpx.URL("http://example.com/v1/api")
483+
new_request = httpx.Request(
484+
"GET",
485+
"http://example.com/v1/api",
486+
headers={
487+
AUTHORIZATION_HEADER: "Bearer token",
488+
"Cookie": "session=SECRET",
489+
"Content-Type": "application/json"
490+
}
491+
)
486492

487-
default_scrub_sensitive_headers(headers, original_url, new_url)
493+
default_scrub_sensitive_headers(new_request, original_url)
488494

489-
assert AUTHORIZATION_HEADER not in headers
490-
assert "Cookie" not in headers
491-
assert "Content-Type" in headers
495+
assert AUTHORIZATION_HEADER not in new_request.headers
496+
assert "Cookie" not in new_request.headers
497+
assert "Content-Type" in new_request.headers
492498

493499

494500
def test_default_scrub_sensitive_headers_keeps_on_same_origin():
495501
"""Test that default scrubber keeps headers when host and scheme are the same"""
496502
from kiota_http.middleware.options.redirect_handler_option import default_scrub_sensitive_headers
497503

498-
headers = httpx.Headers({
499-
AUTHORIZATION_HEADER: "Bearer token",
500-
"Cookie": "session=SECRET",
501-
"Content-Type": "application/json"
502-
})
503504
original_url = httpx.URL("https://example.com/v1/api")
504-
new_url = httpx.URL("https://example.com/v2/api")
505+
new_request = httpx.Request(
506+
"GET",
507+
"https://example.com/v2/api",
508+
headers={
509+
AUTHORIZATION_HEADER: "Bearer token",
510+
"Cookie": "session=SECRET",
511+
"Content-Type": "application/json"
512+
}
513+
)
505514

506-
default_scrub_sensitive_headers(headers, original_url, new_url)
515+
default_scrub_sensitive_headers(new_request, original_url)
507516

508-
assert AUTHORIZATION_HEADER in headers
509-
assert "Cookie" in headers
510-
assert "Content-Type" in headers
517+
assert AUTHORIZATION_HEADER in new_request.headers
518+
assert "Cookie" in new_request.headers
519+
assert "Content-Type" in new_request.headers
511520

512521

513522
def test_default_scrub_sensitive_headers_removes_on_port_change():
514523
"""Test that default scrubber removes Authorization and Cookie when port changes"""
515524
from kiota_http.middleware.options.redirect_handler_option import default_scrub_sensitive_headers
516525

517-
headers = httpx.Headers({
518-
AUTHORIZATION_HEADER: "Bearer token",
519-
"Cookie": "session=SECRET",
520-
"Content-Type": "application/json"
521-
})
522526
original_url = httpx.URL("http://example.org:8080/foo")
523-
new_url = httpx.URL("http://example.org:9090/bar")
527+
new_request = httpx.Request(
528+
"GET",
529+
"http://example.org:9090/bar",
530+
headers={
531+
AUTHORIZATION_HEADER: "Bearer token",
532+
"Cookie": "session=SECRET",
533+
"Content-Type": "application/json"
534+
}
535+
)
524536

525-
default_scrub_sensitive_headers(headers, original_url, new_url)
537+
default_scrub_sensitive_headers(new_request, original_url)
526538

527-
assert AUTHORIZATION_HEADER not in headers
528-
assert "Cookie" not in headers
529-
assert "Content-Type" in headers # Other headers should remain
539+
assert AUTHORIZATION_HEADER not in new_request.headers
540+
assert "Cookie" not in new_request.headers
541+
assert "Content-Type" in new_request.headers # Other headers should remain
530542

531543

532544
def test_default_scrub_sensitive_headers_keeps_on_same_port():
533545
"""Test that default scrubber keeps headers when port is the same"""
534546
from kiota_http.middleware.options.redirect_handler_option import default_scrub_sensitive_headers
535547

536-
headers = httpx.Headers({
537-
AUTHORIZATION_HEADER: "Bearer token",
538-
"Cookie": "session=SECRET",
539-
"Content-Type": "application/json"
540-
})
541548
original_url = httpx.URL("http://example.org:8080/foo")
542-
new_url = httpx.URL("http://example.org:8080/bar")
549+
new_request = httpx.Request(
550+
"GET",
551+
"http://example.org:8080/bar",
552+
headers={
553+
AUTHORIZATION_HEADER: "Bearer token",
554+
"Cookie": "session=SECRET",
555+
"Content-Type": "application/json"
556+
}
557+
)
543558

544-
default_scrub_sensitive_headers(headers, original_url, new_url)
559+
default_scrub_sensitive_headers(new_request, original_url)
545560

546-
assert AUTHORIZATION_HEADER in headers
547-
assert "Cookie" in headers
548-
assert "Content-Type" in headers
561+
assert AUTHORIZATION_HEADER in new_request.headers
562+
assert "Cookie" in new_request.headers
563+
assert "Content-Type" in new_request.headers
549564

550565

551566
def test_default_scrub_sensitive_headers_handles_none_gracefully():
552567
"""Test that default scrubber handles None/empty inputs gracefully"""
553568
from kiota_http.middleware.options.redirect_handler_option import default_scrub_sensitive_headers
554569

555570
# Should not raise exceptions
556-
default_scrub_sensitive_headers(None, httpx.URL(BASE_URL), httpx.URL(BASE_URL))
557-
default_scrub_sensitive_headers(httpx.Headers(), None, httpx.URL(BASE_URL))
558-
default_scrub_sensitive_headers(httpx.Headers(), httpx.URL(BASE_URL), None)
571+
default_scrub_sensitive_headers(None, httpx.URL(BASE_URL))
572+
default_scrub_sensitive_headers(httpx.Request("GET", BASE_URL), None)
559573

560574

561575
def test_custom_scrub_sensitive_headers():
562576
"""Test that custom scrubber can be set on options"""
563-
def custom_scrubber(headers, original_url, new_url):
577+
def custom_scrubber(new_request, original_url):
564578
# Custom logic
565579
pass
566580

0 commit comments

Comments
 (0)