Merge pull request #314 from cipherstash/refactor-name

♻️ refactor: convert Name struct to enum and extract to own module

Author: tobyhede

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

@@ -2,10 +2,7 @@ pub mod column;
 
 use super::{
     format_code::FormatCode,
-    messages::{
-        describe::{Describe, Target},
-        Name,
-    },
+    messages::{describe::Describe, Name, Target},
     Column,
 };
 use crate::{
@@ -264,7 +261,7 @@ impl Context {
             } => self.get_portal_statement(name),
             Describe {
                 ref name,
-                target: Target::PreparedStatement,
+                target: Target::Statement,
             } => self.get_statement(name),
         }
     }
@@ -583,7 +580,7 @@ mod tests {
     use crate::{
         config::LogConfig,
         log,
-        postgresql::messages::{describe::Target, Name},
+        postgresql::messages::{Name, Target},
     };
     use cipherstash_client::IdentifiedBy;
     use eql_mapper::Schema;
@@ -621,15 +618,15 @@ mod tests {
 
         let mut context = Context::new(1, schema);
 
-        let name = Name("name".to_string());
+        let name = Name::from("name");
 
         context.add_statement(name.clone(), statement());
 
         let statement = context.get_statement(&name).unwrap();
 
         let describe = Describe {
             name,
-            target: Target::PreparedStatement,
+            target: Target::Statement,
         };
         context.set_describe(describe);
 
@@ -646,8 +643,8 @@ mod tests {
 
         let mut context = Context::new(1, schema);
 
-        let statement_name = Name("statement".to_string());
-        let portal_name = Name("portal".to_string());
+        let statement_name = Name::from("statement");
+        let portal_name = Name::from("portal");
 
         // Add statement to context
         context.add_statement(statement_name.clone(), statement());
@@ -688,8 +685,8 @@ mod tests {
         let mut context = Context::new(1, schema);
 
         // Create multiple statements
-        let statement_name_1 = Name("statement_1".to_string());
-        let statement_name_2 = Name("statement_2".to_string());
+        let statement_name_1 = Name::from("statement_1");
+        let statement_name_2 = Name::from("statement_2");
 
         // Add statements to context
         context.add_statement(statement_name_1.clone(), statement());
@@ -698,7 +695,7 @@ mod tests {
         // Replicate pipelined execution
         // Add multiple portals with the same name
         // Pointing to different statements
-        let portal_name = Name("portal".to_string());
+        let portal_name = Name::from("portal");
 
         let statement_1 = context.get_statement(&statement_name_1).unwrap();
         context.add_portal(portal_name.clone(), portal(&statement_1));
@@ -737,14 +734,14 @@ mod tests {
 
         let mut context = Context::new(1, schema);
 
-        let statement_name_1 = Name("statement_1".to_string());
+        let statement_name_1 = Name::from("statement_1");
         let portal_name_1 = Name::unnamed();
 
-        let statement_name_2 = Name("statement_2".to_string());
+        let statement_name_2 = Name::from("statement_2");
         let portal_name_2 = Name::unnamed();
 
-        let statement_name_3 = Name("statement_3".to_string());
-        let portal_name_3 = Name("portal_3".to_string());
+        let statement_name_3 = Name::from("statement_3");
+        let portal_name_3 = Name::from("portal_3");
 
         // Add statement to context
         context.add_statement(statement_name_1.clone(), statement());

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

@@ -15,9 +15,10 @@ use crate::log::{CONTEXT, MAPPER, PROTOCOL};
 use crate::postgresql::context::column::Column;
 use crate::postgresql::context::Portal;
 use crate::postgresql::data::literal_from_sql;
+use crate::postgresql::messages::close::Close;
 use crate::postgresql::messages::ready_for_query::ReadyForQuery;
 use crate::postgresql::messages::terminate::Terminate;
-use crate::postgresql::messages::Name;
+use crate::postgresql::messages::{Name, Target};
 use crate::prometheus::{
     CLIENTS_BYTES_RECEIVED_TOTAL, ENCRYPTED_VALUES_TOTAL, ENCRYPTION_DURATION_SECONDS,
     ENCRYPTION_ERROR_TOTAL, ENCRYPTION_REQUESTS_TOTAL, SERVER_BYTES_SENT_TOTAL,
@@ -287,6 +288,9 @@ where
                     return Ok(());
                 }
             }
+            Code::Close => {
+                self.close_handler(&bytes).await?;
+            }
             code => {
                 debug!(target: PROTOCOL,
                     client_id = self.context.client_id,
@@ -322,6 +326,19 @@ where
         Ok(())
     }
 
+    async fn close_handler(&mut self, bytes: &BytesMut) -> Result<(), Error> {
+        let close = Close::try_from(bytes)?;
+        debug!(target: PROTOCOL, client_id = self.context.client_id, ?close);
+        match close.target {
+            Target::Portal => self.context.close_portal(&close.name),
+            Target::Statement => {
+                self.context.close_portal(&close.name);
+                // self.context.close_statement(&close.name);
+            }
+        }
+        Ok(())
+    }
+
     async fn execute_handler(&mut self, bytes: &BytesMut) -> Result<(), Error> {
         let execute = Execute::try_from(bytes)?;
         debug!(target: PROTOCOL, client_id = self.context.client_id, ?execute);

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

@@ -227,10 +227,10 @@ impl TryFrom<&BytesMut> for Bind {
         let _len = cursor.get_i32();
 
         let portal = cursor.read_string()?;
-        let portal = Name(portal);
+        let portal = Name::from(portal);
 
         let prepared_statement = cursor.read_string()?;
-        let prepared_statement = Name(prepared_statement);
+        let prepared_statement = Name::from(prepared_statement);
 
         let num_param_format_codes = cursor.get_i16();
         let mut param_format_codes = Vec::new();

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

@@ -0,0 +1,144 @@
+use crate::error::{Error, ProtocolError};
+use crate::postgresql::protocol::BytesMutReadString;
+use crate::{SIZE_I32, SIZE_U8};
+
+use bytes::{Buf, BufMut, BytesMut};
+use std::convert::TryFrom;
+use std::ffi::CString;
+use std::io::Cursor;
+
+use super::target::Target;
+use super::{FrontendCode, Name};
+
+///
+/// Close b'C' (Frontend) message.
+///
+/// See: <https://www.postgresql.org/docs/current/protocol-message-formats.html>
+///
+///     Byte1('C')
+///     Identifies the message as a Close command.
+///
+///     Int32
+///     Length of message contents in bytes, including self.
+///
+///     Byte1
+///     'S' to close a prepared statement; or 'P' to close a portal.
+///
+///     String
+///     The name of the prepared statement or portal to close (an empty string selects the unnamed prepared statement or portal).
+
+#[derive(Debug, Clone)]
+pub(crate) struct Close {
+    pub target: Target,
+    pub name: Name,
+}
+
+impl TryFrom<&BytesMut> for Close {
+    type Error = Error;
+
+    fn try_from(bytes: &BytesMut) -> Result<Close, Self::Error> {
+        let mut cursor = Cursor::new(bytes);
+        let code = cursor.get_u8();
+
+        if FrontendCode::from(code) != FrontendCode::Close {
+            return Err(ProtocolError::UnexpectedMessageCode {
+                expected: FrontendCode::Close.into(),
+                received: code as char,
+            }
+            .into());
+        }
+
+        let _len = cursor.get_i32(); // read and progress cursor
+        let target = cursor.get_u8();
+        let target = Target::try_from(target)?;
+        let name = cursor.read_string()?;
+        let name = Name::from(name);
+
+        Ok(Close { target, name })
+    }
+}
+
+impl TryFrom<Close> for BytesMut {
+    type Error = Error;
+
+    fn try_from(close: Close) -> Result<BytesMut, Error> {
+        let mut bytes = BytesMut::new();
+
+        let name = CString::new(close.name.as_str())?;
+        let name = name.as_bytes_with_nul();
+
+        let len = SIZE_I32 + SIZE_U8 + name.len();
+
+        bytes.put_u8(FrontendCode::Close.into());
+        bytes.put_i32(len as i32);
+        bytes.put_u8(close.target.into());
+        bytes.put_slice(name);
+
+        Ok(bytes)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{config::LogConfig, log, postgresql::messages::Name};
+    use bytes::BytesMut;
+    use std::convert::TryFrom;
+
+    fn to_message(s: &[u8]) -> BytesMut {
+        BytesMut::from(s)
+    }
+
+    #[test]
+    pub fn test_close_statement() {
+        log::init(LogConfig::default());
+
+        // Close unnamed prepared statement: C\0\0\0\x06S\0
+        let bytes = to_message(b"C\0\0\0\x06S\0");
+        let close = Close::try_from(&bytes).unwrap();
+
+        assert!(matches!(close.target, Target::Statement));
+        assert!(close.name.is_unnamed());
+    }
+
+    #[test]
+    pub fn test_close_portal() {
+        log::init(LogConfig::default());
+
+        // Close unnamed portal: C\0\0\0\x06P\0
+        let bytes = to_message(b"C\0\0\0\x06P\0");
+        let close = Close::try_from(&bytes).unwrap();
+
+        assert!(matches!(close.target, Target::Portal));
+        assert!(close.name.is_unnamed());
+    }
+
+    #[test]
+    pub fn test_close_named_statement() {
+        log::init(LogConfig::default());
+
+        // Close named prepared statement "stmt1": C\0\0\0\x0bSstmt1\0
+        let bytes = to_message(b"C\0\0\0\x0bSstmt1\0");
+        let close = Close::try_from(&bytes).unwrap();
+
+        assert!(matches!(close.target, Target::Statement));
+        assert_eq!(close.name.as_str(), "stmt1");
+        assert!(!close.name.is_unnamed());
+    }
+
+    #[test]
+    pub fn test_close_to_bytes() {
+        log::init(LogConfig::default());
+
+        let close = Close {
+            target: Target::Portal,
+            name: Name::from("portal1"),
+        };
+
+        let bytes = BytesMut::try_from(close).unwrap();
+        let parsed = Close::try_from(&bytes).unwrap();
+
+        assert!(matches!(parsed.target, Target::Portal));
+        assert_eq!(parsed.name.as_str(), "portal1");
+    }
+}

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

@@ -7,6 +7,7 @@ use std::convert::TryFrom;
 use std::ffi::CString;
 use std::io::Cursor;
 
+use super::target::Target;
 use super::{FrontendCode, Name};
 
 ///
@@ -32,29 +33,6 @@ pub(crate) struct Describe {
     pub name: Name,
 }
 
-///
-/// The target of the describe message.
-///
-/// Valid values are PreparedStatment or Portal
-///
-/// A Portal is a parsed statement PLUS any bound parameters
-/// Describe with `Target::Portal` returns the RowDescription describing the result set.
-/// The assuumption is that the parameters are already bound to the portal, so the Describe message is not required to include any parameter information.
-///
-/// Calls to Execute are made on a Portal (not a prepared statement) as execute requires any bound parameters
-///
-/// A PreparedStatement is the parsed statement
-/// Describe with `Target::PreparedStatement` returns a ParameterDescription followed by the RowDescription.
-///
-///
-/// See https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY
-///
-#[derive(Debug, Clone)]
-pub enum Target {
-    Portal,
-    PreparedStatement,
-}
-
 impl TryFrom<&BytesMut> for Describe {
     type Error = Error;
 
@@ -74,7 +52,7 @@ impl TryFrom<&BytesMut> for Describe {
         let target = cursor.get_u8();
         let target = Target::try_from(target)?;
         let name = cursor.read_string()?;
-        let name = Name(name);
+        let name = Name::from(name);
 
         Ok(Describe { target, name })
     }
@@ -86,28 +64,16 @@ impl TryFrom<Describe> for BytesMut {
     fn try_from(describe: Describe) -> Result<BytesMut, Error> {
         let mut bytes = BytesMut::new();
 
-        let name = CString::new(describe.name.0.as_str())?;
+        let name = CString::new(describe.name.as_str())?;
         let name = name.as_bytes_with_nul();
 
         let len = SIZE_I32 + SIZE_U8 + name.len();
 
         bytes.put_u8(FrontendCode::Describe.into());
         bytes.put_i32(len as i32);
-        bytes.put_u8(describe.target as u8);
+        bytes.put_u8(describe.target.into());
         bytes.put_slice(name);
 
         Ok(bytes)
     }
 }
-
-impl TryFrom<u8> for Target {
-    type Error = Error;
-
-    fn try_from(t: u8) -> Result<Target, Error> {
-        match t as char {
-            'S' => Ok(Target::PreparedStatement),
-            'P' => Ok(Target::Portal),
-            t => Err(ProtocolError::UnexpectedDescribeTarget(t).into()),
-        }
-    }
-}

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

@@ -29,7 +29,7 @@ impl TryFrom<&BytesMut> for Execute {
         let _len = cursor.get_i32(); // read and progress cursor
 
         let portal = cursor.read_string()?;
-        let portal = Name(portal);
+        let portal = Name::from(portal);
         let max_rows = cursor.get_i32();
 
         Ok(Execute { portal, max_rows })