From b7730d5e90419ce9fff7ec28acd13a5e5ae0df2e Mon Sep 17 00:00:00 2001
From: YuriyZ <yzabrovarniy@gmail.com>
Date: Fri, 28 Oct 2022 07:39:26 +0300
Subject: [PATCH] feat(jans-auth-server): draft implementation of native sso

Native SSO

#2518
---
 .../as/common/model/session/SessionId.java    |   1 +
 .../model/configuration/AppConfiguration.java |  12 ++
 .../authorize/ws/rs/AuthzRequestService.java  |   4 +
 .../as/server/model/token/IdTokenFactory.java |  44 +++----
 .../token/ws/rs/TokenExchangeService.java     |  76 ++++++++++--
 .../token/ws/rs/TokenRestWebServiceImpl.java  |  17 ++-
 .../ws/rs/TokenRestWebServiceValidator.java   |  24 ++--
 .../ws/rs/AuthzRequestServiceTest.java        |  25 ++++
 .../token/ws/rs/TokenExchangeServiceTest.java | 116 ++++++++++++++++++
 .../server/src/test/resources/testng.xml      |   1 +
 10 files changed, 274 insertions(+), 46 deletions(-)
 create mode 100644 jans-auth-server/server/src/test/java/io/jans/as/server/token/ws/rs/TokenExchangeServiceTest.java

