1212from ...pipeline .stage import Stage
1313from ...util .clock import ASYNC_SYSTEM_CLOCK , SYSTEM_CLOCK , AsyncClock , Clock
1414from .access_token import AccessTokenInfo , TokenRequestOptions
15+ from .challenge import parse_challenges
16+ from .challenge_handler import ChallengeHandler
1517from .credentials import (
1618 AsyncTokenCredential ,
1719 BasicAuthCredential ,
@@ -91,11 +93,23 @@ class BearerTokenPolicy(Policy):
9193 at most once per refresh window even under heavy concurrent send pressure.
9294
9395 The ``on_challenge`` hook is a no-op by default; subclasses override it
94- to handle CAE/claims challenges.
96+ to handle CAE/claims challenges. Alternatively, callers can pass a
97+ ``challenge_handler`` to plug in a scheme-aware handler (e.g.
98+ ``DigestChallengeHandler``); the handler is consulted before
99+ ``on_challenge`` on 401/407 and, if it satisfies the challenge, the
100+ returned ``(name, value)`` pair is stamped on the retried request.
95101 """
96102
97103 STAGE = Stage .AUTH
98- __slots__ = ("_audience" , "_cache" , "_clock" , "_credential" , "_lock" , "_scopes" )
104+ __slots__ = (
105+ "_audience" ,
106+ "_cache" ,
107+ "_challenge_handler" ,
108+ "_clock" ,
109+ "_credential" ,
110+ "_lock" ,
111+ "_scopes" ,
112+ )
99113
100114 def __init__ (
101115 self ,
@@ -104,6 +118,7 @@ def __init__(
104118 cache : TokenCache | None = None ,
105119 audience : str | None = None ,
106120 clock : Clock | None = None ,
121+ challenge_handler : ChallengeHandler | None = None ,
107122 ) -> None :
108123 if not scopes :
109124 raise ValueError ("at least one scope is required" )
@@ -113,21 +128,65 @@ def __init__(
113128 self ._audience = audience
114129 self ._clock : Clock = clock if clock is not None else SYSTEM_CLOCK
115130 self ._lock = threading .Lock ()
131+ self ._challenge_handler = challenge_handler
116132
117133 def send (self , request : Request , ctx : PipelineContext ) -> Response :
118134 request = self ._authorize (request , ctx )
119135 response = self .next .send (request , ctx )
120- if int (response .status ) != 401 :
136+ status = int (response .status )
137+ if status not in (401 , 407 ):
138+ return response
139+ # On 401 the bearer token was rejected: drop it from the cache so the
140+ # ``on_challenge`` fallback path forces a refresh. A 407 means the
141+ # proxy rejected us, not the origin — leave the cached token alone.
142+ if status == 401 :
143+ self ._cache .set (self ._scopes , _expired_token (), self ._audience )
144+ handler_header = self ._apply_challenge_handler (request , response , status )
145+ if handler_header is not None :
146+ request = request .with_header (* handler_header )
147+ response = self .next .send (request , ctx )
148+ if int (response .status ) not in (401 , 407 ):
149+ return response
150+ raise ClientAuthenticationError (response = response )
151+ if status != 401 :
152+ # No handler match on a 407; surface the response unchanged so
153+ # callers can inspect proxy-auth failures themselves.
121154 return response
122- # Drop cached token and ask subclasses to handle the challenge.
123- self ._cache .set (self ._scopes , _expired_token (), self ._audience )
124155 if "WWW-Authenticate" in response .headers and self .on_challenge (request , response ):
125156 request = self ._authorize (request , ctx , force_refresh = True )
126157 response = self .next .send (request , ctx )
127158 if int (response .status ) != 401 :
128159 return response
129160 raise ClientAuthenticationError (response = response )
130161
162+ def _apply_challenge_handler (
163+ self ,
164+ request : Request ,
165+ response : Response ,
166+ status : int ,
167+ ) -> tuple [str , str ] | None :
168+ """Delegate to the configured ``ChallengeHandler``, if any.
169+
170+ Returns the ``(name, value)`` pair to stamp on the retry, or ``None``
171+ when no handler is configured or the handler declines the challenge.
172+ """
173+ if self ._challenge_handler is None :
174+ return None
175+ is_proxy = status == 407
176+ header_name = "Proxy-Authenticate" if is_proxy else "WWW-Authenticate"
177+ raw = response .headers .get (header_name )
178+ if raw is None :
179+ return None
180+ challenges = parse_challenges (raw )
181+ if not challenges or not self ._challenge_handler .can_handle (challenges ):
182+ return None
183+ return self ._challenge_handler .handle (
184+ request .method ,
185+ request .url ,
186+ challenges ,
187+ is_proxy = is_proxy ,
188+ )
189+
131190 def on_challenge (self , request : Request , response : Response ) -> bool :
132191 """Override to handle a ``WWW-Authenticate`` challenge.
133192
0 commit comments