saving the used cookie secret on the context

src/main/java/com/perimeterx/api/PerimeterX.java

@@ -58,6 +58,7 @@
 import com.perimeterx.utils.logger.IPXLogger;
 import com.perimeterx.utils.StringUtils;
 import com.perimeterx.utils.logger.LoggerFactory;
+import edu.emory.mathcs.backport.java.util.Collections;
 
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponseWrapper;
@@ -68,8 +69,10 @@
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.util.Base64;
+import java.util.List;
 
 import static com.perimeterx.utils.Constants.*;
+import static com.perimeterx.utils.PXCommonUtils.cookieKeysToCheck;
 import static java.util.Objects.isNull;
 
 /**
@@ -303,7 +306,9 @@ public boolean isValidTelemetryRequest(HttpServletRequest request, PXContext con
                 return false;
             }
 
-            for (String key : configuration.getCookieKeys()) {
+            List<String> cookieKeyToCheck = cookieKeysToCheck(context, configuration);
+
+            for (String key : cookieKeyToCheck) {
                 final byte[] hmacBytes = HMACUtils.HMACString(timestamp, key);
                 final String generatedHmac = StringUtils.byteArrayToHexString(hmacBytes).toLowerCase();
 
@@ -320,6 +325,7 @@ public boolean isValidTelemetryRequest(HttpServletRequest request, PXContext con
         return false;
     }
 
+
     /**
      * Set activity handler
      *

src/main/java/com/perimeterx/internals/cookie/AbstractPXCookie.java

@@ -19,7 +19,8 @@
 import java.nio.charset.StandardCharsets;
 import java.security.NoSuchAlgorithmException;
 import java.util.Arrays;
-import java.util.List;
+
+import static com.perimeterx.utils.PXCommonUtils.cookieKeysToCheck;
 
 /**
  * Created by nitzangoldfeder on 13/04/2017.
@@ -41,7 +42,6 @@ public abstract class AbstractPXCookie implements PXCookie {
     protected PXConfiguration pxConfiguration;
     protected String pxCookie;
     protected JsonNode decodedCookie;
-    protected List<String> cookieKeys;
     protected String cookieOrig;
 
     public AbstractPXCookie(PXConfiguration pxConfiguration, CookieData cookieData, PXContext context) {
@@ -52,7 +52,6 @@ public AbstractPXCookie(PXConfiguration pxConfiguration, CookieData cookieData,
         this.pxConfiguration = pxConfiguration;
         this.userAgent = cookieData.isMobileToken() ? "" : cookieData.getUserAgent();
         this.ip = cookieData.getIp();
-        this.cookieKeys = pxConfiguration.getCookieKeys();
         this.cookieVersion = cookieData.getCookieVersion();
     }
 
@@ -84,7 +83,7 @@ public boolean deserialize() throws PXCookieDecryptionException {
 
         JsonNode decodedCookie;
         if (this.pxConfiguration.isEncryptionEnabled()) {
-            decodedCookie = this.decrypt();
+            decodedCookie = this.decrypt(context);
         } else {
             decodedCookie = this.decode();
         }
@@ -97,7 +96,7 @@ public boolean deserialize() throws PXCookieDecryptionException {
         return true;
     }
 
-    private JsonNode decrypt() throws PXCookieDecryptionException {
+    private JsonNode decrypt(PXContext context) throws PXCookieDecryptionException {
         final String[] parts = this.pxCookie.split(":");
         if (parts.length != 3) {
             throw new PXCookieDecryptionException("Part length invalid");
@@ -124,7 +123,7 @@ private JsonNode decrypt() throws PXCookieDecryptionException {
         final int dkLen = KEY_LEN + cipher.getBlockSize();
         PBKDF2Parameters p = new PBKDF2Parameters(HMAC_SHA_256, "UTF-8", salt, iterations);
 
-        for (String cookieKey : this.cookieKeys) {
+        for (String cookieKey : cookieKeysToCheck(this.context, this.pxConfiguration)) {
             try {
                 byte[] dk = new PBKDF2Engine(p).deriveKey(cookieKey, dkLen);
                 byte[] key = Arrays.copyOf(dk, KEY_LEN);
@@ -135,7 +134,9 @@ private JsonNode decrypt() throws PXCookieDecryptionException {
                 final byte[] data = cipher.doFinal(encrypted, 0, encrypted.length);
 
                 String decryptedString = new String(data, StandardCharsets.UTF_8);
-                return mapper.readTree(decryptedString);
+                JsonNode result = mapper.readTree(decryptedString);
+                context.setCookieKeyUsed(cookieKey);
+                return result;
             } catch (Exception ignored) {
             }
         }
@@ -162,7 +163,7 @@ public boolean isExpired() {
     }
 
     public boolean isHmacValid(String hmacStr, String cookieHmac) {
-        boolean isValid = this.cookieKeys.stream()
+        boolean isValid = cookieKeysToCheck(this.context, this.pxConfiguration).stream()
                 .anyMatch(cookieKey -> HMACUtils.isHMACValid(hmacStr, cookieHmac, cookieKey, logger));
         if (!isValid) {
             context.logger.debug(LogReason.DEBUG_COOKIE_DECRYPTION_HMAC_FAILED, pxCookie, this.userAgent);

src/main/java/com/perimeterx/internals/cookie/cookieparsers/HeaderParser.java

@@ -16,6 +16,7 @@
 import java.util.*;
 import java.util.stream.Stream;
 
+import static com.perimeterx.utils.PXCommonUtils.cookieKeysToCheck;
 import static java.util.stream.Collectors.toList;
 
 public abstract class HeaderParser {

src/main/java/com/perimeterx/models/PXContext.java

@@ -41,6 +41,7 @@
 import static com.perimeterx.utils.Constants.BREACHED_ACCOUNT_KEY_NAME;
 import static com.perimeterx.utils.Constants.LOGGER_TOKEN_HEADER_NAME;
 import static com.perimeterx.utils.PXCommonUtils.cookieHeadersNames;
+import static com.perimeterx.utils.PXCommonUtils.cookieKeysToCheck;
 
 /**
  * PXContext - Populate relevant data from HttpRequest
@@ -229,6 +230,11 @@ public class PXContext {
     private String pxhdDomain;
     private long enforcerStartTime;
 
+    /**
+     * The cookie key used to decrypt the cookie
+     */
+    private String cookieKeyUsed;
+
     /**
      * The base64 encoded request full url
      */
