fix: fix bug with token re-wrapping

eclipse-tractusx · Mar 7, 2024 · e047550 · e047550
1 parent b92c924
commit e047550
Show file tree

Hide file tree

Showing 8 changed files with 55 additions and 46 deletions.
diff --git a/...in/java/org/eclipse/tractusx/managedidentitywallets/controller/SecureTokenController.java b/...in/java/org/eclipse/tractusx/managedidentitywallets/controller/SecureTokenController.java
@@ -22,7 +22,6 @@
 package org.eclipse.tractusx.managedidentitywallets.controller;
 
 import com.nimbusds.jwt.JWT;
-import com.nimbusds.jwt.JWTClaimsSet;
 import com.nimbusds.jwt.JWTParser;
 import io.swagger.v3.oas.annotations.tags.Tag;
 import jakarta.validation.Valid;
@@ -32,8 +31,6 @@
 import org.apache.commons.lang3.StringUtils;
 import org.eclipse.tractusx.managedidentitywallets.apidocs.SecureTokenControllerApiDoc;
 import org.eclipse.tractusx.managedidentitywallets.constant.StringPool;
-import org.eclipse.tractusx.managedidentitywallets.dao.entity.JtiRecord;
-import org.eclipse.tractusx.managedidentitywallets.dao.repository.JtiRepository;
 import org.eclipse.tractusx.managedidentitywallets.dao.repository.WalletRepository;
 import org.eclipse.tractusx.managedidentitywallets.domain.BusinessPartnerNumber;
 import org.eclipse.tractusx.managedidentitywallets.domain.DID;
@@ -59,10 +56,8 @@
 import org.springframework.web.bind.annotation.RestController;
 
 import java.util.Set;
-import java.util.UUID;
 import java.util.regex.Pattern;
 
-import static org.eclipse.tractusx.managedidentitywallets.utils.TokenParsingUtils.getJtiAccessToken;
 
 @RestController
 @Slf4j
