Adding a more secure way of caching tokens in the STS

rt/ws/security/src/main/java/org/apache/cxf/ws/security/tokenstore/TokenStoreUtils.java

@@ -18,9 +18,28 @@
  */
 package org.apache.cxf.ws.security.tokenstore;
 
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.util.HexFormat;
+
+import org.w3c.dom.Element;
+
+import org.apache.cxf.helpers.DOMUtils;
 import org.apache.cxf.message.Message;
 import org.apache.cxf.service.model.EndpointInfo;
 import org.apache.cxf.ws.security.SecurityConstants;
+import org.apache.wss4j.common.ext.WSSecurityException;
+import org.apache.wss4j.common.saml.SamlAssertionWrapper;
+import org.apache.wss4j.common.token.BinarySecurity;
+import org.apache.wss4j.dom.WSConstants;
+import org.apache.wss4j.dom.message.token.SecurityContextToken;
+import org.apache.wss4j.dom.message.token.UsernameToken;
+import org.apache.xml.security.Init;
+import org.apache.xml.security.c14n.Canonicalizer;
 
 /**
  * Some common functionality
@@ -61,4 +80,135 @@ public static TokenStore getTokenStore(Message message) throws TokenStoreExcepti
             return tokenStore;
         }
     }
+
+    /**
+     * Get a cache key for a signed SAML Assertion, that can be used to store and retrieve the (validated)
+     * Assertion in/from a TokenStore. The key is a SHA-256 digest over the canonicalized SignedInfo
+     * (which binds the signed content of the Assertion) and the SignatureValue. It returns null if the
+     * Assertion is not signed, or if the signature does not conform to the SAML signature profile (in which
+     * case the signature might not cover the Assertion itself, and so the Assertion must not be cached).
+     *
+     * This method must be used to look up a received Assertion in a TokenStore.
+     */
+    public static String getCacheKey(SamlAssertionWrapper assertion) throws WSSecurityException {
+        return getCacheKey(assertion, true);
+    }
+
+    /**
+     * Get a cache key for a signed SAML Assertion - see getCacheKey(SamlAssertionWrapper).
+     * @param validateSignatureProfile whether to check that the signature conforms to the SAML signature
+     *        profile first, returning null if it doesn't. This must be true when looking up a received
+     *        Assertion. It can be false when storing an Assertion that was just signed by the STS, as
+     *        its DOM Element might not be attached to a Document, which the profile validation requires.
+     */
+    public static String getCacheKey(SamlAssertionWrapper assertion, boolean validateSignatureProfile)
+        throws WSSecurityException {
+        byte[] signatureValue = assertion.getSignatureValue();
+        if (signatureValue == null || signatureValue.length == 0) {
+            return null;
+        }
+
+        if (validateSignatureProfile) {
+            try {
+                assertion.validateSignatureAgainstProfile();
+            } catch (WSSecurityException ex) {
+                return null;
+            }
+        }
+
+        Element signedInfo = null;
+        Element assertionElement = assertion.getElement();
+        if (assertionElement == null && assertion.getSamlObject() != null) {
+            assertionElement = assertion.getSamlObject().getDOM();
+        }
+        if (assertionElement != null) {
+            Element signature =
+                DOMUtils.getFirstChildWithName(assertionElement, WSConstants.SIG_NS, WSConstants.SIG_LN);
+            if (signature != null) {
+                signedInfo = DOMUtils.getFirstChildWithName(signature, WSConstants.SIG_NS, "SignedInfo");
+            }
+        }
+        if (signedInfo == null) {
+            return null;
+        }
+
+        try {
+            if (!Init.isInitialized()) {
+                Init.init();
+            }
+            ByteArrayOutputStream signedInfoBytes = new ByteArrayOutputStream();
+            Canonicalizer.getInstance(Canonicalizer.ALGO_ID_C14N_EXCL_OMIT_COMMENTS)
+                .canonicalizeSubtree(signedInfo, signedInfoBytes);
+
+            return computeCacheKey("SAML", signedInfoBytes.toByteArray(), signatureValue);
+        } catch (Exception ex) {
+            throw new WSSecurityException(WSSecurityException.ErrorCode.FAILURE, ex);
+        }
+    }
+
+    /**
+     * Get a cache key for a UsernameToken, that can be used to store and retrieve the (validated)
+     * UsernameToken in/from a TokenStore. The key is a SHA-256 digest over the UsernameToken values.
+     */
+    public static String getCacheKey(UsernameToken usernameToken) throws WSSecurityException {
+        byte[] salt = usernameToken.getSalt();
+        return computeCacheKey("UsernameToken",
+                               toBytes(usernameToken.getName()),
+                               toBytes(usernameToken.getPassword()),
+                               toBytes(usernameToken.getPasswordType()),
+                               toBytes(usernameToken.getNonce()),
+                               toBytes(usernameToken.getCreated()),
+                               salt,
+                               toBytes(Integer.toString(usernameToken.getIteration())));
+    }
+
+    /**
+     * Get a cache key for a BinarySecurityToken, that can be used to store and retrieve the (validated)
+     * BinarySecurityToken in/from a TokenStore. The key is a SHA-256 digest over the token values.
+     */
+    public static String getCacheKey(BinarySecurity binarySecurity) {
+        return computeCacheKey("BinarySecurityToken",
+                               toBytes(binarySecurity.getValueType()),
+                               toBytes(binarySecurity.getEncodingType()),
+                               binarySecurity.getToken());
+    }
+
+    /**
+     * Get a cache key for a SecurityContextToken, that can be used to store and retrieve the (validated)
+     * SecurityContextToken in/from a TokenStore. The key is a SHA-256 digest over the token identifier.
+     */
+    public static String getCacheKey(SecurityContextToken securityContextToken) {
+        String identifier = securityContextToken.getIdentifier();
+        if (identifier == null) {
+            return null;
+        }
+        return computeCacheKey("SecurityContextToken", toBytes(identifier));
+    }
+
+    private static byte[] toBytes(String value) {
+        return value == null ? null : value.getBytes(StandardCharsets.UTF_8);
+    }
+
+    private static String computeCacheKey(String tokenType, byte[]... values) {
+        try {
+            ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+            DataOutputStream out = new DataOutputStream(bytes);
+            out.writeUTF(tokenType);
+            for (byte[] value : values) {
+                // Length-prefix each value so that different combinations can't produce the same input
+                if (value == null) {
+                    out.writeInt(-1);
+                } else {
+                    out.writeInt(value.length);
+                    out.write(value);
+                }
+            }
+            out.flush();
+
+            MessageDigest digest = MessageDigest.getInstance("SHA-256");
+            return HexFormat.of().formatHex(digest.digest(bytes.toByteArray()));
+        } catch (IOException | NoSuchAlgorithmException ex) {
+            throw new IllegalStateException(ex);
+        }
+    }
 }

