feat(OIDC): Add OIDC token exchange endpoint (#8043)

Co-authored-by: flagsmith-engineering[bot] <flagsmith-engineering[bot]@users.noreply.github.com>

Author: khvn26

api/api_keys/views.py

@@ -6,6 +6,9 @@
 from organisations.permissions.permissions import (
     NestedIsOrganisationAdminPermission,
 )
+from trust_relationships.authentication import (
+    TrustRelationshipTokenAuthentication,
+)
 
 from .authentication import MasterAPIKeyAuthentication
 from .models import MasterAPIKey
@@ -18,7 +21,10 @@ def get_authenticators(self) -> list[BaseAuthentication]:
         return [
             authenticator
             for authenticator in super().get_authenticators()
-            if not isinstance(authenticator, MasterAPIKeyAuthentication)
+            if not isinstance(
+                authenticator,
+                (MasterAPIKeyAuthentication, TrustRelationshipTokenAuthentication),
+            )
         ]
 
 

api/app/settings/common.py

@@ -343,6 +343,10 @@
 
 LOGIN_THROTTLE_RATE = env("LOGIN_THROTTLE_RATE", "20/min")
 DCR_THROTTLE_RATE = env("DCR_THROTTLE_RATE", "500/month")
+OIDC_TOKEN_EXCHANGE_THROTTLE_RATE = env("OIDC_TOKEN_EXCHANGE_THROTTLE_RATE", "60/min")
+TRUST_RELATIONSHIP_ACCESS_TOKEN_LIFETIME_SECONDS = env.int(
+    "TRUST_RELATIONSHIP_ACCESS_TOKEN_LIFETIME_SECONDS", default=3600
+)
 SIGNUP_THROTTLE_RATE = env("SIGNUP_THROTTLE_RATE", "10000/min")
 USER_THROTTLE_RATE = env("USER_THROTTLE_RATE", default=None)
 MASTER_API_KEY_THROTTLE_RATE = env("MASTER_API_KEY_THROTTLE_RATE", default=None)
@@ -354,6 +358,7 @@
         "rest_framework.authentication.TokenAuthentication",
         "api_keys.authentication.MasterAPIKeyAuthentication",
         "oauth2_metadata.authentication.OAuth2BearerTokenAuthentication",
+        "trust_relationships.authentication.TrustRelationshipTokenAuthentication",
     ),
     "PAGE_SIZE": 10,
     "UNICODE_JSON": False,
@@ -362,6 +367,7 @@
     "DEFAULT_THROTTLE_RATES": {
         "login": LOGIN_THROTTLE_RATE,
         "dcr_register": DCR_THROTTLE_RATE,
+        "oidc_token_exchange": OIDC_TOKEN_EXCHANGE_THROTTLE_RATE,
         "signup": SIGNUP_THROTTLE_RATE,
         "master_api_key": MASTER_API_KEY_THROTTLE_RATE,
         "mfa_code": "5/min",

api/app/settings/test.py

@@ -21,6 +21,7 @@
 REST_FRAMEWORK["DEFAULT_THROTTLE_RATES"] = {
     "login": "100/min",
     "dcr_register": "100/min",
+    "oidc_token_exchange": "100/min",
     "mfa_code": "5/min",
     "invite": "10/min",
     "signup": "100/min",

api/custom_auth/urls.py

@@ -8,6 +8,7 @@
     FFAdminUserViewSet,
     delete_token,
 )
+from trust_relationships.views import OIDCTokenExchangeView
 
 app_name = "custom_auth"
 
@@ -38,4 +39,9 @@
     path("", include("djoser.urls")),
     path("", include("custom_auth.mfa.trench.urls")),  # MFA
     path("oauth/", include("custom_auth.oauth.urls")),
+    path(
+        "oidc/token/",
+        OIDCTokenExchangeView.as_view(),
+        name="oidc-token-exchange",
+    ),
 ]

api/tests/integration/trust_relationships/conftest.py

@@ -1,7 +1,15 @@
+from typing import Generator
+
+import jwt
 import pytest
+import responses as responses_lib
+from pytest_mock import MockerFixture
 from rest_framework import status
 from rest_framework.test import APIClient
 