@@ -76,8 +71,6 @@ public class SecureTokenController {
 
     private final WalletRepository walletRepo;
 
-    private final JtiRepository jtiRepository;
-
     @InitBinder
     void initBinder(WebDataBinder webDataBinder) {
         webDataBinder.addValidators(new SecureTokenRequestValidator());
@@ -122,17 +115,11 @@ public ResponseEntity<StsTokenResponse> token(
             throw new InvalidSecureTokenRequestException("The provided data could not be used to create and sign a token.");
         }
 
-        // store jti info in repository
-        JWTClaimsSet jwtClaimsSet = responseJwt.getJWTClaimsSet();
-        String jtiValue = getJtiAccessToken(jwtClaimsSet);
-        JtiRecord jtiRecord = JtiRecord.builder().jti(UUID.fromString(jtiValue)).isUsedStatus(false).build();
-        jtiRepository.save(jtiRecord);
-
         // create the response
         log.debug("Preparing StsTokenResponse.");
         StsTokenResponse response = StsTokenResponse.builder()
                 .token(responseJwt.serialize())
-                .expiresAt(jwtClaimsSet.getExpirationTime().getTime())
+                .expiresAt(responseJwt.getJWTClaimsSet().getExpirationTime().getTime())
                 .build();
         return ResponseEntity.status(HttpStatus.CREATED).body(response);
     }

diff --git a/src/main/java/org/eclipse/tractusx/managedidentitywallets/service/PresentationService.java b/src/main/java/org/eclipse/tractusx/managedidentitywallets/service/PresentationService.java
@@ -357,7 +357,9 @@ private JtiRecord getJtiRecord(JWTClaimsSet jwtClaimsSet) {
         String jtiValue = getStringClaim(jwtClaimsSet, JTI);
         JtiRecord jtiRecord = jtiRepository.getByJti(UUID.fromString(jtiValue));
         if (Objects.isNull(jtiRecord)) {
-            throw new BadDataException("Jti record does not exist");
+            JtiRecord jtiToAdd = JtiRecord.builder().jti(UUID.fromString(jtiValue)).isUsedStatus(false).build();
+            jtiRepository.save(jtiToAdd);
+            return jtiToAdd;
         } else if (jtiRecord.isUsedStatus()) {
             throw new BadDataException("The token was already used");
         } else {

diff --git a/...main/java/org/eclipse/tractusx/managedidentitywallets/service/SecureTokenServiceImpl.java b/...main/java/org/eclipse/tractusx/managedidentitywallets/service/SecureTokenServiceImpl.java
@@ -24,7 +24,9 @@
 import com.nimbusds.jwt.JWT;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
+import org.eclipse.tractusx.managedidentitywallets.dao.entity.JtiRecord;
 import org.eclipse.tractusx.managedidentitywallets.dao.entity.WalletKey;
+import org.eclipse.tractusx.managedidentitywallets.dao.repository.JtiRepository;
 import org.eclipse.tractusx.managedidentitywallets.dao.repository.WalletKeyRepository;
 import org.eclipse.tractusx.managedidentitywallets.dao.repository.WalletRepository;
 import org.eclipse.tractusx.managedidentitywallets.domain.BusinessPartnerNumber;
@@ -36,8 +38,12 @@
 import org.eclipse.tractusx.managedidentitywallets.sts.SecureTokenConfigurationProperties;
 
 import java.time.Instant;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.UUID;
+
+import static org.eclipse.tractusx.managedidentitywallets.utils.TokenParsingUtils.getJtiAccessToken;
 
 @Slf4j
 @RequiredArgsConstructor
@@ -51,6 +57,8 @@ public class SecureTokenServiceImpl implements SecureTokenService {
 
     private final SecureTokenConfigurationProperties properties;
 
+    private final JtiRepository jtiRepository;
+
     @Override
     public JWT issueToken(final DID self, final DID partner, final Set<String> scopes) {
         log.debug("'issueToken' using scopes and DID.");
@@ -59,6 +67,7 @@ public JWT issueToken(final DID self, final DID partner, final Set<String> scope
         // as we're signing two tokens.
         Instant expirationTime = Instant.now().plus(properties.tokenDuration());
         JWT accessToken = this.tokenIssuer.createAccessToken(keyPair, self, partner, expirationTime, scopes);
+        checkAndStoreJti(accessToken);
         return this.tokenIssuer.createIdToken(keyPair, self, partner, expirationTime, accessToken);
     }
 
@@ -67,9 +76,19 @@ public JWT issueToken(DID self, DID partner, JWT accessToken) {
         log.debug("'issueToken' using an access_token and DID.");
         KeyPair keyPair = walletKeyRepository.findFirstByWallet_Did(self.toString()).toDto();
         Instant expirationTime = Instant.now().plus(properties.tokenDuration());
+        checkAndStoreJti(accessToken);
         return this.tokenIssuer.createIdToken(keyPair, self, partner, expirationTime, accessToken);
     }
 
+    private void checkAndStoreJti(JWT accessToken) {
+        String jtiValue = getJtiAccessToken(accessToken);
+        JtiRecord jti = jtiRepository.getByJti(UUID.fromString(jtiValue));
+        if (Objects.isNull(jti)) {
+            JtiRecord jtiRecord = JtiRecord.builder().jti(UUID.fromString(jtiValue)).isUsedStatus(false).build();
+            jtiRepository.save(jtiRecord);
+        }
+    }
+
     @Override
     public JWT issueToken(BusinessPartnerNumber self, BusinessPartnerNumber partner, Set<String> scopes) {
         log.debug("'issueToken' using scopes and BPN.");

diff --git a/src/main/java/org/eclipse/tractusx/managedidentitywallets/sts/SecureTokenBeanConfig.java b/src/main/java/org/eclipse/tractusx/managedidentitywallets/sts/SecureTokenBeanConfig.java
@@ -21,6 +21,7 @@
 
 package org.eclipse.tractusx.managedidentitywallets.sts;
 
+import org.eclipse.tractusx.managedidentitywallets.dao.repository.JtiRepository;
 import org.eclipse.tractusx.managedidentitywallets.dao.repository.WalletKeyRepository;
 import org.eclipse.tractusx.managedidentitywallets.dao.repository.WalletRepository;
 import org.eclipse.tractusx.managedidentitywallets.interfaces.SecureTokenIssuer;
@@ -37,9 +38,10 @@ public SecureTokenService secureTokenService(
             WalletKeyRepository keyRepository,
             WalletRepository walletRepository,
             SecureTokenIssuer issuer,
-            SecureTokenConfigurationProperties properties
+            SecureTokenConfigurationProperties properties,
+            JtiRepository jtiRepository
     ) {
-        return new SecureTokenServiceImpl(keyRepository, walletRepository, issuer, properties);
+        return new SecureTokenServiceImpl(keyRepository, walletRepository, issuer, properties, jtiRepository);
     }
 
 }
diff --git a/src/main/java/org/eclipse/tractusx/managedidentitywallets/utils/TokenParsingUtils.java b/src/main/java/org/eclipse/tractusx/managedidentitywallets/utils/TokenParsingUtils.java
@@ -94,18 +94,17 @@ public static String getScope(JWTClaimsSet jwtClaimsSet) {
         }
     }
 
-    public static String getJtiAccessToken(JWTClaimsSet jwtClaimsSet) {
-        Optional<String> token = getAccessToken(jwtClaimsSet);
-        String accessToken = token.orElseThrow(() -> new BadDataException(ACCESS_TOKEN_ERROR));
-        SignedJWT accessTokenJwt = parseToken(accessToken);
-        JWTClaimsSet accessTokenClaimsSet = getClaimsSet(accessTokenJwt);
-        return getStringClaim(accessTokenClaimsSet, JTI);
+    public static String getJtiAccessToken(JWT accessToken) {
+        try {
+            return getStringClaim(accessToken.getJWTClaimsSet(), JTI);
+        } catch (ParseException e) {
+            throw new BadDataException(PARSING_TOKEN_ERROR, e);
+        }
     }
 
     public static String getNonceAccessToken(JWT accessToken) {
         try {
-            JWTClaimsSet claimsSet = accessToken.getJWTClaimsSet();
-            return claimsSet.getStringClaim(NONCE);
+            return accessToken.getJWTClaimsSet().getStringClaim(NONCE);
         } catch (ParseException e) {
             throw new BadDataException(PARSING_TOKEN_ERROR, e);
         }

diff --git a/src/main/resources/db/changelog/changes/create_jti_table.sql b/src/main/resources/db/changelog/changes/create_jti_table.sql
@@ -30,6 +30,7 @@ CREATE TABLE IF NOT EXISTS public.jti
     created_at    timestamp(6) NOT NULL DEFAULT CURRENT_TIMESTAMP,
     modified_at   timestamp(6) NULL,
     modified_from varchar(255) NULL,
-    CONSTRAINT jti_pkey PRIMARY KEY (id)
+    CONSTRAINT jti_pkey PRIMARY KEY (id),
+    CONSTRAINT uk_jti UNIQUE (jti)
 );
 COMMENT ON TABLE public.jti IS 'This table will store jti field statuses';
diff --git a/...ava/org/eclipse/tractusx/managedidentitywallets/controller/SecureTokenControllerTest.java b/...ava/org/eclipse/tractusx/managedidentitywallets/controller/SecureTokenControllerTest.java
@@ -24,7 +24,6 @@
 import org.eclipse.tractusx.managedidentitywallets.ManagedIdentityWalletsApplication;
 import org.eclipse.tractusx.managedidentitywallets.config.MIWSettings;
 import org.eclipse.tractusx.managedidentitywallets.config.TestContextInitializer;
-import org.eclipse.tractusx.managedidentitywallets.dao.repository.JtiRepository;
 import org.eclipse.tractusx.managedidentitywallets.utils.AuthenticationUtils;
 import org.eclipse.tractusx.managedidentitywallets.utils.TestUtils;
 import org.eclipse.tractusx.ssi.lib.did.web.DidWebFactory;
@@ -55,9 +54,6 @@ class SecureTokenControllerTest {
     @Autowired
     private TestRestTemplate testTemplate;
 
-    @Autowired
-    private JtiRepository jtiRepository;
-
     @Test
     void token() {
         // given

diff --git a/src/test/java/org/eclipse/tractusx/managedidentitywallets/vp/PresentationServiceTest.java b/src/test/java/org/eclipse/tractusx/managedidentitywallets/vp/PresentationServiceTest.java
@@ -21,7 +21,6 @@
 
 package org.eclipse.tractusx.managedidentitywallets.vp;
 
-import com.nimbusds.jose.JOSEException;
 import com.nimbusds.jwt.JWT;
 import com.nimbusds.jwt.JWTClaimsSet;
 import com.nimbusds.jwt.JWTParser;
@@ -127,6 +126,23 @@ void createPresentation200ResponseAsJsonLD() {
     }
 
     @SneakyThrows
+    @Test
+    void createPresentation200ResponseNoJtiRecord() {
+        boolean asJwt = true;
+        String bpn = TestUtils.getRandomBpmNumber();
+        String did = generateWalletAndGetDid(bpn);
+        String jtiValue = generateUuid();
+        String accessToken = generateAccessToken(did, did, did, BPN_CREDENTIAL_READ, jtiValue);
+
+        Map<String, Object> presentation = presentationService.createVpWithRequiredScopes(SignedJWT.parse(accessToken), asJwt);
+        String vpAsJwt = String.valueOf(presentation.get(VERIFIABLE_PRESENTATION));
+        JWT jwt = JWTParser.parse(vpAsJwt);
+
+        Assertions.assertNotNull(presentation);
+        Assertions.assertEquals(did, jwt.getJWTClaimsSet().getSubject());
+        Assertions.assertEquals(did, jwt.getJWTClaimsSet().getIssuer());
+    }
+
     @Test
     void createPresentationIncorrectVcTypeResponse() {
         boolean asJwt = true;
@@ -141,7 +157,6 @@ void createPresentationIncorrectVcTypeResponse() {
                 presentationService.createVpWithRequiredScopes(SignedJWT.parse(accessToken), asJwt));
     }
 
-    @SneakyThrows
     @Test
     void createPresentationIncorrectRightsRequested() {
         boolean asJwt = true;
@@ -154,19 +169,6 @@ void createPresentationIncorrectRightsRequested() {
                 presentationService.createVpWithRequiredScopes(SignedJWT.parse(accessToken), asJwt));
     }
 
-    @SneakyThrows
-    @Test
-    void createPresentationIncorrectNoJtiRecord() {
-        boolean asJwt = false;
-        String bpn = TestUtils.getRandomBpmNumber();
-        String did = generateWalletAndGetDid(bpn);
-        String accessToken = generateAccessToken(did, did, did, BPN_CREDENTIAL_READ, generateUuid());
-
-        BadDataException ex = Assertions.assertThrows(BadDataException.class, () -> presentationService.createVpWithRequiredScopes(SignedJWT.parse(accessToken), asJwt));
-        Assertions.assertEquals("Jti record does not exist", ex.getMessage());
-    }
-
-    @SneakyThrows
     @Test
     void createPresentationIncorrectJtiAlreadyUsed() {
         boolean asJwt = false;
@@ -193,7 +195,8 @@ private JtiRecord buildJti(String value, boolean isUsed) {
         return JtiRecord.builder().jti(UUID.fromString(value)).isUsedStatus(isUsed).build();
     }
 
-    private String generateAccessToken(String issUrl, String sub, String aud, String scope, String jwt) throws JOSEException {
+    @SneakyThrows
+    private String generateAccessToken(String issUrl, String sub, String aud, String scope, String jwt)  {
         JWTClaimsSet innerSet = buildClaimsSet(issUrl, sub, aud, TestConstants.NONCE, scope, EXP_VALID_DATE, IAT_VALID_DATE, jwt);
         return buildJWTToken(JWK_INNER, innerSet);
     }