♻️ refactor: replace Proxy with Context in handler and use async_trait

- Replace handler function signature to accept Context<ZeroKms> instead of Proxy + client_id
- Add config access methods to Context for database, TLS, and other settings
- Convert EncryptionService trait from impl Future to async_trait for cleaner syntax
- Update ZeroKms implementation to use async_trait
- Modify startup functions to work with Context instead of Proxy
- Remove services module as functionality is now integrated into Context
- Fix type annotation errors by using concrete Context<ZeroKms> type

This refactoring simplifies the architecture by making Context the primary
interface for handler operations, avoiding dyn trait objects while maintaining
strong typing.

Author: tobyhede

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

@@ -22,30 +22,4 @@ mod tests {
 
         assert!(rows.is_empty());
     }
-
-    #[tokio::test]
-    async fn disable_mapping_disables_schema_reload() {
-        let client = connect_with_tls(PROXY).await;
-
-        let sql = "SET CIPHERSTASH.UNSAFE_DISABLE_MAPPING = true";
-        client.query(sql, &[]).await.unwrap();
-
-        let id = random_id();
-
-        let sql = format!(
-            "CREATE TABLE table_{id} (
-            id bigint,
-            PRIMARY KEY(id)
-        );"
-        );
-
-        let _ = client.execute(&sql, &[]).await.unwrap();
-
-        let sql = "SET CIPHERSTASH.UNSAFE_DISABLE_MAPPING = false";
-        client.query(sql, &[]).await.unwrap();
-
-        let sql = format!("SELECT id FROM table_{id}");
-        let result = client.query(&sql, &[]).await;
-        assert!(result.is_err());
-    }
 }

packages/cipherstash-proxy/src/lib.rs

@@ -9,7 +9,6 @@ pub mod log;
 pub mod postgresql;
 pub mod prometheus;
 pub mod proxy;
-pub mod services;
 pub mod tls;
 
 pub use crate::cli::Args;

packages/cipherstash-proxy/src/main.rs

