From 805d10346af1bdfb233f867bf45a1e601e8c334d Mon Sep 17 00:00:00 2001
From: yimiaoxiehou <yimiaoxiehou@gmail.com>
Date: Mon, 14 Sep 2026 18:31:27 +0800
Subject: [PATCH] feat(elasticsearch): support custom CA and mutual TLS

---
 .../connection/ConnectionDialog.vue           |  24 ++-
 crates/dbx-core/src/connection.rs             |  10 +-
 crates/dbx-core/src/db/easysearch_driver.rs   |  22 ++-
 .../dbx-core/src/db/elasticsearch_driver.rs   | 114 ++++++++++--
 crates/dbx-core/src/db/mod.rs                 | 163 ++++++++++++++++++
 src-tauri/src/commands/connection.rs          |  20 ++-
 6 files changed, 327 insertions(+), 26 deletions(-)

diff --git a/apps/desktop/src/components/connection/ConnectionDialog.vue b/apps/desktop/src/components/connection/ConnectionDialog.vue
index badd87f8c0..c7494b0599 100644
--- a/apps/desktop/src/components/connection/ConnectionDialog.vue
+++ b/apps/desktop/src/components/connection/ConnectionDialog.vue
@@ -4402,12 +4402,30 @@ function connectionConfigForSubmit(id: string, generatedName = "", validatePlugi
     if ((config.client_cert_path && !config.client_key_path) || (!config.client_cert_path && config.client_key_path)) {
       throw new Error(t("connection.etcdClientCertPairRequired"));
     }
-  } else if (form.value.db_type !== "consul" && config.db_type !== "zookeeper") {
+  } else if (form.value.db_type !== "consul" && config.db_type !== "zookeeper" && config.db_type !== "elasticsearch" && config.db_type !== "easysearch") {
     config.etcd_endpoints = undefined;
     config.client_cert_path = undefined;
     config.client_key_path = undefined;
   }
-  if (config.db_type !== "mysql" && config.db_type !== "clickhouse" && config.db_type !== "etcd" && config.db_type !== "consul" && config.db_type !== "starrocks" && config.db_type !== "mongodb" && config.db_type !== "victoriametrics" && config.db_type !== "zookeeper") {
+  if (config.db_type === "elasticsearch" || config.db_type === "easysearch") {
+    config.client_cert_path = config.client_cert_path?.trim() || "";
+    config.client_key_path = config.client_key_path?.trim() || "";
+    if ((config.client_cert_path && !config.client_key_path) || (!config.client_cert_path && config.client_key_path)) {
+      throw new Error(t("connection.etcdClientCertPairRequired"));
+    }
+  }
+  if (
+    config.db_type !== "mysql" &&
+    config.db_type !== "clickhouse" &&
+    config.db_type !== "etcd" &&
+    config.db_type !== "consul" &&
+    config.db_type !== "starrocks" &&
+    config.db_type !== "mongodb" &&
+    config.db_type !== "victoriametrics" &&
+    config.db_type !== "zookeeper" &&
+    config.db_type !== "elasticsearch" &&
+    config.db_type !== "easysearch"
+  ) {
     config.ca_cert_path = undefined;
   } else {
     config.ca_cert_path = config.ca_cert_path?.trim() || "";
@@ -8344,7 +8362,7 @@ function openExternalUrl(url: string) {
                   </label>
                 </div>
 
-                <template v-if="form.db_type === 'etcd' || form.db_type === 'consul' || form.db_type === 'zookeeper'">
+                <template v-if="form.db_type === 'etcd' || form.db_type === 'consul' || form.db_type === 'zookeeper' || form.db_type === 'elasticsearch' || form.db_type === 'easysearch'">
                   <div class="grid grid-cols-4 items-start gap-4">
                     <Label :class="connectionLabelSmallPaddedClass">
                       <span class="inline-flex items-center justify-end gap-1">
diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs
index 7695f1900c..a92d9633e4 100644
--- a/crates/dbx-core/src/connection.rs
+++ b/crates/dbx-core/src/connection.rs
@@ -2468,7 +2468,10 @@ impl AppState {
                     db_config.url_params.as_deref(),
                     db_config.external_config.as_ref(),
                     connect_timeout,
-                );
+                    Some(db_config.ca_cert_path.as_str()),
+                    Some(db_config.client_cert_path.as_str()),
+                    Some(db_config.client_key_path.as_str()),
+                )?;
                 db::elasticsearch_driver::test_connection(&mut client, connect_timeout).await?;
                 PoolKind::Elasticsearch(client)
             }
@@ -2481,7 +2484,10 @@ impl AppState {
                     db_config.url_params.as_deref(),
                     db_config.external_config.as_ref(),
                     connect_timeout,
-                );
+                    Some(db_config.ca_cert_path.as_str()),
+                    Some(db_config.client_cert_path.as_str()),
+                    Some(db_config.client_key_path.as_str()),
+                )?;
                 db::easysearch_driver::test_connection(&mut client, connect_timeout).await?;
                 PoolKind::Easysearch(client)
             }
