From 351b3e8a57c9e0143b7b6f3ed2160dcdce00225e Mon Sep 17 00:00:00 2001
From: Ellie Huxtable <ellie@elliehuxtable.com>
Date: Mon, 18 Sep 2023 08:39:19 +0100
Subject: Add connect timeout and overall timeout (#1238)

* Add connect timeout and overall timeout

* Make it configurable

* Fix test

* Add docs
---
 atuin-client/src/api_client.rs  | 10 +++++++++-
 atuin-client/src/record/sync.rs | 14 ++++++++++++--
 atuin-client/src/settings.rs    |  5 +++++
 atuin-client/src/sync.rs        |  7 ++++++-
 4 files changed, 32 insertions(+), 4 deletions(-)

(limited to 'atuin-client')

diff --git a/atuin-client/src/api_client.rs b/atuin-client/src/api_client.rs
index d2ca339f..b6625a34 100644
--- a/atuin-client/src/api_client.rs
+++ b/atuin-client/src/api_client.rs
@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 use std::env;
+use std::time::Duration;
 
 use eyre::{bail, Result};
 use reqwest::{
@@ -106,7 +107,12 @@ pub async fn latest_version() -> Result<Version> {
 }
 
 impl<'a> Client<'a> {
-    pub fn new(sync_addr: &'a str, session_token: &'a str) -> Result<Self> {
+    pub fn new(
+        sync_addr: &'a str,
+        session_token: &'a str,
+        connect_timeout: u64,
+        timeout: u64,
+    ) -> Result<Self> {
         let mut headers = HeaderMap::new();
         headers.insert(AUTHORIZATION, format!("Token {session_token}").parse()?);
 
@@ -115,6 +121,8 @@ impl<'a> Client<'a> {
             client: reqwest::Client::builder()
                 .user_agent(APP_USER_AGENT)
                 .default_headers(headers)
+                .connect_timeout(Duration::new(connect_timeout, 0))
+                .timeout(Duration::new(timeout, 0))
                 .build()?,
         })
     }
diff --git a/atuin-client/src/record/sync.rs b/atuin-client/src/record/sync.rs
index ad9079ed..56be0638 100644
--- a/atuin-client/src/record/sync.rs
+++ b/atuin-client/src/record/sync.rs
@@ -22,7 +22,12 @@ pub enum Operation {
 }
 
 pub async fn diff(settings: &Settings, store: &mut impl Store) -> Result<(Vec<Diff>, RecordIndex)> {
-    let client = Client::new(&settings.sync_address, &settings.session_token)?;
+    let client = Client::new(
+        &settings.sync_address,
+        &settings.session_token,
+        settings.network_connect_timeout,
+        settings.network_timeout,
+    )?;
 
     let local_index = store.tail_records().await?;
     let remote_index = client.record_index().await?;
@@ -218,7 +223,12 @@ pub async fn sync_remote(
     local_store: &mut impl Store,
     settings: &Settings,
 ) -> Result<(i64, i64)> {
-    let client = Client::new(&settings.sync_address, &settings.session_token)?;
+    let client = Client::new(
+        &settings.sync_address,
+        &settings.session_token,
+        settings.network_connect_timeout,
+        settings.network_timeout,
+    )?;
 
     let mut uploaded = 0;
     let mut downloaded = 0;
diff --git a/atuin-client/src/settings.rs b/atuin-client/src/settings.rs
index e2d22e93..e49ab0cb 100644
--- a/atuin-client/src/settings.rs
+++ b/atuin-client/src/settings.rs
@@ -176,6 +176,9 @@ pub struct Settings {
     pub workspaces: bool,
     pub ctrl_n_shortcuts: bool,
 
+    pub network_connect_timeout: u64,
+    pub network_timeout: u64,
+
     // This is automatically loaded when settings is created. Do not set in
     // config! Keep secrets and settings apart.
     pub session_token: String,
@@ -372,6 +375,8 @@ impl Settings {
             .set_default("workspaces", false)?
             .set_default("ctrl_n_shortcuts", false)?
             .set_default("secrets_filter", true)?
+            .set_default("network_connect_timeout", 5)?
+            .set_default("network_timeout", 30)?
             .add_source(
                 Environment::with_prefix("atuin")
                     .prefix_separator("_")
diff --git a/atuin-client/src/sync.rs b/atuin-client/src/sync.rs
index 439ed6d0..bd2ff474 100644
--- a/atuin-client/src/sync.rs
+++ b/atuin-client/src/sync.rs
@@ -189,7 +189,12 @@ async fn sync_upload(
 }
 
 pub async fn sync(settings: &Settings, force: bool, db: &mut (impl Database + Send)) -> Result<()> {
-    let client = api_client::Client::new(&settings.sync_address, &settings.session_token)?;
+    let client = api_client::Client::new(
+        &settings.sync_address,
+        &settings.session_token,
+        settings.network_connect_timeout,
+        settings.network_timeout,
+    )?;
 
     let key = load_key(settings)?; // encryption key
 
-- 
cgit v1.3.1