diff --git a/netmito/src/client/http.rs b/netmito/src/client/http.rs index dfc7360..cc84fed 100644 --- a/netmito/src/client/http.rs +++ b/netmito/src/client/http.rs @@ -36,6 +36,7 @@ impl MitoHttpClient { credential_path: Option, user: Option, password: Option, + refresh: bool, ) -> crate::error::Result { let client_credential_path = credential_path .as_ref() @@ -56,6 +57,7 @@ impl MitoHttpClient { self.url.clone(), user, password, + refresh, ) .await?; self.credential_path = client_credential_path; diff --git a/netmito/src/client/mod.rs b/netmito/src/client/mod.rs index e77cdf2..9a5ae70 100644 --- a/netmito/src/client/mod.rs +++ b/netmito/src/client/mod.rs @@ -84,25 +84,23 @@ impl MitoClient { pub async fn setup(config: ClientConfig) -> crate::error::Result { tracing::debug!("Client is setting up"); - let refresh = config.refresh; let mut http_client = MitoHttpClient::new(config.coordinator_addr); let username = http_client - .connect(config.credential_path, config.user, config.password) + .connect( + config.credential_path, + config.user, + config.password, + config.refresh, + ) .await?; - let mut client = MitoClient { + Ok(MitoClient { http_client, username, redis_client: None, redis_pubsub_client: None, async_redis_client: None, - }; - - if refresh { - client.refresh_token().await?; - } - - Ok(client) + }) } pub fn http_client(&self) -> &http::MitoHttpClient { diff --git a/netmito/src/service/auth/cred.rs b/netmito/src/service/auth/cred.rs index a10c7d3..0f0f57f 100644 --- a/netmito/src/service/auth/cred.rs +++ b/netmito/src/service/auth/cred.rs @@ -179,6 +179,7 @@ pub async fn get_user_credential( mut url: Url, user: Option, password: Option, + refresh: bool, ) -> crate::error::Result<(String, String)> { // Try to load credential from file let cred_path = cred_path @@ -197,26 +198,38 @@ pub async fn get_user_credential( if cred_path.exists() { if let Ok(mut lines) = read_lines(&cred_path).await { if let Some((username, cred)) = extract_credential(user.as_ref(), &mut lines).await? { - url.set_path("auth"); - let resp = client - .get(url.as_str()) - .bearer_auth(&cred) - .send() - .await - .map_err(|e| { - if e.is_request() && e.is_connect() { - url.set_path(""); - RequestError::ConnectionError(url.to_string()) - } else { - e.into() + let request = if refresh { + url.set_path("refresh"); + client.post(url.as_str()) + } else { + url.set_path("auth"); + client.get(url.as_str()) + }; + let resp = request.bearer_auth(&cred).send().await.map_err(|e| { + if e.is_request() && e.is_connect() { + url.set_path(""); + RequestError::ConnectionError(url.to_string()) + } else { + e.into() + } + })?; + let status = resp.status(); + if status.is_success() { + if refresh { + let resp = resp + .json::() + .await + .map_err(RequestError::from)?; + let token = resp.token; + modify_or_append_credential(&cred_path, &username, &token).await?; + return Ok((username, token)); + } else { + let resp_name = resp.text().await.map_err(RequestError::from)?; + if resp_name == username { + return Ok((username, cred)); } - })?; - if resp.status().is_success() { - let resp_name = resp.text().await.map_err(RequestError::from)?; - if resp_name == username { - return Ok((username, cred)); } - } else if resp.status().is_server_error() { + } else if status.is_server_error() { return Err(ApiError::InternalServerError.into()); } } @@ -224,7 +237,7 @@ pub async fn get_user_credential( } // Local credential not found or invalid, need to login tracing::warn!("Local credential not found or invalid, need to login"); - let req = fill_user_login(user, password, false)?; + let req = fill_user_login(user, password, refresh)?; url.set_path("login"); let resp = client .post(url.as_str()) diff --git a/netmito/src/worker.rs b/netmito/src/worker.rs index 7b1bc69..b9607eb 100644 --- a/netmito/src/worker.rs +++ b/netmito/src/worker.rs @@ -257,6 +257,7 @@ impl MitoWorker { config.coordinator_addr.clone(), config.user.take(), config.password.take(), + false, ) .await?; let mut url = config.coordinator_addr.clone();