[minor] use custom claims overrides so we can directly specify special wlcg/scitoken claims (#7)

Author: dsschult

src/scitoken_issuer/config.py

@@ -1,6 +1,8 @@
 import dataclasses as dc
+import json
 import logging
 import secrets
+from typing import Any
 
 from wipac_dev_tools import from_environment_as_dataclass
 
@@ -28,7 +30,7 @@ class EnvConfig:
     IDP_USERNAME_CLAIM: str = 'preferred_username'
 
     ISSUER_ADDRESS: str = ''
-    AUDIENCE: str = ''
+    CUSTOM_CLAIMS: dict[str, Any] | str = ''
     KEY_TYPE: str = 'RS256'
 
     ACCESS_TOKEN_EXPIRATION: int = 300  # seconds
@@ -68,10 +70,10 @@ def __post_init__(self) -> None:
                 raise ConfigError('Must specify IDP_CLIENT_ID in production')
             if not self.IDP_CLIENT_SECRET:
                 raise ConfigError('Must specify IDP_CLIENT_SECRET in production')
-            if not self.AUDIENCE:
-                raise ConfigError('Must specify AUDIENCE in production')
             if not self.MONGODB_URL:
                 raise ConfigError('Must specify MONGODB_URL in production')
+        if self.CUSTOM_CLAIMS and isinstance(self.CUSTOM_CLAIMS, str):
+            object.__setattr__(self, 'CUSTOM_CLAIMS', json.loads(self.CUSTOM_CLAIMS))
         if self.KEY_TYPE not in DEFAULT_KEY_ALGORITHMS:
             raise ConfigError(f'KEY_TYPE must be one of {DEFAULT_KEY_ALGORITHMS}')
         if self.MONGODB_WRITE_CONCERN < 1:

src/scitoken_issuer/server.py

@@ -450,27 +450,32 @@ async def post(self):
         current_key = await self.state.get_current_key()
         auth = Auth(
             secret=current_key['private_key'],
-            audience=config.ENV.AUDIENCE,
             issuer=config.ENV.ISSUER_ADDRESS,
             algorithm=config.ENV.KEY_TYPE,
             integer_times=True,  # scitokens-cpp can't handle floats
         )
+        access_claims = {
+            'jti': uuid.uuid4().hex,
+            config.ENV.IDP_USERNAME_CLAIM: username,
+            'scope': access_scope,
+        }
+        if config.ENV.CUSTOM_CLAIMS:
+            logger.info('custom claims: %r', config.ENV.CUSTOM_CLAIMS)
+            if not set(access_claims).isdisjoint(config.ENV.CUSTOM_CLAIMS):
+                logger.error('CUSTOM_CLAIMS should not override existing claims')
+                raise OAuthError(500, error="invalid_claims")
+            access_claims.update(config.ENV.CUSTOM_CLAIMS)
         access_token = auth.create_token(
             subject=username,
             expiration=config.ENV.ACCESS_TOKEN_EXPIRATION,
-            payload={
-                'aud': [config.ENV.AUDIENCE],
-                config.ENV.IDP_USERNAME_CLAIM: username,
-                'scope': access_scope,
-                'jti': uuid.uuid4().hex,
-                'wlcg.ver': '1.0',
-            },
+            payload=access_claims,
             headers={'kid': current_key['kid']},
         )
         refresh_token = auth.create_token(
             subject=username,
             expiration=config.ENV.REFRESH_TOKEN_EXPIRATION,
             payload={
+                'jti': uuid.uuid4().hex,
                 'aud': config.ENV.ISSUER_ADDRESS,
                 config.ENV.IDP_USERNAME_CLAIM: username,
                 'idp_username': username,

tests/test_server.py

@@ -126,7 +126,6 @@ async def fn():
             CI_TESTING=True,
             PORT=port,
             ISSUER_ADDRESS='',
-            AUDIENCE='storage',
             POSIX_PATH=str(tmp_path),
             DEVICE_CODE_POLLING_INTERVAL=.1):
 
@@ -307,7 +306,6 @@ async def common(scopes):
                         assert login.call_count == 1
 
                         data = jwt.decode(token, options={"verify_signature": False})
-                        assert data['aud'] == [scitoken_issuer.config.ENV.AUDIENCE]
                         assert data['sub'] == user
                         return data
 
@@ -455,15 +453,19 @@ async def test_scitokens(server, storage, monkeypatch):
     """
     users, _, _ = storage
     user = 'test1'
-    async with server() as address:
-        login = Mock(side_effect=do_device_login(user))
-        monkeypatch.setattr('rest_tools.client.device_client._print_qrcode', login)
-        scopes = ['storage.read:/data/ana/project1/sub1']
-        async with make_client(address, scopes) as rc:
-            token = rc._openid_token()
-            assert token
-            access_data = jwt.decode(token, options={"verify_signature": False})
-            logging.info('access token: %r', access_data)
-            assert isinstance(access_data['exp'], int)
-
-            assert access_data['wlcg.ver'] == '1.0'
+
+    with env(CUSTOM_CLAIMS='{"aud": ["https://wlcg.cern.ch/jwt/v1/any"], "wlcg.ver":"1.0"}'):
+        async with server() as address:
+            login = Mock(side_effect=do_device_login(user))
+            monkeypatch.setattr('rest_tools.client.device_client._print_qrcode', login)
+            scopes = ['storage.read:/data/ana/project1/sub1']
+            async with make_client(address, scopes) as rc:
+                token = rc._openid_token()
+                assert token
+                access_data = jwt.decode(token, options={"verify_signature": False})
+                logging.info('access token: %r', access_data)
+                assert isinstance(access_data['exp'], int)
+
+                assert access_data['wlcg.ver'] == '1.0'
+
+                assert 'https://wlcg.cern.ch/jwt/v1/any' in access_data['aud']