diff --git a/jans-auth-server/common/src/main/java/io/jans/as/common/model/session/SessionId.java b/jans-auth-server/common/src/main/java/io/jans/as/common/model/session/SessionId.java
index 07e50cee392..090c7d8ee9c 100644
--- a/jans-auth-server/common/src/main/java/io/jans/as/common/model/session/SessionId.java
+++ b/jans-auth-server/common/src/main/java/io/jans/as/common/model/session/SessionId.java
@@ -95,6 +95,7 @@ public class SessionId implements Deletable, Serializable {
     @Expiration
     private int ttl;
 
+    @NotNull
     public List<String> getDeviceSecrets() {
         if (deviceSecrets == null) deviceSecrets = new ArrayList<>();
         return deviceSecrets;
diff --git a/jans-auth-server/model/src/main/java/io/jans/as/model/configuration/AppConfiguration.java b/jans-auth-server/model/src/main/java/io/jans/as/model/configuration/AppConfiguration.java
index 7cbe69efe8d..c519bd7b64f 100644
--- a/jans-auth-server/model/src/main/java/io/jans/as/model/configuration/AppConfiguration.java
+++ b/jans-auth-server/model/src/main/java/io/jans/as/model/configuration/AppConfiguration.java
@@ -249,8 +249,12 @@ public class AppConfiguration implements Configuration {
     private Set<AuthorizationRequestCustomParameter> authorizationRequestCustomAllowedParameters;
     private Boolean openidScopeBackwardCompatibility = false;
     private Boolean disableU2fEndpoint = false;
+
+    // Token Exchange
     private Boolean rotateDeviceSecret = false;
+    private Boolean returnDeviceSecretFromAuthzEndpoint = false;
 
+    // DCR
     private Boolean dcrSignatureValidationEnabled = false;
     private String dcrSignatureValidationSharedSecret;
     private String dcrSignatureValidationSoftwareStatementJwksURIClaim;
@@ -346,6 +350,14 @@ public void setAllowAllValueForRevokeEndpoint(Boolean allowAllValueForRevokeEndp
         this.allowAllValueForRevokeEndpoint = allowAllValueForRevokeEndpoint;
     }
 
+    public Boolean getReturnDeviceSecretFromAuthzEndpoint() {
+        return returnDeviceSecretFromAuthzEndpoint;
+    }
+
+    public void setReturnDeviceSecretFromAuthzEndpoint(Boolean returnDeviceSecretFromAuthzEndpoint) {
+        this.returnDeviceSecretFromAuthzEndpoint = returnDeviceSecretFromAuthzEndpoint;
+    }
+
     public Boolean getRotateDeviceSecret() {
         if (rotateDeviceSecret == null) rotateDeviceSecret = false;
         return rotateDeviceSecret;
diff --git a/jans-auth-server/server/src/main/java/io/jans/as/server/authorize/ws/rs/AuthzRequestService.java b/jans-auth-server/server/src/main/java/io/jans/as/server/authorize/ws/rs/AuthzRequestService.java
index f6304b7e3ba..2a6b564f16a 100644
--- a/jans-auth-server/server/src/main/java/io/jans/as/server/authorize/ws/rs/AuthzRequestService.java
+++ b/jans-auth-server/server/src/main/java/io/jans/as/server/authorize/ws/rs/AuthzRequestService.java
@@ -49,6 +49,7 @@
 import jakarta.ws.rs.core.MediaType;
 import jakarta.ws.rs.core.Response;
 import org.apache.commons.lang.ArrayUtils;
+import org.apache.commons.lang.BooleanUtils;
 import org.apache.commons.lang.StringUtils;
 import org.jetbrains.annotations.Nullable;
 import org.json.JSONObject;
@@ -108,6 +109,9 @@ public class AuthzRequestService {
     private RedirectionUriService redirectionUriService;
 
     public void addDeviceSecretToSession(AuthzRequest authzRequest, SessionId sessionId) {
+        if (BooleanUtils.isFalse(appConfiguration.getReturnDeviceSecretFromAuthzEndpoint())) {
+            return;
+        }
         if (!Arrays.asList(authzRequest.getScope().split(" ")).contains(ScopeConstants.DEVICE_SSO)) {
             return;
         }
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 abb95e6e561..4fe604237dc 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
@@ -10,7 +10,9 @@
 import io.jans.as.common.claims.Audience;
 import io.jans.as.common.model.common.User;
 import io.jans.as.common.model.registration.Client;
+import io.jans.as.common.model.session.SessionId;
 import io.jans.as.common.service.AttributeService;
+import io.jans.as.model.authorize.CodeVerifier;
 import io.jans.as.model.configuration.AppConfiguration;
 import io.jans.as.model.exception.InvalidClaimException;
 import io.jans.as.model.jwt.JwtClaimName;
@@ -20,15 +22,7 @@
 import io.jans.as.persistence.model.Scope;
 import io.jans.as.server.model.authorize.Claim;
 import io.jans.as.server.model.authorize.JwtAuthorizationRequest;
-import io.jans.as.server.model.common.AbstractToken;
-import io.jans.as.server.model.common.AccessToken;
-import io.jans.as.server.model.common.AuthorizationCode;
-import io.jans.as.server.model.common.CIBAGrant;
-import io.jans.as.server.model.common.ExecutionContext;
-import io.jans.as.server.model.common.IAuthorizationGrant;
-import io.jans.as.server.model.common.RefreshToken;
-import io.jans.as.common.model.session.SessionId;
-import io.jans.as.server.model.common.UnmodifiableAuthorizationGrant;
+import io.jans.as.server.model.common.*;
 import io.jans.as.server.service.ScopeService;
 import io.jans.as.server.service.SessionIdService;
 import io.jans.as.server.service.external.ExternalAuthenticationService;
@@ -39,26 +33,19 @@
 import io.jans.model.GluuAttribute;
 import io.jans.model.custom.script.conf.CustomScriptConfiguration;
 import io.jans.model.custom.script.type.auth.PersonAuthenticationType;
+import jakarta.ejb.Stateless;
+import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.inject.Inject;
+import jakarta.inject.Named;
 import org.apache.commons.lang.StringUtils;
 import org.apache.logging.log4j.util.Strings;
 import org.json.JSONObject;
 import org.slf4j.Logger;
 
-import jakarta.ejb.Stateless;
-import jakarta.enterprise.context.ApplicationScoped;
-import jakarta.inject.Inject;
-import jakarta.inject.Named;
-import java.util.ArrayList;
-import java.util.Calendar;
-import java.util.Collection;
-import java.util.Date;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.UUID;
+import java.util.*;
 
 import static io.jans.as.model.common.ScopeType.DYNAMIC;
+import static io.jans.as.server.token.ws.rs.TokenExchangeService.DEVICE_SECRET;
 
 /**
  * JSON Web Token (JWT) is a compact token format intended for space constrained
@@ -165,6 +152,8 @@ private void fillClaims(JsonWebResponse jwr,
             jwr.setClaim("sid", session.getOutsideSid());
         }
 
+        addTokenExchangeClaims(jwr, executionContext, session);
+
         if (authorizationGrant.getAcrValues() != null) {
             jwr.setClaim(JwtClaimName.AUTHENTICATION_CONTEXT_CLASS_REFERENCE, authorizationGrant.getAcrValues());
             setAmrClaim(jwr, authorizationGrant.getAcrValues());
@@ -264,6 +253,17 @@ private void fillClaims(JsonWebResponse jwr,
         }
     }
 
+    private void addTokenExchangeClaims(JsonWebResponse jwr, ExecutionContext executionContext, SessionId sessionId) {
+        if (sessionId == null) { // unable to find session
+            return;
+        }
+
+        String deviceSecret = executionContext.getHttpRequest().getParameter(DEVICE_SECRET);
+        if (StringUtils.isNotBlank(deviceSecret) && sessionId.getDeviceSecrets().contains(deviceSecret)) {
+            jwr.setClaim("ds_hash", CodeVerifier.s256(deviceSecret));
+        }
+    }
+
     /**
      * Filters some claims from id_token based on if access_token is issued or not.
      * openid-connect-core-1_0.html Section 5.4
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 1ae725a7923..23e9964e56b 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
@@ -4,6 +4,8 @@
 import io.jans.as.common.model.registration.Client;
 import io.jans.as.common.model.session.SessionId;
 import io.jans.as.common.service.AttributeService;
+import io.jans.as.model.common.GrantType;
+import io.jans.as.model.common.ScopeConstants;
 import io.jans.as.model.configuration.AppConfiguration;
 import io.jans.as.model.token.JsonWebResponse;
 import io.jans.as.server.model.audit.OAuth2AuditLog;
@@ -16,6 +18,7 @@
 import jakarta.inject.Inject;
 import jakarta.inject.Named;
 import jakarta.servlet.http.HttpServletRequest;
+import org.apache.commons.lang.ArrayUtils;
 import org.apache.commons.lang.StringUtils;
 import org.apache.commons.lang3.BooleanUtils;
 import org.json.JSONException;
@@ -37,6 +40,8 @@
 @Named
 public class TokenExchangeService {
 
+    public static final String DEVICE_SECRET = "device_secret";
+
     @Inject
     private Logger log;
 
@@ -61,15 +66,39 @@ public class TokenExchangeService {
     @Inject
     private AttributeService attributeService;
 
-    public String rotateDeviceSecret(SessionId sessionId, String actorToken) {
-        if (BooleanUtils.isFalse(appConfiguration.getRotateDeviceSecret())) {
+    public void rotateDeviceSecretOnRefreshToken(HttpServletRequest httpRequest, AuthorizationGrant refreshGrant, String scope) {
+        if (!scope.contains(ScopeConstants.DEVICE_SSO)) {
+            return;
+        }
+        if (StringUtils.isBlank(refreshGrant.getSessionDn())) {
+            return;
+        }
+        final SessionId sessionId = sessionIdService.getSessionByDn(refreshGrant.getSessionDn());
+        if (sessionId == null) {
+            return;
+        }
+
+        final String deviceSecret = httpRequest.getParameter(DEVICE_SECRET);
+
+        // spec: rotate only if device_secret is not specified
+        if (StringUtils.isBlank(deviceSecret)) {
+            rotateDeviceSecret(sessionId, deviceSecret, true);
+        }
+    }
+
+    public String rotateDeviceSecret(SessionId sessionId, String deviceSecret) {
+        return rotateDeviceSecret(sessionId, deviceSecret, false);
+    }
+
+    public String rotateDeviceSecret(SessionId sessionId, String deviceSecret, boolean forceRotation) {
+        if (BooleanUtils.isFalse(appConfiguration.getRotateDeviceSecret()) && !forceRotation) {
             return null;
         }
 
         String newDeviceSecret = HandleTokenFactory.generateDeviceSecret();
 
         final List<String> deviceSecrets = sessionId.getDeviceSecrets();
-        deviceSecrets.remove(actorToken);
+        deviceSecrets.remove(deviceSecret);
         deviceSecrets.add(newDeviceSecret);
 
         sessionIdService.updateSessionId(sessionId, false);
@@ -85,19 +114,19 @@ public JSONObject processTokenExchange(String scope, Function<JsonWebResponse, V
         final String audience = httpRequest.getParameter("audience");
         final String subjectToken = httpRequest.getParameter("subject_token");
         final String subjectTokenType = httpRequest.getParameter("subject_token_type");
-        final String actorToken = httpRequest.getParameter("actor_token");
+        final String deviceSecret = httpRequest.getParameter("actor_token");
         final String actorTokenType = httpRequest.getParameter("actor_token_type");
 
         tokenRestWebServiceValidator.validateAudience(audience, auditLog);
         tokenRestWebServiceValidator.validateSubjectTokenType(subjectTokenType, auditLog);
         tokenRestWebServiceValidator.validateActorTokenType(actorTokenType, auditLog);
-        tokenRestWebServiceValidator.validateActorToken(actorToken, auditLog);
+        tokenRestWebServiceValidator.validateActorToken(deviceSecret, auditLog);
 
-        final SessionId sessionId = sessionIdService.getSessionByDeviceSecret(actorToken);
-        tokenRestWebServiceValidator.validateSessionForTokenExchange(sessionId, actorToken, auditLog);
+        final SessionId sessionId = sessionIdService.getSessionByDeviceSecret(deviceSecret);
+        tokenRestWebServiceValidator.validateSessionForTokenExchange(sessionId, deviceSecret, auditLog);
         checkNotNull(sessionId); // it checked by validator, added it only to relax static bugs checker
 
-        tokenRestWebServiceValidator.validateSubjectToken(subjectToken, sessionId, auditLog);
+        tokenRestWebServiceValidator.validateSubjectToken(deviceSecret, subjectToken, sessionId, auditLog);
 
         TokenExchangeGrant tokenExchangeGrant = authorizationGrantList.createTokenExchangeGrant(new User(), client);
         tokenExchangeGrant.setSessionDn(sessionId.getDn());
@@ -127,14 +156,14 @@ public JSONObject processTokenExchange(String scope, Function<JsonWebResponse, V
 
         executionContext.getAuditLog().updateOAuth2AuditLog(tokenExchangeGrant, true);
 
-        String rotatedDeviceSecret = rotateDeviceSecret(sessionId, actorToken);
+        String rotatedDeviceSecret = rotateDeviceSecret(sessionId, deviceSecret);
 
         JSONObject jsonObj = new JSONObject();
         try {
             TokenRestWebServiceImpl.fillJsonObject(jsonObj, accessToken, accessToken.getTokenType(), accessToken.getExpiresIn(), reToken, scope, idToken);
             jsonObj.put("issued_token_type", TOKEN_TYPE_ACCESS_TOKEN);
             if (StringUtils.isNotBlank(rotatedDeviceSecret)) {
-                jsonObj.put("device_secret", rotatedDeviceSecret);
+                jsonObj.put(DEVICE_SECRET, rotatedDeviceSecret);
             }
         } catch (JSONException e) {
             log.error(e.getMessage(), e);
@@ -142,4 +171,31 @@ public JSONObject processTokenExchange(String scope, Function<JsonWebResponse, V
 
         return jsonObj;
     }
+
+    public void putNewDeviceSecret(JSONObject jsonObj, String sessionDn, Client client, String scope) {
+        if (!scope.contains(ScopeConstants.DEVICE_SSO)) {
+            return;
+        }
+        if (!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;
+        }
+
+        try {
+            final SessionId sessionId = sessionIdService.getSessionByDn(sessionDn);
+            if (sessionId == null) {
+                log.debug("Unable to find session by dn: {}", sessionDn);
+                return;
+            }
+
+            String newDeviceSecret = HandleTokenFactory.generateDeviceSecret();
+            sessionId.getDeviceSecrets().add(newDeviceSecret);
+
+            sessionIdService.updateSessionId(sessionId, false);
+
+            jsonObj.put("device_token", newDeviceSecret);
+        } catch (Exception e) {
+            log.error("Failed to generate device_secret", e);
+        }
+    }
 }
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 9bfa3f17e6c..f94da0bbd51 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
@@ -194,9 +194,7 @@ public Response requestAccessToken(String grantType, String code,
 
             if (gt == GrantType.AUTHORIZATION_CODE) {
                 return processAuthorizationCode(code, scope, codeVerifier, sessionIdObj, executionContext);
-            }
-
-            if (gt == GrantType.REFRESH_TOKEN) {
+            } else if (gt == GrantType.REFRESH_TOKEN) {
                 return processRefreshTokenGrant(scope, refreshToken, idTokenPreProcessing, executionContext);
             } else if (gt == GrantType.CLIENT_CREDENTIALS) {
                 return processClientGredentials(scope, request, auditLog, client, idTokenPreProcessing, executionContext);
@@ -368,6 +366,8 @@ private Response processRefreshTokenGrant(String scope, String refreshToken, Fun
             grantService.removeByCode(refreshToken); // remove refresh token after access token and id_token is created.
         }
 
+        tokenExchangeService.rotateDeviceSecretOnRefreshToken(executionContext.getHttpRequest(), authorizationGrant, scope);
+
         auditLog.updateOAuth2AuditLog(authorizationGrant, true);
 
         return response(Response.ok().entity(getJSonResponse(accToken,
@@ -426,8 +426,15 @@ private Response processAuthorizationCode(String code, String scope, String code
 
         grantService.removeAuthorizationCode(authorizationCodeGrant.getAuthorizationCode().getCode());
 
-        final String entity = getJSonResponse(accToken, accToken.getTokenType(), accToken.getExpiresIn(), reToken, scope, idToken);
-        return response(Response.ok().entity(entity), executionContext.getAuditLog());
+        JSONObject jsonObj = new JSONObject();
+        try {
+            fillJsonObject(jsonObj, accToken, accToken.getTokenType(), accToken.getExpiresIn(), reToken, scope, idToken);
+            tokenExchangeService.putNewDeviceSecret(jsonObj, authorizationCodeGrant.getSessionDn(), client, scope);
+        } catch (JSONException e) {
+            log.error(e.getMessage(), e);
+        }
+
+        return response(Response.ok().entity(jsonObj.toString()), executionContext.getAuditLog());
     }
 
     @Nullable
diff --git a/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenRestWebServiceValidator.java b/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenRestWebServiceValidator.java
index 5513814fed9..793ca976ad9 100644
--- a/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenRestWebServiceValidator.java
+++ b/jans-auth-server/server/src/main/java/io/jans/as/server/token/ws/rs/TokenRestWebServiceValidator.java
@@ -4,6 +4,7 @@
 import io.jans.as.common.model.registration.Client;
 import io.jans.as.common.model.session.SessionId;
 import io.jans.as.common.model.session.SessionIdState;
+import io.jans.as.model.authorize.CodeVerifier;
 import io.jans.as.model.common.GrantType;
 import io.jans.as.model.configuration.AppConfiguration;
 import io.jans.as.model.crypto.AbstractCryptoProvider;
@@ -218,11 +219,10 @@ public void validateSessionForTokenExchange(SessionId session, String actorToken
         }
     }
 
-    public Jwt validateSubjectToken(String subjectToken, SessionId sidSession, OAuth2AuditLog auditLog) {
+    public void validateSubjectToken(String subjectToken, String deviceSecret, SessionId sidSession, OAuth2AuditLog auditLog) {
         try {
             final Jwt jwt = Jwt.parse(subjectToken);
-            validateSubjectTokenSignature(sidSession, jwt, auditLog);
-            return jwt;
+            validateSubjectTokenSignature(deviceSecret, sidSession, jwt, auditLog);
         } catch (InvalidJwtException e) {
             log.error("Unable to parse subject_token as JWT.", e);
             throw new WebApplicationException(response(error(400, TokenErrorResponseType.INVALID_REQUEST, "Unable to parse subject_token as JWT."), auditLog));
@@ -234,7 +234,7 @@ public Jwt validateSubjectToken(String subjectToken, SessionId sidSession, OAuth
         }
     }
 
-    private void validateSubjectTokenSignature(SessionId sidSession, Jwt jwt, OAuth2AuditLog auditLog) throws InvalidJwtException, CryptoProviderException {
+    private void validateSubjectTokenSignature(String deviceSecret, SessionId sidSession, Jwt jwt, OAuth2AuditLog auditLog) throws InvalidJwtException, CryptoProviderException {
         // verify jwt signature if we can't find it in db
         if (!cryptoProvider.verifySignature(jwt.getSigningInput(), jwt.getEncodedSignature(), jwt.getHeader().getKeyId(),
                 null, null, jwt.getHeader().getSignatureAlgorithm())) {
@@ -242,12 +242,18 @@ private void validateSubjectTokenSignature(SessionId sidSession, Jwt jwt, OAuth2
             throw new WebApplicationException(response(error(400, TokenErrorResponseType.INVALID_REQUEST, "subject_token signature verification failed."), auditLog));
         }
 
-        final String sidClaim = jwt.getClaims().getClaimAsString("sid");
-        if (sidSession != null && org.apache.commons.lang.StringUtils.equals(sidSession.getOutsideSid(), sidClaim)) {
-            return;
+        final String sid = jwt.getClaims().getClaimAsString("sid");
+        if (StringUtils.isBlank(sid) || !StringUtils.equals(sidSession.getOutsideSid(), sid)) {
+            log.error("sid claim from subject_token does not match to session sid.");
+            throw new WebApplicationException(response(error(400, TokenErrorResponseType.INVALID_REQUEST, "sid claim from subject_token does not match to session sid."), auditLog));
+        }
+
+        final String dsHash = jwt.getClaims().getClaimAsString("ds_hash");
+        if (StringUtils.isBlank(dsHash) || !dsHash.equals(CodeVerifier.s256(deviceSecret))) {
+            final String msg = "ds_hash claim from subject_token does not match to hash of device_secret";
+            log.error(msg);
+            throw new WebApplicationException(response(error(400, TokenErrorResponseType.INVALID_REQUEST, msg), auditLog));
         }
-        log.error("sid claim from subject_token does not match to session sid.");
-        throw new WebApplicationException(response(error(400, TokenErrorResponseType.INVALID_REQUEST, "sid claim from subject_token does not match to session sid."), auditLog));
     }
 
     public void validateAudience(String audience, OAuth2AuditLog auditLog) {
diff --git a/jans-auth-server/server/src/test/java/io/jans/as/server/authorize/ws/rs/AuthzRequestServiceTest.java b/jans-auth-server/server/src/test/java/io/jans/as/server/authorize/ws/rs/AuthzRequestServiceTest.java
index c05f46dfd38..d5b10badb44 100644
--- a/jans-auth-server/server/src/test/java/io/jans/as/server/authorize/ws/rs/AuthzRequestServiceTest.java
+++ b/jans-auth-server/server/src/test/java/io/jans/as/server/authorize/ws/rs/AuthzRequestServiceTest.java
@@ -24,6 +24,7 @@
 import org.testng.annotations.Test;
 
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertTrue;
 
@@ -69,8 +70,28 @@ public class AuthzRequestServiceTest {
     @Mock
     private RedirectionUriService redirectionUriService;
 
+    @Test
+    public void addDeviceSecretToSession_withoutUnabledConfiguration_shouldNotGenerateDeviceSecret() {
+        when(appConfiguration.getReturnDeviceSecretFromAuthzEndpoint()).thenReturn(false);
+
+        Client client = new Client();
+        client.setGrantTypes(new GrantType[] { GrantType.AUTHORIZATION_CODE, GrantType.TOKEN_EXCHANGE});
+
+        AuthzRequest authzRequest = new AuthzRequest();
+        authzRequest.setScope("openid device_sso");
+        authzRequest.setRedirectUriResponse(new RedirectUriResponse(mock(RedirectUri.class), "", mock(HttpServletRequest.class), mock(ErrorResponseFactory.class)));
+        authzRequest.setClient(client);
+
+        SessionId sessionId = new SessionId();
+
+        authzRequestService.addDeviceSecretToSession(authzRequest, sessionId);
+        assertTrue(sessionId.getDeviceSecrets().isEmpty());
+    }
+
     @Test
     public void addDeviceSecretToSession_withoutDeviceSsoScope_shouldNotGenerateDeviceSecret() {
+        when(appConfiguration.getReturnDeviceSecretFromAuthzEndpoint()).thenReturn(true);
+
         Client client = new Client();
         client.setGrantTypes(new GrantType[] { GrantType.AUTHORIZATION_CODE, GrantType.TOKEN_EXCHANGE});
 
@@ -87,6 +108,8 @@ public void addDeviceSecretToSession_withoutDeviceSsoScope_shouldNotGenerateDevi
 
     @Test
     public void addDeviceSecretToSession_withDeviceSsoScope_shouldGenerateDeviceSecret() {
+        when(appConfiguration.getReturnDeviceSecretFromAuthzEndpoint()).thenReturn(true);
+
         Client client = new Client();
         client.setGrantTypes(new GrantType[] { GrantType.AUTHORIZATION_CODE, GrantType.TOKEN_EXCHANGE});
 
@@ -105,6 +128,8 @@ public void addDeviceSecretToSession_withDeviceSsoScope_shouldGenerateDeviceSecr
 
     @Test
     public void addDeviceSecretToSession_withClientWithoutTokenExchangeGrantType_shouldNotGenerateDeviceSecret() {
+        when(appConfiguration.getReturnDeviceSecretFromAuthzEndpoint()).thenReturn(true);
+
         Client client = new Client();
         client.setGrantTypes(new GrantType[] { GrantType.AUTHORIZATION_CODE});
 
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
new file mode 100644
index 00000000000..6303467cf01
--- /dev/null
+++ b/jans-auth-server/server/src/test/java/io/jans/as/server/token/ws/rs/TokenExchangeServiceTest.java
@@ -0,0 +1,116 @@
+package io.jans.as.server.token.ws.rs;
+
+import io.jans.as.common.model.registration.Client;
+import io.jans.as.common.model.session.SessionId;
+import io.jans.as.model.common.GrantType;
+import io.jans.as.model.configuration.AppConfiguration;
+import io.jans.as.model.crypto.AbstractCryptoProvider;
+import io.jans.as.model.error.ErrorResponseFactory;
+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;
+import org.slf4j.Logger;
+import org.testng.annotations.Listeners;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.*;
+
+/**
+ * @author Yuriy Z
+ */
+@Listeners(MockitoTestNGListener.class)
+public class TokenExchangeServiceTest {
+
+    @Mock
+    private Logger log;
+
+    @Mock
+    private AppConfiguration appConfiguration;
+
+    @Mock
+    private ApplicationAuditLogger applicationAuditLogger;
+
+    @Mock
+    private ErrorResponseFactory errorResponseFactory;
+
+    @Mock
+    private AbstractCryptoProvider cryptoProvider;
+
+    @Mock
+    private SessionIdService sessionIdService;
+
+    @InjectMocks
+    private TokenExchangeService tokenExchangeService;
+
+    @Test
+    public void putNewDeviceSecret_whenScopeDeviceSSOIsNotPresent_shouldNotGenerateDeviceSecret() {
+        SessionId sessionId = new SessionId();
+        Client client = new Client();
+        client.setGrantTypes(new GrantType[] {GrantType.AUTHORIZATION_CODE, GrantType.TOKEN_EXCHANGE});
+
+        final JSONObject jsonObj = new JSONObject();
+        tokenExchangeService.putNewDeviceSecret(jsonObj, "sessionDn", client, "openid");
+
+        assertTrue(sessionId.getDeviceSecrets().isEmpty());
+        assertFalse(jsonObj.has("device_token"));
+    }
+
+
+    @Test
+    public void putNewDeviceSecret_whenTokenExchangeGrantIsNotPresent_shouldNotGenerateDeviceSecret() {
+        SessionId sessionId = new SessionId();
+        Client client = new Client();
+        client.setGrantTypes(new GrantType[] {GrantType.AUTHORIZATION_CODE});
+
+        final JSONObject jsonObj = new JSONObject();
+        tokenExchangeService.putNewDeviceSecret(jsonObj, "sessionDn", client, "openid device_sso");
+
+        assertTrue(sessionId.getDeviceSecrets().isEmpty());
+        assertFalse(jsonObj.has("device_token"));
+    }
+
+    @Test
+    public void putNewDeviceSecret_whenAllConditionsMet_shouldGenerateDeviceSecret() {
+        SessionId sessionId = new SessionId();
+        Client client = new Client();
+        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");
+
+        assertTrue(sessionId.getDeviceSecrets().contains(jsonObj.getString("device_token")));
+    }
+
+    @Test
+    public void rotateDeviceSecret_whenRotationIsDisabled_shouldNotRotateDeviceSecret() {
+        when(appConfiguration.getRotateDeviceSecret()).thenReturn(false);
+
+        SessionId sessionId = new SessionId();
+        sessionId.getDeviceSecrets().add("a");
+
+        final String newDeviceSecret = tokenExchangeService.rotateDeviceSecret(sessionId, "a");
+        assertNull(newDeviceSecret);
+        assertTrue(sessionId.getDeviceSecrets().contains("a"));
+        assertEquals(sessionId.getDeviceSecrets().size(), 1);
+    }
+
+    @Test
+    public void rotateDeviceSecret_whenRotationIsEnabled_shouldRotateDeviceSecret() {
+        when(appConfiguration.getRotateDeviceSecret()).thenReturn(true);
+
+        SessionId sessionId = new SessionId();
+        sessionId.getDeviceSecrets().add("a");
+
+        final String newDeviceSecret = tokenExchangeService.rotateDeviceSecret(sessionId, "a");
+        assertTrue(StringUtils.isNotBlank(newDeviceSecret) && !"a".equals(sessionId.getDeviceSecrets().get(0)));
+        assertFalse(sessionId.getDeviceSecrets().contains("a"));
+        assertEquals(sessionId.getDeviceSecrets().size(), 1);
+    }
+}
diff --git a/jans-auth-server/server/src/test/resources/testng.xml b/jans-auth-server/server/src/test/resources/testng.xml
index 1cd1dd9e674..976129a0509 100644
--- a/jans-auth-server/server/src/test/resources/testng.xml
+++ b/jans-auth-server/server/src/test/resources/testng.xml
@@ -19,6 +19,7 @@
             <class name="io.jans.as.server.service.external.ExternalAuthenticationServiceTest" />
             <class name="io.jans.as.server.servlet.OpenIdConfigurationTest" />
 
+            <class name="io.jans.as.server.token.ws.rs.TokenExchangeServiceTest" />
             <class name="io.jans.as.server.token.ws.rs.TokenRestWebServiceValidatorTest" />
             <class name="io.jans.as.server.ws.rs.stat.MonthsTest" />
             <class name="io.jans.as.server.authorize.ws.rs.AuthorizeRestWebServiceValidatorTest" />