Merge pull request #180 from cipherstash/optional-connection-timeout

fix: make connection timeout optional

Author: auxesis

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

@@ -19,8 +19,7 @@ pub struct DatabaseConfig {
     #[serde(deserialize_with = "protected_string_deserializer")]
     password: Protected<String>,
 
-    #[serde(default = "DatabaseConfig::default_connection_timeout")]
-    pub connection_timeout: u64,
+    pub connection_timeout: Option<u64>,
 
     #[serde(default)]
     pub with_tls_verification: bool,
@@ -41,11 +40,6 @@ impl DatabaseConfig {
         5432
     }
 
-    // 5 minutes
-    pub const fn default_connection_timeout() -> u64 {
-        1000 * 60 * 5
-    }
-
     pub const fn default_config_reload_interval() -> u64 {
         60
     }
@@ -73,8 +67,8 @@ impl DatabaseConfig {
         self.password.to_owned().risky_unwrap()
     }
 
-    pub fn connection_timeout(&self) -> Duration {
-        Duration::from_millis(self.connection_timeout)
+    pub fn connection_timeout(&self) -> Option<Duration> {
+        self.connection_timeout.map(Duration::from_millis)
     }
 
     pub fn server_name(&self) -> Result<ServerName, Error> {

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

@@ -14,6 +14,7 @@ use tokio::{
 use tokio_postgres::Client;
 use tracing::{debug, error, info, warn};
 
+const TCP_USER_TIMEOUT: Duration = Duration::from_secs(10);
 const TCP_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(5);
 const TCP_KEEPALIVE_TIME: Duration = Duration::from_secs(5);
 const TCP_KEEPALIVE_RETRIES: u32 = 5;
@@ -120,6 +121,17 @@ pub fn configure(stream: &TcpStream) {
         );
     });
 
+    #[cfg(target_os = "linux")]
+    match sock_ref.set_tcp_user_timeout(Some(TCP_USER_TIMEOUT)) {
+        Ok(_) => (),
+        Err(err) => {
+            warn!(
+                msg = "Error configuring tcp_user_timeout for connection",
+                error = err.to_string()
+            );
+        }
+    }
+
     match sock_ref.set_keepalive(true) {
         Ok(_) => {
             let params = &TcpKeepalive::new()

packages/cipherstash-proxy/src/error.rs

@@ -2,9 +2,8 @@ use crate::{postgresql::Column, Identifier};
 use bytes::BytesMut;
 use cipherstash_client::encryption;
 use metrics_exporter_prometheus::BuildError;
-use std::io;
+use std::{io, time::Duration};
 use thiserror::Error;
-use tokio::time::error::Elapsed;
 
 const ERROR_DOC_BASE_URL: &str = "https://github.com/cipherstash/proxy/blob/main/docs/errors.md";
 
@@ -25,8 +24,8 @@ pub enum Error {
     #[error("Connection closed by client")]
     ConnectionClosed,
 
-    #[error("Connection timed out")]
-    ConnectionTimeout(#[from] Elapsed),
+    #[error("Connection timed out after {} ms", duration.as_secs())]
+    ConnectionTimeout { duration: Duration },
 
     #[error("Error creating connection")]
     DatabaseConnection,

packages/cipherstash-proxy/src/main.rs

@@ -111,8 +111,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                                     Error::CancelRequest => {
                                         info!(msg = "Database connection closed after cancel request");
                                     }
-                                    Error::ConnectionTimeout(_) => {
-                                        warn!(msg = "Database connection timeout");
+                                    Error::ConnectionTimeout{..} => {
+                                        warn!(msg = "Database connection timeout", error = err.to_string());
                                     }
                                     _ => {
                                         error!(msg = "Database connection error", error = err.to_string());

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

@@ -81,7 +81,7 @@ where
     pub async fn rewrite(&mut self) -> Result<(), Error> {
         let connection_timeout = self.encrypt.config.database.connection_timeout();
 
-        let (code, mut bytes) = protocol::read_message_with_timeout(
+        let (code, mut bytes) = protocol::read_message(
             &mut self.server_reader,
             self.context.client_id,
             connection_timeout,

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

@@ -76,7 +76,7 @@ where
 
     pub async fn rewrite(&mut self) -> Result<(), Error> {
         let connection_timeout = self.encrypt.config.database.connection_timeout();
-        let (code, mut bytes) = protocol::read_message_with_timeout(
+        let (code, mut bytes) = protocol::read_message(
             &mut self.client_reader,
             self.context.client_id,
             connection_timeout,

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

@@ -70,7 +70,7 @@ pub async fn handler(
     );
 
     loop {
-        let startup_message = startup::read_message_with_timeout(
+        let startup_message = startup::read_message(
             &mut client_stream,
             encrypt.config.database.connection_timeout(),
         )
@@ -125,8 +125,7 @@ pub async fn handler(
 
         let connection_timeout = encrypt.config.database.connection_timeout();
         let (_code, bytes) =
-            protocol::read_message_with_timeout(&mut client_stream, client_id, connection_timeout)
-                .await?;
+            protocol::read_message(&mut client_stream, client_id, connection_timeout).await?;
 
         let password_message = PasswordMessage::try_from(&bytes)?;
 

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

@@ -16,6 +16,7 @@ pub const CODE_INVALID_PASSWORD: &str = "28P01";
 pub const CODE_RAISE_EXCEPTION: &str = "P0001";
 pub const CODE_SYNTAX_ERROR: &str = "42601";
 pub const CODE_INVALID_TEXT_REPRESENTATION: &str = "22P02";
+pub const CODE_IDLE_SESSION_TIMEOUT: &str = "57P05";
 
 ///
 /// ErrorResponse (B)
@@ -58,6 +59,29 @@ pub enum ErrorResponseCode {
 }
 
 impl ErrorResponse {
+    pub fn connection_timeout() -> Self {
+        Self {
+            fields: vec![
+                Field {
+                    code: ErrorResponseCode::Severity,
+                    value: "FATAL".to_string(),
+                },
+                Field {
+                    code: ErrorResponseCode::SeverityLegacy,
+                    value: "FATAL".to_string(),
+                },
+                Field {
+                    code: ErrorResponseCode::Code,
+                    value: CODE_IDLE_SESSION_TIMEOUT.to_string(),
+                },
+                Field {
+                    code: ErrorResponseCode::Message,
+                    value: "Connection timeout".to_string(),
+                },
+            ],
+        }
+    }
+
     pub fn invalid_password(message: &str) -> Self {
         Self {
             fields: vec![

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

@@ -81,18 +81,37 @@ pub async fn read_auth_message<S: AsyncRead + Unpin>(
     Authentication::try_from(&bytes)
 }
 
+///
+/// Reads a Postgres message from client with an optional timeout
+///
+/// Timeout values are in config
+///
+///
+pub async fn read_message<S: AsyncRead + Unpin>(
+    mut stream: S,
+    client_id: i32,
+    connection_timeout: Option<Duration>,
+) -> Result<(Code, BytesMut), Error> {
+    match connection_timeout {
+        Some(duration) => read_message_with_timeout(stream, client_id, duration).await,
+        None => read(&mut stream, client_id).await,
+    }
+}
+
 ///
 /// Reads a Postgres message from client with a timeout
 ///
 /// Timeout values are in config
 ///
 ///
-pub async fn read_message_with_timeout<S: AsyncRead + Unpin>(
+async fn read_message_with_timeout<S: AsyncRead + Unpin>(
     mut stream: S,
     client_id: i32,
-    connection_timeout: Duration,
+    duration: Duration,
 ) -> Result<(Code, BytesMut), Error> {
-    timeout(connection_timeout, read_message(&mut stream, client_id)).await?
+    timeout(duration, read(&mut stream, client_id))
+        .await
+        .map_err(|_| Error::ConnectionTimeout { duration })?
 }
 
 ///
@@ -102,7 +121,7 @@ pub async fn read_message_with_timeout<S: AsyncRead + Unpin>(
 /// Byte is then passed as `code` to this function to preserve the message structure
 ///
 ///
-pub async fn read_message<S: AsyncRead + Unpin>(
+async fn read<S: AsyncRead + Unpin>(
     mut stream: S,
     client_id: i32,
 ) -> Result<(Code, BytesMut), Error> {

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

@@ -46,14 +46,35 @@ pub async fn with_tls(stream: AsyncStream, config: &TandemConfig) -> Result<Asyn
     }
 }
 
-pub async fn read_message_with_timeout<C>(
-    client: &mut C,
-    connection_timeout: Duration,
-) -> Result<StartupMessage, Error>
-where
-    C: AsyncRead + Unpin,
-{
-    timeout(connection_timeout, read_message(client)).await?
+///
+/// Reads a Postgres startup message from client with an optional timeout
+///
+/// Timeout values are in config
+///
+///
+pub async fn read_message<S: AsyncRead + Unpin>(
+    mut stream: S,
+    connection_timeout: Option<Duration>,
+) -> Result<StartupMessage, Error> {
+    match connection_timeout {
+        Some(duration) => read_message_with_timeout(stream, duration).await,
+        None => read(&mut stream).await,
+    }
+}
+
+///
+/// Reads a Postgres message from client with a timeout
+///
+/// Timeout values are in config
+///
+///
+async fn read_message_with_timeout<S: AsyncRead + Unpin>(
+    mut stream: S,
+    duration: Duration,
+) -> Result<StartupMessage, Error> {
+    timeout(duration, read(&mut stream))
+        .await
+        .map_err(|_| Error::ConnectionTimeout { duration })?
 }
 
 ///
@@ -62,7 +83,7 @@ where
 ///
 ///
 ///
-async fn read_message<C>(client: &mut C) -> Result<StartupMessage, Error>
+async fn read<C>(client: &mut C) -> Result<StartupMessage, Error>
 where
     C: AsyncRead + Unpin,
 {