Enhancing Pluggable Features of Token Manager

closes #21182

server-spi/src/main/java/org/keycloak/services/clientpolicy/ClientPolicyEvent.java

@@ -40,6 +40,7 @@ public enum ClientPolicyEvent {
     TOKEN_REFRESH,
     TOKEN_REFRESH_RESPONSE,
     TOKEN_REVOKE,
+    TOKEN_REVOKE_RESPONSE,
     TOKEN_INTROSPECT,
     USERINFO_REQUEST,
     LOGOUT_REQUEST,

services/src/main/java/org/keycloak/protocol/oidc/TokenManager.java

@@ -1018,6 +1018,8 @@ public class TokenManager {
 
         String stateHash;
 
+        private AccessTokenResponse response;
+
         public AccessTokenResponseBuilder(RealmModel realm, ClientModel client, EventBuilder event, KeycloakSession session,
                                           UserSessionModel userSession, ClientSessionContext clientSessionCtx) {
             this.realm = realm;
@@ -1152,6 +1154,8 @@ public class TokenManager {
         }
 
         public AccessTokenResponse build() {
+            if (response != null) return response;
+
             if (accessToken != null) {
                 event.detail(Details.TOKEN_ID, accessToken.getId());
             }
@@ -1214,7 +1218,8 @@ public class TokenManager {
             res.setScope(responseScope);
             event.detail(Details.SCOPE, responseScope);
 
-            return res;
+            response = res;
+            return response;
         }
 
 

services/src/main/java/org/keycloak/protocol/oidc/endpoints/LogoutEndpoint.java

@@ -512,6 +512,7 @@ public class LogoutEndpoint {
 
         try {
             session.clientPolicy().triggerOnEvent(new LogoutRequestContext(form));
+            refreshToken = form.getFirst(OAuth2Constants.REFRESH_TOKEN);
         } catch (ClientPolicyException cpe) {
             throw new CorsErrorResponseException(cors, cpe.getError(), cpe.getErrorDetail(), cpe.getErrorStatus());
         }

services/src/main/java/org/keycloak/protocol/oidc/endpoints/TokenEndpoint.java

@@ -497,6 +497,7 @@ public class TokenEndpoint {
         } else {
             res = responseBuilder.build();
         }
+
         event.success();
 
         return cors.builder(Response.ok(res).type(MediaType.APPLICATION_JSON_TYPE)).build();
@@ -528,6 +529,7 @@ public class TokenEndpoint {
 
         try {
             session.clientPolicy().triggerOnEvent(new TokenRefreshContext(formParams));
+            refreshToken = formParams.getFirst(OAuth2Constants.REFRESH_TOKEN);
         } catch (ClientPolicyException cpe) {
             event.error(cpe.getError());
             throw new CorsErrorResponseException(cors, cpe.getError(), cpe.getErrorDetail(), cpe.getErrorStatus());

services/src/main/java/org/keycloak/protocol/oidc/endpoints/TokenIntrospectionEndpoint.java

@@ -46,8 +46,8 @@ import jakarta.ws.rs.core.Response.Status;
  */
 public class TokenIntrospectionEndpoint {
 
-    private static final String PARAM_TOKEN_TYPE_HINT = "token_type_hint";
+    public static final String PARAM_TOKEN_TYPE_HINT = "token_type_hint";
-    private static final String PARAM_TOKEN = "token";
+    public static final String PARAM_TOKEN = "token";
 
     private final KeycloakSession session;
 
@@ -100,6 +100,7 @@ public class TokenIntrospectionEndpoint {
 
         try {
             session.clientPolicy().triggerOnEvent(new TokenIntrospectContext(formParams));
+            token = formParams.getFirst(PARAM_TOKEN);
         } catch (ClientPolicyException cpe) {
             throw throwErrorResponseException(Errors.INVALID_REQUEST, cpe.getErrorDetail(), Status.BAD_REQUEST);
         }

services/src/main/java/org/keycloak/protocol/oidc/endpoints/TokenRevocationEndpoint.java

@@ -50,6 +50,7 @@ import org.keycloak.representations.AccessToken;
 import org.keycloak.services.CorsErrorResponseException;
 import org.keycloak.services.clientpolicy.ClientPolicyException;
 import org.keycloak.services.clientpolicy.context.TokenRevokeContext;
+import org.keycloak.services.clientpolicy.context.TokenRevokeResponseContext;
 import org.keycloak.services.managers.UserSessionCrossDCManager;
 import org.keycloak.services.managers.UserSessionManager;
 import org.keycloak.services.resources.Cors;
@@ -59,7 +60,7 @@ import org.keycloak.util.TokenUtil;
  * @author <a href="mailto:yoshiyuki.tabata.jy@hitachi.com">Yoshiyuki Tabata</a>
  */
 public class TokenRevocationEndpoint {
-    private static final String PARAM_TOKEN = "token";
+    public static final String PARAM_TOKEN = "token";
 
     private final KeycloakSession session;
 
@@ -120,6 +121,13 @@ public class TokenRevocationEndpoint {
 
         event.success();
 
+        try {
+            session.clientPolicy().triggerOnEvent(new TokenRevokeResponseContext(formParams));
+        } catch (ClientPolicyException cpe) {
+            event.error(cpe.getError());
+            throw new CorsErrorResponseException(cors, cpe.getError(), cpe.getErrorDetail(), cpe.getErrorStatus());
+        }
+
         session.getProvider(SecurityHeadersProvider.class).options().allowEmptyContentType();
         return cors.builder(Response.ok()).build();
     }

services/src/main/java/org/keycloak/protocol/oidc/endpoints/UserInfoEndpoint.java

@@ -100,7 +100,7 @@ public class UserInfoEndpoint {
     private final RealmModel realm;
     private final OAuth2Error error;
     private Cors cors;
-    private String authorization;
+    private TokenForUserInfo tokenForUserInfo = new TokenForUserInfo();
 
     public UserInfoEndpoint(KeycloakSession session, org.keycloak.protocol.oidc.TokenManager tokenManager) {
         this.session = session;
@@ -163,7 +163,7 @@ public class UserInfoEndpoint {
         cors.allowAllOrigins();
 
         try {
-            session.clientPolicy().triggerOnEvent(new UserInfoRequestContext(authorization));
+            session.clientPolicy().triggerOnEvent(new UserInfoRequestContext(tokenForUserInfo));
         } catch (ClientPolicyException cpe) {
             throw error.error(cpe.getError()).errorDescription(cpe.getErrorDetail()).status(cpe.getErrorStatus()).build();
         }
@@ -172,7 +172,7 @@ public class UserInfoEndpoint {
                 .event(EventType.USER_INFO_REQUEST)
                 .detail(Details.AUTH_METHOD, Details.VALIDATE_ACCESS_TOKEN);
 
-        if (authorization == null) {
+        if (tokenForUserInfo.getToken() == null) {
             event.error(Errors.INVALID_TOKEN);
             throw error.unauthorized();
         }
@@ -180,7 +180,7 @@ public class UserInfoEndpoint {
         AccessToken token;
         ClientModel clientModel = null;
         try {
-            TokenVerifier<AccessToken> verifier = TokenVerifier.create(authorization, AccessToken.class).withDefaultChecks()
+            TokenVerifier<AccessToken> verifier = TokenVerifier.create(tokenForUserInfo.getToken(), AccessToken.class).withDefaultChecks()
                     .realmUrl(Urls.realmIssuer(session.getContext().getUri().getBaseUri(), realm.getName()));
 
             SignatureVerifierContext verifierContext = session.getProvider(SignatureProvider.class, verifier.getHeader().getAlgorithm().name()).verifier(verifier.getHeader().getKeyId());
@@ -417,11 +417,24 @@ public class UserInfoEndpoint {
 
     private void authorization(String accessToken) {
         if (accessToken != null) {
-            if (authorization == null) {
+            if (tokenForUserInfo.getToken() == null) {
-                authorization = accessToken;
+                tokenForUserInfo.setToken(accessToken);
             } else {
                 throw error.cors(cors.allowAllOrigins()).invalidRequest("More than one method used for including an access token");
             }
         }
     }
+
+    public static class TokenForUserInfo {
+
+        private String token;
+
+        public String getToken() {
+            return token;
+        }
+
+        public void setToken(String token) {
+            this.token = token;
+        }
+    }
 }

services/src/main/java/org/keycloak/services/clientpolicy/context/LogoutRequestContext.java

@@ -17,7 +17,6 @@
 
 package org.keycloak.services.clientpolicy.context;
 
-import jakarta.ws.rs.core.MultivaluedHashMap;
 import jakarta.ws.rs.core.MultivaluedMap;
 
 import org.keycloak.services.clientpolicy.ClientPolicyContext;

services/src/main/java/org/keycloak/services/clientpolicy/context/TokenRevokeResponseContext.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2020 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.keycloak.services.clientpolicy.context;
+
+import jakarta.ws.rs.core.MultivaluedMap;
+
+import org.keycloak.services.clientpolicy.ClientPolicyContext;
+import org.keycloak.services.clientpolicy.ClientPolicyEvent;
+
+/**
+ * @author <a href="mailto:takashi.norimatsu.ws@hitachi.com">Takashi Norimatsu</a>
+ */
+public class TokenRevokeResponseContext implements ClientPolicyContext {
+
+    private final MultivaluedMap<String, String> params;
+
+    public TokenRevokeResponseContext(MultivaluedMap<String, String> params) {
+        this.params = params;
+    }
+
+    @Override
+    public ClientPolicyEvent getEvent() {
+        return ClientPolicyEvent.TOKEN_REVOKE_RESPONSE;
+    }
+
+    public MultivaluedMap<String, String> getParams() {
+        return params;
+    }
+
+}

services/src/main/java/org/keycloak/services/clientpolicy/context/UserInfoRequestContext.java

@@ -17,6 +17,7 @@
 
 package org.keycloak.services.clientpolicy.context;
 
+import org.keycloak.protocol.oidc.endpoints.UserInfoEndpoint;
 import org.keycloak.services.clientpolicy.ClientPolicyContext;
 import org.keycloak.services.clientpolicy.ClientPolicyEvent;
 
@@ -25,10 +26,10 @@ import org.keycloak.services.clientpolicy.ClientPolicyEvent;
  */
 public class UserInfoRequestContext implements ClientPolicyContext {
 
-    private final String tokenString;
+    private UserInfoEndpoint.TokenForUserInfo tokenForUserInfo;
 
-    public UserInfoRequestContext(String tokenString) {
+    public UserInfoRequestContext(UserInfoEndpoint.TokenForUserInfo tokenForUserInfo) {
-        this.tokenString = tokenString;
+        this.tokenForUserInfo = tokenForUserInfo;
     }
 
     @Override
@@ -36,8 +37,8 @@ public class UserInfoRequestContext implements ClientPolicyContext {
         return ClientPolicyEvent.USERINFO_REQUEST;
     }
 
-    public String getTokenString() {
+    public UserInfoEndpoint.TokenForUserInfo getTokenForUserInfo() {
-        return tokenString;
+        return tokenForUserInfo;
     }
 
 }

services/src/main/java/org/keycloak/services/clientpolicy/executor/HolderOfKeyEnforcerExecutor.java

@@ -143,7 +143,7 @@ public class HolderOfKeyEnforcerExecutor implements ClientPolicyExecutorProvider
     }
 
     private void checkUserInfo(UserInfoRequestContext context, HttpRequest request) throws ClientPolicyException {
-        String encodedAccessToken = context.getTokenString();
+        String encodedAccessToken = context.getTokenForUserInfo().getToken();
 
         AccessToken accessToken = session.tokens().decode(encodedAccessToken, AccessToken.class);
         if (accessToken == null) {

testsuite/integration-arquillian/servers/auth-server/services/testsuite-providers/src/main/java/org/keycloak/testsuite/services/clientpolicy/executor/TestEnhancedPluggableTokenManagerExecutor.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2023 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.keycloak.testsuite.services.clientpolicy.executor;
+
+import org.jboss.logging.Logger;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.protocol.oidc.TokenManager;
+import org.keycloak.representations.idm.ClientPolicyExecutorConfigurationRepresentation;
+import org.keycloak.services.clientpolicy.ClientPolicyContext;
+import org.keycloak.services.clientpolicy.ClientPolicyEvent;
+import org.keycloak.services.clientpolicy.ClientPolicyException;
+import org.keycloak.services.clientpolicy.context.TokenResponseContext;
+import org.keycloak.services.clientpolicy.executor.ClientPolicyExecutorProvider;
+
+/**
+ * @author <a href="mailto:takashi.norimatsu.ws@hitachi.com">Takashi Norimatsu</a>
+ */
+public class TestEnhancedPluggableTokenManagerExecutor implements ClientPolicyExecutorProvider<ClientPolicyExecutorConfigurationRepresentation> {
+
+    private static final Logger logger = Logger.getLogger(TestEnhancedPluggableTokenManagerExecutor.class);
+
+    protected final KeycloakSession session;
+
+    public TestEnhancedPluggableTokenManagerExecutor(KeycloakSession session) {
+        this.session = session;
+    }
+
+    @Override
+    public String getProviderId() {
+        return TestEnhancedPluggableTokenManagerExecutorFactory.PROVIDER_ID;
+    }
+
+    @Override
+    public void executeOnEvent(ClientPolicyContext context) throws ClientPolicyException {
+        ClientPolicyEvent event = context.getEvent();
+
+        if (event.equals(ClientPolicyEvent.TOKEN_RESPONSE)) {
+            TokenResponseContext tokenResponseContext = (TokenResponseContext)context;
+            dropSubClaimAndBuildTokenResponse(tokenResponseContext.getAccessTokenResponseBuilder());
+        }
+    }
+
+    private void dropSubClaimAndBuildTokenResponse(TokenManager.AccessTokenResponseBuilder builder) throws ClientPolicyException {
+        builder.getAccessToken().subject(null);
+        builder.build();
+    }
+
+}

testsuite/integration-arquillian/servers/auth-server/services/testsuite-providers/src/main/java/org/keycloak/testsuite/services/clientpolicy/executor/TestEnhancedPluggableTokenManagerExecutorFactory.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2023 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.keycloak.testsuite.services.clientpolicy.executor;
+
+import java.util.Collections;
+import java.util.List;
+
+import org.keycloak.Config.Scope;
+import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.KeycloakSessionFactory;
+import org.keycloak.provider.ProviderConfigProperty;
+import org.keycloak.services.clientpolicy.executor.ClientPolicyExecutorProvider;
+import org.keycloak.services.clientpolicy.executor.ClientPolicyExecutorProviderFactory;
+
+/**
+ * @author <a href="mailto:takashi.norimatsu.ws@hitachi.com">Takashi Norimatsu</a>
+ */
+public class TestEnhancedPluggableTokenManagerExecutorFactory implements ClientPolicyExecutorProviderFactory {
+
+    public static final String PROVIDER_ID = "test-enhanced-token-mgr";
+
+    @Override
+    public ClientPolicyExecutorProvider create(KeycloakSession session) {
+        return new TestEnhancedPluggableTokenManagerExecutor(session);
+    }
+
+    @Override
+    public void init(Scope config) {
+    }
+
+    @Override
+    public void postInit(KeycloakSessionFactory factory) {
+    }
+
+    @Override
+    public void close() {
+    }
+
+    @Override
+    public String getId() {
+        return PROVIDER_ID;
+    }
+
+    @Override
+    public String getHelpText() {
+        return "NA";
+    }
+
+    @Override
+    public List<ProviderConfigProperty> getConfigProperties() {
+        return Collections.emptyList();
+    }
+
+}

testsuite/integration-arquillian/servers/auth-server/services/testsuite-providers/src/main/resources/META-INF/services/org.keycloak.services.clientpolicy.executor.ClientPolicyExecutorProviderFactory

@@ -1 +1,2 @@
-org.keycloak.testsuite.services.clientpolicy.executor.TestRaiseExceptionExecutorFactory
+org.keycloak.testsuite.services.clientpolicy.executor.TestRaiseExceptionExecutorFactory
+org.keycloak.testsuite.services.clientpolicy.executor.TestEnhancedPluggableTokenManagerExecutorFactory

testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/client/policies/ClientPoliciesExtendedEventTest.java

@@ -46,6 +46,7 @@ import org.keycloak.models.Constants;
 import org.keycloak.models.OAuth2DeviceConfig;
 import org.keycloak.models.utils.KeycloakModelUtils;
 import org.keycloak.protocol.oidc.OIDCConfigAttributes;
+import org.keycloak.representations.AccessToken;
 import org.keycloak.representations.RefreshToken;
 import org.keycloak.representations.idm.ClientRepresentation;
 import org.keycloak.representations.idm.CredentialRepresentation;
@@ -64,6 +65,8 @@ import org.keycloak.testsuite.pages.ErrorPage;
 import org.keycloak.testsuite.pages.LogoutConfirmPage;
 import org.keycloak.testsuite.pages.OAuth2DeviceVerificationPage;
 import org.keycloak.testsuite.pages.OAuthGrantPage;
+import org.keycloak.testsuite.services.clientpolicy.executor.TestEnhancedPluggableTokenManagerExecutor;
+import org.keycloak.testsuite.services.clientpolicy.executor.TestEnhancedPluggableTokenManagerExecutorFactory;
 import org.keycloak.testsuite.services.clientpolicy.executor.TestRaiseExceptionExecutorFactory;
 import org.keycloak.testsuite.util.ClientBuilder;
 import org.keycloak.testsuite.util.ClientPoliciesUtil.ClientPoliciesBuilder;
@@ -597,4 +600,44 @@ public class ClientPoliciesExtendedEventTest extends AbstractClientPoliciesTest
         assertTrue(errorPage.isCurrent());
         assertEquals("Exception thrown intentionally", errorPage.getError());
     }
+
+    @Test
+    public void testEnhancedPluggableTokenManagerForTokenResponse() throws Exception {
+        // register a confidential client
+        String clientId = generateSuffixedName(CLIENT_NAME);
+        String clientSecret = "secret";
+        createClientByAdmin(clientId, (ClientRepresentation clientRep) -> {
+            clientRep.setSecret(clientSecret);
+            clientRep.setPublicClient(Boolean.FALSE);
+            clientRep.setBearerOnly(Boolean.FALSE);
+        });
+
+        // register profiles
+        String json = (new ClientProfilesBuilder()).addProfile(
+                (new ClientProfileBuilder()).createProfile(PROFILE_NAME, "Den Forste Profilen")
+                        .addExecutor(TestEnhancedPluggableTokenManagerExecutorFactory.PROVIDER_ID, null)
+                        .toRepresentation()
+        ).toString();
+        updateProfiles(json);
+
+        // register policies
+        json = (new ClientPoliciesBuilder()).addPolicy(
+                (new ClientPolicyBuilder()).createPolicy(POLICY_NAME, "La Primera Plitica", Boolean.TRUE)
+                        .addCondition(ClientAccessTypeConditionFactory.PROVIDER_ID,
+                                createClientAccessTypeConditionConfig(Arrays.asList(ClientAccessTypeConditionFactory.TYPE_CONFIDENTIAL)))
+                        .addProfile(PROFILE_NAME)
+                        .toRepresentation()
+        ).toString();
+        updatePolicies(json);
+
+        oauth.clientId(clientId);
+        oauth.doLogin(TEST_USER_NAME, TEST_USER_PASSWORD);
+
+        events.expectLogin().client(clientId).assertEvent();
+        String code = oauth.getCurrentQuery().get(OAuth2Constants.CODE);
+        OAuthClient.AccessTokenResponse response = oauth.doAccessTokenRequest(code, clientSecret);
+        assertEquals(200, response.getStatusCode());
+        AccessToken token = oauth.verifyToken(response.getAccessToken());
+        assertNull(token.getSubject());
+    }
 }