add --backend-connection-timeout arg

demarches-simplifiees · Sep 27, 2023 · 300bc04 · 300bc04
1 parent 11658be
commit 300bc04
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/src/args.rs b/src/args.rs
@@ -6,7 +6,7 @@ DS encryption proxy.
 Usage:
   ds_proxy encrypt <input-file> <output-file> [--password-file=<password-file>] [--salt=<salt>] [--chunk-size=<chunk-size>] [--keyring-file=<keyring-file>]
   ds_proxy decrypt <input-file> <output-file> [--password-file=<password-file>] [--salt=<salt>] [--chunk-size=<chunk-size>] [--keyring-file=<keyring-file>]
-  ds_proxy proxy [--address=<address>] [--password-file=<password-file>] [--salt=<salt>] [--chunk-size=<chunk-size>] [--upstream-url=<upstream-url>] [--local-encryption-directory=<local-encryption-directory>] [--keyring-file=<keyring-file>] [--aws-access-key=<aws-access-key>] [--aws-secret-key=<aws-secret-key>] [--aws-region=<aws-region>]
+  ds_proxy proxy [--address=<address>] [--password-file=<password-file>] [--salt=<salt>] [--chunk-size=<chunk-size>] [--upstream-url=<upstream-url>] [--local-encryption-directory=<local-encryption-directory>] [--keyring-file=<keyring-file>] [--aws-access-key=<aws-access-key>] [--aws-secret-key=<aws-secret-key>] [--aws-region=<aws-region>] [--backend-connection-timeout=<backend-connection-timeout>]
   ds_proxy add-key [--password-file=<password-file>] [--salt=<salt>] [--keyring-file=<keyring-file>]
   ds_proxy (-h | --help)
   ds_proxy --version
@@ -30,6 +30,7 @@ pub struct Args {
     pub flag_aws_access_key: Option<String>,
     pub flag_aws_secret_key: Option<String>,
     pub flag_aws_region: Option<String>,
+    pub flag_backend_connection_timeout: Option<u64>,
     pub cmd_encrypt: bool,
     pub cmd_decrypt: bool,
     pub cmd_proxy: bool,

diff --git a/src/config.rs b/src/config.rs
@@ -5,6 +5,7 @@ use std::env;
 use std::net::{SocketAddr, ToSocketAddrs};
 use std::path::PathBuf;
 use url::Url;
+use std::time::Duration;
 
 // match nginx default (proxy_buffer_size in ngx_stream_proxy_module)
 pub const DEFAULT_CHUNK_SIZE: usize = 16 * 1024;
@@ -42,6 +43,7 @@ pub struct HttpConfig {
     pub aws_access_key: Option<String>,
     pub aws_secret_key: Option<String>,
     pub aws_region: Option<String>,
+    pub backend_connection_timeout: Duration,
 }
 
 #[derive(Debug, Clone)]
@@ -149,6 +151,16 @@ impl Config {
             }
             .unwrap();
 
+            let backend_connection_timeout = match&args.flag_backend_connection_timeout {
+                Some(timeout_u64) => Duration::from_secs(*timeout_u64),
+                None => match env::var("BACKEND_CONNECTION_TIMEOUT") {
+                    Ok(timeout_string) => Duration::from_secs(timeout_string.parse().expect("BACKEND_CONNECTION_TIMEOUT is not a u64")),
+                    _ => Duration::from_secs(1)
+                }
+            };
+
+            log::info!("backend_connection_timeout: {:?}", backend_connection_timeout);
+
             Config::Http(HttpConfig {
                 keyring,
                 chunk_size,
@@ -158,6 +170,7 @@ impl Config {
                 aws_access_key: args.flag_aws_access_key.clone(),
                 aws_secret_key: args.flag_aws_secret_key.clone(),
                 aws_region: args.flag_aws_region.clone(),
+                backend_connection_timeout,
             })
         }
     }
@@ -322,6 +335,7 @@ mod tests {
             aws_access_key: None,
             aws_secret_key: None,
             aws_region: None,
+            backend_connection_timeout: Duration::from_secs(1)
         }
     }
 }
diff --git a/src/http/proxy.rs b/src/http/proxy.rs
@@ -12,7 +12,6 @@ use actix_web::dev::Service;
 use futures::FutureExt;
 use std::time::Duration;
 
-const CONNECT_TIMEOUT: Duration = Duration::from_secs(1);
 const RESPONSE_TIMEOUT: Duration = Duration::from_secs(30);
 
 #[actix_web::main]
@@ -24,7 +23,7 @@ pub async fn main(config: HttpConfig) -> std::io::Result<()> {
             .app_data(Data::new(
                 awc::Client::builder()
                     .connector(
-                        awc::Connector::new().timeout(CONNECT_TIMEOUT), // max time to connect to remote host including dns name resolution
+                        awc::Connector::new().timeout(config.backend_connection_timeout), // max time to connect to remote host including dns name resolution
                     )
                     .timeout(RESPONSE_TIMEOUT) // the total time before a response must be received
                     .finish(),