@@ -82,7 +82,11 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             },
             _ = sighup() => {
                 info!(msg = "Received SIGHUP. Reloading configuration");
-                (listener, proxy) = reload_config(listener, &args, proxy).await;
+                (listener, proxy) = match reload_config(listener, &args).await {
+                    Ok((listener, proxy)) => (listener, proxy),
+                    Err(_) => todo!(),
+                };
+
                 info!(msg = "Reloaded configuration");
             },
             _ = sigterm() => {
@@ -91,16 +95,15 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             },
             Ok(client_stream) = AsyncStream::accept(&listener) => {
 
-                    let proxy = proxy.clone();
-
                     client_id += 1;
 
+                    let context = proxy.context(client_id);
+
                     tracker.spawn(async move {
-                        let proxy = proxy.clone();
 
                         gauge!(CLIENTS_ACTIVE_CONNECTIONS).increment(1);
 
-                        match pg::handler(client_stream, proxy, client_id).await {
+                        match pg::handler(client_stream,context).await {
                             Ok(_) => (),
                             Err(err) => {
 
@@ -261,15 +264,15 @@ async fn sighup() -> std::io::Result<()> {
     Ok(())
 }
 
-async fn reload_config(listener: TcpListener, args: &Args, proxy: Proxy) -> (TcpListener, Proxy) {
+async fn reload_config(listener: TcpListener, args: &Args) -> Result<(TcpListener, Proxy), Error> {
     let new_config = match TandemConfig::load(args) {
         Ok(config) => config,
         Err(err) => {
             warn!(
                 msg = "Configuration could not be reloaded: {}",
                 error = err.to_string()
             );
-            return (listener, proxy);
+            return Err(err);
         }
     };
 
@@ -278,8 +281,8 @@ async fn reload_config(listener: TcpListener, args: &Args, proxy: Proxy) -> (Tcp
     // Explicit drop needed here to free the network resources before binding if using the same address & port
     std::mem::drop(listener);
 
-    (
+    Ok((
         connect::bind_with_retry(&new_proxy.config.server).await,
         new_proxy,
-    )
+    ))
 }

packages/cipherstash-proxy/src/postgresql/backend.rs

@@ -19,7 +19,7 @@ use crate::prometheus::{
     DECRYPTION_ERROR_TOTAL, DECRYPTION_REQUESTS_TOTAL, ROWS_ENCRYPTED_TOTAL,
     ROWS_PASSTHROUGH_TOTAL, ROWS_TOTAL, SERVER_BYTES_RECEIVED_TOTAL,
 };
-use crate::proxy::Proxy;
+use crate::proxy::EncryptionService;
 use bytes::BytesMut;
 use metrics::{counter, histogram};
 use std::time::Instant;
@@ -70,25 +70,25 @@ use tracing::{debug, error, info, warn};
 /// - `RowDescription`: Result column metadata (modified for encrypted columns)
 /// - `ParameterDescription`: Parameter metadata (modified for encrypted parameters)
 /// - `ReadyForQuery`: Session ready state (triggers schema reload if needed)
-pub struct Backend<R>
+pub struct Backend<R, S>
 where
     R: AsyncRead + Unpin,
+    S: EncryptionService,
 {
     /// Sender for outgoing messages to client
     client_sender: Sender,
     /// Reader for incoming messages from server
     server_reader: R,
-    /// Encryption service for column decryption
-    proxy: Proxy,
     /// Session context with portal and statement metadata
-    context: Context,
+    context: Context<S>,
     /// Buffer for batching DataRow messages before decryption
     buffer: MessageBuffer,
 }
 
-impl<R> Backend<R>
+impl<R, S> Backend<R, S>
 where
     R: AsyncRead + Unpin,
+    S: EncryptionService,
 {
     /// Creates a new Backend instance.
     ///
@@ -98,12 +98,11 @@ where
     /// * `server_reader` - Stream for reading messages from the PostgreSQL server
     /// * `encrypt` - Encryption service for handling column decryption
     /// * `context` - Session context shared with the frontend
-    pub fn new(client_sender: Sender, server_reader: R, proxy: Proxy, context: Context) -> Self {
+    pub fn new(client_sender: Sender, server_reader: R, context: Context<S>) -> Self {
         let buffer = MessageBuffer::new();
         Backend {
             client_sender,
             server_reader,
-            proxy,
             context,
             buffer,
         }
@@ -150,19 +149,17 @@ where
     /// Returns `Ok(())` on successful message processing, or an `Error` if a fatal
     /// error occurs that should terminate the connection.
     pub async fn rewrite(&mut self) -> Result<(), Error> {
-        let connection_timeout = self.proxy.config.database.connection_timeout();
-
         let (code, mut bytes) = protocol::read_message(
             &mut self.server_reader,
             self.context.client_id,
-            connection_timeout,
+            self.context.connection_timeout(),
         )
         .await?;
 
         let sent: u64 = bytes.len() as u64;
         counter!(SERVER_BYTES_RECEIVED_TOTAL).increment(sent);
 
-        if self.proxy.is_passthrough() {
+        if self.context.is_passthrough() {
             debug!(target: DEVELOPMENT,
                 client_id = self.context.client_id,
                 msg = "Passthrough enabled"
@@ -250,7 +247,7 @@ where
                     msg = "ReadyForQuery"
                 );
                 if self.context.schema_changed() {
-                    self.proxy.reload_schema().await;
+                    self.context.reload_schema().await;
                 }
             }
 
@@ -450,16 +447,12 @@ where
         );
 
         // Decrypt CipherText -> Plaintext
-        let plaintexts = self
-            .proxy
-            .decrypt(keyset_id, ciphertexts)
-            .await
-            .inspect_err(|_| {
-                counter!(DECRYPTION_ERROR_TOTAL).increment(1);
-            })?;
+        let plaintexts = self.context.decrypt(ciphertexts).await.inspect_err(|_| {
+            counter!(DECRYPTION_ERROR_TOTAL).increment(1);
+        })?;
 
         // Avoid the iter calculation if we can
-        if self.proxy.config.prometheus_enabled() {
+        if self.context.prometheus_enabled() {
             let decrypted_count =
                 plaintexts
                     .iter()
@@ -655,9 +648,10 @@ where
 }
 
 /// Implementation of PostgreSQL error handling for the Backend component.
-impl<R> PostgreSqlErrorHandler for Backend<R>
+impl<R, S> PostgreSqlErrorHandler for Backend<R, S>
 where
     R: AsyncRead + Unpin,
+    S: EncryptionService,
 {
     fn client_sender(&mut self) -> &mut Sender {
         &mut self.client_sender

packages/cipherstash-proxy/src/postgresql/column_mapper.rs

@@ -3,7 +3,7 @@ use crate::{
     error::{EncryptError, Error},
     log::MAPPER,
     postgresql::Column,
-    services::SchemaService,
+    proxy::EncryptConfig,
 };
 use eql_mapper::{EqlTerm, TableColumn, TypeCheckedStatement};
 use postgres_types::Type;
@@ -14,13 +14,13 @@ use tracing::{debug, warn};
 /// and mapping them to encryption configurations.
 #[derive(Clone)]
 pub struct ColumnMapper {
-    schema_service: Arc<dyn SchemaService>,
+    encrypt_config: Arc<EncryptConfig>,
 }
 
 impl ColumnMapper {
     /// Create a new ColumnProcessor with the given schema service and client ID
-    pub fn new(schema_service: Arc<dyn SchemaService>) -> Self {
-        Self { schema_service }
+    pub fn new(encrypt_config: Arc<EncryptConfig>) -> Self {
+        Self { encrypt_config }
     }
 
     /// Maps typed statement projection columns to an Encrypt column configuration
@@ -127,7 +127,7 @@ impl ColumnMapper {
         identifier: Identifier,
         eql_term: &EqlTerm,
     ) -> Result<Option<Column>, Error> {
-        match self.schema_service.get_column_config(&identifier) {
+        match self.encrypt_config.get_column_config(&identifier) {
             Some(config) => {
                 debug!(
                     target: MAPPER,