From 328c5edf4d82f3cffd91dd8ffd530b8644f80a7c Mon Sep 17 00:00:00 2001
From: Avinash Sajjanshetty <avi@turso.tech>
Date: Mon, 25 Aug 2025 02:17:53 +0530
Subject: [PATCH] Add `PRAGMA cipher` to allow setting cipher algo

---
 core/lib.rs                | 12 ++++++++++++
 core/pragma.rs             |  4 ++++
 core/storage/encryption.rs | 24 ++++++++++++++++++++++++
 core/translate/pragma.rs   | 17 ++++++++++++++++-
 parser/src/ast.rs          |  4 ++++
 5 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/core/lib.rs b/core/lib.rs
index e8e629887..3a8c7ead1 100644
--- a/core/lib.rs
+++ b/core/lib.rs
@@ -41,6 +41,7 @@ pub mod numeric;
 mod numeric;
 
 use crate::incremental::view::ViewTransactionState;
+use crate::storage::encryption::CipherMode;
 use crate::translate::optimizer::optimize_plan;
 use crate::translate::pragma::TURSO_CDC_DEFAULT_TABLE_NAME;
 #[cfg(all(feature = "fs", feature = "conn_raw_api"))]
@@ -455,6 +456,7 @@ impl Database {
             metrics: RefCell::new(ConnectionMetrics::new()),
             is_nested_stmt: Cell::new(false),
             encryption_key: RefCell::new(None),
+            encryption_cipher: RefCell::new(None),
         });
         let builtin_syms = self.builtin_syms.borrow();
         // add built-in extensions symbols to the connection to prevent having to load each time
@@ -886,6 +888,7 @@ pub struct Connection {
     /// Generally this is only true for ParseSchema.
     is_nested_stmt: Cell<bool>,
     encryption_key: RefCell<Option<EncryptionKey>>,
+    encryption_cipher: RefCell<Option<CipherMode>>,
 }
 
 impl Connection {
@@ -1961,6 +1964,15 @@ impl Connection {
         let pager = self.pager.borrow();
         pager.set_encryption_context(&key);
     }
+
+    pub fn set_encryption_cipher(&self, cipher: CipherMode) {
+        tracing::trace!("setting encryption cipher for connection");
+        self.encryption_cipher.replace(Some(cipher));
+    }
+
+    pub fn get_encryption_cipher_mode(&self) -> Option<CipherMode> {
+        self.encryption_cipher.borrow().clone()
+    }
 }
 
 pub struct Statement {
diff --git a/core/pragma.rs b/core/pragma.rs
index 0ae48b97a..e006963c0 100644
--- a/core/pragma.rs
+++ b/core/pragma.rs
@@ -111,6 +111,10 @@ pub fn pragma_for(pragma: &PragmaName) -> Pragma {
             PragmaFlags::Result0 | PragmaFlags::SchemaReq | PragmaFlags::NoColumns1,
             &["hexkey"],
         ),
+        EncryptionCipher => Pragma::new(
+            PragmaFlags::Result0 | PragmaFlags::SchemaReq | PragmaFlags::NoColumns1,
+            &["cipher"],
+        ),
     }
 }
 
diff --git a/core/storage/encryption.rs b/core/storage/encryption.rs
index e8b7f1181..17343f1c0 100644
--- a/core/storage/encryption.rs
+++ b/core/storage/encryption.rs
@@ -134,6 +134,30 @@ pub enum CipherMode {
     Aegis256,
 }
 
+impl TryFrom<&str> for CipherMode {
+    type Error = LimboError;
+
+    fn try_from(s: &str) -> Result<Self, Self::Error> {
+        match s.to_lowercase().as_str() {
+            "aes256gcm" | "aes-256-gcm" | "aes_256_gcm" => Ok(CipherMode::Aes256Gcm),
+            "aegis256" | "aegis-256" | "aegis_256" => Ok(CipherMode::Aegis256),
+            _ => Err(LimboError::InvalidArgument(format!(
+                "Unknown cipher name: {}",
+                s
+            ))),
+        }
+    }
+}
+
+impl std::fmt::Display for CipherMode {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            CipherMode::Aes256Gcm => write!(f, "aes256gcm"),
+            CipherMode::Aegis256 => write!(f, "aegis256"),
+        }
+    }
+}
+
 impl CipherMode {
     /// Every cipher requires a specific key size. For 256-bit algorithms, this is 32 bytes.
     /// For 128-bit algorithms, it would be 16 bytes, etc.
diff --git a/core/translate/pragma.rs b/core/translate/pragma.rs
index b53c7af5b..43abf02e7 100644
--- a/core/translate/pragma.rs
+++ b/core/translate/pragma.rs
@@ -10,7 +10,7 @@ use turso_parser::ast::{PragmaName, QualifiedName};
 use super::integrity_check::translate_integrity_check;
 use crate::pragma::pragma_for;
 use crate::schema::Schema;
-use crate::storage::encryption::EncryptionKey;
+use crate::storage::encryption::{CipherMode, EncryptionKey};
 use crate::storage::pager::AutoVacuumMode;
 use crate::storage::pager::Pager;
 use crate::storage::sqlite3_ondisk::CacheSize;
@@ -318,6 +318,12 @@ fn update_pragma(
             connection.set_encryption_key(key);
             Ok((program, TransactionMode::None))
         }
+        PragmaName::EncryptionCipher => {
+            let value = parse_string(&value)?;
+            let cipher = CipherMode::try_from(value.as_str())?;
+            connection.set_encryption_cipher(cipher);
+            Ok((program, TransactionMode::None))
+        }
     }
 }
 
@@ -589,6 +595,15 @@ fn query_pragma(
             program.add_pragma_result_column(pragma.to_string());
             Ok((program, TransactionMode::None))
         }
+        PragmaName::EncryptionCipher => {
+            if let Some(cipher) = connection.get_encryption_cipher_mode() {
+                let register = program.alloc_register();
+                program.emit_string8(cipher.to_string(), register);
+                program.emit_result_row(register, 1);
+                program.add_pragma_result_column(pragma.to_string());
+            }
+            Ok((program, TransactionMode::None))
+        }
     }
 }
 
diff --git a/parser/src/ast.rs b/parser/src/ast.rs
index 0c9de857c..a73e4ed36 100644
--- a/parser/src/ast.rs
+++ b/parser/src/ast.rs
@@ -1211,6 +1211,10 @@ pub enum PragmaName {
     AutoVacuum,
     /// `cache_size` pragma
     CacheSize,
+    /// encryption cipher algorithm name for encrypted databases
+    #[strum(serialize = "cipher")]
+    #[cfg_attr(feature = "serde", serde(rename = "cipher"))]
+    EncryptionCipher,
     /// List databases
     DatabaseList,
     /// Encoding - only support utf8