Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ def __init__(
self.client.headers["User-Agent"] = requests.utils.default_user_agent(
f"foundry-dev-tools/{__version__}/python-requests"
)
self.token_provider.set_expiration(self)

if self.config.rich_traceback:
from rich.traceback import install
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,22 +3,24 @@
from __future__ import annotations

import base64
import json
import time
from functools import cached_property
import uuid
from typing import TYPE_CHECKING, ClassVar

import palantir_oauth_client
import requests
from requests.structures import CaseInsensitiveDict

from foundry_dev_tools.config.config_types import Host
from foundry_dev_tools.errors.config import TokenProviderConfigError
from foundry_dev_tools.errors.config import InvalidOrExpiredJWTTokenError, TokenProviderConfigError
from foundry_dev_tools.errors.handling import ErrorHandlingConfig, raise_foundry_api_error
from foundry_dev_tools.errors.multipass import ClientAuthenticationFailedError
from foundry_dev_tools.utils.config import entry_point_fdt_token_provider

if TYPE_CHECKING:
from foundry_dev_tools.config.config_types import FoundryOAuthGrantType, Token
from foundry_dev_tools.config.context import FoundryContext


class TokenProvider:
Expand Down Expand Up @@ -56,10 +58,18 @@ def requests_auth_handler(self, r: requests.PreparedRequest) -> requests.Prepare
def set_requests_session(self, session: requests.Session) -> None:
"""No-op by default."""

def set_expiration(self, ctx: FoundryContext) -> None:
"""No-op by default.

Used for jwt expiration retrieval.
"""


class JWTTokenProvider(TokenProvider):
"""Provides Host and Token."""

__expiration: int | None = None

def __init__(self, host: Host | str, jwt: Token) -> None:
"""Initialize the JWTTokenProvider.

Expand All @@ -70,9 +80,39 @@ def __init__(self, host: Host | str, jwt: Token) -> None:
super().__init__(host)
self._jwt = jwt

@cached_property
def _decode_token_id(self) -> str:
split_jwt = self._jwt.split(".")
if len(split_jwt) < 2:
raise InvalidOrExpiredJWTTokenError
try:
claim_part = split_jwt[1] + "==" # add padding, which is omitted by jwt
claim_json = base64.b64decode(claim_part)
claims = json.loads(claim_json)
jti = base64.b64decode(claims["jti"])
return str(uuid.UUID(bytes=jti))
except Exception as e:
raise InvalidOrExpiredJWTTokenError from e

def set_expiration(self, ctx: FoundryContext) -> None:
"""This method decodes the jwt token and retrieves its expiration."""
if self.__expiration is not None:
return
token_id = self._decode_token_id()
try:
tokens = ctx.multipass.get_tokens() # retrieve all tokens the user has generated
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the performance penalty of this call?
Wouldn't this call fail in case the token used to make this request is already expired?

Copy link
Collaborator Author

@jonas-w jonas-w Jul 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes if this call fails, as mentioned in the PR, the same exception will be thrown (this call is inside a try except block, which on except throws the InvalidOrExpiredJWTTokenError).

This will only be called on creation of the FoundryContext, so it will only increase startup time and not decrease performance for every api request, and for me it takes around 300-500ms for this init step (which is the same time it takes for every other api call with my internet connection)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this could theoretically be cached, so per jwt it would only need to do this request one time, and then save it into a json file, which holds all expiration times for each token id

for token in tokens:
if token["tokenId"] == token_id: # search for this jwt token
# save expiration time, which will be used in the token() property
self.__expiration = time.time() + token["expires_in"] - 15 # subtract 15 seconds, to be safe
break
except: # noqa: E722
raise InvalidOrExpiredJWTTokenError from None

@property
def token(self) -> Token:
"""Returns the token supplied when creating this Provider."""
if self.__expiration is not None and self.__expiration <= time.time():
raise InvalidOrExpiredJWTTokenError
return self._jwt


Expand Down
10 changes: 10 additions & 0 deletions libs/foundry-dev-tools/src/foundry_dev_tools/errors/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,3 +37,13 @@ class MissingFoundryHostError(TokenProviderConfigError):

def __init__(self) -> None:
super().__init__("A domain is missing in your credentials configuration.")


class InvalidOrExpiredJWTTokenError(FoundryConfigError):
"""Error if jwt is either expired or invalid."""

def __init__(self) -> None:
super().__init__(
"The JWT Token you provided in the config, is either expired or invalid.\n"
"Please generate a new JWT Token and update your configuration."
)
Loading