Merge pull request #198 from cipherstash/terminate-on-close

fix: ensure terminate message is sent to server on connection close

Author: tobyhede

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

@@ -15,6 +15,7 @@ use crate::postgresql::context::column::Column;
 use crate::postgresql::context::Portal;
 use crate::postgresql::data::literal_from_sql;
 use crate::postgresql::messages::error_response::ErrorResponse;
+use crate::postgresql::messages::terminate::Terminate;
 use crate::postgresql::messages::Name;
 use crate::prometheus::{
     CLIENTS_BYTES_RECEIVED_TOTAL, ENCRYPTED_VALUES_TOTAL, ENCRYPTION_DURATION_SECONDS,
@@ -220,6 +221,13 @@ where
         Ok(())
     }
 
+    pub async fn terminate(&mut self) -> Result<(), Error> {
+        debug!(target: PROTOCOL, msg = "Terminate server connection");
+        let bytes = Terminate::message();
+        self.write_to_server(bytes).await?;
+        Ok(())
+    }
+
     async fn describe_handler(&mut self, bytes: &BytesMut) -> Result<(), Error> {
         let describe = Describe::try_from(bytes)?;
         debug!(target: PROTOCOL, client_id = self.context.client_id, ?describe);

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

@@ -251,7 +251,13 @@ pub async fn handler(
 
     let client_to_server = async {
         loop {
-            frontend.rewrite().await?;
+            let result = frontend.rewrite().await;
+            // Ensure the connection is terminated if the client closes the connection
+            // The client ConnectionClosed error is triggered before the terminate message is passed through
+            if matches!(result, Err(Error::ConnectionClosed)) {
+                frontend.terminate().await?
+            }
+            result?;
         }
         // Unreachable, but helps the compiler understand the return type
         // TODO: extract into a function or something with type

packages/cipherstash-proxy/src/postgresql/messages/mod.rs

@@ -10,6 +10,7 @@ pub mod param_description;
 pub mod parse;
 pub mod query;
 pub mod row_description;
+pub mod terminate;
 
 pub const NULL: i32 = -1;
 
@@ -24,6 +25,7 @@ pub enum FrontendCode {
     SASLInitialResponse,
     SASLResponse,
     Sync,
+    Terminate,
     Unknown(char),
 }
 
@@ -72,6 +74,7 @@ impl From<char> for FrontendCode {
             #[allow(unreachable_patterns)]
             'p' => FrontendCode::SASLResponse, // Uses same char, here for completeness
             'S' => FrontendCode::Sync,
+            'X' => FrontendCode::Terminate,
             _ => FrontendCode::Unknown(code),
         }
     }
@@ -89,6 +92,7 @@ impl From<FrontendCode> for u8 {
             FrontendCode::SASLInitialResponse => b'p',
             FrontendCode::SASLResponse => b'p',
             FrontendCode::Sync => b'S',
+            FrontendCode::Terminate => b'X',
             FrontendCode::Unknown(c) => c as u8,
         }
     }
@@ -106,6 +110,7 @@ impl From<FrontendCode> for char {
             FrontendCode::SASLInitialResponse => 'p',
             FrontendCode::SASLResponse => 'p',
             FrontendCode::Sync => 'S',
+            FrontendCode::Terminate => 'X',
             FrontendCode::Unknown(c) => c,
         }
     }

packages/cipherstash-proxy/src/postgresql/messages/terminate.rs

@@ -0,0 +1,16 @@
+use bytes::{BufMut, BytesMut};
+
+use super::FrontendCode;
+
+pub struct Terminate;
+
+impl Terminate {
+    pub fn message() -> BytesMut {
+        let mut bytes = BytesMut::new();
+
+        bytes.put_u8(FrontendCode::Terminate.into());
+        bytes.put_i32(4);
+
+        bytes
+    }
+}