Merge pull request #341 from cipherstash/encryption-sanity-checks

tobyhede · web-flow · commit f843f1f2b2ff · 2025-12-17T08:49:57.000+11:00
Encryption sanity checks
diff --git a/packages/cipherstash-proxy-integration/src/common.rs b/packages/cipherstash-proxy-integration/src/common.rs
@@ -306,6 +306,9 @@ pub async fn insert_jsonb() -> Value {
 
     insert(&sql, &[&id, &encrypted_jsonb]).await;
 
+    // Verify encryption actually occurred
+    assert_encrypted_jsonb(id, &encrypted_jsonb).await;
+
     encrypted_jsonb
 }
 
@@ -321,9 +324,72 @@ pub async fn insert_jsonb_for_search() {
 
         let sql = "INSERT INTO encrypted (id, encrypted_jsonb) VALUES ($1, $2)";
         insert(sql, &[&id, &encrypted_jsonb]).await;
+
+        // Verify encryption actually occurred for each row
+        assert_encrypted_jsonb(id, &encrypted_jsonb).await;
+    }
+}
+
+/// Verifies that a text value was actually encrypted in the database.
+/// Queries directly (bypassing proxy) and asserts stored value differs from plaintext.
+pub async fn assert_encrypted_text(id: i64, column: &str, plaintext: &str) {
+    let sql = format!("SELECT {}::text FROM encrypted WHERE id = $1", column);
+    let stored: Vec<String> = query_direct_by(&sql, &id).await;
+
+    assert_eq!(stored.len(), 1, "Expected exactly one row");
+    let stored_text = &stored[0];
+
+    assert_ne!(
+        stored_text, plaintext,
+        "ENCRYPTION FAILED for {}: Stored value matches plaintext! Data was not encrypted.",
+        column
+    );
+}
+
+/// Verifies that a JSONB value was actually encrypted in the database.
+/// Queries directly (bypassing proxy) and asserts stored value differs from plaintext.
+pub async fn assert_encrypted_jsonb(id: i64, plaintext: &Value) {
+    let sql = "SELECT encrypted_jsonb::text FROM encrypted WHERE id = $1";
+    let stored: Vec<String> = query_direct_by(sql, &id).await;
+
+    assert_eq!(stored.len(), 1, "Expected exactly one row");
+    let stored_text = &stored[0];
+
+    let plaintext_str = plaintext.to_string();
+    assert_ne!(
+        stored_text, &plaintext_str,
+        "ENCRYPTION FAILED for encrypted_jsonb: Stored value matches plaintext! Data was not encrypted."
+    );
+
+    // Additional verification: the encrypted format should be different structure
+    if let Ok(stored_json) = serde_json::from_str::<Value>(stored_text) {
+        assert_ne!(
+            stored_json, *plaintext,
+            "ENCRYPTION FAILED for encrypted_jsonb: Stored JSON structure matches plaintext!"
+        );
     }
 }
 
+/// Verifies that a numeric value was actually encrypted in the database.
+/// Queries directly (bypassing proxy) and asserts stored value differs from plaintext.
+pub async fn assert_encrypted_numeric<T>(id: i64, column: &str, plaintext: T)
+where
+    T: std::fmt::Display + std::str::FromStr + PartialEq,
+{
+    let sql = format!("SELECT {}::text FROM encrypted WHERE id = $1", column);
+    let stored: Vec<String> = query_direct_by(&sql, &id).await;
+
+    assert_eq!(stored.len(), 1, "Expected exactly one row");
+    let stored_text = &stored[0];
+
+    let plaintext_str = plaintext.to_string();
+    assert_ne!(
+        stored_text, &plaintext_str,
+        "ENCRYPTION FAILED for {}: Stored value matches plaintext! Data was not encrypted.",
+        column
+    );
+}
+
 ///
 /// Configure the client TLS settings.
 /// These are the settings for connecting to the database with TLS.