diff --git a/crates/dbx-core/src/db/easysearch_driver.rs b/crates/dbx-core/src/db/easysearch_driver.rs
index 79527c60f8..382438c3b2 100644
--- a/crates/dbx-core/src/db/easysearch_driver.rs
+++ b/crates/dbx-core/src/db/easysearch_driver.rs
@@ -20,10 +20,24 @@ impl EasysearchClient {
         url_params: Option<&str>,
         external_config: Option<&Value>,
         timeout: Duration,
-    ) -> Self {
-        Self {
-            inner: EsClient::from_config(url, username, password, tls_enabled, url_params, external_config, timeout),
-        }
+        ca_cert_path: Option<&str>,
+        client_cert_path: Option<&str>,
+        client_key_path: Option<&str>,
+    ) -> Result<Self, String> {
+        Ok(Self {
+            inner: EsClient::from_config(
+                url,
+                username,
+                password,
+                tls_enabled,
+                url_params,
+                external_config,
+                timeout,
+                ca_cert_path,
+                client_cert_path,
+                client_key_path,
+            )?,
+        })
     }
 }
 
diff --git a/crates/dbx-core/src/db/elasticsearch_driver.rs b/crates/dbx-core/src/db/elasticsearch_driver.rs
index bf61e2b330..54acfce871 100644
--- a/crates/dbx-core/src/db/elasticsearch_driver.rs
+++ b/crates/dbx-core/src/db/elasticsearch_driver.rs
@@ -10,7 +10,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
-use super::{http_client_builder, with_connection_timeout};
+use super::{apply_tls_certificates, http_client_builder, with_connection_timeout};
 use crate::db::document_result::DocumentQueryResult;
 use crate::types::QueryResult;
 
@@ -83,9 +83,14 @@ impl EsClient {
             "/".to_string(),
             false,
             None,
+            None,
+            None,
+            None,
         )
+        .expect("failed to build Elasticsearch HTTP client")
     }
 
