feat: enhance session management with token refresh logic and error handling (7d5c9298) · Commits · OOP / code / portal

src/portal_service/api/dependencies.py

+16 −3

Changes for src/portal_service/api/dependencies.py: 16 added lines, 3 removed lines.
Original line number	Diff line number	Diff line
		@@ -2,6 +2,7 @@ from collections.abc import AsyncGenerator
		from datetime import timedelta
		from typing import Annotated, cast

		import structlog
		from fastapi import Depends, HTTPException, Request, status
		from sqlalchemy import text
		from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
		@@ -15,11 +16,13 @@ from portal_service.app_state import AppState
		from portal_service.application.login_service import LoginService
		from portal_service.application.session_service import ActiveSession, SessionService
		from portal_service.config import OIDCSettings, Settings, get_settings
		from portal_service.domain.errors import InvalidSessionError
		from portal_service.domain.errors import IdentityProviderUnavailableError, InvalidSessionError
		from portal_service.domain.ports.database.login_attempts import LoginAttemptRepository
		from portal_service.domain.ports.database.login_sessions import LoginSessionRepository
		from portal_service.domain.ports.identity_provider import IdentityProvider

		logger: structlog.BoundLogger = structlog.get_logger(__name__)


		def get_app_state(request: Request) -> AppState:
		return cast(AppState, request.app.state)
		@@ -93,10 +96,12 @@ def get_login_service(

		def get_session_service(
		sessions: Annotated[LoginSessionRepository, Depends(get_login_session_repo)],
		identity_provider: Annotated[IdentityProvider, Depends(get_identity_provider)],
		settings: SettingsDep,
		) -> SessionService:
		return SessionService(
		sessions=sessions,
		identity_provider=identity_provider,
		session_ttl=timedelta(seconds=settings.session.max_age),
		)

		@@ -111,10 +116,18 @@ async def get_current_session(
		if not cookie:
		raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Not authenticated")
		try:
		return await service.load(cookie)
		# Refreshed here, so every route that needs the user also gets a usable token.
		return await service.ensure_fresh(await service.load(cookie))
		except InvalidSessionError:
		# Same answer for every reason: malformed, unknown, expired or undecryptable.
		# Same answer for every reason: malformed, unknown, expired, undecryptable, or
		# refused by Keycloak when refreshing.
		raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Not authenticated") from None
		except IdentityProviderUnavailableError as exc:
		# Keycloak is unreachable: the session may well be fine, so this is not a 401.
		logger.error("could not refresh the session", error=str(exc))
		raise HTTPException(
		status.HTTP_503_SERVICE_UNAVAILABLE, "Authentication service unavailable"
		) from exc


		CurrentSessionDep = Annotated[ActiveSession, Depends(get_current_session)]

src/portal_service/application/session_service.py

+51 −2

Changes for src/portal_service/application/session_service.py: 51 added lines, 2 removed lines.
Original line number	Diff line number	Diff line
		@@ -13,9 +13,10 @@ from portal_service.application.session_crypto import (
		decrypt_tokens,
		encrypt_tokens,
		)
		from portal_service.domain.errors import InvalidSessionError
		from portal_service.domain.errors import InvalidGrantError, InvalidSessionError
		from portal_service.domain.models import IssuedTokens, LoginSession, SessionTokens
		from portal_service.domain.ports.database.login_sessions import LoginSessionRepository
		from portal_service.domain.ports.identity_provider import IdentityProvider


		@dataclass(frozen=True)
		@@ -35,9 +36,18 @@ class ActiveSession:
		credentials: SessionCredentials = field(repr=False)


		_REFRESH_LEEWAY = timedelta(seconds=30)


		class SessionService:
		def __init__(self, sessions: LoginSessionRepository, session_ttl: timedelta) -> None:
		def __init__(
		self,
		sessions: LoginSessionRepository,
		identity_provider: IdentityProvider,
		session_ttl: timedelta,
		) -> None:
		self._sessions = sessions
		self._identity_provider = identity_provider
		self._session_ttl = session_ttl

		async def create(self, subject: str, issued: IssuedTokens) -> CreatedSession:
		@@ -74,3 +84,42 @@ class SessionService:
		tokens=decrypt_tokens(row.encrypted_tokens, credentials),
		credentials=credentials,
		)

		async def ensure_fresh(self, active: ActiveSession) -> ActiveSession:
		"""
		Returns a session whose access token is usable, refreshing it if needed.

		Refreshing is also how revocation reaches us: Keycloak refuses a refresh once
		the user is disabled, their password is reset, or the session is ended there.
		We drop the session when that happens, so the next request is a clean logout.
		"""
		now = datetime.now(UTC)
		if active.session.access_token_expires_at - _REFRESH_LEEWAY > now:
		return active
		if active.tokens.refresh_token is None:
		raise InvalidSessionError("Session has no refresh token")

		try:
		issued = await self._identity_provider.refresh(active.tokens.refresh_token)
		except InvalidGrantError as exc:
		await self._sessions.delete_by_id(active.session.id)
		raise InvalidSessionError("Keycloak no longer accepts this session") from exc

		# Keycloak may return only some of the tokens; keep what it did not replace.
		tokens = SessionTokens(
		access_token=issued.tokens.access_token,
		refresh_token=issued.tokens.refresh_token or active.tokens.refresh_token,
		id_token=issued.tokens.id_token or active.tokens.id_token,
		)
		row = await self._sessions.save(
		LoginSession(
		id=active.session.id,
		subject=active.session.subject,
		access_token_expires_at=now + timedelta(seconds=issued.expires_in),
		# Same key as before: it lives only in the cookie the browser still holds.
		encrypted_tokens=encrypt_tokens(tokens, active.credentials),
		# The session's own deadline is unchanged by a refresh.
		expires_at=active.session.expires_at,
		)
		)
		return ActiveSession(session=row, tokens=tokens, credentials=active.credentials)

tests/test_me_route.py

+12 −12

Changes for tests/test_me_route.py: 12 added lines, 12 removed lines.
Original line number	Diff line number	Diff line
		@@ -12,10 +12,10 @@ import pytest
		from fastapi import FastAPI
		from fastapi.testclient import TestClient

		from portal_service.api.dependencies import get_login_session_repo
		from portal_service.api.dependencies import get_identity_provider, get_login_session_repo
		from portal_service.application.session_crypto import SessionCredentials, encrypt_tokens
		from portal_service.domain.models import LoginSession, SessionTokens
		from tests.fakes import FakeLoginSessionRepository
		from tests.fakes import FakeIdentityProvider, FakeLoginSessionRepository

		SESSION_COOKIE = "oop_session"

		@@ -25,11 +25,20 @@ def sessions() -> FakeLoginSessionRepository:
		return FakeLoginSessionRepository()


		@pytest.fixture
		def identity_provider() -> FakeIdentityProvider:
		return FakeIdentityProvider()


		@pytest.fixture
		def client(
		app_jwt: FastAPI, patch_jwks: None, sessions: FakeLoginSessionRepository
		app_jwt: FastAPI,
		patch_jwks: None,
		sessions: FakeLoginSessionRepository,
		identity_provider: FakeIdentityProvider,
		) -> Iterator[TestClient]:
		app_jwt.dependency_overrides[get_login_session_repo] = lambda: sessions
		app_jwt.dependency_overrides[get_identity_provider] = lambda: identity_provider
		yield TestClient(app_jwt)
		app_jwt.dependency_overrides.clear()

		@@ -117,15 +126,6 @@ def test_an_expired_session_is_not_a_user(
		assert client.get("/api/auth/me").status_code == 401


		def test_an_expired_access_token_is_not_yet_a_user(
		client: TestClient, sessions: FakeLoginSessionRepository, make_token: Callable[..., str]
		) -> None:
		# Refreshing it is the next step; until then the session cannot answer.
		_open_session(client, sessions, make_token(expires_in=-10))

		assert client.get("/api/auth/me").status_code == 401


		def test_a_tampered_cookie_key_is_not_a_user(
		client: TestClient, sessions: FakeLoginSessionRepository, make_token: Callable[..., str]
		) -> None:

tests/test_session_refresh.py

0 → 100644

+221 −0

Changes for tests/test_session_refresh.py: 221 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		"""
		Refreshing the access token.

		Access tokens live minutes, sessions live hours, so a session refreshes itself.
		Refreshing is also how revocation reaches us: Keycloak refuses once the user is
		disabled or the session ended there.
		"""

		from collections.abc import Callable, Iterator
		from datetime import UTC, datetime, timedelta

		import pytest
		from fastapi import FastAPI
		from fastapi.testclient import TestClient

		from portal_service.api.dependencies import get_identity_provider, get_login_session_repo
		from portal_service.application.session_crypto import (
		SessionCredentials,
		decrypt_tokens,
		encrypt_tokens,
		)
		from portal_service.domain.errors import IdentityProviderUnavailableError, InvalidGrantError
		from portal_service.domain.models import IssuedTokens, LoginSession, SessionTokens
		from tests.fakes import FakeIdentityProvider, FakeLoginSessionRepository

		SESSION_COOKIE = "oop_session"


		@pytest.fixture
		def sessions() -> FakeLoginSessionRepository:
		return FakeLoginSessionRepository()


		@pytest.fixture
		def identity_provider() -> FakeIdentityProvider:
		return FakeIdentityProvider()


		@pytest.fixture
		def client(
		app_jwt: FastAPI,
		patch_jwks: None,
		sessions: FakeLoginSessionRepository,
		identity_provider: FakeIdentityProvider,
		) -> Iterator[TestClient]:
		app_jwt.dependency_overrides[get_login_session_repo] = lambda: sessions
		app_jwt.dependency_overrides[get_identity_provider] = lambda: identity_provider
		yield TestClient(app_jwt)
		app_jwt.dependency_overrides.clear()


		def _open_session(
		client: TestClient,
		sessions: FakeLoginSessionRepository,
		access_token: str,
		*,
		access_token_expires_in: timedelta,
		refresh_token: str \| None = "the-refresh-token",
		) -> SessionCredentials:
		credentials = SessionCredentials.generate()
		now = datetime.now(UTC)
		sessions.rows[credentials.session_id_hash] = LoginSession(
		id=credentials.session_id_hash,
		subject="user-123",
		access_token_expires_at=now + access_token_expires_in,
		encrypted_tokens=encrypt_tokens(
		SessionTokens(
		access_token=access_token, refresh_token=refresh_token, id_token="the-id-token"
		),
		credentials,
		),
		expires_at=now + timedelta(hours=8),
		)
		client.cookies.set(SESSION_COOKIE, credentials.to_cookie())
		return credentials


		def _issued(make_token: Callable[..., str], **claims: object) -> IssuedTokens:
		return IssuedTokens(
		tokens=SessionTokens(
		access_token=make_token(**claims),
		refresh_token="a-newer-refresh-token",
		id_token="a-newer-id-token",
		),
		expires_in=300,
		)


		def test_a_valid_token_is_not_refreshed(
		client: TestClient,
		sessions: FakeLoginSessionRepository,
		identity_provider: FakeIdentityProvider,
		make_token: Callable[..., str],
		) -> None:
		_open_session(client, sessions, make_token(), access_token_expires_in=timedelta(minutes=5))

		assert client.get("/api/auth/me").status_code == 200
		assert identity_provider.refreshed == []


		def test_an_expired_token_is_refreshed_and_the_user_answers(
		client: TestClient,
		sessions: FakeLoginSessionRepository,
		identity_provider: FakeIdentityProvider,
		make_token: Callable[..., str],
		) -> None:
		_open_session(
		client, sessions, make_token(expires_in=-10), access_token_expires_in=-timedelta(seconds=10)
		)
		identity_provider.issued = _issued(make_token, roles=["operator"])

		body = client.get("/api/auth/me").json()

		assert identity_provider.refreshed == ["the-refresh-token"]
		assert body["roles"] == ["operator"]


		def test_a_token_about_to_expire_is_refreshed_early(
		client: TestClient,
		sessions: FakeLoginSessionRepository,
		identity_provider: FakeIdentityProvider,
		make_token: Callable[..., str],
		) -> None:
		# Still valid, but not for long enough to hand to another service.
		_open_session(client, sessions, make_token(), access_token_expires_in=timedelta(seconds=10))
		identity_provider.issued = _issued(make_token)

		assert client.get("/api/auth/me").status_code == 200
		assert identity_provider.refreshed == ["the-refresh-token"]


		def test_the_refreshed_tokens_are_stored_under_the_same_cookie(
		client: TestClient,
		sessions: FakeLoginSessionRepository,
		identity_provider: FakeIdentityProvider,
		make_token: Callable[..., str],
		) -> None:
		credentials = _open_session(
		client, sessions, make_token(expires_in=-10), access_token_expires_in=-timedelta(seconds=10)
		)
		identity_provider.issued = _issued(make_token)

		client.get("/api/auth/me")

		row = sessions.rows[credentials.session_id_hash]
		stored = decrypt_tokens(row.encrypted_tokens, credentials)
		assert stored.access_token == identity_provider.issued.tokens.access_token
		assert stored.refresh_token == "a-newer-refresh-token"
		assert row.access_token_expires_at > datetime.now(UTC)


		def test_a_refresh_token_keycloak_did_not_replace_is_kept(
		client: TestClient,
		sessions: FakeLoginSessionRepository,
		identity_provider: FakeIdentityProvider,
		make_token: Callable[..., str],
		) -> None:
		credentials = _open_session(
		client, sessions, make_token(expires_in=-10), access_token_expires_in=-timedelta(seconds=10)
		)
		identity_provider.issued = IssuedTokens(
		tokens=SessionTokens(access_token=make_token()), expires_in=300
		)

		client.get("/api/auth/me")

		stored = decrypt_tokens(
		sessions.rows[credentials.session_id_hash].encrypted_tokens, credentials
		)
		assert stored.refresh_token == "the-refresh-token"
		assert stored.id_token == "the-id-token"


		def test_a_revoked_session_is_dropped(
		client: TestClient,
		sessions: FakeLoginSessionRepository,
		identity_provider: FakeIdentityProvider,
		make_token: Callable[..., str],
		) -> None:
		# Keycloak refuses the refresh: the user was disabled, or logged out there.
		credentials = _open_session(
		client, sessions, make_token(expires_in=-10), access_token_expires_in=-timedelta(seconds=10)
		)
		identity_provider.error = InvalidGrantError("session not active")

		assert client.get("/api/auth/me").status_code == 401
		assert credentials.session_id_hash not in sessions.rows


		def test_keycloak_being_down_keeps_the_session(
		client: TestClient,
		sessions: FakeLoginSessionRepository,
		identity_provider: FakeIdentityProvider,
		make_token: Callable[..., str],
		) -> None:
		credentials = _open_session(
		client, sessions, make_token(expires_in=-10), access_token_expires_in=-timedelta(seconds=10)
		)
		identity_provider.error = IdentityProviderUnavailableError("token endpoint unreachable")

		assert client.get("/api/auth/me").status_code == 503
		# The session may be perfectly good; we could not ask.
		assert credentials.session_id_hash in sessions.rows


		def test_a_session_without_a_refresh_token_cannot_be_renewed(
		client: TestClient,
		sessions: FakeLoginSessionRepository,
		identity_provider: FakeIdentityProvider,
		make_token: Callable[..., str],
		) -> None:
		_open_session(
		client,
		sessions,
		make_token(expires_in=-10),
		access_token_expires_in=-timedelta(seconds=10),
		refresh_token=None,
		)

		assert client.get("/api/auth/me").status_code == 401
		assert identity_provider.refreshed == []