diff --git a/packages/cipherstash-proxy-integration/src/encryption_sanity.rs b/packages/cipherstash-proxy-integration/src/encryption_sanity.rs
@@ -0,0 +1,231 @@
+//! Encryption sanity checks - verify data is actually encrypted.
+//!
+//! These tests insert data through the proxy, then query DIRECTLY from the database
+//! (bypassing the proxy) to verify the stored value is encrypted (differs from plaintext).
+//!
+//! This catches silent mapping failures where data passes through unencrypted.
+
+#[cfg(test)]
+mod tests {
+    use crate::common::{
+        assert_encrypted_jsonb, assert_encrypted_numeric, assert_encrypted_text, clear,
+        connect_with_tls, random_id, trace, PROXY,
+    };
+    use chrono::NaiveDate;
+
+    #[tokio::test]
+    async fn text_encryption_sanity_check() {
+        trace();
+        clear().await;
+
+        let id = random_id();
+        let plaintext = "hello world";
+
+        // Insert through proxy (should encrypt)
+        let client = connect_with_tls(PROXY).await;
+        let sql = "INSERT INTO encrypted (id, encrypted_text) VALUES ($1, $2)";
+        client.query(sql, &[&id, &plaintext]).await.unwrap();
+
+        // Verify encryption occurred
+        assert_encrypted_text(id, "encrypted_text", plaintext).await;
+
+        // Round-trip: query through proxy should decrypt back to original
+        let sql = "SELECT encrypted_text FROM encrypted WHERE id = $1";
+        let rows = client.query(sql, &[&id]).await.unwrap();
+        assert_eq!(rows.len(), 1, "Expected exactly one row for round-trip");
+        let decrypted: String = rows[0].get(0);
+        assert_eq!(
+            decrypted, plaintext,
+            "DECRYPTION FAILED: Round-trip value doesn't match original!"
+        );
+    }
+
+    #[tokio::test]
+    async fn jsonb_encryption_sanity_check() {
+        trace();
+        clear().await;
+
+        let id = random_id();
+        let plaintext_json = serde_json::json!({"key": "value", "number": 42});
+
+        // Insert through proxy (should encrypt)
+        let client = connect_with_tls(PROXY).await;
+        let sql = "INSERT INTO encrypted (id, encrypted_jsonb) VALUES ($1, $2)";
+        client.query(sql, &[&id, &plaintext_json]).await.unwrap();
+
+        // Verify encryption occurred
+        assert_encrypted_jsonb(id, &plaintext_json).await;
+
+        // Round-trip: query through proxy should decrypt back to original
+        let sql = "SELECT encrypted_jsonb FROM encrypted WHERE id = $1";
+        let rows = client.query(sql, &[&id]).await.unwrap();
+        assert_eq!(rows.len(), 1, "Expected exactly one row for round-trip");
+        let decrypted: serde_json::Value = rows[0].get(0);
+        assert_eq!(
+            decrypted, plaintext_json,
+            "DECRYPTION FAILED: Round-trip value doesn't match original!"
+        );
+    }
+
+    #[tokio::test]
+    async fn float8_encryption_sanity_check() {
+        trace();
+        clear().await;
+
+        let id = random_id();
+        let plaintext: f64 = 123.456;
+
+        // Insert through proxy (should encrypt)
+        let client = connect_with_tls(PROXY).await;
+        let sql = "INSERT INTO encrypted (id, encrypted_float8) VALUES ($1, $2)";
+        client.query(sql, &[&id, &plaintext]).await.unwrap();
+
+        // Verify encryption occurred
+        assert_encrypted_numeric(id, "encrypted_float8", plaintext).await;
+
+        // Round-trip: query through proxy should decrypt back to original
+        let sql = "SELECT encrypted_float8 FROM encrypted WHERE id = $1";
+        let rows = client.query(sql, &[&id]).await.unwrap();
+        assert_eq!(rows.len(), 1, "Expected exactly one row for round-trip");
+        let decrypted: f64 = rows[0].get(0);
+        assert!(
+            (decrypted - plaintext).abs() < f64::EPSILON,
+            "DECRYPTION FAILED: Round-trip value doesn't match original!"
+        );
+    }
+
+    #[tokio::test]
+    async fn bool_encryption_sanity_check() {
+        trace();
+        clear().await;
+
+        let id = random_id();
+        let plaintext: bool = true;
+
+        // Insert through proxy (should encrypt)
+        let client = connect_with_tls(PROXY).await;
+        let sql = "INSERT INTO encrypted (id, encrypted_bool) VALUES ($1, $2)";
+        client.query(sql, &[&id, &plaintext]).await.unwrap();
+
+        // Verify encryption occurred
+        assert_encrypted_text(id, "encrypted_bool", &plaintext.to_string()).await;
+
+        // Round-trip: query through proxy should decrypt back to original
+        let sql = "SELECT encrypted_bool FROM encrypted WHERE id = $1";
+        let rows = client.query(sql, &[&id]).await.unwrap();
+        assert_eq!(rows.len(), 1, "Expected exactly one row for round-trip");
+        let decrypted: bool = rows[0].get(0);
+        assert_eq!(
+            decrypted, plaintext,
+            "DECRYPTION FAILED: Round-trip value doesn't match original!"
+        );
+    }
+
+    #[tokio::test]
+    async fn date_encryption_sanity_check() {
+        trace();
+        clear().await;
+
+        let id = random_id();
+        let plaintext = NaiveDate::from_ymd_opt(2024, 6, 15).unwrap();
+
+        // Insert through proxy (should encrypt)
+        let client = connect_with_tls(PROXY).await;
+        let sql = "INSERT INTO encrypted (id, encrypted_date) VALUES ($1, $2)";
+        client.query(sql, &[&id, &plaintext]).await.unwrap();
+
+        // Verify encryption occurred
+        assert_encrypted_text(id, "encrypted_date", &plaintext.to_string()).await;
+
+        // Round-trip: query through proxy should decrypt back to original
+        let sql = "SELECT encrypted_date FROM encrypted WHERE id = $1";
+        let rows = client.query(sql, &[&id]).await.unwrap();
+        assert_eq!(rows.len(), 1, "Expected exactly one row for round-trip");
+        let decrypted: NaiveDate = rows[0].get(0);
+        assert_eq!(
+            decrypted, plaintext,
+            "DECRYPTION FAILED: Round-trip value doesn't match original!"
+        );
+    }
+
+    #[tokio::test]
+    async fn int2_encryption_sanity_check() {
+        trace();
+        clear().await;
+
+        let id = random_id();
+        let plaintext: i16 = 42;
+
+        // Insert through proxy (should encrypt)
+        let client = connect_with_tls(PROXY).await;
+        let sql = "INSERT INTO encrypted (id, encrypted_int2) VALUES ($1, $2)";
+        client.query(sql, &[&id, &plaintext]).await.unwrap();
+
+        // Verify encryption occurred
+        assert_encrypted_numeric(id, "encrypted_int2", plaintext).await;
+
+        // Round-trip: query through proxy should decrypt back to original
+        let sql = "SELECT encrypted_int2 FROM encrypted WHERE id = $1";
+        let rows = client.query(sql, &[&id]).await.unwrap();
+        assert_eq!(rows.len(), 1, "Expected exactly one row for round-trip");
+        let decrypted: i16 = rows[0].get(0);
+        assert_eq!(
+            decrypted, plaintext,
+            "DECRYPTION FAILED: Round-trip value doesn't match original!"
+        );
+    }
+
+    #[tokio::test]
+    async fn int4_encryption_sanity_check() {
+        trace();
+        clear().await;
+
+        let id = random_id();
+        let plaintext: i32 = 12345;
+
+        // Insert through proxy (should encrypt)
+        let client = connect_with_tls(PROXY).await;
+        let sql = "INSERT INTO encrypted (id, encrypted_int4) VALUES ($1, $2)";
+        client.query(sql, &[&id, &plaintext]).await.unwrap();
+
+        // Verify encryption occurred
+        assert_encrypted_numeric(id, "encrypted_int4", plaintext).await;
+
+        // Round-trip: query through proxy should decrypt back to original
+        let sql = "SELECT encrypted_int4 FROM encrypted WHERE id = $1";
+        let rows = client.query(sql, &[&id]).await.unwrap();
+        assert_eq!(rows.len(), 1, "Expected exactly one row for round-trip");
+        let decrypted: i32 = rows[0].get(0);
+        assert_eq!(
+            decrypted, plaintext,
+            "DECRYPTION FAILED: Round-trip value doesn't match original!"
+        );
+    }
+
+    #[tokio::test]
+    async fn int8_encryption_sanity_check() {
+        trace();
+        clear().await;
+
+        let id = random_id();
+        let plaintext: i64 = 9876543210;
+
+        // Insert through proxy (should encrypt)
+        let client = connect_with_tls(PROXY).await;
+        let sql = "INSERT INTO encrypted (id, encrypted_int8) VALUES ($1, $2)";
+        client.query(sql, &[&id, &plaintext]).await.unwrap();
+
+        // Verify encryption occurred
+        assert_encrypted_numeric(id, "encrypted_int8", plaintext).await;
+
+        // Round-trip: query through proxy should decrypt back to original
+        let sql = "SELECT encrypted_int8 FROM encrypted WHERE id = $1";
+        let rows = client.query(sql, &[&id]).await.unwrap();
+        assert_eq!(rows.len(), 1, "Expected exactly one row for round-trip");
+        let decrypted: i64 = rows[0].get(0);
+        assert_eq!(
+            decrypted, plaintext,
+            "DECRYPTION FAILED: Round-trip value doesn't match original!"
+        );
+    }
+}
diff --git a/packages/cipherstash-proxy-integration/src/lib.rs b/packages/cipherstash-proxy-integration/src/lib.rs
@@ -2,6 +2,7 @@ mod common;
 mod decrypt;
 mod disable_mapping;
 mod empty_result;
