Fix Device Code session propagation and OIDC claims

openam-oauth2/src/main/java/org/forgerock/oauth2/core/DeviceCode.java

@@ -21,8 +21,10 @@
 import static org.forgerock.openam.utils.CollectionUtils.newList;
 import static org.forgerock.openam.utils.Time.*;
 
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Date;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -115,6 +117,21 @@ public void setResourceOwnerId(String resourceOwnerId) {
         setStringProperty(OAuth2Constants.CoreTokenParams.USERNAME, resourceOwnerId);
     }
 
+
+
+    public void setAuthModules(String authModules) {
+    	setStringProperty(AUTH_MODULES, authModules);
+    }
+
+
+    /**
+     * Get the auth modules string.
+     * @return The pipe-separated list of auth modules.
+     */
+    public String getAuthModules() {
+        return getStringProperty(AUTH_MODULES);
+    }
+
     /**
      * Gets the Client ID parameter.
      * @return The Client ID.
@@ -138,6 +155,14 @@ public String getNonce() {
     public String getAcrValues() {
         return getStringProperty(OAuth2Constants.Params.ACR_VALUES);
     }
+
+    /**
+     * Sets the ACR Values for device code object.
+     */
+    public void setAcrValues(String acrValues) {
+        setStringProperty(OAuth2Constants.Params.ACR_VALUES, acrValues);
+    }
+
 
     /**
      * Gets the Code Challenge Method parameter.
@@ -338,6 +363,7 @@ public boolean isAuthorized() {
         return Boolean.valueOf(getStringProperty("AUTHORIZED"));
     }
 
+
     /**
      * {@inheritDoc}
      */

openam-oauth2/src/main/java/org/forgerock/oauth2/core/DeviceCodeGrantTypeHandler.java

@@ -85,22 +85,31 @@ protected AccessToken handle(OAuth2Request request, ClientRegistration client,
 
         String clientId = client.getClientId();
         DeviceCode deviceCode = tokenStore.readDeviceCode(clientId, code, request);
-
+        
         if (deviceCode == null ||
                 !clientId.equals(deviceCode.getClientId()) ||
                 !request.getParameter(REALM).equals(deviceCode.getRealm())) {
             throw new AuthorizationDeclinedException();
         }
 
+        AccessToken accessToken;
         try {
             if (deviceCode.isAuthorized()) {
                 String grantType = request.getParameter(OAuth2Constants.Params.GRANT_TYPE);
                 Set<String> scope = deviceCode.getScope();
                 String resourceOwnerId = deviceCode.getResourceOwnerId();
                 String validatedClaims = providerSettings.validateRequestedClaims(
                         deviceCode.getStringProperty(OAuth2Constants.Custom.CLAIMS));
-                return generateAccessToken(providerSettings, grantType, clientId, resourceOwnerId, scope,
-                        validatedClaims, request);
+                final String nonce = deviceCode.getNonce();
+
+                accessToken = generateAccessToken(providerSettings, grantType, clientId, resourceOwnerId, scope,
+                        validatedClaims, nonce, request);
+
+                providerSettings.additionalDataToReturnFromTokenEndpoint(
+                        accessToken,
+                        request);
+
+                return accessToken;
             }
 
             if (deviceCode.getExpiryTime() < currentTimeMillis()) {
@@ -130,9 +139,9 @@ protected AccessToken handle(OAuth2Request request, ClientRegistration client,
     }
 
     private AccessToken generateAccessToken(OAuth2ProviderSettings providerSettings, String grantType, String clientId,
-            String resourceOwnerId, Set<String> scope, String validatedClaims, OAuth2Request request)
+            String resourceOwnerId, Set<String> scope, String validatedClaims, String nonce, OAuth2Request request)
             throws ServerException, NotFoundException {
         return accessTokenGenerator.generateAccessToken(providerSettings, grantType, clientId, resourceOwnerId, null,
-                scope, validatedClaims, null, null, request);
+                scope, validatedClaims, null, nonce, request);
     }
 }

openam-oauth2/src/main/java/org/forgerock/oauth2/core/ResourceOwnerSessionValidator.java

