From aa12c66c25e68eefec22c213e8f228fd0699d8ce Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@windmill.dev>
Date: Wed, 20 May 2026 06:48:10 +0000
Subject: [PATCH] feat(snowflake): derive public key from private key when
 omitted (WIN-1959) (#9251)

* feat(snowflake): derive public key from private key when omitted (WIN-1959)

Snowflake key-pair auth needs a SHA256 fingerprint of the public key for
the JWT iss claim, but the public key is mathematically derivable from
the RSA private key. Other tools (e.g. Power BI) only require the
private key, so requiring users to supply both is redundant. When
public_key is missing, fall back to deriving it from private_key (PKCS#8
or PKCS#1 PEM) instead of erroring out.

Fixes WIN-1959

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(snowflake): treat empty public_key/private_key as missing

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 backend/Cargo.lock                            |  1 +
 backend/windmill-worker/Cargo.toml            |  3 +-
 .../windmill-worker/src/snowflake_executor.rs | 48 ++++++++++++++++---
 3 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/backend/Cargo.lock b/backend/Cargo.lock
index 1a3efc3e11..b8b0ab8625 100644
--- a/backend/Cargo.lock
+++ b/backend/Cargo.lock
@@ -15667,6 +15667,7 @@ dependencies = [
  "regex",
  "reqwest 0.13.1",
  "reqwest-middleware",
+ "rsa",
  "rust_decimal",
  "serde",
  "serde_json",
diff --git a/backend/windmill-worker/Cargo.toml b/backend/windmill-worker/Cargo.toml
index 5ebcb0b370..2e75ac3632 100644
--- a/backend/windmill-worker/Cargo.toml
+++ b/backend/windmill-worker/Cargo.toml
@@ -13,7 +13,7 @@ default = []
 private = ["windmill-worker-volumes/private", "windmill-queue/private", "windmill-common/private", "windmill-dep-map/private", "windmill-runtime-nativets?/private"]
 mcp = ["windmill-ai/mcp", "dep:windmill-mcp"]
 prometheus = ["dep:prometheus", "windmill-common/prometheus"]
-enterprise = ["windmill-queue/enterprise", "windmill-git-sync/enterprise", "windmill-common/enterprise", "windmill-worker-volumes/enterprise", "windmill-runtime-nativets?/enterprise", "dep:pem", "dep:tokio-util", "dep:opentelemetry-proto", "dep:prost", "dep:hudsucker", "dep:rcgen", "dep:hyper-http-proxy", "dep:hyper-tls", "dep:hyper-util"]
+enterprise = ["windmill-queue/enterprise", "windmill-git-sync/enterprise", "windmill-common/enterprise", "windmill-worker-volumes/enterprise", "windmill-runtime-nativets?/enterprise", "dep:pem", "dep:rsa", "dep:tokio-util", "dep:opentelemetry-proto", "dep:prost", "dep:hudsucker", "dep:rcgen", "dep:hyper-http-proxy", "dep:hyper-tls", "dep:hyper-util"]
 mssql = ["dep:tiberius"]
 mssql-kerberos = ["mssql", "tiberius/integrated-auth-gssapi"]  # Linux/Unix integrated auth
 mssql-winauth = ["mssql", "tiberius/winauth"]  # Windows integrated auth
@@ -112,6 +112,7 @@ jsonwebtoken.workspace = true
 sha2.workspace = true
 hmac.workspace = true
 pem = { workspace = true, optional = true }
+rsa = { workspace = true, optional = true }
 urlencoding.workspace = true
 nix.workspace = true
 bytes.workspace = true
diff --git a/backend/windmill-worker/src/snowflake_executor.rs b/backend/windmill-worker/src/snowflake_executor.rs
index 839609d9da..35051aee98 100644
--- a/backend/windmill-worker/src/snowflake_executor.rs
+++ b/backend/windmill-worker/src/snowflake_executor.rs
@@ -630,14 +630,50 @@ pub async fn do_snowflake(
         )
         .to_uppercase();
 
-        let public_key = match database.public_key.as_deref() {
-            Some(key) => pem::parse(key.as_bytes()).map_err(|e| {
-                Error::ExecutionErr(format!("Failed to parse public key: {}", e.to_string()))
-            })?,
-            None => return Err(Error::ExecutionErr("Public key is missing".to_string())),
+        let public_key_der: Vec<u8> = match database
+            .public_key
+            .as_deref()
+            .map(str::trim)
+            .filter(|s| !s.is_empty())
+        {
+            Some(key) => pem::parse(key.as_bytes())
+                .map_err(|e| Error::ExecutionErr(format!("Failed to parse public key: {e}")))?
+                .into_contents(),
+            None => {
+                // Derive the public key from the private key — RSA private keys
+                // contain the public components (n, e).
+                use rsa::pkcs8::{DecodePrivateKey, EncodePublicKey};
+                let pk_pem = database
+                    .private_key
+                    .as_deref()
+                    .map(str::trim)
+                    .filter(|s| !s.is_empty())
+                    .ok_or_else(|| {
+                        Error::ExecutionErr(
+                            "Either public_key or private_key must be provided".to_string(),
+                        )
+                    })?;
+                let rsa_priv = rsa::RsaPrivateKey::from_pkcs8_pem(pk_pem)
+                    .or_else(|_| {
+                        use rsa::pkcs1::DecodeRsaPrivateKey;
+                        rsa::RsaPrivateKey::from_pkcs1_pem(pk_pem)
+                    })
+                    .map_err(|e| {
+                        Error::ExecutionErr(format!(
+                            "Failed to parse private key to derive public key: {e}"
+                        ))
+                    })?;
+                let rsa_pub = rsa::RsaPublicKey::from(&rsa_priv);
+                rsa_pub
+                    .to_public_key_der()
+                    .map_err(|e| {
+                        Error::ExecutionErr(format!("Failed to encode derived public key: {e}"))
+                    })?
+                    .to_vec()
+            }
         };
         let mut public_key_hash = Sha256::new();
-        public_key_hash.update(public_key.contents());
+        public_key_hash.update(&public_key_der);
 
         let public_key_fp = engine::general_purpose::STANDARD.encode(public_key_hash.finalize());