From 1a6ec768d45332fcc74b5e86791a39184d7c4ed0 Mon Sep 17 00:00:00 2001
From: nanhe <papaya.dyn@alibaba-inc.com>
Date: Tue, 24 Dec 2024 15:11:55 +0800
Subject: [PATCH] feat: support adding multiple certificates to the keystore

---
 .../tea/okhttp/OkHttpClientBuilder.java       | 43 +++++++++++++++++--
 .../tea/okhttp/OkHttpClientBuilderTest.java   |  8 ++++
 2 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/src/main/java/com/aliyun/tea/okhttp/OkHttpClientBuilder.java b/src/main/java/com/aliyun/tea/okhttp/OkHttpClientBuilder.java
index ab4ab51..85df524 100644
--- a/src/main/java/com/aliyun/tea/okhttp/OkHttpClientBuilder.java
+++ b/src/main/java/com/aliyun/tea/okhttp/OkHttpClientBuilder.java
@@ -16,13 +16,17 @@
 import java.security.SecureRandom;
 import java.security.cert.Certificate;
 import java.security.cert.CertificateFactory;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 public class OkHttpClientBuilder {
     private static final String charset = "UTF-8";
     private final OkHttpClient.Builder builder;
+    public static final String PEM_BEGIN = "-----BEGIN CERTIFICATE-----";
+    public static final String PEM_END = "-----END CERTIFICATE-----";
 
     public OkHttpClientBuilder() {
         builder = new OkHttpClient().newBuilder();
@@ -84,6 +88,33 @@ public OkHttpClientBuilder connectionPool(Map<String, Object> map) {
         return this;
     }
 
+    private List<String> splitPemCertificates(String pemData) throws Exception {
+        List<String> certificates = new ArrayList<String>();
+
+        if (null != pemData && pemData.contains(PEM_BEGIN) && pemData.contains(PEM_END)) {
+            StringBuilder sb = null;
+            BufferedReader reader = new BufferedReader(new StringReader(pemData));
+
+            String line;
+            while ((line = reader.readLine()) != null) {
+                if (line.contains(PEM_BEGIN)) {
+                    sb = new StringBuilder();
+                    sb.append(PEM_BEGIN).append('\n');
+                } else if (null != sb && line.contains(PEM_END)) {
+                    sb.append(PEM_END).append('\n');
+                    certificates.add(sb.toString());
+                    sb = null;
+                } else if (null != sb) {
+                    sb.append(line).append('\n');
+                }
+            }
+        } else if (null != pemData) {
+            certificates.add(pemData);
+        }
+
+        return certificates;
+    }
+
     public OkHttpClientBuilder certificate(Map<String, Object> map) {
         try {
             if (null != map.get("ignoreSSL") && Boolean.parseBoolean(String.valueOf(map.get("ignoreSSL")))) {
@@ -107,12 +138,16 @@ public OkHttpClientBuilder certificate(Map<String, Object> map) {
                 KeyStore trustStore = KeyStore.getInstance(KeyStore.getDefaultType());
                 trustStore.load(null);
                 String ca = String.valueOf(map.get("ca"));
+                List<String> pemCerts = splitPemCertificates(ca);
                 CertificateFactory certFactory = CertificateFactory.getInstance("X.509");
-                Certificate certificate;
-                try (InputStream is = new ByteArrayInputStream(ca.getBytes(charset))) {
-                    certificate = certFactory.generateCertificate(is);
+                int certIndex = 0;
+                // Process each certificate and add to the keystore
+                for (String pemCert : pemCerts) {
+                    try (InputStream is = new ByteArrayInputStream(pemCert.getBytes(charset))) {
+                        Certificate certificate = certFactory.generateCertificate(is);
+                        trustStore.setCertificateEntry("ca" + certIndex++, certificate);
+                    }
                 }
-                trustStore.setCertificateEntry("server-ca", certificate);
                 TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
                 trustManagerFactory.init(trustStore);
                 X509TrustManager trustManager = (X509TrustManager) trustManagerFactory.getTrustManagers()[0];
diff --git a/src/test/java/com/aliyun/tea/okhttp/OkHttpClientBuilderTest.java b/src/test/java/com/aliyun/tea/okhttp/OkHttpClientBuilderTest.java
index 293076f..4bff8f5 100644
--- a/src/test/java/com/aliyun/tea/okhttp/OkHttpClientBuilderTest.java
+++ b/src/test/java/com/aliyun/tea/okhttp/OkHttpClientBuilderTest.java
@@ -94,6 +94,14 @@ public void certificateTest() throws IOException {
             Assert.assertTrue(e.getMessage().contains("Unable to initialize"));
         }
 
+        map.put("ca", "-----BEGIN CERTIFICATE-----\nwrong ca-----END CERTIFICATE-----\n\n-----BEGIN CERTIFICATE-----\nwrong ca-----END CERTIFICATE-----");
+        try {
+            new OkHttpClientBuilder().certificate(map);
+            Assert.fail();
+        } catch (TeaException e) {
+            Assert.assertTrue(e.getMessage().contains("Unable to initialize"));
+        }
+
         map.put("ca", null);
         new OkHttpClientBuilder().certificate(map);
         map.put("ca", "");