From 45689a2ad3dbb91f7cd6c61a215c7d36fa968bf4 Mon Sep 17 00:00:00 2001
From: Jonathan Hess <hessjc@google.com>
Date: Mon, 3 Mar 2025 09:58:34 -0700
Subject: [PATCH] fix: Add GraalVM configuration for the JDK's DNS Resolver.

This fixes the GraalVM distribution so that it can resolve DNS Names.
---
 .../com/google/cloud/sql/core/JndiDnsResolver.java   | 10 +++++++++-
 .../cloud/sql/nativeimage/CloudSqlFeature.java       |  6 ++++++
 .../proxy-config.json                                | 12 ++++++++++++
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/com/google/cloud/sql/core/JndiDnsResolver.java b/core/src/main/java/com/google/cloud/sql/core/JndiDnsResolver.java
index 2dfe5cb94..aa24bab63 100644
--- a/core/src/main/java/com/google/cloud/sql/core/JndiDnsResolver.java
+++ b/core/src/main/java/com/google/cloud/sql/core/JndiDnsResolver.java
@@ -18,7 +18,9 @@
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Hashtable;
 import java.util.stream.Collectors;
+import javax.naming.Context;
 import javax.naming.NameNotFoundException;
 import javax.naming.NamingException;
 import javax.naming.directory.Attribute;
@@ -51,13 +53,19 @@ class JndiDnsResolver implements DnsResolver {
    * @throws javax.naming.NameNotFoundException when the domain name did not resolve.
    */
   @Override
+  @SuppressWarnings("JdkObsolete")
   public Collection<String> resolveTxt(String domainName)
       throws javax.naming.NameNotFoundException {
     try {
       // Notice: This is old Java 1.2 style code. It uses the ancient JNDI DNS Provider api.
       // See https://docs.oracle.com/javase/7/docs/technotes/guides/jndi/jndi-dns.html
+
+      // Explicitly reference the JNDI DNS classes. This is required for GraalVM.
+      Hashtable contextProps = new Hashtable<>();
+      contextProps.put(Context.INITIAL_CONTEXT_FACTORY, "com.sun.jndi.dns.DnsContextFactory");
+      contextProps.put(Context.OBJECT_FACTORIES, "com.sun.jndi.url.dns.dnsURLContextFactory");
       Attribute attr =
-          new InitialDirContext()
+          new InitialDirContext(contextProps)
               .getAttributes(jndiPrefix + domainName, new String[] {"TXT"})
               .get("TXT");
       // attr.getAll() returns a Vector containing strings, one for each record returned by dns.
diff --git a/core/src/main/java/com/google/cloud/sql/nativeimage/CloudSqlFeature.java b/core/src/main/java/com/google/cloud/sql/nativeimage/CloudSqlFeature.java
index 6028c6cde..29a6928de 100644
--- a/core/src/main/java/com/google/cloud/sql/nativeimage/CloudSqlFeature.java
+++ b/core/src/main/java/com/google/cloud/sql/nativeimage/CloudSqlFeature.java
@@ -45,6 +45,8 @@ final class CloudSqlFeature implements Feature {
   private static final String POSTGRES_SOCKET_CLASS = "com.google.cloud.sql.postgres.SocketFactory";
 
   private static final String MYSQL_SOCKET_CLASS = "com.google.cloud.sql.mysql.SocketFactory";
+  private static final String JNDI_DNS_FACTORY = "com.sun.jndi.dns.DnsContextFactory";
+  private static final String JNDI_DNS_OBJECT_FACTORY = "com.sun.jndi.url.dns.dnsURLContextFactory";
 
   @Override
   public void beforeAnalysis(BeforeAnalysisAccess access) {
@@ -55,6 +57,10 @@ public void beforeAnalysis(BeforeAnalysisAccess access) {
     // The Core Cloud SQL Socket
     NativeImageUtils.registerClassForReflection(access, CLOUD_SQL_SOCKET_CLASS);
 
+    // The JNDI DNS factory for looking up DNS names.
+    NativeImageUtils.registerClassForReflection(access, JNDI_DNS_FACTORY);
+    NativeImageUtils.registerClassForReflection(access, JNDI_DNS_OBJECT_FACTORY);
+
     // Register Hikari configs if used with Cloud SQL.
     if (access.findClassByName("com.zaxxer.hikari.HikariConfig") != null) {
       NativeImageUtils.registerClassForReflection(access, "com.zaxxer.hikari.HikariConfig");
diff --git a/core/src/main/resources/META-INF/native-image/com.google.cloud.sql/cloud-sql-jdbc-socket-factory-parent/proxy-config.json b/core/src/main/resources/META-INF/native-image/com.google.cloud.sql/cloud-sql-jdbc-socket-factory-parent/proxy-config.json
index 45c06e418..209e1620f 100644
--- a/core/src/main/resources/META-INF/native-image/com.google.cloud.sql/cloud-sql-jdbc-socket-factory-parent/proxy-config.json
+++ b/core/src/main/resources/META-INF/native-image/com.google.cloud.sql/cloud-sql-jdbc-socket-factory-parent/proxy-config.json
@@ -18,5 +18,17 @@
       "jnr.enxio.channels.Native$LibC",
       "jnr.ffi.provider.LoadedLibrary"
     ]
+  },
+  {
+    "name": "com.sun.jndi.dns.DnsContextFactory",
+    "interfaces": [
+      "javax.naming.spi.InitialContextFactory"
+    ]
+  },
+  {
+    "name": "com.sun.jndi.url.dns.dnsURLContextFactory",
+    "interfaces": [
+      "javax.naming.spi.ObjectFactory"
+    ]
   }
 ]