+    #[allow(clippy::too_many_arguments)]
     fn new_with_mode(
         url: &str,
         username: Option<&str>,
@@ -96,19 +101,24 @@ impl EsClient {
         connectivity_check_path: String,
         connectivity_check_disabled: bool,
         index_grouping: Option<Regex>,
-    ) -> Self {
+        ca_cert_path: Option<&str>,
+        client_cert_path: Option<&str>,
+        client_key_path: Option<&str>,
+    ) -> Result<Self, String> {
         let base_url = url.trim_end_matches('/').to_string();
         let auth = match (username, password) {
             (Some(u), Some(p)) if !u.is_empty() => Some((u.to_string(), p.to_string())),
             _ => None,
         };
         let mut builder = http_client_builder(timeout).danger_accept_invalid_certs(accept_invalid_certs);
+        builder = apply_tls_certificates(builder, ca_cert_path, client_cert_path, client_key_path)?;
         if let Some(addrs) = elasticsearch_localhost_resolve_addrs(&base_url, connectivity_check_disabled) {
             builder = builder.resolve_to_addrs("localhost", &addrs);
         }
-        let http = builder.build().unwrap_or_else(|_| HttpClient::new());
+        let http =
+            builder.build().map_err(|error| format!("Failed to initialize Elasticsearch HTTP client: {error}"))?;
         let fallback_base_urls = elasticsearch_base_url_fallbacks(&base_url);
-        Self {
+        Ok(Self {
             http,
             base_url,
             fallback_base_urls,
@@ -118,7 +128,7 @@ impl EsClient {
             connectivity_check_disabled,
             index_grouping,
             pit_search_supported: Arc::new(AtomicBool::new(true)),
-        }
+        })
     }
 
     pub fn from_config(
@@ -129,7 +139,10 @@ impl EsClient {
         url_params: Option<&str>,
         external_config: Option<&Value>,
         timeout: Duration,
-    ) -> Self {
+        ca_cert_path: Option<&str>,
+        client_cert_path: Option<&str>,
+        client_key_path: Option<&str>,
+    ) -> Result<Self, String> {
         let kibana_base_path = elasticsearch_kibana_base_path(external_config);
         let transport_mode = if kibana_base_path.is_some() {
             ElasticsearchTransportMode::KibanaProxy
@@ -150,6 +163,9 @@ impl EsClient {
             connectivity_check_path,
             connectivity_check_disabled,
             index_grouping,
+            ca_cert_path,
+            client_cert_path,
+            client_key_path,
         )
     }
 
@@ -3536,7 +3552,11 @@ mod tests {
             None,
             Some(&json!({ "connectivityCheckDisabled": "yes" })),
             Duration::from_secs(2),
-        );
+            None,
+            None,
+            None,
+        )
+        .expect("failed to build Elasticsearch test client");
         super::test_connection(&mut client, Duration::from_secs(2)).await.unwrap();
         let response = client.get("/_cluster/health").send().await.unwrap();
 
@@ -3554,7 +3574,11 @@ mod tests {
             Some("sslmode=disable"),
             None,
             Duration::from_secs(1),
-        );
+            None,
+            None,
+            None,
+        )
+        .expect("failed to build Elasticsearch test client");
 
         assert_eq!(client.base_url, "https://localhost:9200");
         assert_eq!(client.fallback_base_urls, vec!["https://127.0.0.1:9200"]);
@@ -3589,7 +3613,11 @@ mod tests {
                 "connectivityCheckPath": "GET pro-logs-*/_search"
             })),
             Duration::from_secs(1),
-        );
+            None,
+            None,
+            None,
+        )
+        .expect("failed to build Elasticsearch test client");
         assert_eq!(client.connectivity_check_path, "/pro-logs-*/_search");
     }
 
@@ -3622,7 +3650,11 @@ mod tests {
             None,
             Some(&json!({ "connectivityCheckPath": "/pro-logs-*/_search" })),
             Duration::from_secs(2),
-        );
+            None,
+            None,
+            None,
+        )
+        .expect("failed to build Elasticsearch test client");
         super::test_connection(&mut client, Duration::from_secs(2)).await.unwrap();
         server.await.unwrap();
     }
@@ -3638,7 +3670,11 @@ mod tests {
             None,
             Some(&external_config),
             Duration::from_secs(1),
-        );
+            None,
+            None,
+            None,
+        )
+        .expect("failed to build Elasticsearch test client");
 
         assert_eq!(client.base_url, "https://localhost:5601/kibana/s/analytics");
         assert_eq!(client.fallback_base_urls, vec!["https://127.0.0.1:5601/kibana/s/analytics"]);
@@ -5146,7 +5182,11 @@ mod tests {
             None,
             Some(&external_config),
             Duration::from_secs(1),
-        );
+            None,
+            None,
+            None,
+        )
+        .expect("failed to build Elasticsearch test client");
         let result =
             super::execute_rest_query(&client, "DELETE /missing/_doc/1?refresh=true\n{\"reason\":\"cleanup\"}")
                 .await
@@ -5499,7 +5539,11 @@ mod tests {
             None,
             Some(&external_config),
             Duration::from_secs(20),
-        );
+            None,
+            None,
+            None,
+        )
+        .expect("failed to build Elasticsearch test client");
         super::test_connection(&mut client, Duration::from_secs(20)).await.unwrap();
 
         let index = format!("dbx-kibana-proxy-{}", uuid::Uuid::new_v4().simple());
