From 22f67996381d44271cd0032fcd6391d3f786f626 Mon Sep 17 00:00:00 2001
From: YuriyZ <yzabrovarniy@gmail.com>
Date: Mon, 31 Oct 2022 22:11:47 +0200
Subject: [PATCH] fix(jans-auth-server): native sso - return device secret if
 device_sso scope is present #2790

https://github.com/JanssenProject/jans/issues/2518
---
 .../java/io/jans/as/client/TokenResponse.java |  9 +++++
 .../io/jans/as/client/client/Asserter.java    | 17 ++++++++
 .../assertbuilders/JwtAssertBuilder.java      | 11 +++++
 .../TokenResponseAssertBuilder.java           | 11 ++++-
 .../server/model/common/ExecutionContext.java |  9 +++++
 .../as/server/model/token/IdTokenFactory.java |  6 ++-
 .../token/ws/rs/TokenExchangeService.java     | 12 +++---
 .../token/ws/rs/TokenRestWebServiceImpl.java  |  6 ++-
 .../token/ws/rs/TokenExchangeServiceTest.java | 40 +++++++++----------
 9 files changed, 92 insertions(+), 29 deletions(-)

diff --git a/jans-auth-server/client/src/main/java/io/jans/as/client/TokenResponse.java b/jans-auth-server/client/src/main/java/io/jans/as/client/TokenResponse.java
index 9324868a6e3..5eafd82f934 100644
--- a/jans-auth-server/client/src/main/java/io/jans/as/client/TokenResponse.java
+++ b/jans-auth-server/client/src/main/java/io/jans/as/client/TokenResponse.java
@@ -30,6 +30,7 @@ public class TokenResponse extends BaseResponseWithErrors<TokenErrorResponseType
     private String refreshToken;
     private String scope;
     private String idToken;
+    private String deviceToken;
 
     public TokenResponse() {
     }
@@ -75,12 +76,20 @@ public void injectDataFromJson(String json) {
                 if (jsonObj.has("id_token")) {
                     setIdToken(jsonObj.getString("id_token"));
                 }
+                setDeviceToken(jsonObj.optString("device_token"));
             } catch (JSONException e) {
                 LOG.error(e.getMessage(), e);
             }
         }
     }
 
