From 10137ffaa44282ccf3424e9361c0c286c86781d9 Mon Sep 17 00:00:00 2001
From: Avinash Sajjanshetty <avi@turso.tech>
Date: Wed, 17 Sep 2025 21:58:27 +0530
Subject: [PATCH] run whopper with encryption if arg is passed

---
 Cargo.lock         |  1 +
 whopper/Cargo.toml |  4 ++-
 whopper/main.rs    | 75 +++++++++++++++++++++++++++++++++++++++++-----
 3 files changed, 72 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d2e512cf4..9e2bb330d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4338,6 +4338,7 @@ version = "0.2.0-pre.3"
 dependencies = [
  "anyhow",
  "clap",
+ "hex",
  "memmap2",
  "rand 0.9.2",
  "rand_chacha 0.9.0",
diff --git a/whopper/Cargo.toml b/whopper/Cargo.toml
index 0695ebcf4..8056fbc43 100644
--- a/whopper/Cargo.toml
+++ b/whopper/Cargo.toml
@@ -25,6 +25,8 @@ tracing = { workspace = true }
 tracing-subscriber = { workspace = true, features = ["env-filter"] }
 turso_core = { workspace = true, features = ["simulator"]}
 turso_parser = { workspace = true }
+hex = "0.4.3"
 
 [features]
-checksum = ["turso_core/checksum"]
\ No newline at end of file
+checksum = ["turso_core/checksum"]
+encryption = ["turso_core/encryption"]
\ No newline at end of file
diff --git a/whopper/main.rs b/whopper/main.rs
index ef6f1788e..ba155ca7e 100644
--- a/whopper/main.rs
+++ b/whopper/main.rs
@@ -14,7 +14,9 @@ use std::cell::RefCell;
 use std::sync::Arc;
 use tracing::trace;
 use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};
-use turso_core::{Connection, Database, IO, Statement};
+use turso_core::{
+    CipherMode, Connection, Database, DatabaseOpts, EncryptionOpts, IO, OpenFlags, Statement,
+};
 use turso_parser::ast::SortOrder;
 
 mod io;
@@ -36,6 +38,9 @@ struct Args {
     /// Enable MVCC (Multi-Version Concurrency Control)
     #[arg(long)]
     enable_mvcc: bool,
+    /// Enable database encryption
+    #[arg(long)]
+    enable_encryption: bool,
 }
 
 struct SimulatorConfig {
@@ -74,6 +79,17 @@ struct Stats {
     integrity_checks: usize,
 }
 
+fn may_be_set_encryption(
+    conn: Arc<Connection>,
+    opts: &Option<EncryptionOpts>,
+) -> anyhow::Result<Arc<Connection>> {
+    if let Some(opts) = opts {
+        conn.pragma_update("cipher", format!("'{}'", opts.cipher.clone()))?;
+        conn.pragma_update("hexkey", format!("'{}'", opts.hexkey.clone()))?;
+    }
+    Ok(conn)
+}
+
 fn main() -> anyhow::Result<()> {
     init_logger();
 
@@ -109,14 +125,35 @@ fn main() -> anyhow::Result<()> {
 
     let db_path = format!("whopper-{}-{}.db", seed, std::process::id());
 
-    let db = match Database::open_file(io.clone(), &db_path, args.enable_mvcc, true) {
-        Ok(db) => db,
-        Err(e) => {
-            return Err(anyhow::anyhow!("Database open failed: {}", e));
+    let encryption_opts = if args.enable_encryption {
+        let opts = random_encryption_config(&mut rng);
+        println!("cipher = {}, key = {}", opts.cipher, opts.hexkey);
+        Some(opts)
+    } else {
+        None
+    };
+
+    let db = {
+        let opts = DatabaseOpts::new()
+            .with_mvcc(args.enable_mvcc)
+            .with_indexes(true);
+
+        match Database::open_file_with_flags(
+            io.clone(),
+            &db_path,
+            OpenFlags::default(),
+            opts,
+            encryption_opts.clone(),
+        ) {
+            Ok(db) => db,
+            Err(e) => {
+                return Err(anyhow::anyhow!("Database open failed: {}", e));
+            }
         }
     };
+
     let boostrap_conn = match db.connect() {
-        Ok(conn) => conn,
+        Ok(conn) => may_be_set_encryption(conn, &encryption_opts)?,
         Err(e) => {
             return Err(anyhow::anyhow!("Connection failed: {}", e));
         }
@@ -146,7 +183,7 @@ fn main() -> anyhow::Result<()> {
     let mut fibers = Vec::new();
     for i in 0..config.max_connections {
         let conn = match db.connect() {
-            Ok(conn) => conn,
+            Ok(conn) => may_be_set_encryption(conn, &encryption_opts)?,
             Err(e) => {
                 return Err(anyhow::anyhow!(
                     "Failed to create fiber connection {}: {}",
@@ -323,6 +360,30 @@ fn create_initial_schema(rng: &mut ChaCha8Rng) -> Vec<Create> {
     schema
 }
 
+fn random_encryption_config(rng: &mut ChaCha8Rng) -> EncryptionOpts {
+    let cipher_modes = [
+        CipherMode::Aes128Gcm,
+        CipherMode::Aes256Gcm,
+        CipherMode::Aegis256,
+        CipherMode::Aegis128L,
+        CipherMode::Aegis128X2,
+        CipherMode::Aegis128X4,
+        CipherMode::Aegis256X2,
+        CipherMode::Aegis256X4,
+    ];
+
+    let cipher_mode = cipher_modes[rng.random_range(0..cipher_modes.len())];
+
+    let key_size = cipher_mode.required_key_size();
+    let mut key = vec![0u8; key_size];
+    rng.fill_bytes(&mut key);
+
+    EncryptionOpts {
+        cipher: cipher_mode.to_string(),
+        hexkey: hex::encode(&key),
+    }
+}
+
 fn perform_work(
     fiber_idx: usize,
     rng: &mut ChaCha8Rng,