@@ -5590,4 +5634,48 @@ mod tests {
 
         assert_eq!(serde_json::to_string(&doc).unwrap(), r#"{"z":1,"a":2}"#);
     }
+
+    // Integration test against a real TLS-enabled Elasticsearch. Gated by env vars
+    // so it stays a no-op in CI:
+    //   DBX_ELASTICSEARCH_TLS_TEST_URL            e.g. https://localhost:9200
+    //   DBX_ELASTICSEARCH_TLS_TEST_CA_CERT_PATH   path to the server's CA cert (PEM)
+    // Mirrors the ZooKeeper TLS integration test: tls_enabled=true loads the custom
+    // CA (and any client cert) through apply_tls_certificates, then connects over HTTPS.
+    #[tokio::test]
+    async fn tls_integration_connects_with_custom_ca() {
+        let url = match std::env::var("DBX_ELASTICSEARCH_TLS_TEST_URL") {
+            Ok(u) => u,
+            Err(_) => {
+                eprintln!("skipping ES TLS integration test: DBX_ELASTICSEARCH_TLS_TEST_URL not set");
+                return;
+            }
+        };
+        let ca_cert_path = match std::env::var("DBX_ELASTICSEARCH_TLS_TEST_CA_CERT_PATH") {
+            Ok(c) => c,
+            Err(_) => {
+                eprintln!("skipping ES TLS integration test: DBX_ELASTICSEARCH_TLS_TEST_CA_CERT_PATH not set");
+                return;
+            }
+        };
+        // Optional basic-auth credentials (required when the server has security enabled).
+        let username = std::env::var("DBX_ELASTICSEARCH_TLS_TEST_USERNAME").ok();
+        let password = std::env::var("DBX_ELASTICSEARCH_TLS_TEST_PASSWORD").ok();
+        let timeout = std::time::Duration::from_secs(15);
+        let mut client = EsClient::from_config(
+            &url,
+            username.as_deref(),
+            password.as_deref(),
+            true,
+            None,
+            None,
+            timeout,
+            Some(&ca_cert_path),
+            None,
+            None,
+        )
+        .expect("failed to build Elasticsearch TLS client");
+        super::test_connection(&mut client, timeout)
+            .await
+            .expect("Elasticsearch TLS connection with custom CA should succeed");
+    }
 }
diff --git a/crates/dbx-core/src/db/mod.rs b/crates/dbx-core/src/db/mod.rs
index 62aabbbb9a..c788de078b 100644
--- a/crates/dbx-core/src/db/mod.rs
+++ b/crates/dbx-core/src/db/mod.rs
@@ -54,7 +54,9 @@ pub mod vector_driver;
 pub mod victoriametrics_driver;
 pub mod wkb;
 
+use crate::path_utils::expand_tilde;
 use reqwest::ClientBuilder;
+use reqwest::{Certificate, Identity};
 use std::fmt;
 use std::future::Future;
 use std::time::Duration;
@@ -131,6 +133,64 @@ pub fn http_client_builder(timeout: Duration) -> ClientBuilder {
     reqwest::Client::builder().connect_timeout(timeout).no_proxy()
 }
 
