Add EnvVar as a parameter

crates/criticalup-cli/src/commands/auth_remove.rs

@@ -3,12 +3,14 @@
 
 use crate::errors::Error;
 use crate::Context;
-use criticalup_core::state::State;
+use criticalup_core::state::{EnvVars, State};
 
 pub(crate) async fn run(ctx: &Context) -> Result<(), Error> {
     let state = State::load(&ctx.config).await?;
 
-    if state.authentication_token(None).await.is_some() {
+    let env_vars = EnvVars::default();
+
+    if state.authentication_token(None, &env_vars).await.is_some() {
         state.set_authentication_token(None);
         state.persist().await?;
     }

crates/criticalup-core/src/download_server_client.rs

@@ -15,6 +15,7 @@ use reqwest_middleware::{ClientBuilder, ClientWithMiddleware, RequestBuilder};
 use reqwest_retry::policies::ExponentialBackoff;
 use reqwest_retry::RetryTransientMiddleware;
 use serde::Deserialize;
+use crate::state;
 
 const CLIENT_MAX_RETRIES: u32 = 5;
 
@@ -124,9 +125,11 @@ impl DownloadServerClient {
             None
         };
 
+        let env_vars = state::EnvVars::default();
+
         let header = self
             .state
-            .authentication_token(path_to_token_file)
+            .authentication_token(path_to_token_file, &env_vars)
             .await
             .as_ref()
             .and_then(|token| HeaderValue::from_str(&format!("Bearer {}", token.unseal())).ok());

crates/criticalup-core/src/state.rs

@@ -60,20 +60,21 @@ impl State {
     pub async fn authentication_token(
         &self,
         token_path: Option<&str>,
+        env_vars: &EnvVars,
     ) -> Option<AuthenticationToken> {
         match token_path {
             Some(token_path) => {
                 let token_path = std::path::Path::new(token_path);
                 if token_path.exists() {
                     match tokio::fs::read_to_string(token_path).await {
                         Ok(token) => Some(AuthenticationToken(token.to_string().trim().into())),
-                        Err(_) => self.authentication_token_inner(),
+                        Err(_) => self.authentication_token_inner(env_vars),
                     }
                 } else {
-                    self.authentication_token_inner()
+                    self.authentication_token_inner(env_vars)
                 }
             }
-            None => self.authentication_token_inner(),
+            None => self.authentication_token_inner(env_vars),
         }
     }
 
@@ -82,14 +83,15 @@ impl State {
     /// Attempts to read from:
     ///  1. The `CRITICALUP_TOKEN_ENV_VAR_NAME` environment
     ///  2. The state
-    fn authentication_token_inner(&self) -> Option<AuthenticationToken> {
-        match std::env::var(CRITICALUP_TOKEN_ENV_VAR_NAME) {
-            Ok(token_from_env) => Some(AuthenticationToken(token_from_env)),
-            Err(_) => {
+    fn authentication_token_inner(&self, env_vars: &EnvVars) -> Option<AuthenticationToken> {
+        match &env_vars.criticalup_token {
+            Some(e) => Some(AuthenticationToken(e.to_string())),
+            None => {
                 let borrowed = self.inner.borrow();
                 borrowed.repr.authentication_token.clone()
             }
         }
+
     }
 
     pub fn set_authentication_token(&self, token: Option<AuthenticationToken>) {
@@ -374,6 +376,18 @@ impl StateInstallation {
     }
 }
 
+pub struct EnvVars {
+    criticalup_token: Option<String>,
+}
+
+impl Default for EnvVars {
+    fn default() -> Self {
+        EnvVars {
+            criticalup_token: std::env::var(CRITICALUP_TOKEN_ENV_VAR_NAME).ok()
+        }
+    }
+}
+
 #[derive(Clone, Serialize, Deserialize)]
 #[cfg_attr(test, derive(PartialEq, Eq))]
 pub struct AuthenticationToken(String);
@@ -420,6 +434,7 @@ impl std::fmt::Debug for AuthenticationToken {
 #[cfg(test)]
 mod tests {
     use crate::test_utils::TestEnvironment;
+    use crate::state::EnvVars;
 
     use super::*;
 
@@ -431,6 +446,8 @@ mod tests {
         }}
     }
 
+    const ENV_VARS: EnvVars = EnvVars { criticalup_token: None };
+
     #[tokio::test]
     async fn test_load_state_without_existing_file() {
         let test = TestEnvironment::prepare().await;
@@ -460,7 +477,7 @@ mod tests {
         let state = State::load(test_env.config()).await.unwrap();
         assert_eq!(
             Some(AuthenticationToken("hello".into())),
-            state.authentication_token(None).await
+            state.authentication_token(None, &ENV_VARS).await
         );
     }
 
@@ -751,7 +768,7 @@ mod tests {
         state.set_authentication_token(None);
 
         //  Make sure the state file has authentication token as None.
-        assert_eq!(state.authentication_token(None).await, None);
+        assert_eq!(state.authentication_token(None, &ENV_VARS).await, None);
 
         let file_token_content = "my-awesome-token-from-file";
         let token_name = "CRITICALUP_TOKEN";
@@ -763,7 +780,7 @@ mod tests {
         std::fs::write(secrets_dir.join(token_name), file_token_content).unwrap();
         let token = test_env
             .state()
-            .authentication_token(Some(secrets_dir.join(token_name).to_str().unwrap()))
+            .authentication_token(Some(secrets_dir.join(token_name).to_str().unwrap()), &ENV_VARS)
             .await;
         assert_eq!(Some(AuthenticationToken(file_token_content.into())), token)
     }
@@ -774,12 +791,12 @@ mod tests {
         let state = test_env.state();
 
         state.set_authentication_token(None);
-        assert_eq!(None, state.authentication_token(None).await);
+        assert_eq!(None, state.authentication_token(None, &ENV_VARS).await);
 
         state.set_authentication_token(Some(AuthenticationToken("hello world".into())));
         assert_eq!(
             Some(AuthenticationToken("hello world".into())),
-            state.authentication_token(None).await
+            state.authentication_token(None, &ENV_VARS).await
         );
     }
 
@@ -794,7 +811,7 @@ mod tests {
         test_env.state().persist().await.unwrap();
 
         let new_state = State::load(test_env.config()).await.unwrap();
-        assert_eq!(Some(token), new_state.authentication_token(None).await);
+        assert_eq!(Some(token), new_state.authentication_token(None, &ENV_VARS).await);
     }
 
     #[tokio::test]
@@ -1009,4 +1026,24 @@ mod tests {
             unused_installations
         )
     }
+
+    #[tokio::test]
+    async fn test_set_authn_token_env_var() {
+        let test_env = TestEnvironment::with().state().prepare().await;
+        let state = test_env.state();
+
+        let token = "my_awesome_token".to_string();
+
+        let env_vars = EnvVars {
+            criticalup_token: Some(token.clone()),
+        };
+
+        state.set_authentication_token(None);
+        assert_eq!(Some(AuthenticationToken(token)), state.authentication_token(None, &env_vars).await);
+
+        let env_vars = EnvVars {
+            criticalup_token: None,
+        };
+        assert_eq!(None, state.authentication_token(None, &env_vars).await);
+    }
 }