+from tests.test_helpers import OIDCIssuerStub
+from trust_relationships.oidc import get_jwks_client
+
 
 @pytest.fixture()
 def trust_relationship(
@@ -21,3 +29,35 @@ def trust_relationship(
     )
     assert response.status_code == status.HTTP_201_CREATED
     return response.json()["id"]  # type: ignore[no-any-return]
+
+
+@pytest.fixture(autouse=True)
+def clear_jwks_client_cache() -> Generator[None, None, None]:
+    get_jwks_client.cache_clear()
+    yield
+    get_jwks_client.cache_clear()
+
+
+@pytest.fixture()
+def oidc_issuer(
+    responses: responses_lib.RequestsMock,
+    mocker: MockerFixture,
+) -> OIDCIssuerStub:
+    issuer = OIDCIssuerStub("https://token.actions.githubusercontent.com")
+    # Rejections short-circuiting before discovery are expected.
+    responses.assert_all_requests_are_fired = False
+    responses.add(
+        responses_lib.GET,
+        "https://token.actions.githubusercontent.com/.well-known/openid-configuration",
+        json={
+            "jwks_uri": "https://token.actions.githubusercontent.com/.well-known/jwks"
+        },
+    )
+    mocker.patch.object(jwt.PyJWKClient, "fetch_data", return_value=issuer.jwks)
+    return issuer
+
+
+@pytest.fixture()
+def machine_client() -> APIClient:
+    # A client that is never force-authenticated
+    return APIClient()

api/tests/integration/trust_relationships/test_token_exchange.py

@@ -0,0 +1,166 @@
+from rest_framework import status
+from rest_framework.test import APIClient
+
+from tests.test_helpers import OIDCIssuerStub
+
+
+def test_token_exchange__matching_token__returns_usable_access_token(
+    machine_client: APIClient,
+    organisation: int,
+    trust_relationship: int,
+    oidc_issuer: OIDCIssuerStub,
+) -> None:
+    # Given
+    token = oidc_issuer.sign_token(
+        aud="https://github.com/Flagsmith",
+        sub="repo:Flagsmith/flagsmith:ref:refs/heads/main",
+        repository="Flagsmith/flagsmith",
+    )
+
+    # When
+    response = machine_client.post(
+        "/api/v1/auth/oidc/token/", data={"token": token}, format="json"
+    )
+
+    # Then
+    assert response.status_code == status.HTTP_200_OK
+    response_json = response.json()
+    assert response_json["token_type"] == "Bearer"
+    assert response_json["expires_in"] == 3600
+
+    # And the minted token authenticates against the admin API
+    machine_client.credentials(
+        HTTP_AUTHORIZATION=f"Bearer {response_json['access_token']}"
+    )
+    organisations_response = machine_client.get("/api/v1/organisations/")
+    assert organisations_response.status_code == status.HTTP_200_OK
+    assert organisations_response.json()["results"][0]["id"] == organisation
+
+
+def test_token_exchange__deleted_trust_relationship__access_token_rejected(
+    machine_client: APIClient,
+    admin_client: APIClient,
+    organisation: int,
+    trust_relationship: int,
+    oidc_issuer: OIDCIssuerStub,
+) -> None:
+    # Given
+    token = oidc_issuer.sign_token(
+        aud="https://github.com/Flagsmith",
+        sub="repo:Flagsmith/flagsmith:ref:refs/heads/main",
+        repository="Flagsmith/flagsmith",
+    )
+    access_token = machine_client.post(
+        "/api/v1/auth/oidc/token/", data={"token": token}, format="json"
+    ).json()["access_token"]
+    admin_client.delete(
+        f"/api/v1/organisations/{organisation}"
+        f"/trust-relationships/{trust_relationship}/"
+    )
+
+    # When
+    machine_client.credentials(HTTP_AUTHORIZATION=f"Bearer {access_token}")
+    response = machine_client.get("/api/v1/organisations/")
+
+    # Then
+    assert response.status_code == status.HTTP_401_UNAUTHORIZED
+
+
+def test_token_exchange__minted_token__cannot_manage_machine_credentials(
+    machine_client: APIClient,
+    organisation: int,
+    trust_relationship: int,
+    oidc_issuer: OIDCIssuerStub,
+) -> None:
+    # Given
+    token = oidc_issuer.sign_token(
+        aud="https://github.com/Flagsmith",
+        sub="repo:Flagsmith/flagsmith:ref:refs/heads/main",
+        repository="Flagsmith/flagsmith",
+    )
+    access_token = machine_client.post(
+        "/api/v1/auth/oidc/token/", data={"token": token}, format="json"
+    ).json()["access_token"]
+    machine_client.credentials(HTTP_AUTHORIZATION=f"Bearer {access_token}")
+
+    # When
+    master_api_keys_response = machine_client.post(
+        f"/api/v1/organisations/{organisation}/master-api-keys/",
+        data={"name": "sneaky", "organisation": organisation},
+    )
+    trust_relationships_response = machine_client.get(
+        f"/api/v1/organisations/{organisation}/trust-relationships/"
+    )
+
+    # Then
+    assert master_api_keys_response.status_code == status.HTTP_401_UNAUTHORIZED
+    assert trust_relationships_response.status_code == status.HTTP_401_UNAUTHORIZED
+
+
+def test_token_exchange__no_matching_trust_relationship__returns_401(
+    machine_client: APIClient,
+    trust_relationship: int,
+    oidc_issuer: OIDCIssuerStub,
+) -> None:
+    # Given
+    token = oidc_issuer.sign_token(
+        aud="https://github.com/Flagsmith",
+        sub="repo:SomeoneElse/repo:ref:refs/heads/main",
+        repository="SomeoneElse/repo",
+    )
+
+    # When
+    response = machine_client.post(
+        "/api/v1/auth/oidc/token/", data={"token": token}, format="json"
+    )
+
+    # Then
+    assert response.status_code == status.HTTP_401_UNAUTHORIZED
+
+
+def test_token_exchange__multi_audience_token_matching_multiple__returns_401(
+    machine_client: APIClient,
+    admin_client: APIClient,
+    organisation: int,
+    trust_relationship: int,
+    oidc_issuer: OIDCIssuerStub,
+) -> None:
+    # Given: a second trust relationship under a different audience, and a
+    # token listing both audiences
+    create_response = admin_client.post(
+        f"/api/v1/organisations/{organisation}/trust-relationships/",
+        data={
+            "name": "GitHub Actions (CI)",
+            "issuer": "https://token.actions.githubusercontent.com",
+            "audience": "flagsmith-ci",
+            "is_admin": True,
+        },
+        format="json",
+    )
+    assert create_response.status_code == status.HTTP_201_CREATED
+    token = oidc_issuer.sign_token(
+        aud=["https://github.com/Flagsmith", "flagsmith-ci"],
+        sub="repo:Flagsmith/flagsmith:ref:refs/heads/main",
+        repository="Flagsmith/flagsmith",
+    )
+
+    # When
+    response = machine_client.post(
+        "/api/v1/auth/oidc/token/", data={"token": token}, format="json"
+    )
+
+    # Then
+    assert response.status_code == status.HTTP_401_UNAUTHORIZED
+    assert response.json()["detail"] == (
+        "Token audience matches multiple trust relationships."
+    )
+
+
+def test_token_exchange__missing_token__returns_400(
+    machine_client: APIClient,
+) -> None:
+    # Given / When
+    response = machine_client.post("/api/v1/auth/oidc/token/", data={}, format="json")
+
+    # Then
+    assert response.status_code == status.HTTP_400_BAD_REQUEST

api/tests/integration/trust_relationships/test_viewset.py

@@ -33,16 +33,17 @@ def test_create_trust_relationship__valid_data__returns_backing_key_details(
     ]
 
 
-def test_create_trust_relationship__trailing_slash_issuer__returns_normalised_issuer(
+def test_create_trust_relationship__trailing_slash_issuer__returns_issuer_verbatim(
     admin_client: APIClient,
     organisation: int,
 ) -> None:
-    # Given
+    # Given: an issuer that ends in a slash, as Auth0's does. `iss` is matched
+    # by exact string at exchange time, so the slash must survive the round trip.
     url = f"/api/v1/organisations/{organisation}/trust-relationships/"
     data = {
-        "name": "GitHub Actions",
-        "issuer": "https://token.actions.githubusercontent.com/",
-        "audience": "https://github.com/Flagsmith",
+        "name": "Auth0",
+        "issuer": "https://tenant.eu.auth0.com/",
+        "audience": "https://api.flagsmith.com",
         "is_admin": True,
     }
 
@@ -51,7 +52,7 @@ def test_create_trust_relationship__trailing_slash_issuer__returns_normalised_is
 
     # Then
     assert response.status_code == status.HTTP_201_CREATED
-    assert response.json()["issuer"] == "https://token.actions.githubusercontent.com"
+    assert response.json()["issuer"] == "https://tenant.eu.auth0.com/"
 
 
 def test_create_trust_relationship__http_issuer__returns_400(

api/tests/test_helpers.py

@@ -1,5 +1,9 @@
+import json
+from datetime import datetime, timedelta, timezone
 from typing import Any
 
+import jwt
+from cryptography.hazmat.primitives.asymmetric import rsa
 from flag_engine.segments.types import ConditionOperator
 
 
@@ -32,3 +36,40 @@ def generate_segment_data(
             }
         ],
     }
+
+
+class OIDCIssuerStub:
+    """An RSA keypair posing as an OIDC issuer: signs tokens and serves the
+    matching JWKS in tests.
+    """
+
+    KEY_ID = "test-key"
+
+    def __init__(self, issuer: str) -> None:
+        self.issuer = issuer
+        self._private_key = rsa.generate_private_key(
+            public_exponent=65537, key_size=2048
+        )
+
+    @property
+    def jwks(self) -> dict[str, Any]:
+        jwk = json.loads(
+            jwt.algorithms.RSAAlgorithm.to_jwk(self._private_key.public_key())
+        )
+        jwk["kid"] = self.KEY_ID
+        return {"keys": [jwk]}
+
+    def sign_token(self, expires_in_seconds: int = 300, **claims: Any) -> str:
+        now = datetime.now(tz=timezone.utc)
+        payload: dict[str, Any] = {
+            "iss": self.issuer,
+            "iat": now,
+            "exp": now + timedelta(seconds=expires_in_seconds),
+            **claims,
+        }
+        return jwt.encode(
+            payload,
+            self._private_key,
+            algorithm="RS256",
+            headers={"kid": self.KEY_ID},
+        )