From a0f86a3e38a09a5ae616252c3b8d64cb5982643d Mon Sep 17 00:00:00 2001
From: harinath <harinath@harinaths-MacBook-Pro.local>
Date: Thu, 1 May 2025 14:48:02 -0700
Subject: [PATCH] Add PEMKeyManager to handle PEM based certs and keys.

---
 .../main/java/org/postgresql/PGProperty.java  |   7 +
 .../java/org/postgresql/ssl/LibPQFactory.java |  35 +++-
 .../org/postgresql/ssl/PEMKeyManager.java     | 168 ++++++++++++++++++
 .../test/ssl/PEMKeyManagerTest.java           |  58 ++++++
 4 files changed, 262 insertions(+), 6 deletions(-)
 create mode 100644 pgjdbc/src/main/java/org/postgresql/ssl/PEMKeyManager.java
 create mode 100644 pgjdbc/src/test/java/org/postgresql/test/ssl/PEMKeyManagerTest.java

diff --git a/pgjdbc/src/main/java/org/postgresql/PGProperty.java b/pgjdbc/src/main/java/org/postgresql/PGProperty.java
index 740f4e59..b31fc5a5 100644
--- a/pgjdbc/src/main/java/org/postgresql/PGProperty.java
+++ b/pgjdbc/src/main/java/org/postgresql/PGProperty.java
@@ -830,6 +830,13 @@ public enum PGProperty {
       "",
       "Factory class to instantiate factories for XML processing"),
 
+    /**
+     * Algorithm for the PEM key.
+     */
+    PEM_KEY_ALGORITHM(
+            "pemKeyAlgorithm",
+            "RSA",
+            "Algorithm of the PEM key"),
   ;
 
   private final String name;
diff --git a/pgjdbc/src/main/java/org/postgresql/ssl/LibPQFactory.java b/pgjdbc/src/main/java/org/postgresql/ssl/LibPQFactory.java
index 33f16e0d..e51b981e 100644
--- a/pgjdbc/src/main/java/org/postgresql/ssl/LibPQFactory.java
+++ b/pgjdbc/src/main/java/org/postgresql/ssl/LibPQFactory.java
@@ -70,16 +70,20 @@ public class LibPQFactory extends WrappedFactory {
     return cbh;
   }
 