+mod encryption_sanity;
 mod extended_protocol_error_messages;
 mod insert;
 mod map_concat;
diff --git a/packages/cipherstash-proxy-integration/src/multitenant/set_keyset_name.rs b/packages/cipherstash-proxy-integration/src/multitenant/set_keyset_name.rs
@@ -115,6 +115,12 @@ mod tests {
         let result = client.simple_query(&sql).await;
         assert!(result.is_ok());
 
+        // SET TENANT_1 WITHOUT QUOTES
+        // VALID AS LONG AS NAME IS A VALID PG IDENTIFIER
+        let sql = format!("SET CIPHERSTASH.KEYSET_NAME = {tenant_keyset_name_1}");
+        let result = client.simple_query(&sql).await;
+        assert!(result.is_ok());
+
         // INSERT
         let tenant_1_id = random_id();
         let encrypted_text = "hello";
@@ -372,8 +378,8 @@ mod tests {
 
         // Test cases that should potentially fail or be handled gracefully
         let invalid_cases = vec![
-            format!("SET CIPHERSTASH.KEYSET_NAME = {tenant_keyset_name_1}"), // unquoted string
-            format!("SET CIPHERSTASH.KEYSET_NAME = NULL"),                   // null value
+            format!("SET CIPHERSTASH.KEYSET_NAME = test-1"), // unquoted string that is NOT a valid pg Identifier
+            format!("SET CIPHERSTASH.KEYSET_NAME = NULL"),   // null value
         ];
 
         for invalid_sql in invalid_cases {
diff --git a/packages/cipherstash-proxy/src/postgresql/context/mod.rs b/packages/cipherstash-proxy/src/postgresql/context/mod.rs
@@ -446,17 +446,20 @@ where
         }) = statement
         {
             if variable == &*SQL_SETTING_NAME_KEYSET_NAME {
-                if let Some(Expr::Value(ValueWithSpan { value, .. })) = values.first() {
-                    let keyset_name = match value {
-                        Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => s.clone(),
-                        Value::Number(n, _) => n.to_string(),
-                        _ => {
-                            let err = EncryptError::KeysetNameCouldNotBeSet;
-                            warn!(target: CONTEXT, client_id = self.client_id, msg = err.to_string());
-                            return Ok(None);
+                // Try to extract keyset name from Value (quoted string/number) or Identifier (unquoted)
+                let keyset_name = match values.first() {
+                    Some(Expr::Value(ValueWithSpan { value, .. })) => match value {
+                        Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => {
+                            Some(s.clone())
                         }
-                    };
-
+                        Value::Number(n, _) => Some(n.to_string()),
+                        _ => None,
+                    },
+                    Some(Expr::Identifier(ident)) => Some(ident.value.clone()),
+                    _ => None,
+                };
+
+                if let Some(keyset_name) = keyset_name {
                     debug!(target: CONTEXT, client_id = self.client_id, msg = "Set KeysetName", ?keyset_name);
 
                     let identifier = KeysetIdentifier(IdentifiedBy::Name(keyset_name.into()));
@@ -469,7 +472,7 @@ where
                 } else {
                     let err = EncryptError::KeysetNameCouldNotBeSet;
                     warn!(target: CONTEXT, client_id = self.client_id, msg = err.to_string());
-                    // We let the database handle any syntax errors to avoid complexifying the fronted flow (more)
+                    // We let the database handle any syntax errors to avoid complexifying the frontend flow
                 }
             }
         }
