import typing
from dataclasses import dataclass
from threading import RLock
import httpx
from ..logger import getLogger
from .types import SalesforceLogin, SalesforceToken, TokenRefreshCallback
from .origins import SalesforceOriginError, https_origin, trusted_origin
from ..request_policy import request_policy
LOGGER = getLogger("auth")
SESSION_REVISION_EXTENSION = "sf_toolkit.session_revision"
[docs]
@dataclass(frozen=True, repr=False)
class SessionSnapshot:
revision: int
token: SalesforceToken | None
[docs]
class SalesforceAuth(httpx.Auth):
login: SalesforceLogin | None
callback: TokenRefreshCallback | None
token: SalesforceToken | None
[docs]
def __init__(
self,
login: SalesforceLogin | None = None,
session_token: SalesforceToken | None = None,
callback: TokenRefreshCallback | None = None,
allowed_resource_origins: typing.Iterable[str | httpx.URL] = (),
*,
cookie_callback: typing.Callable[[httpx.Request], None] | None = None,
session_callback: typing.Callable[[SalesforceToken | None, typing.Any], None]
| None = None,
):
self.login = login
self._session_lock = RLock()
self._notification_lock = RLock()
self._token = session_token
self._session_revision = 0
self.callback = callback
self.session_callback = session_callback
self.cookie_callback = cookie_callback
self.allowed_resource_origins = frozenset(
trusted_origin(url) for url in allowed_resource_origins
)
@property
def token(self) -> SalesforceToken | None:
with self._session_lock:
return self._token
@token.setter
def token(self, token: SalesforceToken | None) -> None:
# Explicit assignment also keeps client routing coherent, but preserves
# the legacy setter's behavior of not notifying persistence callbacks.
self.install_session(token, notify=False)
[docs]
def session_snapshot(self) -> SessionSnapshot:
with self._session_lock:
return SessionSnapshot(self._session_revision, self._token)
[docs]
def install_session(
self,
token: SalesforceToken | None,
*,
expected: SessionSnapshot | None = None,
userinfo: typing.Any = None,
notify: bool = True,
) -> SessionSnapshot:
if token is not None:
https_origin(token.instance)
with self._session_lock:
if expected is not None and expected.revision != self._session_revision:
return SessionSnapshot(self._session_revision, self._token)
# This callback is the client's small internal routing/cache update,
# not a user callback. It validates before changing any state.
if self.session_callback is not None:
self.session_callback(token, userinfo)
self._token = token
self._session_revision += 1
installed = SessionSnapshot(self._session_revision, token)
if notify and token is not None and self.callback is not None:
# Serialize notifications outside the state lock. A stale queued
# notification cannot subsequently revert persisted session state.
with self._notification_lock:
if self.session_snapshot().revision == installed.revision:
self.callback(token)
return installed
[docs]
def check_resource_url(self, url: httpx.URL | str) -> None:
"""Fence resource credentials; login exchanges use their own endpoints."""
self._check_resource_url(url, self.token)
def _check_resource_url(
self, url: httpx.URL | str, token: SalesforceToken | None
) -> None:
if token is None:
raise SalesforceOriginError("A Salesforce session is required")
expected = https_origin(token.instance)
actual = https_origin(url)
if actual != expected and actual not in self.allowed_resource_origins:
raise SalesforceOriginError(
"Refusing authenticated request to an untrusted Salesforce origin"
)
def _authorize(self, request: httpx.Request, token: SalesforceToken) -> None:
self._check_resource_url(request.url, token)
request.headers["Authorization"] = f"Bearer {token.token}"
@staticmethod
def _expired_session(response: httpx.Response) -> bool:
if response.status_code == 403:
return response.text == "Bad_OAuth_Token"
if response.status_code != 401:
return False
try:
errors = response.json()
except ValueError:
return False
return isinstance(errors, list) and any(
isinstance(error, dict)
and (error.get("errorCode") or error.get("errorDetails"))
== "INVALID_SESSION_ID"
for error in errors
)
[docs]
def sync_auth_flow(self, request: httpx.Request):
flow = self.auth_flow(request)
try:
current = next(flow)
while True:
response = yield current
if current is not request or response.status_code in {401, 403}:
response.read()
try:
current = flow.send(response)
except StopIteration:
return
finally:
flow.close()
[docs]
async def async_auth_flow(self, request: httpx.Request):
flow = self.auth_flow(request)
try:
current = next(flow)
while True:
response = yield current
if current is not request or response.status_code in {401, 403}:
await response.aread()
try:
current = flow.send(response)
except StopIteration:
return
finally:
flow.close()
[docs]
def auth_flow(
self, request: httpx.Request
) -> typing.Generator[httpx.Request, httpx.Response, None]:
new_token: SalesforceToken
policy = request_policy(request)
initial = self.session_snapshot()
if initial.token is None:
assert self.login is not None, "No login method provided"
try:
login_flow = self.login()
login_request = next(login_flow)
while True:
if login_request is not None:
login_response = yield login_request
login_request = login_flow.send(login_response)
else:
login_request = next(login_flow)
except StopIteration as login_result:
new_token = login_result.value
self.install_session(SalesforceToken(*new_token), expected=initial)
assert self.token is not None, "Failed to perform initial login"
authorized = self.session_snapshot()
request_token = authorized.token
if request.url.is_relative_url:
absolute_url = (
request_token.instance.raw_path + request.url.raw_path.lstrip(b"/")
)
request.url = request_token.instance.copy_with(raw_path=absolute_url)
request._prepare({**request.headers})
self._authorize(request, request_token)
request.extensions[SESSION_REVISION_EXTENSION] = authorized.revision
response = yield request
if policy.allow_auth_replay and self.login and self._expired_session(response):
# The shared token may have changed while this response was pending.
# Fence migration against the token that actually authorized it.
previous_origin = https_origin(request_token.instance)
expected = authorized
if self.session_snapshot().revision == expected.revision:
try:
login_flow = self.login()
login_request = next(login_flow)
while True:
if login_request is not None:
login_response = yield login_request
login_request = login_flow.send(login_response)
else:
login_request = next(login_flow)
except StopIteration as login_result:
new_token = login_result.value
try:
https_origin(new_token.instance)
except SalesforceOriginError:
raise SalesforceOriginError(
"Cannot replay after an unsafe Salesforce instance change"
) from None
# Do not overwrite another refresh completed during login.
self.install_session(new_token, expected=expected)
renewed = self.session_snapshot()
new_token = renewed.token
new_origin = https_origin(new_token.instance)
if (
new_origin != previous_origin
and https_origin(request.url) == previous_origin
):
# Only the previous authenticated instance may migrate. Explicit
# alternate origins stay put, and path/query bytes are unchanged.
request.url = request.url.copy_with(
scheme=new_token.instance.scheme,
host=new_token.instance.host,
port=new_token.instance.port,
)
request.headers["Host"] = request.url.netloc.decode("ascii")
# Cookies were selected for the original URL when this request
# was built. Never carry them across an instance-origin change.
request.headers.pop("Cookie", None)
if self.cookie_callback is not None:
self.cookie_callback(request)
self._authorize(request, new_token)
request.extensions[SESSION_REVISION_EXTENSION] = renewed.revision
response = yield request