@@ -199,7 +199,8 @@ public ResourceOwner validate(OAuth2Request request) throws ResourceOwnerAuthent
                     throw new LoginRequiredException();
                 }
             } else if (OAuth2Constants.TokenEndpoint.PASSWORD.equals(request.getParameter(GRANT_TYPE))
-                    || OAuth2Constants.TokenEndpoint.CLIENT_CREDENTIALS.equals(request.getParameter(GRANT_TYPE))) {
+                    || OAuth2Constants.TokenEndpoint.CLIENT_CREDENTIALS.equals(request.getParameter(GRANT_TYPE))
+                    || OAuth2Constants.TokenEndpoint.DEVICE_CODE.equals(request.getParameter(GRANT_TYPE))) {
                 // If we're doing password grant type, the SSOToken will have been created and deleted again within
                 // OpenAMResourceOwnerAuthenticator. The request will not have a session, and so the token will have
                 // been null from the attempted creation in L148.

openam-oauth2/src/main/java/org/forgerock/oauth2/restlet/DeviceCodeVerificationResource.java

@@ -27,6 +27,7 @@
 import java.util.Map;
 import java.util.Set;
 
+import com.iplanet.sso.SSOException;
 import com.iplanet.sso.SSOToken;
 import org.forgerock.oauth2.core.AuthorizationService;
 import org.forgerock.oauth2.core.ClientRegistration;
@@ -71,6 +72,7 @@
 import org.restlet.routing.Router;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import com.sun.identity.authentication.util.ISAuthConstants;
 
 /**
  * A restlet resource for user codes
@@ -128,6 +130,7 @@ public DeviceCodeVerificationResource(XUIState xuiState, @Named("OAuth2Router")
     @Post
     public Representation verify(Representation body) throws ServerException, NotFoundException,
             InvalidGrantException, OAuth2RestletException {
+
         final Request restletRequest = getRequest();
         OAuth2Request request = requestFactory.create(restletRequest);
 
@@ -167,7 +170,13 @@ public Representation verify(Representation body) throws ServerException, NotFou
                         saveConsent(request);
                     }
                     if (consentGiven) {
+
                         ResourceOwner resourceOwner = resourceOwnerSessionValidator.validate(request);
+                        deviceCode.setAcrValues(getAuthenticationContextClassReferenceFromRequest(request));
+                        SSOToken token = resourceOwnerSessionValidator.getResourceOwnerSession(request);
+                        if (token != null) {
+                        	populateAuthenticationInfo(deviceCode, token);
+                        }
                         deviceCode.setResourceOwnerId(resourceOwner.getId());
                         deviceCode.setAuthorized(true);
                         tokenStore.updateDeviceCode(deviceCode, request);
@@ -179,10 +188,16 @@ public Representation verify(Representation body) throws ServerException, NotFou
                 }
             } else {
                 ResourceOwner resourceOwner = resourceOwnerSessionValidator.validate(request);
+                deviceCode.setAcrValues(getAuthenticationContextClassReferenceFromRequest(request));
+                SSOToken token = resourceOwnerSessionValidator.getResourceOwnerSession(request);
+                if (token != null) {
+                	populateAuthenticationInfo(deviceCode, token);
+                }
                 deviceCode.setResourceOwnerId(resourceOwner.getId());
                 deviceCode.setAuthorized(true);
                 tokenStore.updateDeviceCode(deviceCode, request);
             }
+
         } catch (IllegalArgumentException e) {
             if (e.getMessage().contains("client_id")) {
                 throw new OAuth2RestletException(400, "invalid_request", e.getMessage(),
@@ -299,4 +314,21 @@ private TemplateFactory getTemplateFactory(Context context) {
     protected void doCatch(Throwable throwable) {
         exceptionHandler.handle(throwable, getContext(), getRequest(), getResponse());
     }
+
+
+    private void populateAuthenticationInfo(DeviceCode deviceCode, SSOToken token) {
+        if (token == null) {
+            return;
+        }
+
+        try {
+            deviceCode.setAuthModules(token.getProperty(ISAuthConstants.AUTH_TYPE));
+        } catch (SSOException e) {
+            logger.warn("Could not get list of auth modules from authentication", e);
+        }
+    }
+
+    private String getAuthenticationContextClassReferenceFromRequest(OAuth2Request request) {
+        return (String) request.getRequest().getAttributes().get(OAuth2Constants.JWTTokenParams.ACR);
+    }
 }

openam-oauth2/src/main/java/org/forgerock/openam/oauth2/StatefulTokenStore.java

@@ -407,6 +407,8 @@ private List<String> getAMRFromAuthModules(OAuth2Request request, OAuth2Provider
         String authModules;
         if (request.getToken(AuthorizationCode.class) != null) {
             authModules = request.getToken(AuthorizationCode.class).getAuthModules();
+        } else if (request.getToken(DeviceCode.class) != null) {
+            authModules = request.getToken(DeviceCode.class).getAuthModules();
         } else if (request.getToken(RefreshToken.class) != null) {
             authModules = request.getToken(RefreshToken.class).getAuthModules();
         } else {
@@ -432,6 +434,8 @@ private List<String> getAMRFromAuthModules(OAuth2Request request, OAuth2Provider
     private String getAuthenticationContextClassReference(OAuth2Request request) {
         if (request.getToken(AuthorizationCode.class) != null) {
             return request.getToken(AuthorizationCode.class).getAuthenticationContextClassReference();
+        } else if(request.getToken(DeviceCode.class) != null){
+        	return request.getToken(DeviceCode.class).getAcrValues();
         } else if (request.getToken(RefreshToken.class) != null) {
             return request.getToken(RefreshToken.class).getAuthenticationContextClassReference();
         } else {
@@ -897,7 +901,7 @@ public DeviceCode createDeviceCode(Set<String> scope, ResourceOwner resourceOwne
             Integer maxAge, String claims, OAuth2Request request, String codeChallenge, String codeChallengeMethod)
             throws ServerException, NotFoundException {
 
-        logger.message("DefaultOAuthTokenStoreImpl::Creating Authorization code");
+        logger.message("DefaultOAuthTokenStoreImpl::Creating Device code");
 
         final OAuth2ProviderSettings providerSettings = providerSettingsFactory.get(request);
         final String deviceCode = UUID.randomUUID().toString();