@@ -394,7 +400,7 @@ private void parseCookies(HttpServletRequest request, boolean isMobileToken) {
             setVidAndPxhd(cookies);
             tokens.addAll(headerParser.createRawCookieDataList(cookieHeaders));
             this.tokens = tokens;
-            DataEnrichmentCookie deCookie = headerParser.getRawDataEnrichmentCookie(this.tokens, this.pxConfiguration.getCookieKeys());
+            DataEnrichmentCookie deCookie = headerParser.getRawDataEnrichmentCookie(this.tokens, cookieKeysToCheck(this, this.pxConfiguration));
             this.pxde = deCookie.getJsonPayload();
             this.pxdeVerified = deCookie.isValid();
         }

src/main/java/com/perimeterx/utils/PXCommonUtils.java

@@ -1,5 +1,6 @@
 package com.perimeterx.utils;
 
+import com.perimeterx.models.PXContext;
 import com.perimeterx.models.configuration.PXConfiguration;
 import org.apache.http.Header;
 import org.apache.http.HttpHeaders;
@@ -19,6 +20,14 @@
 public class PXCommonUtils {
     final static int FIRST_PARTY_CONNECTION_TIMEOUT_IN_MSEC = 4000;
 
+    public static List<String> cookieKeysToCheck(PXContext context, PXConfiguration configuration) {
+        if (context.getCookieKeyUsed() == null) {
+            return configuration.getCookieKeys();
+        }
+
+        return Collections.singletonList(context.getCookieKeyUsed());
+    }
+
     public static List<Header> getDefaultHeaders(String authToken) {
         Header contentType = new BasicHeader(HttpHeaders.CONTENT_TYPE, "application/json");
         Header authorization = new BasicHeader(HttpHeaders.AUTHORIZATION, "Bearer " + authToken);