rt/ws/security/src/main/java/org/apache/cxf/ws/security/trust/STSTokenValidator.java

@@ -20,7 +20,6 @@
 package org.apache.cxf.ws.security.trust;
 
 import java.io.IOException;
-import java.util.Arrays;
 import java.util.List;
 
 import javax.security.auth.callback.Callback;
@@ -86,23 +85,20 @@ public Credential validateWithSTS(Credential credential, Message message) throws
         try {
             SecurityToken token = new SecurityToken();
             Element tokenElement = null;
-            int hash = 0;
+            String cacheKey = null;
             if (credential.getSamlAssertion() != null) {
                 SamlAssertionWrapper assertion = credential.getSamlAssertion();
-                byte[] signatureValue = assertion.getSignatureValue();
-                if (signatureValue != null && signatureValue.length > 0) {
-                    hash = Arrays.hashCode(signatureValue);
-                }
-                tokenElement = credential.getSamlAssertion().getElement();
+                cacheKey = TokenStoreUtils.getCacheKey(assertion);
+                tokenElement = assertion.getElement();
             } else if (credential.getUsernametoken() != null) {
                 tokenElement = credential.getUsernametoken().getElement();
-                hash = credential.getUsernametoken().hashCode();
+                cacheKey = TokenStoreUtils.getCacheKey(credential.getUsernametoken());
             } else if (credential.getBinarySecurityToken() != null) {
                 tokenElement = credential.getBinarySecurityToken().getElement();
-                hash = credential.getBinarySecurityToken().hashCode();
+                cacheKey = TokenStoreUtils.getCacheKey(credential.getBinarySecurityToken());
             } else if (credential.getSecurityContextToken() != null) {
                 tokenElement = credential.getSecurityContextToken().getElement();
-                hash = credential.getSecurityContextToken().hashCode();
+                cacheKey = TokenStoreUtils.getCacheKey(credential.getSecurityContextToken());
             }
             token.setToken(tokenElement);
 
@@ -112,8 +108,8 @@ public Credential validateWithSTS(Credential credential, Message message) throws
                 if (ts == null) {
                     ts = tokenStore;
                 }
-                if (ts != null && hash != 0) {
-                    SecurityToken transformedToken = getTransformedToken(ts, hash);
+                if (ts != null && cacheKey != null) {
+                    SecurityToken transformedToken = getTransformedToken(ts, cacheKey);
                     if (transformedToken != null && !transformedToken.isExpired()) {
                         SamlAssertionWrapper assertion = new SamlAssertionWrapper(transformedToken.getToken());
                         credential.setPrincipal(new SAMLTokenPrincipalImpl(assertion));
@@ -122,7 +118,6 @@ public Credential validateWithSTS(Credential credential, Message message) throws
                     }
                 }
             }
-            token.setTokenHash(hash);
 
             STSClient c = stsClient;
             if (c == null) {
@@ -156,10 +151,10 @@ public Credential validateWithSTS(Credential credential, Message message) throws
                     SamlAssertionWrapper assertion = new SamlAssertionWrapper(returnedToken.getToken());
                     credential.setTransformedToken(assertion);
                     credential.setPrincipal(new SAMLTokenPrincipalImpl(assertion));
-                    if (!disableCaching && hash != 0 && ts != null) {
+                    if (!disableCaching && cacheKey != null && ts != null) {
                         ts.add(returnedToken);
                         token.setTransformedTokenIdentifier(returnedToken.getId());
-                        ts.add(Integer.toString(hash), token);
+                        ts.add(cacheKey, token);
                     }
                 }
                 return credential;
@@ -195,9 +190,9 @@ protected boolean isValidatedLocally(Credential credential, RequestData data)
         return false;
     }
 
-    private SecurityToken getTransformedToken(TokenStore ts, int hash) {
-        SecurityToken recoveredToken = ts.getToken(Integer.toString(hash));
-        if (recoveredToken != null && recoveredToken.getTokenHash() == hash) {
+    private SecurityToken getTransformedToken(TokenStore ts, String cacheKey) {
+        SecurityToken recoveredToken = ts.getToken(cacheKey);
+        if (recoveredToken != null) {
             String transformedTokenId = recoveredToken.getTransformedTokenIdentifier();
             if (transformedTokenId != null) {
                 return ts.getToken(transformedTokenId);