+/// Augments an HTTP client builder with custom TLS material for drivers that
+/// need to trust a non-system CA or present a client certificate (mTLS).
+///
+/// * `ca_cert_path` adds an extra root certificate (PEM bundle) alongside the
+///   system roots, so self-signed server certificates are trusted.
+/// * `client_cert_path` + `client_key_path` (required together) are concatenated
+///   into a single PEM identity, as expected by reqwest's rustls backend.
+///
+/// Returns the augmented builder, or a descriptive error if a configured
+/// certificate/key file cannot be read or parsed. Callers surface that error
+/// instead of silently falling back to the system trust store, so a typo in a
+/// cert path fails the connection loudly rather than being ignored.
+pub(crate) fn apply_tls_certificates(
+    mut builder: ClientBuilder,
+    ca_cert_path: Option<&str>,
+    client_cert_path: Option<&str>,
+    client_key_path: Option<&str>,
+) -> Result<ClientBuilder, String> {
+    // The desktop build unifies reqwest with both `default-tls` and `rustls-tls`,
+    // and `TlsBackend` defaults to native-tls there, which rejects the PEM identity
+    // below with `incompatible TLS identity type`. Pin rustls like the Consul and
+    // etcd-metrics clients so custom CA roots and PEM client identities work in
+    // every target.
+    builder = builder.use_rustls_tls();
+    if let Some(ca) = ca_cert_path.filter(|path| !path.is_empty()) {
+        let ca = expand_tilde(ca);
+        let contents = std::fs::read(&ca)
+            .map_err(|error| format!("Failed to read Elasticsearch CA certificate at {ca}: {error}"))?;
+        let certificates = Certificate::from_pem_bundle(&contents)
+            .map_err(|error| format!("Failed to parse Elasticsearch CA certificate at {ca}: {error}"))?;
+        for certificate in certificates {
+            builder = builder.add_root_certificate(certificate);
+        }
+    }
+    if let (Some(cert_path), Some(key_path)) =
+        (client_cert_path.filter(|path| !path.is_empty()), client_key_path.filter(|path| !path.is_empty()))
+    {
+        let cert_path = expand_tilde(cert_path);
+        let key_path = expand_tilde(key_path);
+        let mut pem = std::fs::read(&cert_path)
+            .map_err(|error| format!("Failed to read Elasticsearch client certificate at {cert_path}: {error}"))?;
+        // reqwest's rustls backend concatenates the cert and key as two PEM
+        // blocks; ensure they are newline-separated so parsing succeeds even
+        // when the cert file does not end with a trailing newline.
+        if !pem.ends_with(b"\n") {
+            pem.push(b'\n');
+        }
+        pem.extend(
+            std::fs::read(&key_path)
+                .map_err(|error| format!("Failed to read Elasticsearch client key at {key_path}: {error}"))?,
+        );
+        let identity = Identity::from_pem(&pem)
+            .map_err(|error| format!("Failed to parse Elasticsearch client identity: {error}"))?;
+        builder = builder.identity(identity);
+    }
+    Ok(builder)
+}
+
 pub(crate) const JS_MAX_SAFE_INTEGER: i64 = 9_007_199_254_740_991;
 
 pub fn safe_i64_to_json(v: i64) -> serde_json::Value {
@@ -294,3 +354,106 @@ mod tests {
         assert_eq!(json_value_for_js(value), expected);
     }
 }
