From 589346acba3faeb3d1071a431cacb1f9e0010ec3 Mon Sep 17 00:00:00 2001
From: Quentin BEY <quentin.bey@mail.numerique.gouv.fr>
Date: Tue, 28 Jan 2025 11:33:36 +0100
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8(oidc)=20store=20refresh=20token=20in?=
 =?UTF-8?q?=20session?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This code is not very clean but it reduces the footprint
to add refresh token storage in session.
This could be cleaned up once the PR it merged:
https://github.com/mozilla/mozilla-django-oidc/pull/377
---
 src/backend/core/authentication/backends.py   | 43 ++++++++++++++-
 .../tests/authentication/test_backends.py     | 52 +++++++++++++++++++
 2 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/src/backend/core/authentication/backends.py b/src/backend/core/authentication/backends.py
index 1a99555e..bcc81d80 100644
--- a/src/backend/core/authentication/backends.py
+++ b/src/backend/core/authentication/backends.py
@@ -17,6 +17,12 @@ from core.models import DuplicateEmailError, User
 logger = logging.getLogger(__name__)
 
 
+def store_oidc_refresh_token(session, refresh_token):
+    """Store the OIDC refresh token in the session if enabled in settings."""
+    if import_from_settings("OIDC_STORE_REFRESH_TOKEN", False):
+        session["oidc_refresh_token"] = refresh_token
+
+
 def store_tokens(session, access_token, id_token, refresh_token):
     """Store tokens in the session if enabled in settings."""
     if import_from_settings("OIDC_STORE_ACCESS_TOKEN", False):
@@ -25,8 +31,7 @@ def store_tokens(session, access_token, id_token, refresh_token):
     if import_from_settings("OIDC_STORE_ID_TOKEN", False):
         session["oidc_id_token"] = id_token
 
-    if import_from_settings("OIDC_STORE_REFRESH_TOKEN", False):
-        session["oidc_refresh_token"] = refresh_token
+    store_oidc_refresh_token(session, refresh_token)
 
 
 class OIDCAuthenticationBackend(MozillaOIDCAuthenticationBackend):
@@ -36,6 +41,40 @@ class OIDCAuthenticationBackend(MozillaOIDCAuthenticationBackend):
     in the User and Identity models, and handles signed and/or encrypted UserInfo response.
     """
 
+    def __init__(self, *args, **kwargs):
+        """
+        Initialize the OIDC Authentication Backend.
+
+        Adds an internal attribute to store the token_info dictionary.
+        The purpose of `self._token_info` is to not duplicate code from
+        the original `authenticate` method.
+        This won't be needed after https://github.com/mozilla/mozilla-django-oidc/pull/377
+        is merged.
+        """
+        super().__init__(*args, **kwargs)
+        self._token_info = None
+
+    def get_token(self, payload):
+        """
+        Return token object as a dictionary.
+
+        Store the value to extract the refresh token in the `authenticate` method.
+        """
+        self._token_info = super().get_token(payload)
+        return self._token_info
+
+    def authenticate(self, request, **kwargs):
+        """Authenticates a user based on the OIDC code flow."""
+        user = super().authenticate(request, **kwargs)
+
+        if user is not None:
+            # Then the user successfully authenticated
+            store_oidc_refresh_token(
+                request.session, self._token_info.get("refresh_token")
+            )
+
+        return user
+
     def get_userinfo(self, access_token, id_token, payload):
         """Return user details dictionary.
 
diff --git a/src/backend/core/tests/authentication/test_backends.py b/src/backend/core/tests/authentication/test_backends.py
index 8bd47cab..360fd376 100644
--- a/src/backend/core/tests/authentication/test_backends.py
+++ b/src/backend/core/tests/authentication/test_backends.py
@@ -547,3 +547,55 @@ def test_authentication_verify_claims_success(django_assert_num_queries, monkeyp
     assert user.full_name == "Doe"
     assert user.short_name is None
     assert user.email == "john.doe@example.com"
+
+
+@responses.activate
+def test_authentication_session_tokens(
+    django_assert_num_queries, monkeypatch, rf, settings
+):
+    """
+    Test that the session contains oidc_refresh_token and oidc_access_token after authentication.
+    """
+    settings.OIDC_OP_TOKEN_ENDPOINT = "http://oidc.endpoint.test/token"
+    settings.OIDC_OP_USER_ENDPOINT = "http://oidc.endpoint.test/userinfo"
+    settings.OIDC_OP_JWKS_ENDPOINT = "http://oidc.endpoint.test/jwks"
+    settings.OIDC_STORE_ACCESS_TOKEN = True
+    settings.OIDC_STORE_REFRESH_TOKEN = True
+
+    klass = OIDCAuthenticationBackend()
+    request = rf.get("/some-url", {"state": "test-state", "code": "test-code"})
+    request.session = {}
+
+    def verify_token_mocked(*args, **kwargs):
+        return {"sub": "123", "email": "test@example.com"}
+
+    monkeypatch.setattr(OIDCAuthenticationBackend, "verify_token", verify_token_mocked)
+
+    responses.add(
+        responses.POST,
+        re.compile(settings.OIDC_OP_TOKEN_ENDPOINT),
+        json={
+            "access_token": "test-access-token",
+            "refresh_token": "test-refresh-token",
+        },
+        status=200,
+    )
+
+    responses.add(
+        responses.GET,
+        re.compile(settings.OIDC_OP_USER_ENDPOINT),
+        json={"sub": "123", "email": "test@example.com"},
+        status=200,
+    )
+
+    with django_assert_num_queries(6):
+        user = klass.authenticate(
+            request,
+            code="test-code",
+            nonce="test-nonce",
+            code_verifier="test-code-verifier",
+        )
+
+    assert user is not None
+    assert request.session["oidc_access_token"] == "test-access-token"
+    assert request.session["oidc_refresh_token"] == "test-refresh-token"