11from __future__ import annotations
22
3+ import hashlib
34from collections .abc import Callable , Iterable
45from typing import Any , ClassVar
56
@@ -298,6 +299,55 @@ def authenticate(self, request: Request):
298299class UserAuthTokenAuthentication (StandardAuthentication ):
299300 token_name = b"bearer"
300301
302+ def _find_or_update_token_by_hash (self , token_str : str ) -> ApiToken | ApiTokenReplica :
303+ """
304+ Find token by hash or update token's hash value if only found via plaintext.
305+
306+ 1. Hash provided plaintext token.
307+ 2. Perform lookup based on hashed value.
308+ 3. If found, return the token.
309+ 4. If not found, search for the token based on its plaintext value.
310+ 5. If found, update the token's hashed value and return the token.
311+ 6. If not found via hash or plaintext value, raise AuthenticationFailed
312+
313+ Returns `ApiTokenReplica` if running in REGION silo or
314+ `ApiToken` if running in CONTROL silo.
315+ """
316+
317+ hashed_token = hashlib .sha256 (token_str .encode ()).hexdigest ()
318+
319+ if SiloMode .get_current_mode () == SiloMode .REGION :
320+ try :
321+ # Try to find the token by its hashed value first
322+ return ApiTokenReplica .objects .get (hashed_token = hashed_token )
323+ except ApiTokenReplica .DoesNotExist :
324+ try :
325+ # If we can't find it by hash, use the plaintext string
326+ return ApiTokenReplica .objects .get (token = token_str )
327+ except ApiTokenReplica .DoesNotExist :
328+ # If the token does not exist by plaintext either, it is not a valid token
329+ raise AuthenticationFailed ("Invalid token" )
330+ else :
331+ try :
332+ # Try to find the token by its hashed value first
333+ return ApiToken .objects .select_related ("user" , "application" ).get (
334+ hashed_token = hashed_token
335+ )
336+ except ApiToken .DoesNotExist :
337+ try :
338+ # If we can't find it by hash, use the plaintext string
339+ api_token = ApiToken .objects .select_related ("user" , "application" ).get (
340+ token = token_str
341+ )
342+ except ApiToken .DoesNotExist :
343+ # If the token does not exist by plaintext either, it is not a valid token
344+ raise AuthenticationFailed ("Invalid token" )
345+ else :
346+ # Update it with the hashed value if found by plaintext
347+ api_token .hashed_token = hashed_token
348+ api_token .save (update_fields = ["hashed_token" ])
349+ return api_token
350+
301351 def accepts_auth (self , auth : list [bytes ]) -> bool :
302352 if not super ().accepts_auth (auth ):
303353 return False
@@ -320,26 +370,16 @@ def authenticate_token(self, request: Request, token_str: str) -> tuple[Any, Any
320370 application_is_inactive = False
321371
322372 if not token :
323- if SiloMode .get_current_mode () == SiloMode .REGION :
324- try :
325- atr = token = ApiTokenReplica .objects .get (token = token_str )
326- except ApiTokenReplica .DoesNotExist :
327- raise AuthenticationFailed ("Invalid token" )
328- user = user_service .get_user (user_id = atr .user_id )
329- application_is_inactive = not atr .application_is_active
330- else :
331- try :
332- at = token = (
333- ApiToken .objects .filter (token = token_str )
334- .select_related ("user" , "application" )
335- .get ()
336- )
337- except ApiToken .DoesNotExist :
338- raise AuthenticationFailed ("Invalid token" )
339- user = at .user
373+ token = self ._find_or_update_token_by_hash (token_str )
374+ if isinstance (token , ApiTokenReplica ): # we're running as a REGION silo
375+ user = user_service .get_user (user_id = token .user_id )
376+ application_is_inactive = not token .application_is_active
377+ else : # the token returned is an ApiToken from the CONTROL silo
378+ user = token .user
340379 application_is_inactive = (
341- at .application is not None and not at .application .is_active
380+ token .application is not None and not token .application .is_active
342381 )
382+
343383 elif isinstance (token , SystemToken ):
344384 user = token .user
345385
@@ -389,9 +429,9 @@ def authenticate_token(self, request: Request, token_str: str) -> tuple[Any, Any
389429 raise AuthenticationFailed ("Invalid org token" )
390430 else :
391431 try :
392- token = OrgAuthToken .objects .filter (
432+ token = OrgAuthToken .objects .get (
393433 token_hashed = token_hashed , date_deactivated__isnull = True
394- ). get ()
434+ )
395435 except OrgAuthToken .DoesNotExist :
396436 raise AuthenticationFailed ("Invalid org token" )
397437
0 commit comments