Add SET command for run-time configuration of keyset_id

docs/errors.md

@@ -15,6 +15,9 @@
 
 - Encrypt errors:
   - [Column could not be encrypted](#encrypt-column-could-not-be-encrypted)
+  - [Column could not be encrypted](#encrypt-column-could-not-be-encrypted)
+  - [Could not decrypt data for keyset](#encrypt-encrypt-could-not-decrypt-data-for-keyset)
+  - [KeysetId could not be set](#encrypt-keyset-id-could-not-be-set)
   - [Plaintext could not be encoded](#encrypt-plaintext-could-not-be-encoded)
   - [Unknown column](#encrypt-unknown-column)
   - [Unknown table](#encrypt-unknown-table)
@@ -272,9 +275,68 @@ The most likely cause is network access to the ZeroKMS service.
 
 
 
+
+<!-- ---------------------------------------------------------------------------------------------------- -->
+
+
+## KeysetId could not be set <a id='encrypt-keyset-id-could-not-be-set'></a>
+
+A keyset_id could not be set using the `SET CIPHERSTASH.KEYSET_ID` command.
+
+
+### Error message
+
+```
+A keyset_id could not be set using `SET CIPHERSTASH.KEYSET_ID`
+```
+
+### How to Fix
+
+1. Check the syntax of the `SET CIPHERSTASH.KEYSET_ID` command. The `keyset_id` value should be in single quotes.
+2. Check that the provided `keyset_id` is a valid UUID.
+2. Check that the value is being set as a literal. The PostgreSQL `SET` statement does not support parameterised querying.
+
+
+```
+   SET [ SESSION ] CIPHERSTASH.KEYSET_ID { TO | = } '{keyset_id}'
+```
+
 <!-- ---------------------------------------------------------------------------------------------------- -->
 
 
+## Could not decrypt data for keyset <a id='encrypt-could-not-decrypt-data-for-keyset'></a>
+
+The data belonging to the active `keyset_id` could not be decrypted.
+
+
+### Error message
+
+```
+Could not decrypt data for keyset '{keyset_id}'
+```
+
+### Notes
+
+This error is caused because the active `keyset_id` does not match the `keyset_id` of the data being decrypted.
+
+Each encrypted record belong to a `keyset` with a unique identifier (the `keyset_id`).
+
+The proxy encrypts data in the currently active `keyset`.
+
+A single default keyset can be defined in the proxy configuration or the `SET CIPHERSTASH.KEYSET_ID` statement can be used to dynamically set the active `keyset` at runtime.
+
+If the `keyset_id` of the record does not match the current `keyset_id` the data cannot be decrypted.
+
+### How to Fix
+
+1. Check that the `keyset_id` in the configuration matches the encrypted records.
+2. If using the `SET CIPHERSTASH.KEYSET_ID` statement, check that this `keyset_id`matches the encrypted records.
+is a valid UUID.
+3. Check that the configured `client` has been granted access to to the `keyset_id`.
+
+
+<!-- ---------------------------------------------------------------------------------------------------- -->
+
 
 ## Plaintext could not be encoded <a id='encrypt-plaintext-could-not-be-encoded'></a>
 
@@ -312,8 +374,10 @@ For example:
 2. Check that the configuration has not changed.
 3. Check [EQL](https://github.com/cipherstash/encrypt-query-language).
 
+
 <!-- ---------------------------------------------------------------------------------------------------- -->
 
+
 ## Unknown Column <a id='encrypt-unknown-column'></a>
 
 The column has an encrypted type (PostgreSQL `eql_v2_encrypted` type ) with no encryption configuration.

packages/cipherstash-proxy-integration/src/map_ore_index_where.rs

@@ -73,15 +73,30 @@ mod tests {
 
         // GT: given [1, 3], `> 1` returns [3]
         let sql = format!("SELECT {col_name} FROM encrypted WHERE {col_name} > $1");
-        test_ore_op(&client, col_name, &sql, &[&low], &[high.clone()]).await;
+
+        test_ore_op(
+            &client,
+            col_name,
+            &sql,
+            &[&low],
+            std::slice::from_ref(&high),
+        )
+        .await;
 
         // GT 2nd case: given [1, 3], `> 3` returns []
         let sql = format!("SELECT {col_name} FROM encrypted WHERE {col_name} > $1");
         test_ore_op::<T>(&client, col_name, &sql, &[&high], &[]).await;
 
         // LT: given [1, 3], `< 3` returns [1]
         let sql = format!("SELECT {col_name} FROM encrypted WHERE {col_name} < $1");
-        test_ore_op(&client, col_name, &sql, &[&high], &[low.clone()]).await;
+        test_ore_op(
+            &client,
+            col_name,
+            &sql,
+            &[&high],
+            std::slice::from_ref(&low),
+        )
+        .await;
 
         // LT 2nd case: given [1, 3], `< 3` returns []
         let sql = format!("SELECT {col_name} FROM encrypted WHERE {col_name} < $1");
@@ -116,11 +131,18 @@ mod tests {
 
         // EQ: given [1, 3], `= 1` returns [1]
         let sql = format!("SELECT {col_name} FROM encrypted WHERE {col_name} = $1");
-        test_ore_op(&client, col_name, &sql, &[&low], &[low.clone()]).await;
+        test_ore_op(&client, col_name, &sql, &[&low], std::slice::from_ref(&low)).await;
 
         // NEQ: given [1, 3], `<> 3` returns [1]
         let sql = format!("SELECT {col_name} FROM encrypted WHERE {col_name} <> $1");
-        test_ore_op(&client, col_name, &sql, &[&high], &[low.clone()]).await;
+        test_ore_op(
+            &client,
+            col_name,
+            &sql,
+            &[&high],
+            std::slice::from_ref(&low),
+        )
+        .await;
     }
 
     /// Runs the query and checks the returned results match the expected results.

packages/cipherstash-proxy/src/config/database.rs

@@ -73,7 +73,7 @@ impl DatabaseConfig {
         self.connection_timeout.map(Duration::from_millis)
     }
 
-    pub fn server_name(&self) -> Result<ServerName, Error> {
+    pub fn server_name(&self) -> Result<ServerName<'_>, Error> {
         let name = ServerName::try_from(self.host.as_str()).map_err(|_| {
             ConfigError::InvalidServerName {
                 name: self.host.to_owned(),

packages/cipherstash-proxy/src/encrypt/mod.rs

@@ -18,7 +18,8 @@ use cipherstash_client::{
         PlaintextTarget, Queryable, ReferencedPendingPipeline,
     },
     schema::ColumnConfig,
-    ConsoleConfig, CtsConfig, ZeroKMSConfig,
+    zerokms::ClientKey,
+    ConsoleConfig, CtsConfig, ZeroKMS, ZeroKMSConfig,
 };
 use cipherstash_client::{
     config::{ConfigError, ZeroKMSConfigWithClientKey},
@@ -28,6 +29,7 @@ use config::EncryptConfigManager;
 use schema::SchemaManager;
 use std::{sync::Arc, vec};
 use tracing::{debug, info, warn};
+use uuid::Uuid;
 
 /// SQL Statement for loading encrypt configuration from database
 const ENCRYPT_CONFIG_QUERY: &str = include_str!("./sql/select_config.sql");
@@ -40,24 +42,27 @@ const AGGREGATE_QUERY: &str = include_str!("./sql/select_aggregates.sql");
 
 type ScopedCipher = encryption::ScopedCipher<AutoRefresh<ServiceCredentials>>;
 
+type ZerokmsClient = ZeroKMS<AutoRefresh<ServiceCredentials>, ClientKey>;
+
 ///
 /// All of the things required for Encrypt-as-a-Product
 ///
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct Encrypt {
     pub config: TandemConfig,
-    cipher: Arc<ScopedCipher>,
     pub encrypt_config: EncryptConfigManager,
     pub schema: SchemaManager,
     /// The EQL version installed in the database or `None` if it was not present
     pub eql_version: Option<String>,
+    zerokms_client: Arc<ZerokmsClient>,
 }
 
 impl Encrypt {
     pub async fn init(config: TandemConfig) -> Result<Encrypt, Error> {
-        let cipher = Arc::new(init_cipher(&config).await?);
+        let zerokms_client = init_zerokms_client(&config)?;
+
         let encrypt_config = EncryptConfigManager::init(&config.database).await?;
-        // TODO: populate EqlTraitImpls based on config
+        // TODO: populate EqlTraitImpls based in config
         let schema = SchemaManager::init(&config.database).await?;
 
         let eql_version = {
@@ -80,23 +85,44 @@ impl Encrypt {
 
         Ok(Encrypt {
             config,
-            cipher,
+            zerokms_client: Arc::new(zerokms_client),
             encrypt_config,
             schema,
             eql_version,
         })
     }
 
+    /// Initialize cipher using the stored zerokms_config
+    pub async fn init_cipher(&self, keyset_id: Option<Uuid>) -> Result<ScopedCipher, Error> {
+        let zerokms_client = self.zerokms_client.clone();
+        match ScopedCipher::init(zerokms_client, keyset_id).await {
+            Ok(cipher) => {
+                debug!(target: ENCRYPT, msg = "Initialized ZeroKMS ScopedCipher");
+                Ok(cipher)
+            }
+            Err(err) => {
+                debug!(target: ENCRYPT, msg = "Error initializing ZeroKMS ScopedCipher", error = err.to_string());
+                Err(err.into())
+            }
+        }
+    }
+
     ///
     /// Encrypt `Plaintexts` using the `Column` configuration
     ///
     ///
     pub async fn encrypt(
         &self,
+        keyset_id: Option<Uuid>,
         plaintexts: Vec<Option<Plaintext>>,
         columns: &[Option<Column>],
     ) -> Result<Vec<Option<eql::EqlEncrypted>>, Error> {
-        let mut pipeline = ReferencedPendingPipeline::new(self.cipher.clone());
+        let keyset_id = keyset_id.or(self.config.encrypt.default_keyset_id);
+        debug!(target: ENCRYPT, ?keyset_id);
+
+        let cipher = Arc::new(self.init_cipher(keyset_id).await?);
+
+        let mut pipeline = ReferencedPendingPipeline::new(cipher.clone());
         let mut index_term_plaintexts = vec![None; columns.len()];
 
         for (idx, item) in plaintexts.into_iter().zip(columns.iter()).enumerate() {
@@ -145,7 +171,7 @@ impl Encrypt {
                     let index = column.config.clone().into_ste_vec_index().unwrap();
                     let op = QueryOp::SteVecSelector;
 
-                    let index_term = (index, plaintext).build_queryable(self.cipher.clone(), op)?;
+                    let index_term = (index, plaintext).build_queryable(cipher.clone(), op)?;
 
                     encrypted = Some(to_eql_encrypted_from_index_term(
                         index_term,
@@ -170,8 +196,14 @@ impl Encrypt {
     ///
     pub async fn decrypt(
         &self,
+        keyset_id: Option<Uuid>,
         ciphertexts: Vec<Option<eql::EqlEncrypted>>,
     ) -> Result<Vec<Option<Plaintext>>, Error> {
+        let keyset_id = keyset_id.or(self.config.encrypt.default_keyset_id);
+        debug!(target: ENCRYPT, ?keyset_id);
+
+        let cipher = Arc::new(self.init_cipher(keyset_id).await?);
+
         // Create a mutable vector to hold the decrypted results
         let mut results = vec![None; ciphertexts.len()];
 
@@ -182,8 +214,24 @@ impl Encrypt {
             .filter_map(|(idx, eql)| Some((idx, eql?.body.ciphertext.unwrap())))
             .collect::<_>();
 
-        // Decrypt the ciphertexts
-        let decrypted = self.cipher.decrypt(encrypted).await?;
+        let decrypted = cipher.decrypt(encrypted).await.map_err(|e| -> Error {
+            match &e {
+                cipherstash_client::zerokms::Error::Decrypt(_) => {
+                    let error_msg = e.to_string();
+                    if error_msg.contains("Failed to retrieve key") {
+                        EncryptError::CouldNotDecryptDataForKeyset {
+                            keyset_id: keyset_id
+                                .map(|id| id.to_string())
+                                .unwrap_or("default_keyset".to_string()),
+                        }
+                        .into()
+                    } else {
+                        e.into()
+                    }
+                }
+                _ => e.into(),
+            }
+        })?;
 
         // Merge the decrypted values as plaintext into their original indexed positions
         for (idx, decrypted) in indices.into_iter().zip(decrypted) {
@@ -213,6 +261,15 @@ impl Encrypt {
     }
 }
 
+fn init_zerokms_client(
+    config: &TandemConfig,
+) -> Result<ZeroKMS<AutoRefresh<ServiceCredentials>, ClientKey>, ConfigError> {
+    let zerokms_config = build_zerokms_config(config)?;
+
+    Ok(zerokms_config
+        .create_client_with_credentials(AutoRefresh::new(zerokms_config.credentials())))
+}
+
 fn build_zerokms_config(config: &TandemConfig) -> Result<ZeroKMSConfigWithClientKey, ConfigError> {
     let console_config = ConsoleConfig::builder().with_env().build()?;
 
@@ -243,24 +300,6 @@ fn build_zerokms_config(config: &TandemConfig) -> Result<ZeroKMSConfigWithClient
     builder.build_with_client_key()
 }
 
-async fn init_cipher(config: &TandemConfig) -> Result<ScopedCipher, Error> {
-    let zerokms_config = build_zerokms_config(config)?;
-
-    let zerokms_client = zerokms_config
-        .create_client_with_credentials(AutoRefresh::new(zerokms_config.credentials()));
-
-    match ScopedCipher::init(Arc::new(zerokms_client), config.encrypt.default_keyset_id).await {
-        Ok(cipher) => {
-            debug!(target: ENCRYPT, msg = "Initialized ZeroKMS ScopedCipher");
-            Ok(cipher)
-        }
-        Err(err) => {
-            debug!(target: ENCRYPT, msg =  "Error initializing ZeroKMS ScopedCipher", error = err.to_string());
-            Err(err.into())
-        }
-    }
-}
-
 fn to_eql_encrypted_from_index_term(
     index_term: IndexTerm,
     identifier: &Identifier,

packages/cipherstash-proxy/src/error.rs

@@ -227,9 +227,21 @@ pub enum EncryptError {
     #[error("Column configuration for column '{column}' in table '{table}' does not match the encrypted column. For help visit {}#encrypt-column-config-mismatch", ERROR_DOC_BASE_URL)]
     ColumnConfigurationMismatch { table: String, column: String },
 
+    #[error(
+        "Could not decrypt data for keyset '{keyset_id}'. For help visit {}#encrypt-could-not-decrypt-data-for-keyset",
+        ERROR_DOC_BASE_URL
+    )]
+    CouldNotDecryptDataForKeyset { keyset_id: String },
+
     #[error("InvalidIndexTerm")]
     InvalidIndexTerm,
 
+    #[error(
+        "A keyset_id could not be set using `SET CIPHERSTASH.KEYSET_ID`. For help visit {}#encrypt-keyset-id-could-not-be-set",
+        ERROR_DOC_BASE_URL
+    )]
+    KeysetIdCouldNotBeSet,
+
     /// This should in practice be unreachable
     #[error("Missing encrypt configuration for column type `{plaintext_type}`. For help visit {}#encrypt-missing-encrypt-configuration", ERROR_DOC_BASE_URL)]
     MissingEncryptConfiguration { plaintext_type: String },