+
+#[cfg(test)]
+mod tls_tests {
+    use super::{apply_tls_certificates, http_client_builder};
+    use std::path::Path;
+    use std::process::Command;
+    use std::time::Duration;
+
+    /// Generates a throwaway CA + client certificate/key (PEM) via openssl and
+    /// returns their paths. If openssl is unavailable the files stay empty, which
+    /// makes `apply_tls_certificates` surface a read error instead of degrading.
+    fn gen_certs(dir: &Path) -> (std::path::PathBuf, std::path::PathBuf, std::path::PathBuf) {
+        let ca_key = dir.join("ca.key");
+        let ca_pem = dir.join("ca.pem");
+        let csr = dir.join("client.csr");
+        let client_key = dir.join("client.key");
+        let client_crt = dir.join("client.crt");
+        // Extensions force an X.509 v3 client cert; some toolchains (e.g.
+        // OpenSSL 3.0 on CI) emit v1 without one, which rustls rejects.
+        let client_ext = dir.join("client.ext");
+        let _ = std::fs::write(&client_ext, "basicConstraints=critical,CA:FALSE\n");
+        let _ = Command::new("openssl")
+            .args([
+                "req",
+                "-x509",
+                "-newkey",
+                "rsa:2048",
+                "-nodes",
+                "-keyout",
+                ca_key.to_str().unwrap(),
+                "-out",
+                ca_pem.to_str().unwrap(),
+                "-days",
+                "1",
+                "-subj",
+                "/CN=dbx-test-ca",
+            ])
+            .output();
+        let _ = Command::new("openssl")
+            .args([
+                "req",
+                "-newkey",
+                "rsa:2048",
+                "-nodes",
+                "-keyout",
+                client_key.to_str().unwrap(),
+                "-out",
+                csr.to_str().unwrap(),
+                "-subj",
+                "/CN=dbx-test-client",
+            ])
+            .output();
+        let _ = Command::new("openssl")
+            .args([
+                "x509",
+                "-req",
+                "-in",
+                csr.to_str().unwrap(),
+                "-CA",
+                ca_pem.to_str().unwrap(),
+                "-CAkey",
+                ca_key.to_str().unwrap(),
+                "-CAcreateserial",
+                "-out",
+                client_crt.to_str().unwrap(),
+                "-days",
+                "1",
+                "-extfile",
+                client_ext.to_str().unwrap(),
+            ])
+            .output();
+        (ca_pem, client_crt, client_key)
+    }
+
+    #[test]
+    fn apply_tls_certificates_builds_with_ca_and_client_cert() {
+        let dir = std::env::temp_dir().join(format!("dbx-tls-test-{}", std::process::id()));
+        std::fs::create_dir_all(&dir).unwrap();
+        let (ca, client_crt, client_key) = gen_certs(&dir);
+
+        let built = apply_tls_certificates(
+            http_client_builder(Duration::from_secs(5)),
+            ca.to_str(),
+            client_crt.to_str(),
+            client_key.to_str(),
+        )
+        .expect("client should build with CA + client cert")
+        .build();
+        assert!(built.is_ok(), "built client should be usable with CA + client cert: {:?}", built.err());
+
+        // A misconfigured CA path must now surface a descriptive error instead of
+        // silently degrading to the system trust store.
+        let missing = apply_tls_certificates(
+            http_client_builder(Duration::from_secs(5)),
+            Some("/nonexistent/ca.pem"),
+            None,
+            None,
+        );
+        assert!(missing.is_err(), "missing CA path must surface a read error");
+
+        let _ = std::fs::remove_dir_all(&dir);
+    }
+}
diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs
index 5f5cb23814..d1fdf57cd6 100644
--- a/src-tauri/src/commands/connection.rs
+++ b/src-tauri/src/commands/connection.rs
@@ -1409,7 +1409,10 @@ async fn test_connection_with_info_inner(
                     config.url_params.as_deref(),
                     config.external_config.as_ref(),
                     connect_timeout,
-                );
+                    Some(config.ca_cert_path.as_str()),
+                    Some(config.client_cert_path.as_str()),
+                    Some(config.client_key_path.as_str()),
+                )?;
                 db::elasticsearch_driver::test_connection(&mut client, connect_timeout)
                     .await
                     .map(|_| "Connection successful".to_string())
@@ -1423,7 +1426,10 @@ async fn test_connection_with_info_inner(
                     config.url_params.as_deref(),
                     config.external_config.as_ref(),
                     connect_timeout,
-                );
+                    Some(config.ca_cert_path.as_str()),
+                    Some(config.client_cert_path.as_str()),
+                    Some(config.client_key_path.as_str()),
+                )?;
                 db::easysearch_driver::test_connection(&mut client, connect_timeout)
                     .await
                     .map(|_| "Connection successful".to_string())
@@ -1854,7 +1860,10 @@ pub async fn connect_db(
                 db_config.url_params.as_deref(),
                 db_config.external_config.as_ref(),
                 connect_timeout,
-            );
+                Some(db_config.ca_cert_path.as_str()),
+                Some(db_config.client_cert_path.as_str()),
+                Some(db_config.client_key_path.as_str()),
+            )?;
             db::elasticsearch_driver::test_connection(&mut client, connect_timeout).await?;
             PoolKind::Elasticsearch(client)
         }
@@ -1867,7 +1876,10 @@ pub async fn connect_db(
                 db_config.url_params.as_deref(),
                 db_config.external_config.as_ref(),
                 connect_timeout,
-            );
+                Some(db_config.ca_cert_path.as_str()),
+                Some(db_config.client_cert_path.as_str()),
+                Some(db_config.client_key_path.as_str()),
+            )?;
             db::easysearch_driver::test_connection(&mut client, connect_timeout).await?;
             PoolKind::Easysearch(client)
         }