From c691742810d365f1f932cf667d700a812ee564a2 Mon Sep 17 00:00:00 2001
From: Yuri Mizushima <equanz@apache.org>
Date: Thu, 28 Nov 2024 05:35:08 +0900
Subject: [PATCH] [fix][client] Fixed an issue where a cert chain could not be
 used in TLS authentication (#23644)

(cherry picked from commit e236d61c650bcaef271bc529ec9023edd8c04b01)
(cherry picked from commit e0e195678cc10c8aea579e2874cbe0f945c3b850)
---
 .../pulsar/common/util/KeyManagerProxy.java   | 19 +++----
 .../common/util/KeyManagerProxyTest.java      | 51 +++++++++++++++++++
 2 files changed, 61 insertions(+), 9 deletions(-)
 create mode 100644 pulsar-common/src/test/java/org/apache/pulsar/common/util/KeyManagerProxyTest.java

diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/util/KeyManagerProxy.java b/pulsar-common/src/main/java/org/apache/pulsar/common/util/KeyManagerProxy.java
index 171f5acd2bf75..6b6b0492b4525 100644
--- a/pulsar-common/src/main/java/org/apache/pulsar/common/util/KeyManagerProxy.java
+++ b/pulsar-common/src/main/java/org/apache/pulsar/common/util/KeyManagerProxy.java
@@ -30,11 +30,14 @@
 import java.security.Principal;
 import java.security.PrivateKey;
 import java.security.UnrecoverableKeyException;
+import java.security.cert.Certificate;
 import java.security.cert.CertificateException;
 import java.security.cert.CertificateFactory;
 import java.security.cert.X509Certificate;
+import java.util.List;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 import javax.net.ssl.KeyManagerFactory;
 import javax.net.ssl.SSLEngine;
 import javax.net.ssl.X509ExtendedKeyManager;
@@ -88,18 +91,16 @@ private void updateKeyManager()
             return;
         }
 
-        X509Certificate certificate;
-        PrivateKey privateKey = null;
-        KeyStore keyStore;
-        try (InputStream publicCertStream = new FileInputStream(certFile.getFileName());
-                InputStream privateKeyStream = new FileInputStream(keyFile.getFileName())) {
+        final KeyStore keyStore;
+        try (InputStream publicCertStream = new FileInputStream(certFile.getFileName())) {
             final CertificateFactory cf = CertificateFactory.getInstance("X.509");
-            certificate = (X509Certificate) cf.generateCertificate(publicCertStream);
+            final List<X509Certificate> certificateList = cf.generateCertificates(publicCertStream)
+                    .stream().map(o -> (X509Certificate) o).collect(Collectors.toList());
             keyStore = KeyStore.getInstance("JKS");
-            String alias = certificate.getSubjectX500Principal().getName();
-            privateKey = SecurityUtility.loadPrivateKeyFromPemFile(keyFile.getFileName());
+            final String alias = certificateList.get(0).getSubjectX500Principal().getName();
+            final PrivateKey privateKey = SecurityUtility.loadPrivateKeyFromPemFile(keyFile.getFileName());
             keyStore.load(null);
-            keyStore.setKeyEntry(alias, privateKey, KEYSTORE_PASSWORD, new X509Certificate[] { certificate });
+            keyStore.setKeyEntry(alias, privateKey, KEYSTORE_PASSWORD, certificateList.toArray(new Certificate[0]));
         } catch (IOException | KeyManagementException e) {
             throw new IllegalArgumentException(e);
         }
diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/util/KeyManagerProxyTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/util/KeyManagerProxyTest.java
new file mode 100644
index 0000000000000..5542f0b22ac95
--- /dev/null
+++ b/pulsar-common/src/test/java/org/apache/pulsar/common/util/KeyManagerProxyTest.java
@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.pulsar.common.util;
+
+import static org.testng.Assert.assertEquals;
+import com.google.common.io.Resources;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import lombok.Cleanup;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+public class KeyManagerProxyTest {
+
+    @DataProvider(name = "certDataProvider")
+    public static Object[][] caDataProvider() {
+        return new Object[][]{
+                {"ca/multiple-ca.pem", 2},
+                {"ca/single-ca.pem", 1}
+        };
+    }
+
+    @Test(dataProvider = "certDataProvider")
+    public void testLoadCert(String path, int certCount) {
+        final String certFilePath = Resources.getResource(path).getPath();
+        // This key is not paired with certs, but this is not a problem as the key is not used in this test
+        final String keyFilePath = Resources.getResource("ssl/my-ca/client-key.pem").getPath();
+        @Cleanup("shutdownNow")
+        final ScheduledExecutorService scheduledExecutor = Executors.newSingleThreadScheduledExecutor();
+
+        final KeyManagerProxy keyManager = new KeyManagerProxy(certFilePath, keyFilePath, 60, scheduledExecutor);
+        assertEquals(keyManager.getCertificateChain("cn=test1").length, certCount);
+    }
+}