+    public String getDeviceToken() {
+        return deviceToken;
+    }
+
+    public void setDeviceToken(String deviceToken) {
+        this.deviceToken = deviceToken;
+    }
 
     /**
      * Returns the access token issued by the authorization server.
diff --git a/jans-auth-server/client/src/test/java/io/jans/as/client/client/Asserter.java b/jans-auth-server/client/src/test/java/io/jans/as/client/client/Asserter.java
index 3114e504e0e..0e12409d355 100644
--- a/jans-auth-server/client/src/test/java/io/jans/as/client/client/Asserter.java
+++ b/jans-auth-server/client/src/test/java/io/jans/as/client/client/Asserter.java
@@ -9,6 +9,7 @@
 import io.jans.as.client.OpenIdConfigurationResponse;
 import io.jans.as.client.RegisterResponse;
 import io.jans.as.model.register.RegisterRequestParam;
+import org.apache.commons.lang.StringUtils;
 
 import java.util.Arrays;
 
@@ -24,6 +25,22 @@ public class Asserter {
     private Asserter() {
     }
 
+    public static void assertBlank(String value) {
+        assertTrue(StringUtils.isNotBlank(value));
+    }
+
+    public static void assertBlank(String value, String message) {
+        assertTrue(StringUtils.isBlank(value), message);
+    }
+
+    public static void assertNotBlank(String value) {
+        assertTrue(StringUtils.isNotBlank(value));
+    }
+
+    public static void assertNotBlank(String value, String message) {
+        assertTrue(StringUtils.isNotBlank(value), message);
+    }
+
     public static void assertRegisterResponseClaimsNotNull(RegisterResponse response, RegisterRequestParam... claimsToVerify) {
         if (response == null || claimsToVerify == null) {
             return;
diff --git a/jans-auth-server/client/src/test/java/io/jans/as/client/client/assertbuilders/JwtAssertBuilder.java b/jans-auth-server/client/src/test/java/io/jans/as/client/client/assertbuilders/JwtAssertBuilder.java
index e13a1f26938..c22fa2d1c75 100644
--- a/jans-auth-server/client/src/test/java/io/jans/as/client/client/assertbuilders/JwtAssertBuilder.java
+++ b/jans-auth-server/client/src/test/java/io/jans/as/client/client/assertbuilders/JwtAssertBuilder.java
@@ -16,6 +16,7 @@
 import java.util.List;
 
 import static io.jans.as.client.BaseTest.clientEngine;
+import static io.jans.as.client.client.Asserter.assertNotBlank;
 import static org.testng.Assert.*;
 
 public class JwtAssertBuilder extends BaseAssertBuilder {
@@ -28,6 +29,7 @@ public class JwtAssertBuilder extends BaseAssertBuilder {
     private boolean notNullAuthenticationMethodReferences;
     private boolean notNullClaimsAddressdata;
     private boolean checkMemberOfClaimNoEmpty;
+    private boolean notBlankDsHash;
     private String[] claimsPresence;
     private String[] claimsNoPresence;
 
@@ -53,6 +55,11 @@ public JwtAssertBuilder notNullAccesTokenHash() {
         return this;
     }
 
+    public JwtAssertBuilder notBlankDsHash() {
+        notBlankDsHash = true;
+        return this;
+    }
+
     public JwtAssertBuilder notNullAuthenticationTime() {
         this.notNullAuthenticationTime = true;
         return this;
@@ -181,6 +188,10 @@ public void check() {
             assertTrue(jwt.getClaims().getClaimAsStringList("member_of").size() > 1);
         }
 
+        if (notBlankDsHash) {
+            assertNotBlank(jwt.getClaims().getClaimAsString("ds_hash"), "ds_hash claim is not present");
+        }
+
         if (notNullClaimsAddressdata) {
             assertNotNullClaim(JwtClaimName.ADDRESS_STREET_ADDRESS);
             assertNotNullClaim(JwtClaimName.ADDRESS_COUNTRY);
diff --git a/jans-auth-server/client/src/test/java/io/jans/as/client/client/assertbuilders/TokenResponseAssertBuilder.java b/jans-auth-server/client/src/test/java/io/jans/as/client/client/assertbuilders/TokenResponseAssertBuilder.java
index fff82e2ac8e..7d0e56b0e21 100644
--- a/jans-auth-server/client/src/test/java/io/jans/as/client/client/assertbuilders/TokenResponseAssertBuilder.java
+++ b/jans-auth-server/client/src/test/java/io/jans/as/client/client/assertbuilders/TokenResponseAssertBuilder.java
@@ -1,9 +1,9 @@
 package io.jans.as.client.client.assertbuilders;
 
 import io.jans.as.client.TokenResponse;
-import io.jans.as.client.client.AssertBuilder;
 import io.jans.as.model.token.TokenErrorResponseType;
 
+import static io.jans.as.client.client.Asserter.assertNotBlank;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
 
@@ -12,6 +12,7 @@ public class TokenResponseAssertBuilder extends BaseAssertBuilder {
     private TokenResponse response;
     private int status = 200;
     private boolean notNullRefreshToken;
+    private boolean notBlankDeviceToken;
     private boolean notNullIdToken;
     private boolean notNullScope;
     private boolean nullRefreshToken;
@@ -53,6 +54,11 @@ public TokenResponseAssertBuilder notNullRefreshToken() {
         return this;
     }
 
+    public TokenResponseAssertBuilder notBlankDeviceToken() {
+        this.notBlankDeviceToken = true;
+        return this;
+    }
+
     public TokenResponseAssertBuilder notNullIdToken() {
         this.notNullIdToken = true;
         return this;
@@ -89,6 +95,9 @@ public void check() {
             if (notNullRefreshToken) {
                 assertNotNull(response.getRefreshToken(), "The refresh token is null");
             }
+            if (notBlankDeviceToken) {
+                assertNotBlank(response.getDeviceToken(), "The device token is blank");
+            }
         } else {
             assertEquals(response.getStatus(), status, "Unexpected HTTP status response: " + response.getEntity());
             assertNotNull(response.getEntity(), "The entity is null");
diff --git a/jans-auth-server/server/src/main/java/io/jans/as/server/model/common/ExecutionContext.java b/jans-auth-server/server/src/main/java/io/jans/as/server/model/common/ExecutionContext.java
index 68f6dca4615..f4b289e6c1a 100644
--- a/jans-auth-server/server/src/main/java/io/jans/as/server/model/common/ExecutionContext.java
+++ b/jans-auth-server/server/src/main/java/io/jans/as/server/model/common/ExecutionContext.java
@@ -46,6 +46,7 @@ public class ExecutionContext {
 
     private String dpop;
     private String certAsPem;
+    private String deviceSecret;
 
     private String nonce;
     private String state;
@@ -71,6 +72,14 @@ public ExecutionContext(HttpServletRequest httpRequest, HttpServletResponse http
         this.httpResponse = httpResponse;
     }
 
+    public String getDeviceSecret() {
+        return deviceSecret;
+    }
+
+    public void setDeviceSecret(String deviceSecret) {
+        this.deviceSecret = deviceSecret;
+    }
+
     @NotNull
     public Map<String, String> getAttributes() {
         return attributes;
diff --git a/jans-auth-server/server/src/main/java/io/jans/as/server/model/token/IdTokenFactory.java b/jans-auth-server/server/src/main/java/io/jans/as/server/model/token/IdTokenFactory.java
index 4fe604237dc..fe1cd7a09bf 100644
--- a/jans-auth-server/server/src/main/java/io/jans/as/server/model/token/IdTokenFactory.java
+++ b/jans-auth-server/server/src/main/java/io/jans/as/server/model/token/IdTokenFactory.java
@@ -258,7 +258,11 @@ private void addTokenExchangeClaims(JsonWebResponse jwr, ExecutionContext execut
             return;
         }
 
-        String deviceSecret = executionContext.getHttpRequest().getParameter(DEVICE_SECRET);
+        String deviceSecret = executionContext.getDeviceSecret();
+        if (StringUtils.isBlank(deviceSecret)) {
+            deviceSecret = executionContext.getHttpRequest().getParameter(DEVICE_SECRET);
+        }
+
         if (StringUtils.isNotBlank(deviceSecret) && sessionId.getDeviceSecrets().contains(deviceSecret)) {
             jwr.setClaim("ds_hash", CodeVerifier.s256(deviceSecret));
         }
diff --git a/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenExchangeService.java b/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenExchangeService.java
index 93c168f2bee..d05667158fa 100644
--- a/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenExchangeService.java
+++ b/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenExchangeService.java
@@ -172,20 +172,21 @@ public JSONObject processTokenExchange(String scope, Function<JsonWebResponse, V
         return jsonObj;
     }
 
-    public void putNewDeviceSecret(JSONObject jsonObj, String sessionDn, Client client, String scope) {
+    public String createNewDeviceSecret(String sessionDn, Client client, String scope) {
         if (StringUtils.isBlank(scope) || !scope.contains(ScopeConstants.DEVICE_SSO)) {
-            return;
+            log.debug("Skip device secret. No device_sso scope.");
+            return null;
         }
         if (client == null || !ArrayUtils.contains(client.getGrantTypes(), GrantType.TOKEN_EXCHANGE)) {
             log.debug("Skip device secret. Scope has {} value but client does not have Token Exchange Grant Type enabled ('urn:ietf:params:oauth:grant-type:token-exchange')", ScopeConstants.DEVICE_SSO);
-            return;
+            return null;
         }
 
         try {
             final SessionId sessionId = sessionIdService.getSessionByDn(sessionDn);
             if (sessionId == null) {
                 log.debug("Unable to find session by dn: {}", sessionDn);
-                return;
+                return null;
             }
 
             String newDeviceSecret = HandleTokenFactory.generateDeviceSecret();
@@ -193,9 +194,10 @@ public void putNewDeviceSecret(JSONObject jsonObj, String sessionDn, Client clie
 
             sessionIdService.updateSessionId(sessionId, false);
 
-            jsonObj.put("device_token", newDeviceSecret);
+            return newDeviceSecret;
         } catch (Exception e) {
             log.error("Failed to generate device_secret", e);
         }
+        return null;
     }
 }
diff --git a/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenRestWebServiceImpl.java b/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenRestWebServiceImpl.java
index f94da0bbd51..2492860fe6e 100644
--- a/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenRestWebServiceImpl.java
+++ b/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenRestWebServiceImpl.java
@@ -398,6 +398,7 @@ private Response processAuthorizationCode(String code, String scope, String code
         scope = authorizationCodeGrant.checkScopesPolicy(scope);
 
         AccessToken accToken = authorizationCodeGrant.createAccessToken(executionContext); // create token after scopes are checked
+        final String deviceSecret = tokenExchangeService.createNewDeviceSecret(authorizationCodeGrant.getSessionDn(), client, authorizationCodeGrant.getScopesAsString());
 
         IdToken idToken = null;
         if (authorizationCodeGrant.getScopes().contains(OPENID)) {
@@ -414,6 +415,7 @@ private Response processAuthorizationCode(String code, String scope, String code
 
             ExternalUpdateTokenContext context = new ExternalUpdateTokenContext(executionContext.getHttpRequest(), authorizationCodeGrant, client, appConfiguration, attributeService);
 
+            executionContext.setDeviceSecret(deviceSecret);
             executionContext.setIncludeIdTokenClaims(includeIdTokenClaims);
             executionContext.setPreProcessing(JwrService.wrapWithSidFunction(authorizationCodePreProcessing, sessionIdObj != null ? sessionIdObj.getOutsideSid() : null));
             executionContext.setPostProcessor(externalUpdateTokenService.buildModifyIdTokenProcessor(context));
@@ -429,7 +431,9 @@ private Response processAuthorizationCode(String code, String scope, String code
         JSONObject jsonObj = new JSONObject();
         try {
             fillJsonObject(jsonObj, accToken, accToken.getTokenType(), accToken.getExpiresIn(), reToken, scope, idToken);
-            tokenExchangeService.putNewDeviceSecret(jsonObj, authorizationCodeGrant.getSessionDn(), client, scope);
+            if (StringUtils.isNotBlank(deviceSecret)) {
+                jsonObj.put("device_token", deviceSecret);
+            }
         } catch (JSONException e) {
             log.error(e.getMessage(), e);
         }
diff --git a/jans-auth-server/server/src/test/java/io/jans/as/server/token/ws/rs/TokenExchangeServiceTest.java b/jans-auth-server/server/src/test/java/io/jans/as/server/token/ws/rs/TokenExchangeServiceTest.java
index a8ae3158c39..5abd94ec9b9 100644
--- a/jans-auth-server/server/src/test/java/io/jans/as/server/token/ws/rs/TokenExchangeServiceTest.java
+++ b/jans-auth-server/server/src/test/java/io/jans/as/server/token/ws/rs/TokenExchangeServiceTest.java
@@ -9,7 +9,6 @@
 import io.jans.as.server.audit.ApplicationAuditLogger;
 import io.jans.as.server.service.SessionIdService;
 import org.apache.commons.lang.StringUtils;
-import org.json.JSONObject;
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
 import org.mockito.testng.MockitoTestNGListener;
@@ -17,6 +16,8 @@
 import org.testng.annotations.Listeners;
 import org.testng.annotations.Test;
 
+import static io.jans.as.server.util.TestUtil.assertEmpty;
+import static io.jans.as.server.util.TestUtil.assertNotEmpty;
 import static org.mockito.Mockito.when;
 import static org.testng.Assert.*;
 
@@ -48,56 +49,53 @@ public class TokenExchangeServiceTest {
     private TokenExchangeService tokenExchangeService;
 
     @Test
-    public void putNewDeviceSecret_whenScopeIsNull_shouldNotGenerateDeviceSecretAndShouldNotThrowNPE() {
+    public void createNewDeviceSecret_whenScopeIsNull_shouldNotGenerateDeviceSecretAndShouldNotThrowNPE() {
         SessionId sessionId = new SessionId();
         Client client = new Client();
-        client.setGrantTypes(new GrantType[] {GrantType.AUTHORIZATION_CODE, GrantType.TOKEN_EXCHANGE});
+        client.setGrantTypes(new GrantType[]{GrantType.AUTHORIZATION_CODE, GrantType.TOKEN_EXCHANGE});
 
-        final JSONObject jsonObj = new JSONObject();
-        tokenExchangeService.putNewDeviceSecret(jsonObj, "sessionDn", client, null);
+        String deviceSecret = tokenExchangeService.createNewDeviceSecret("sessionDn", client, null);
 
+        assertEmpty(deviceSecret);
         assertTrue(sessionId.getDeviceSecrets().isEmpty());
-        assertFalse(jsonObj.has("device_token"));
     }
 
     @Test
-    public void putNewDeviceSecret_whenScopeDeviceSSOIsNotPresent_shouldNotGenerateDeviceSecret() {
+    public void createNewDeviceSecret_whenScopeDeviceSSOIsNotPresent_shouldNotGenerateDeviceSecret() {
         SessionId sessionId = new SessionId();
         Client client = new Client();
-        client.setGrantTypes(new GrantType[] {GrantType.AUTHORIZATION_CODE, GrantType.TOKEN_EXCHANGE});
+        client.setGrantTypes(new GrantType[]{GrantType.AUTHORIZATION_CODE, GrantType.TOKEN_EXCHANGE});
 
-        final JSONObject jsonObj = new JSONObject();
-        tokenExchangeService.putNewDeviceSecret(jsonObj, "sessionDn", client, "openid");
+        final String deviceSecret = tokenExchangeService.createNewDeviceSecret("sessionDn", client, "openid");
 
+        assertEmpty(deviceSecret);
         assertTrue(sessionId.getDeviceSecrets().isEmpty());
-        assertFalse(jsonObj.has("device_token"));
     }
 
     @Test
-    public void putNewDeviceSecret_whenTokenExchangeGrantIsNotPresent_shouldNotGenerateDeviceSecret() {
+    public void createNewDeviceSecret_whenTokenExchangeGrantIsNotPresent_shouldNotGenerateDeviceSecret() {
         SessionId sessionId = new SessionId();
         Client client = new Client();
-        client.setGrantTypes(new GrantType[] {GrantType.AUTHORIZATION_CODE});
+        client.setGrantTypes(new GrantType[]{GrantType.AUTHORIZATION_CODE});
 
-        final JSONObject jsonObj = new JSONObject();
-        tokenExchangeService.putNewDeviceSecret(jsonObj, "sessionDn", client, "openid device_sso");
+        final String deviceSecret = tokenExchangeService.createNewDeviceSecret("sessionDn", client, "openid device_sso");
 
+        assertEmpty(deviceSecret);
         assertTrue(sessionId.getDeviceSecrets().isEmpty());
-        assertFalse(jsonObj.has("device_token"));
     }
 
     @Test
-    public void putNewDeviceSecret_whenAllConditionsMet_shouldGenerateDeviceSecret() {
+    public void createNewDeviceSecret_whenAllConditionsMet_shouldGenerateDeviceSecret() {
         SessionId sessionId = new SessionId();
         Client client = new Client();
-        client.setGrantTypes(new GrantType[] {GrantType.AUTHORIZATION_CODE, GrantType.TOKEN_EXCHANGE});
+        client.setGrantTypes(new GrantType[]{GrantType.AUTHORIZATION_CODE, GrantType.TOKEN_EXCHANGE});
 
         when(sessionIdService.getSessionByDn("sessionDn")).thenReturn(sessionId);
 
-        final JSONObject jsonObj = new JSONObject();
-        tokenExchangeService.putNewDeviceSecret(jsonObj, "sessionDn", client, "openid device_sso");
+        final String deviceSecret = tokenExchangeService.createNewDeviceSecret("sessionDn", client, "openid device_sso");
 
-        assertTrue(sessionId.getDeviceSecrets().contains(jsonObj.getString("device_token")));
+        assertNotEmpty(deviceSecret);
+        assertTrue(sessionId.getDeviceSecrets().contains(deviceSecret));
     }
 
     @Test