+  private String getCertFilePath(String defaultdir, Properties info) {
+      // Load the client's certificate and key
+      String sslcertfile = PGProperty.SSL_CERT.getOrDefault(info);
+      if (sslcertfile == null) { // Fall back to default
+          defaultfile = true;
+          sslcertfile = defaultdir + "postgresql.crt";
+      }
+      return sslcertfile;
+  }
   private void initPk8(
       @UnderInitialization(WrappedFactory.class) LibPQFactory this,
       String sslkeyfile, String defaultdir, Properties info) throws  PSQLException {
 
-    // Load the client's certificate and key
-    String sslcertfile = PGProperty.SSL_CERT.getOrDefault(info);
-    if (sslcertfile == null) { // Fall back to default
-      defaultfile = true;
-      sslcertfile = defaultdir + "postgresql.crt";
-    }
+    String sslcertfile = getCertFilePath(defaultdir, info);
 
     // If the properties are empty, give null to prevent client key selection
     km = new LazyKeyManager(("".equals(sslcertfile) ? null : sslcertfile),
@@ -92,6 +96,20 @@ public class LibPQFactory extends WrappedFactory {
     km = new PKCS12KeyManager(sslkeyfile, getCallbackHandler(info));
   }
 
+  private void initPEM(@UnderInitialization(WrappedFactory.class) LibPQFactory this,
+                       String sslKeyFile,
+                       String defaultdir,
+                       Properties info) throws PSQLException {
+    try {
+        String sslCertFile = getCertFilePath(defaultdir, info);
+        String algorithm = PGProperty.PEM_KEY_ALGORITHM.getOrDefault(info);
+        km = new PEMKeyManager(sslKeyFile, sslCertFile, algorithm);
+    } catch (Exception ex) {
+        ex.printStackTrace(System.out);
+        throw new PSQLException(GT.tr("Could not initial PEM files."), PSQLState.CONNECTION_FAILURE, ex);
+    }
+  }
+
   /**
    * @param info the connection parameters The following parameters are used:
    *        sslmode,sslcert,sslkey,sslrootcert,sslhostnameverifier,sslpasswordcallback,sslpassword
@@ -119,6 +137,8 @@ public class LibPQFactory extends WrappedFactory {
 
       if (sslkeyfile.endsWith(".p12") || sslkeyfile.endsWith(".pfx")) {
         initP12(sslkeyfile, info);
+      } else if (sslkeyfile.endsWith(".key") || sslkeyfile.endsWith(".pem")) {
+        initPEM(sslkeyfile, defaultdir, info);
       } else {
         initPk8(sslkeyfile, defaultdir, info);
       }
@@ -209,6 +229,9 @@ public class LibPQFactory extends WrappedFactory {
       if (km instanceof PKCS12KeyManager) {
         ((PKCS12KeyManager) km).throwKeyManagerException();
       }
+      if (km instanceof PEMKeyManager) {
+          ((PEMKeyManager) km).throwKeyManagerException();
+      }
     }
   }
 
diff --git a/pgjdbc/src/main/java/org/postgresql/ssl/PEMKeyManager.java b/pgjdbc/src/main/java/org/postgresql/ssl/PEMKeyManager.java
new file mode 100644
index 00000000..41eb0e93
--- /dev/null
+++ b/pgjdbc/src/main/java/org/postgresql/ssl/PEMKeyManager.java
@@ -0,0 +1,168 @@
+package org.postgresql.ssl;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.postgresql.util.GT;
+import org.postgresql.util.PSQLException;
+import org.postgresql.util.PSQLState;
+
+import javax.net.ssl.X509KeyManager;
+import javax.security.auth.x500.X500Principal;
+import java.io.InputStream;
+import java.net.Socket;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.security.*;
+import java.security.cert.Certificate;
+import java.security.cert.CertificateFactory;
+import java.security.cert.X509Certificate;
+import java.security.spec.PKCS8EncodedKeySpec;
+import java.util.*;
+
+public class PEMKeyManager implements X509KeyManager {
+
+    // private final CallbackHandler cbh;
+
+    private @Nullable PSQLException error;
+
+    private final String keyFilePath;
+    private final String certFilePath;
+
+
+    private final String keyAlgorithm;
+
+    public PEMKeyManager(String pemKeyPath, String pemCertsPath, String keyAlgorithm) {
+        this.keyFilePath = pemKeyPath;
+        this.certFilePath = pemCertsPath;
+        this.keyAlgorithm = keyAlgorithm;
+    }
+
+    /**
+     * getCertificateChain and getPrivateKey cannot throw exceptions, therefore any exception is stored
+     * in {@link #error} and can be raised by this method.
+     *
+     * @throws PSQLException if any exception is stored in {@link #error} and can be raised
+     */
+    public void throwKeyManagerException() throws PSQLException {
+        if (error != null) {
+            throw error;
+        }
+    }
+
+    @Override
+    public @Nullable PrivateKey getPrivateKey(String s) {
+        try {
+            List<String> lines = Files.readAllLines(Paths.get(keyFilePath));
+            StringBuilder keyContent = new StringBuilder();
+            for (String line : lines) {
+                if (!line.contains("BEGIN PRIVATE KEY") &&
+                    !line.contains("BEGIN RSA PRIVATE KEY") &&
+                    !line.contains("END PRIVATE KEY")&&
+                    !line.contains("END RSA PRIVATE KEY")) {
+                    keyContent.append(line.trim());
+                }
+            }
+
+            byte[] privateKeyDERBytes = Base64.getDecoder().decode(keyContent.toString());
+            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyDERBytes);
+            KeyFactory kf = KeyFactory.getInstance(this.keyAlgorithm);
+
+            return kf.generatePrivate(keySpec);
+        } catch (Exception e) {
+            error = new PSQLException(GT.tr("Could not load the private key"), PSQLState.CONNECTION_FAILURE, e);
+        }
+        return null;
+    }
+
+    @Override
+    public X509Certificate @Nullable [] getCertificateChain(String alias) {
+        try (InputStream inStream = Files.newInputStream(Paths.get(this.certFilePath))) {
+            CertificateFactory cf = CertificateFactory.getInstance("X.509");
+
+            Collection<? extends java.security.cert.Certificate> certs = cf.generateCertificates(inStream);
+            List<X509Certificate> certChain = new ArrayList<>();
+
+            for (Certificate cert : certs) {
+                if (cert instanceof X509Certificate) {
+                    certChain.add((X509Certificate) cert);
+                }
+            }
+
+            return certChain.toArray(new X509Certificate[0]);
+        } catch (Exception e) {
+            error = new PSQLException(GT.tr("Could not load cert chain"), PSQLState.CONNECTION_FAILURE, e);
+        }
+        return null;
+    }
+
+    @Override
+    public String @Nullable [] getClientAliases(String keyType, Principal @Nullable [] principals) {
+        String alias = chooseClientAlias(new String[]{keyType}, principals, (Socket) null);
+        return alias == null ? null : new String[]{alias};
+    }
+
+    @Override
+    public @Nullable String chooseClientAlias(String[] keyType, Principal @Nullable [] principals,
+                                              @Nullable Socket socket) {
+        if (principals == null || principals.length == 0) {
+            // Postgres 8.4 and earlier do not send the list of accepted certificate authorities
+            // to the client. See BUG #5468. We only hope, that our certificate will be accepted.
+            return "user";
+        } else {
+            // Sending a wrong certificate makes the connection rejected, even, if clientcert=0 in
+            // pg_hba.conf.
+            // therefore we only send our certificate, if the issuer is listed in issuers
+            X509Certificate[] certchain = getCertificateChain("user");
+            if (certchain == null) {
+                return null;
+            } else {
+                X509Certificate cert = certchain[certchain.length - 1];
+                X500Principal ourissuer = cert.getIssuerX500Principal();
+                String certKeyType = cert.getPublicKey().getAlgorithm();
+                boolean keyTypeFound = false;
+                boolean found = false;
+                if (keyType != null && keyType.length > 0) {
+                    for (String kt : keyType) {
+                        if (kt.equalsIgnoreCase(certKeyType)) {
+                            keyTypeFound = true;
+                        }
+                    }
+                } else {
+                    // If no key types were passed in, assume we don't care
+                    // about checking that the cert uses a particular key type.
+                    keyTypeFound = true;
+                }
+                if (keyTypeFound) {
+                    for (Principal issuer : principals) {
+                        if (ourissuer.equals(issuer)) {
+                            found = keyTypeFound;
+                        }
+                    }
+                }
+                return found ? "user" : null;
+            }
+        }
+    }
+
+    @Override
+    public String @Nullable [] getServerAliases(String s, Principal @Nullable [] principals) {
+        return new String[]{};
+    }
+
+    @Override
+    public @Nullable String chooseServerAlias(String s, Principal @Nullable [] principals,
+                                              @Nullable Socket socket) {
+        // we are not a server
+        return null;
+    }
+}
diff --git a/pgjdbc/src/test/java/org/postgresql/test/ssl/PEMKeyManagerTest.java b/pgjdbc/src/test/java/org/postgresql/test/ssl/PEMKeyManagerTest.java
new file mode 100644
index 00000000..3d66326b
--- /dev/null
+++ b/pgjdbc/src/test/java/org/postgresql/test/ssl/PEMKeyManagerTest.java
@@ -0,0 +1,58 @@
+package org.postgresql.test.ssl;
+
+import org.junit.jupiter.api.Test;
+import org.postgresql.PGProperty;
+import org.postgresql.ssl.PEMKeyManager;
+import org.postgresql.test.TestUtil;
+
+import javax.security.auth.x500.X500Principal;
+import java.sql.Connection;
+import java.util.Properties;
+
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class PEMKeyManagerTest {
+
+    @Test
+    void TestGoodClientPEM() throws Exception {
+        TestUtil.assumeSslTestsEnabled();
+
+        Properties props = new Properties();
+        props.put(TestUtil.DATABASE_PROP, "hostssldb");
+        PGProperty.SSL_MODE.set(props, "prefer");
+        PGProperty.SSL_KEY.set(props, TestUtil.getSslTestCertPath("goodclient.key"));
+        PGProperty.SSL_CERT.set(props, TestUtil.getSslTestCertPath("goodclient.crt"));
+        PGProperty.PEM_KEY_ALGORITHM.set(props, "RSA");
+
+        try (Connection conn = TestUtil.openDB(props)) {
+            boolean sslUsed = TestUtil.queryForBoolean(conn, "SELECT ssl_is_used()");
+            assertTrue(sslUsed, "SSL should be in use");
+        }
+    }
+    @Test
+    void TestChooseClientAlias() {
+        String sslKeyFile = TestUtil.getSslTestCertPath("goodclient.key");
+        String sslCertFile = TestUtil.getSslTestCertPath("goodclient.crt");
+        PEMKeyManager pemKeyManager = new PEMKeyManager(sslKeyFile, sslCertFile, "RSA");
+
+        X500Principal testPrincipal = new X500Principal("CN=root certificate, O=PgJdbc test, ST=CA, C=US");
+        X500Principal[] issuers = new X500Principal[]{testPrincipal};
+
+        String validKeyType = pemKeyManager.chooseClientAlias(new String[]{"RSA"}, issuers, null);
+        assertNotNull(validKeyType);
+
+        String ignoresCase = pemKeyManager.chooseClientAlias(new String[]{"rsa"}, issuers, null);
+        assertNotNull(ignoresCase);
+
+        String invalidKeyType = pemKeyManager.chooseClientAlias(new String[]{"EC"}, issuers, null);
+        assertNull(invalidKeyType);
+
+        String containsValidKeyType = pemKeyManager.chooseClientAlias(new String[]{"EC", "RSA"}, issuers, null);
+        assertNotNull(containsValidKeyType);
+
+        String ignoresBlank = pemKeyManager.chooseClientAlias(new String[]{}, issuers, null);
+        assertNotNull(ignoresBlank);
+    }
+}
-- 

