Merge pull request #1 from NOLAI/multiple-auths

Support multiple auths and errors

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "paas-client"
-version = "0.5.8"
+version = "0.6.0"
 authors = [
     "Job Doesburg <[email protected]>",
     "Julian van der Horst <[email protected]"
@@ -36,6 +36,8 @@ chrono = "0.4.39"
 clap = { version = "4.5.27", optional = true }
 tokio = { version = "1.0.0", features = ["rt", "rt-multi-thread", "macros"] }
 base64 = "0.22.1"
+async-trait = "0.1.86"
+thiserror = "2.0.11"
 
 [dev-dependencies]
 tokio = { version = "1", features = ["full"] }

README.md

@@ -49,12 +49,11 @@ let config = PseudonymServiceConfig {
     ],
 };
 
-let auth_tokens = AuthTokens(HashMap::from([
-    ("test_system_1".to_string(), "test_token_1".to_string()),
-    ("test_system_2".to_string(), "test_token_2".to_string()),
+let auths = SystemAuths::from_auths(HashMap::from([
+    ("test_system_1".to_string(), BearerTokenAuth::new("test_token_1".to_string())),
+    ("test_system_2".to_string(), BearerTokenAuth::new("test_token_2".to_string())),
 ]));
 
-let mut service = PseudonymService::new(config, auth_tokens);
 
 let encrypted_pseudonym = EncryptedPseudonym::from_base64("nr3FRadpFFGCFksYgrloo5J2V9j7JJWcUeiNBna66y78lwMia2-l8He4FfJPoAjuHCpH-8B0EThBr8DS3glHJw==").unwrap();
 let sessions = EncryptionContexts(HashMap::from([
@@ -65,7 +64,7 @@ let sessions = EncryptionContexts(HashMap::from([
 let domain_from = PseudonymizationDomain::from("domain1");
 let domain_to = PseudonymizationDomain::from("domain2");
 
-let mut service = PseudonymService::new(config, auth_tokens);
-let result = service.pseudonymize(&encrypted_pseudonym, &sessions, &domain_from, &domain_to).await;
-let pseudonym = service.decrypt(result).await;
+let mut service = PseudonymService::new(config, &auths).expect("Failed to create service");
+let result = service.pseudonymize(&encrypted_pseudonym, &sessions, &domain_from, &domain_to).await.expect("Failed to pseudonymize");
+let pseudonym = service.decrypt(result).await.expect("Failed to decrypt");
 ```

src/auth.rs

@@ -1,23 +1,97 @@
-use crate::transcryptor_client::AuthToken;
+use async_trait::async_trait;
 use paas_api::status::SystemId;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
-use std::str::FromStr;
+use std::error::Error;
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct AuthTokens(pub HashMap<SystemId, AuthToken>);
+#[derive(Debug, thiserror::Error)]
+pub enum AuthError {
+    #[error("Failed to get token: {0}")]
+    TokenError(String),
+    #[error("Failed to set header: {0}")]
+    HeaderError(#[from] reqwest::header::InvalidHeaderValue),
+}
+
+#[async_trait]
+pub trait Auth {
+    fn token_type(&self) -> &str;
+    async fn token(&self) -> Result<String, Box<dyn Error + Send + Sync>>;
+
+    async fn authenticate(
+        &self,
+        headers: &mut reqwest::header::HeaderMap,
+    ) -> Result<(), Box<dyn Error + Send + Sync>> {
+        let auth_value = format!("{} {}", self.token_type(), self.token().await?);
+        headers.insert(
+            reqwest::header::AUTHORIZATION,
+            reqwest::header::HeaderValue::from_str(&auth_value)?,
+        );
+        Ok(())
+    }
+}
 
-impl AuthTokens {
-    pub fn get(&self, system_id: &SystemId) -> Option<&AuthToken> {
-        self.0.get(system_id)
+#[async_trait]
+pub(crate) trait RequestBuilderExt {
+    async fn with_auth<A: Auth + Send + Sync + ?Sized>(
+        self,
+        auth: &A,
+    ) -> Result<reqwest::RequestBuilder, AuthError>;
+}
+
+#[async_trait]
+impl RequestBuilderExt for reqwest::RequestBuilder {
+    async fn with_auth<A: Auth + Send + Sync + ?Sized>(
+        self,
+        auth: &A,
+    ) -> Result<reqwest::RequestBuilder, AuthError> {
+        let auth_value = format!(
+            "{} {}",
+            auth.token_type(),
+            auth.token()
+                .await
+                .map_err(|e| AuthError::TokenError(e.to_string()))?
+        );
+        Ok(self.header(reqwest::header::AUTHORIZATION, auth_value))
     }
 }
 
-impl FromStr for AuthTokens {
-    type Err = serde_json::Error;
+#[derive(Clone, Serialize, Deserialize)]
+pub struct BearerTokenAuth {
+    token: String,
+}
 
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        let map: HashMap<SystemId, AuthToken> = serde_json::from_str(s)?;
-        Ok(Self(map))
+impl BearerTokenAuth {
+    pub fn new(token: String) -> BearerTokenAuth {
+        BearerTokenAuth { token }
+    }
+}
+
+#[async_trait]
+impl Auth for BearerTokenAuth {
+    fn token_type(&self) -> &str {
+        "Bearer"
+    }
+
+    async fn token(&self) -> Result<String, Box<dyn Error + Send + Sync>> {
+        Ok(self.token.clone())
+    }
+}
+
+pub struct SystemAuths(pub HashMap<SystemId, Box<dyn Auth + Send + Sync>>);
+
+impl SystemAuths {
+    pub fn new(auths: HashMap<SystemId, Box<dyn Auth + Send + Sync>>) -> Self {
+        Self(auths)
+    }
+    pub fn from_auths<A: Auth + Send + Sync + 'static>(auths: HashMap<SystemId, A>) -> Self {
+        Self::new(
+            auths
+                .into_iter()
+                .map(|(id, auth)| (id, Box::new(auth) as Box<dyn Auth + Send + Sync>))
+                .collect(),
+        )
+    }
+    pub fn get(&self, system_id: &SystemId) -> Option<&(dyn Auth + Send + Sync)> {
+        self.0.get(system_id).map(|auth| auth.as_ref())
     }
 }

src/bin/commands/encrypt.rs

@@ -12,7 +12,7 @@ pub fn command() -> Command {
     )
 }
 
-pub async fn execute(matches: &clap::ArgMatches, service: &mut PseudonymService) {
+pub async fn execute(matches: &clap::ArgMatches, service: &mut PseudonymService<'_>) {
     let pseudonym_string = matches
         .get_one::<String>("pseudonym")
         .expect("pseudonym is required");
@@ -21,7 +21,10 @@ pub async fn execute(matches: &clap::ArgMatches, service: &mut PseudonymService)
 
     let rng = &mut OsRng;
 
-    let (encrypted, sessions) = service.encrypt(&pseudonym, rng).await;
+    let (encrypted, sessions) = service
+        .encrypt(&pseudonym, rng)
+        .await
+        .expect("Failed to encrypt");
 
     println!("Encrypted pseudonym: {}", encrypted.as_base64());
     println!("Sessions: {}", sessions.encode());

src/bin/commands/pseudonymize.rs

@@ -46,7 +46,7 @@ pub fn command() -> Command {
         )
 }
 
-pub async fn execute(matches: &clap::ArgMatches, service: &mut PseudonymService) {
+pub async fn execute(matches: &clap::ArgMatches, service: &mut PseudonymService<'_>) {
     let encrypted_pseudonym_str = matches
         .get_one::<String>("encrypted_pseudonym")
         .expect("encrypted_pseudonym is required");
@@ -75,13 +75,14 @@ pub async fn execute(matches: &clap::ArgMatches, service: &mut PseudonymService)
 
     let result = service
         .pseudonymize(&encrypted_pseudonym, &sessions, &domain_from, &domain_to)
-        .await;
+        .await
+        .expect("Failed to pseudonymize");
 
     if matches.get_flag("no_decrypt") {
         eprint!("Transcryption returned: ");
         println!("{}", &result.as_base64());
     } else {
-        let pseudonym = service.decrypt(&result).await;
+        let pseudonym = service.decrypt(&result).await.expect("Failed to decrypt");
         eprint!("Decrypted pseudonym: ");
         println!("{}", &pseudonym.encode_as_hex());
     }

src/bin/paascli.rs

@@ -2,10 +2,12 @@ mod commands;
 
 use clap::{Arg, Command};
 use libpep::high_level::keys::{SessionPublicKey, SessionSecretKey};
-use paas_client::auth::AuthTokens;
+use paas_api::status::SystemId;
+use paas_client::auth::{BearerTokenAuth, SystemAuths};
 use paas_client::pseudonym_service::{PseudonymService, PseudonymServiceConfig};
 use paas_client::sessions::EncryptionContexts;
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 use std::fs;
 
 #[derive(Serialize, Deserialize)]
@@ -60,21 +62,28 @@ async fn main() {
         .get_one::<String>("tokens")
         .expect("tokens path is required");
     let tokens_contents = fs::read_to_string(tokens_path).expect("Failed to read tokens file");
-    let tokens: AuthTokens =
+    let tokens_map: HashMap<SystemId, BearerTokenAuth> =
         serde_json::from_str(&tokens_contents).expect("Failed to parse tokens");
+    let auths: SystemAuths = SystemAuths::from_auths(tokens_map);
 
     // Restore the service from the state dump if it exists
     let state_path = matches.get_one::<String>("state_path");
     let mut service = if let Some(path) = state_path {
-        if let Ok(contents) = fs::read_to_string(path) {
-            let dump: PseudonymServiceDump =
-                serde_json::from_str(&contents).expect("Failed to deserialize service state");
-            PseudonymService::restore(config, tokens, dump.sessions, dump.session_keys)
-        } else {
-            PseudonymService::new(config, tokens)
+        match fs::read_to_string(path) {
+            Ok(contents) => {
+                let dump: PseudonymServiceDump = serde_json::from_str(&contents)
+                    .expect("Failed to deserialize service state from file");
+
+                PseudonymService::restore(config, &auths, dump.sessions, dump.session_keys)
+                    .expect("Failed to restore service from state")
+            }
+            Err(e) => {
+                eprintln!("Failed to read state file: {}, creating new service", e);
+                PseudonymService::new(config, &auths).expect("Failed to create new service")
+            }
         }
     } else {
-        PseudonymService::new(config, tokens)
+        PseudonymService::new(config, &auths).expect("Failed to create new service")
     };
 
     // Execute the subcommand
@@ -92,7 +101,7 @@ async fn main() {
 
     // Write the state dump to the file
     if let Some(path) = state_path {
-        let (sessions, session_keys) = service.dump();
+        let (sessions, session_keys) = service.dump().expect("Failed to dump state");
         let dump = PseudonymServiceDump {
             sessions,
             session_keys,