- Introduce filterKeyStoreByAlias method in KeyStoreUtil to creat…

…e a new KeyStore containing only the specified key and its certificate chain. - Update `StaticSSLContext` to use the filtered KeyStore for initializing the KeyManagerFactory. - Add unit tests for key alias filtering in `KeyStoreUtilTest` and SSL context behavior in `SSLContextTest`. - Include a new truststore for validating second key alias scenarios.
membrane · Sep 11, 2024 · 00d2ebe · 00d2ebe
1 parent 5142b22
commit 00d2ebe
Show file tree

Hide file tree

Showing 5 changed files with 96 additions and 5 deletions.
diff --git a/core/src/main/java/com/predic8/membrane/core/security/KeyStoreUtil.java b/core/src/main/java/com/predic8/membrane/core/security/KeyStoreUtil.java
@@ -19,12 +19,42 @@
 
 import java.io.IOException;
 import java.security.*;
+import java.security.cert.Certificate;
 import java.security.cert.CertificateEncodingException;
 import java.security.cert.CertificateException;
 import java.util.Enumeration;
 import java.util.Optional;
 
 public class KeyStoreUtil {
+    /**
+     * Filters a KeyStore by a specific alias, creating a new KeyStore
+     * containing only the key and certificate chain associated with
+     * the provided alias.
+     *
+     * @param ks       the original KeyStore to filter
+     * @param keyPass  the password for accessing the key
+     * @param keyAlias the alias of the key to filter
+     * @return a new KeyStore containing only the key and certificate
+     * chain associated with the specified alias
+     * @throws KeyStoreException         if the KeyStore cannot be initialized
+     * @throws IOException               if there is an I/O error during the operation
+     * @throws NoSuchAlgorithmException  if the algorithm for recovering
+     *                                   the key cannot be found
+     * @throws CertificateException      if any of the certificates in the
+     *                                   chain are invalid
+     * @throws UnrecoverableKeyException if the key cannot be recovered
+     *                                   using the given password
+     */
+    public static KeyStore filterKeyStoreByAlias(KeyStore ks, char[] keyPass, String keyAlias) throws KeyStoreException, IOException, NoSuchAlgorithmException, CertificateException, UnrecoverableKeyException {
+        KeyStore filteredKeyStore = KeyStore.getInstance(ks.getType());
+        filteredKeyStore.load(null, keyPass);
+
+        Key key = ks.getKey(keyAlias, keyPass);
+        Certificate[] certificateChain = ks.getCertificateChain(keyAlias);
+        filteredKeyStore.setKeyEntry(keyAlias, key, keyPass, certificateChain);
+        return filteredKeyStore;
+    }
+
     /**
      * Generates an SHA-256 digest of the certificate associated with the given alias in the KeyStore.
      *

diff --git a/core/src/main/java/com/predic8/membrane/core/transport/ssl/StaticSSLContext.java b/core/src/main/java/com/predic8/membrane/core/transport/ssl/StaticSSLContext.java
@@ -19,6 +19,7 @@
 import com.predic8.membrane.core.resolver.ResolverMap;
 import com.predic8.membrane.core.transport.TrustManagerWrapper;
 import com.predic8.membrane.core.transport.http2.Http2TlsSupport;
+import org.apache.commons.lang3.NotImplementedException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -85,12 +86,15 @@ public StaticSSLContext(SSLParser sslParser, ResolverMap resourceResolver, Strin
                 if (sslParser.getKeyStore().getType() != null)
                     keyStoreType = sslParser.getKeyStore().getType();
                 KeyStore ks = openKeyStore(sslParser.getKeyStore(), "PKCS12", keyPass, resourceResolver, baseLocation);
-                kmf = KeyManagerFactory.getInstance(algorihm);
-                kmf.init(ks, keyPass);
 
                 String paramAlias = sslParser.getKeyStore().getKeyAlias();
                 String keyAlias = (paramAlias != null) ? aliasOrThrow(ks, paramAlias) : firstAliasOrThrow(ks);
 
+                KeyStore filteredKeyStore = filterKeyStoreByAlias(ks, keyPass, keyAlias);
+
+                kmf = KeyManagerFactory.getInstance(algorihm);
+                kmf.init(filteredKeyStore, keyPass);
+
                 dnsNames = extractDnsNames(ks.getCertificate(keyAlias));
                 List<Certificate> certs = Arrays.asList(ks.getCertificateChain(keyAlias));
                 validUntil = getMinimumValidity(certs);

diff --git a/core/src/test/java/com/predic8/membrane/core/security/KeyStoreUtilTest.java b/core/src/test/java/com/predic8/membrane/core/security/KeyStoreUtilTest.java
@@ -26,13 +26,16 @@
 import java.security.KeyStoreException;
 import java.security.NoSuchAlgorithmException;
 import java.security.NoSuchProviderException;
+import java.security.UnrecoverableKeyException;
 import java.security.cert.CertificateEncodingException;
 import java.security.cert.CertificateException;
 import java.util.Optional;
 
+import static com.predic8.membrane.core.security.KeyStoreUtil.filterKeyStoreByAlias;
+import static com.predic8.membrane.core.security.KeyStoreUtil.getAndLoadKeyStore;
 import static org.junit.jupiter.api.Assertions.*;
 
-class KeyStoreUtilsTest {
+class KeyStoreUtilTest {
 
     private static Router router;
     private static java.security.KeyStore keyStore;
@@ -47,14 +50,26 @@ static void setUp() throws Exception {
         sslParser.setKeyStore(new KeyStore());
         sslParser.getKeyStore().setLocation("classpath:/alias-keystore.p12");
         sslParser.getKeyStore().setKeyPassword(KEYSTORE_PASSWORD);
-        keyStore = KeyStoreUtil.getAndLoadKeyStore(sslParser.getKeyStore(), router.getResolverMap(), router.getBaseLocation(), "PKCS12", KEYSTORE_PASSWORD.toCharArray());
+        keyStore = getAndLoadKeyStore(sslParser.getKeyStore(), router.getResolverMap(), router.getBaseLocation(), "PKCS12", KEYSTORE_PASSWORD.toCharArray());
     }
 
     @AfterAll
     static void tearDown() {
         router.stop();
     }
 
+    @Test
+    void testFilterKeyStoreByAlias() throws KeyStoreException, NoSuchProviderException, IOException, NoSuchAlgorithmException, CertificateException, UnrecoverableKeyException {
+        KeyStore store = new KeyStore();
+        store.setLocation("classpath:/alias-keystore.p12");
+        store.setKeyPassword(KEYSTORE_PASSWORD);
+        java.security.KeyStore loadedKeyStore = getAndLoadKeyStore(store, router.getResolverMap(), router.getBaseLocation(), "PKCS12", KEYSTORE_PASSWORD.toCharArray());
+        assertNotNull(loadedKeyStore);
+        assertEquals(2, loadedKeyStore.size());
+        java.security.KeyStore filteredKeyStore = filterKeyStoreByAlias(loadedKeyStore, "secret".toCharArray(), "key1");
+        assertEquals(1, filteredKeyStore.size());
+    }
+
     @Test
     void testGetDigest() throws CertificateEncodingException, KeyStoreException, NoSuchAlgorithmException {
         String digest = KeyStoreUtil.getDigest(keyStore, ALIAS);
@@ -66,7 +81,7 @@ void testGetAndLoadKeyStore() throws KeyStoreException, NoSuchProviderException,
         KeyStore store = new KeyStore();
         store.setLocation("classpath:/alias-keystore.p12");
         store.setKeyPassword(KEYSTORE_PASSWORD);
-        java.security.KeyStore loadedKeyStore = KeyStoreUtil.getAndLoadKeyStore(store, router.getResolverMap(), router.getBaseLocation(), "PKCS12", KEYSTORE_PASSWORD.toCharArray());
+        java.security.KeyStore loadedKeyStore = getAndLoadKeyStore(store, router.getResolverMap(), router.getBaseLocation(), "PKCS12", KEYSTORE_PASSWORD.toCharArray());
         assertNotNull(loadedKeyStore);
         assertTrue(loadedKeyStore.size() > 0);
         assertTrue(loadedKeyStore.containsAlias(ALIAS));

diff --git a/core/src/test/java/com/predic8/membrane/core/transport/ssl/SSLContextTest.java b/core/src/test/java/com/predic8/membrane/core/transport/ssl/SSLContextTest.java
@@ -18,6 +18,7 @@
 import com.predic8.membrane.core.config.security.KeyStore;
 import com.predic8.membrane.core.config.security.SSLParser;
 import com.predic8.membrane.core.config.security.TrustStore;
+import org.apache.commons.lang3.NotImplementedException;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -97,6 +98,47 @@ public void simpleConfig() {
 		});
 	}
 
+	@Test
+	public void selectFirstKeyAndPass() throws Exception {
+		SSLContext server = cb().withKeyStore("classpath:/alias-keystore.p12").byKeyAlias("key1").build();
+		SSLContext client = cb().withTrustStore("classpath:/alias-truststore.p12").build();
+		testCombination(server, client);
+	}
+
+	@Test
+	public void selectFirstKeyTrustFail() {
+		assertThrows(SocketException.class, () -> {
+			SSLContext server = cb().withKeyStore("classpath:/alias-keystore.p12").byKeyAlias("key2").build();
+			SSLContext client = cb().withTrustStore("classpath:/alias-truststore.p12").build();
+			testCombination(server, client);
+		});
+	}
+
+	@Test
+	public void selectSecondKeyAndPass() throws Exception {
+		SSLContext server = cb().withKeyStore("classpath:/alias-keystore.p12").byKeyAlias("key2").build();
+		SSLContext client = cb().withTrustStore("classpath:/alias-truststore2.p12").build();
+		testCombination(server, client);
+	}
+
+	@Test
+	public void selectSecondKeyTrustFail() {
+		assertThrows(SocketException.class, () -> {
+			SSLContext server = cb().withKeyStore("classpath:/alias-keystore.p12").byKeyAlias("key1").build();
+			SSLContext client = cb().withTrustStore("classpath:/alias-truststore2.p12").build();
+			testCombination(server, client);
+		});
+	}
+
+	@Test
+	public void invalidAlias() {
+		assertThrows(RuntimeException.class, () -> {
+			SSLContext server = cb().withKeyStore("classpath:/alias-keystore.p12").byKeyAlias("key999").build();
+			SSLContext client = cb().withTrustStore("classpath:/alias-truststore.p12").build();
+			testCombination(server, client);
+		});
+	}
+
 	@Test
 	public void serverKeyOnlyWithoutClientTrust() {
 		assertThrows(Exception.class, () -> {

diff --git a/core/src/test/resources/alias-truststore2.p12 b/core/src/test/resources/alias-truststore2.p12