Add cookies to the retried request when performing digest authenticat…

…ion. (#2846)

* Add cookies from the response to the retried request

* Conditionally add cookies from the response

* Fix failing auth module tests

* Fix linting error

* Add tests to check set cookies from server

httpx/_auth.py

@@ -8,7 +8,7 @@
 from urllib.request import parse_http_list
 
 from ._exceptions import ProtocolError
-from ._models import Request, Response
+from ._models import Cookies, Request, Response
 from ._utils import to_bytes, to_str, unquote
 
 if typing.TYPE_CHECKING:  # pragma: no cover
@@ -217,6 +217,8 @@ def auth_flow(self, request: Request) -> typing.Generator[Request, Response, Non
         request.headers["Authorization"] = self._build_auth_header(
             request, self._last_challenge
         )
+        if response.cookies:
+            Cookies(response.cookies).set_cookie_header(request=request)
         yield request
 
     def _parse_challenge(

tests/test_auth.py

@@ -54,7 +54,7 @@ def test_digest_auth_with_401():
         "WWW-Authenticate": 'Digest realm="...", qop="auth", nonce="...", opaque="..."'
     }
     response = httpx.Response(
-        content=b"Auth required", status_code=401, headers=headers
+        content=b"Auth required", status_code=401, headers=headers, request=request
     )
     request = flow.send(response)
     assert request.headers["Authorization"].startswith("Digest")
@@ -79,7 +79,7 @@ def test_digest_auth_with_401_nonce_counting():
         "WWW-Authenticate": 'Digest realm="...", qop="auth", nonce="...", opaque="..."'
     }
     response = httpx.Response(
-        content=b"Auth required", status_code=401, headers=headers
+        content=b"Auth required", status_code=401, headers=headers, request=request
     )
     first_request = flow.send(response)
     assert first_request.headers["Authorization"].startswith("Digest")
@@ -101,3 +101,42 @@ def test_digest_auth_with_401_nonce_counting():
     response = httpx.Response(content=b"Hello, world!", status_code=200)
     with pytest.raises(StopIteration):
         flow.send(response)
+
+
+def set_cookies(request: httpx.Request) -> httpx.Response:
+    headers = {
+        "Set-Cookie": "session=.session_value...",
+        "WWW-Authenticate": 'Digest realm="...", qop="auth", nonce="...", opaque="..."',
+    }
+    if request.url.path == "/auth":
+        return httpx.Response(
+            content=b"Auth required", status_code=401, headers=headers
+        )
+    else:
+        raise NotImplementedError()  # pragma: no cover
+
+
+def test_digest_auth_setting_cookie_in_request():
+    url = "https://www.example.com/auth"
+    client = httpx.Client(transport=httpx.MockTransport(set_cookies))
+    request = client.build_request("GET", url)
+
+    auth = httpx.DigestAuth(username="user", password="pass")
+    flow = auth.sync_auth_flow(request)
+    request = next(flow)
+    assert "Authorization" not in request.headers
+
+    response = client.get(url)
+    assert len(response.cookies) > 0
+    assert response.cookies["session"] == ".session_value..."
+
+    request = flow.send(response)
+    assert request.headers["Authorization"].startswith("Digest")
+    assert request.headers["Cookie"] == "session=.session_value..."
+
+    # No other requests are made.
+    response = httpx.Response(
+        content=b"Hello, world!", status_code=200, request=request
+    )
+    with pytest.raises(StopIteration):
+        flow.send(response)