kbs_protocol: use thiserror for error handling

Signed-off-by: Xynnn007 <xynnn@linux.alibaba.com>

attestation-agent/kbs_protocol/Cargo.toml

@@ -19,6 +19,7 @@ resource_uri.path = "../deps/resource_uri"
 serde.workspace = true
 serde_json.workspace = true
 sha2.workspace = true
+thiserror.workspace = true
 url.workspace = true
 zeroize.workspace = true
 

attestation-agent/kbs_protocol/src/api.rs

@@ -3,7 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use anyhow::Result;
+use crate::Result;
 use async_trait::async_trait;
 pub use resource_uri::ResourceUri;
 

attestation-agent/kbs_protocol/src/client/rcar_client.rs

@@ -3,7 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use anyhow::{bail, Context, Result};
+use anyhow::{bail, Context};
 use async_trait::async_trait;
 use kbs_types::{Attestation, Challenge, ErrorInformation, Request, Response};
 use log::{debug, warn};
@@ -19,6 +19,7 @@ use crate::{
     evidence_provider::EvidenceProvider,
     keypair::TeeKeyPair,
     token_provider::Token,
+    Error, Result,
 };
 
 #[derive(Deserialize, Debug, Clone)]
@@ -35,10 +36,14 @@ impl KbsClient<Box<dyn EvidenceProvider>> {
     pub async fn get_token(&mut self) -> Result<(Token, TeeKeyPair)> {
         if let Some(token) = &self.token {
             if token.check_valid().is_err() {
-                self.rcar_handshake().await?;
+                self.rcar_handshake()
+                    .await
+                    .map_err(|e| Error::RcarHandshake(e.to_string()))?;
             }
         } else {
-            self.rcar_handshake().await?;
+            self.rcar_handshake()
+                .await
+                .map_err(|e| Error::RcarHandshake(e.to_string()))?;
         }
 
         assert!(self.token.is_some());
@@ -53,7 +58,7 @@ impl KbsClient<Box<dyn EvidenceProvider>> {
     ///
     /// Note: if RCAR succeeds, the http client will record the cookie with the kbs server,
     /// which means that this client can be then used to retrieve resources.
-    async fn rcar_handshake(&mut self) -> Result<()> {
+    async fn rcar_handshake(&mut self) -> anyhow::Result<()> {
         let auth_endpoint = format!("{}/{KBS_PREFIX}/auth", self.kbs_host_url);
         let tee = match &self._tee {
             ClientTee::Unitialized => bail!("tee not initialized"),
@@ -133,7 +138,8 @@ impl KbsClient<Box<dyn EvidenceProvider>> {
             .provider
             .get_evidence(ehd)
             .await
-            .context("Get TEE evidence failed")?;
+            .context("Get TEE evidence failed")
+            .map_err(|e| Error::GetEvidence(e.to_string()))?;
 
         Ok(tee_evidence)
     }
@@ -155,39 +161,62 @@ impl KbsClientCapabilities for KbsClient<Box<dyn EvidenceProvider>> {
         for attempt in 1..=KBS_GET_RESOURCE_MAX_ATTEMPT {
             debug!("KBS client: trying to request KBS, attempt {attempt}");
 
-            let res = self.http_client.get(&remote_url).send().await?;
+            let res = self
+                .http_client
+                .get(&remote_url)
+                .send()
+                .await
+                .map_err(|e| Error::HttpError(format!("get failed: {e}")))?;
 
             match res.status() {
                 reqwest::StatusCode::OK => {
-                    let response = res.json::<Response>().await?;
-                    let payload_data = self.tee_key.decrypt_response(response)?;
+                    let response = res
+                        .json::<Response>()
+                        .await
+                        .map_err(|e| Error::KbsResponseDeserializationFailed(e.to_string()))?;
+                    let payload_data = self
+                        .tee_key
+                        .decrypt_response(response)
+                        .map_err(|e| Error::DecryptResponseFailed(e.to_string()))?;
                     return Ok(payload_data);
                 }
                 reqwest::StatusCode::UNAUTHORIZED => {
                     warn!(
                         "Authenticating with KBS failed. Perform a new RCAR handshake: {:#?}",
-                        res.json::<ErrorInformation>().await?
+                        res.json::<ErrorInformation>()
+                            .await
+                            .map_err(|e| Error::KbsResponseDeserializationFailed(e.to_string()))?,
                     );
-                    self.rcar_handshake().await?;
+                    self.rcar_handshake()
+                        .await
+                        .map_err(|e| Error::RcarHandshake(e.to_string()))?;
 
                     continue;
                 }
                 reqwest::StatusCode::NOT_FOUND => {
-                    bail!(
+                    let errorinfo = format!(
                         "KBS resource Not Found (Error 404): {:#?}",
-                        res.json::<ErrorInformation>().await?
-                    )
+                        res.json::<ErrorInformation>()
+                            .await
+                            .map_err(|e| Error::KbsResponseDeserializationFailed(e.to_string()))?
+                    );
+
+                    return Err(Error::ResourceNotFound(errorinfo));
                 }
                 _ => {
-                    bail!(
+                    let errorinfo = format!(
                         "KBS Server Internal Failed, Response: {:#?}",
-                        res.json::<ErrorInformation>().await?
-                    )
+                        res.json::<ErrorInformation>()
+                            .await
+                            .map_err(|e| Error::KbsResponseDeserializationFailed(e.to_string()))?
+                    );
+
+                    return Err(Error::KbsInternalError(errorinfo));
                 }
             }
         }
 
-        bail!("Get resource failed. Unauthorized.")
+        Err(Error::UnAuthorized)
     }
 }
 
@@ -266,4 +295,4 @@ mod test {
         println!("Get token : {token:?}");
         println!("Get key: {key:?}");
     }
-}
+}

attestation-agent/kbs_protocol/src/client/token_client.rs

@@ -3,7 +3,6 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use anyhow::{bail, Result};
 use async_trait::async_trait;
 use kbs_types::{ErrorInformation, Response};
 use log::{debug, warn};
@@ -13,11 +12,16 @@ use crate::{
     api::KbsClientCapabilities,
     client::{KbsClient, KBS_GET_RESOURCE_MAX_ATTEMPT, KBS_PREFIX},
     token_provider::TokenProvider,
+    Error, Result,
 };
 
 impl KbsClient<Box<dyn TokenProvider>> {
     async fn update_token(&mut self) -> Result<()> {
-        let (token, teekey) = self.provider.get_token().await?;
+        let (token, teekey) = self
+            .provider
+            .get_token()
+            .await
+            .map_err(|e| Error::GetTokenFailed(e.to_string()))?;
         self.token = Some(token);
         self.tee_key = teekey;
         Ok(())
@@ -44,38 +48,53 @@ impl KbsClientCapabilities for KbsClient<Box<dyn TokenProvider>> {
                 .get(&remote_url)
                 .bearer_auth(&token.content)
                 .send()
-                .await?;
+                .await
+                .map_err(|e| Error::HttpError(format!("get failed: {e}")))?;
 
             match res.status() {
                 reqwest::StatusCode::OK => {
-                    let response = res.json::<Response>().await?;
-                    let payload_data = self.tee_key.decrypt_response(response)?;
+                    let response = res
+                        .json::<Response>()
+                        .await
+                        .map_err(|e| Error::KbsResponseDeserializationFailed(e.to_string()))?;
+                    let payload_data = self
+                        .tee_key
+                        .decrypt_response(response)
+                        .map_err(|e| Error::DecryptResponseFailed(e.to_string()))?;
                     return Ok(payload_data);
                 }
                 reqwest::StatusCode::UNAUTHORIZED => {
                     warn!(
                         "Authenticating with KBS failed. Get a new token from the token provider: {:#?}",
-                        res.json::<ErrorInformation>().await?
+                        res.json::<ErrorInformation>().await.map_err(|e| Error::KbsResponseDeserializationFailed(e.to_string()))?
                     );
                     self.update_token().await?;
 
                     continue;
                 }
                 reqwest::StatusCode::NOT_FOUND => {
-                    bail!(
+                    let errorinfo = format!(
                         "KBS resource Not Found (Error 404): {:#?}",
-                        res.json::<ErrorInformation>().await?
-                    )
+                        res.json::<ErrorInformation>()
+                            .await
+                            .map_err(|e| Error::KbsResponseDeserializationFailed(e.to_string()))?
+                    );
+
+                    return Err(Error::ResourceNotFound(errorinfo));
                 }
                 _ => {
-                    bail!(
+                    let errorinfo = format!(
                         "KBS Server Internal Failed, Response: {:#?}",
-                        res.json::<ErrorInformation>().await?
-                    )
+                        res.json::<ErrorInformation>()
+                            .await
+                            .map_err(|e| Error::KbsResponseDeserializationFailed(e.to_string()))?
+                    );
+
+                    return Err(Error::KbsInternalError(errorinfo));
                 }
             }
         }
 
-        bail!("Get resource failed. Unauthorized.")
+        Err(Error::UnAuthorized)
     }
 }

attestation-agent/kbs_protocol/src/error.rs

@@ -0,0 +1,47 @@
+// Copyright (c) 2023 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+use thiserror::Error;
+
+pub type Result<T> = std::result::Result<T, Error>;
+
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error("decrypt KBS response body failed: {0}")]
+    DecryptResponseFailed(String),
+
+    #[error("get key pair failed: {0}")]
+    GenerateKeyPairFailed(String),
+
+    #[error("get evidence failed: {0}")]
+    GetEvidence(String),
+
+    #[error("get token failed: {0}")]
+    GetTokenFailed(String),
+
+    #[error("get tee type failed: {0}")]
+    GetTeeTypeFailed(String),
+
+    #[error("http request failed: {0}")]
+    HttpError(String),
+
+    #[error("KBS internal error: {0}")]
+    KbsInternalError(String),
+
+    #[error("deserialize http response failed: {0}")]
+    KbsResponseDeserializationFailed(String),
+
+    #[error("Native Evidence Provider error: {0}")]
+    NativeEvidenceProvider(String),
+
+    #[error("RCAR handshake failed: {0}")]
+    RcarHandshake(String),
+
+    #[error("KBS resource not found: {0}")]
+    ResourceNotFound(String),
+
+    #[error("request unautorized")]
+    UnAuthorized,
+}

attestation-agent/kbs_protocol/src/evidence_provider/mock.rs

@@ -3,12 +3,13 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use anyhow::*;
 use async_trait::async_trait;
 use kbs_types::Tee;
 
 use super::EvidenceProvider;
 
+use crate::Result;
+
 #[derive(Default)]
 pub struct MockedEvidenceProvider {}
 

attestation-agent/kbs_protocol/src/evidence_provider/mod.rs

@@ -9,7 +9,7 @@ pub use native::*;
 pub mod mock;
 pub use mock::*;
 
-use anyhow::*;
+use crate::Result;
 use async_trait::async_trait;
 use kbs_types::Tee;
 

attestation-agent/kbs_protocol/src/evidence_provider/native.rs

@@ -3,30 +3,39 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use anyhow::*;
 use async_trait::async_trait;
 use attester::{detect_tee_type, BoxedAttester};
 use kbs_types::Tee;
 
 use super::EvidenceProvider;
 
+use crate::{Error, Result};
+
 pub struct NativeEvidenceProvider(BoxedAttester);
 
 impl NativeEvidenceProvider {
     pub fn new() -> Result<Self> {
-        let tee =
-            detect_tee_type().ok_or_else(|| anyhow!("failed to get a supported Tee type."))?;
-        Ok(Self(tee.try_into()?))
+        let tee = detect_tee_type()
+            .ok_or_else(|| Error::GetTeeTypeFailed("no supported Tee type detected.".into()))?
+            .try_into()
+            .map_err(|e| {
+                Error::NativeEvidenceProvider(format!("failed to initialize tee driver: {e}"))
+            })?;
+        Ok(Self(tee))
     }
 }
 
 #[async_trait]
 impl EvidenceProvider for NativeEvidenceProvider {
     async fn get_evidence(&self, runtime_data: Vec<u8>) -> Result<String> {
-        self.0.get_evidence(runtime_data).await
+        self.0
+            .get_evidence(runtime_data)
+            .await
+            .map_err(|e| Error::GetEvidence(e.to_string()))
     }
 
     async fn get_tee_type(&self) -> Result<Tee> {
-        detect_tee_type().ok_or_else(|| anyhow!("failed to get a supported Tee type."))
+        detect_tee_type()
+            .ok_or_else(|| Error::GetTeeTypeFailed("no supported Tee type detected.".into()))
     }
 }

attestation-agent/kbs_protocol/src/lib.rs

@@ -73,11 +73,13 @@
 pub mod api;
 pub mod builder;
 pub mod client;
+pub mod error;
 pub mod evidence_provider;
 pub mod keypair;
 pub mod token_provider;
 
 pub use api::*;
 pub use builder::KbsClientBuilder;
+pub use error::{Error, Result};
 pub use keypair::TeeKeyPair;
 pub use token_provider::Token;

attestation-agent/kbs_protocol/src/token_provider/mod.rs

@@ -22,7 +22,7 @@ pub trait TokenProvider: Send + Sync {
     /// Get token provisioned by Kbs and the tee pub key
     ///
     /// The returned value is a (Token, Private key) pair.
-    async fn get_token(&self) -> Result<(Token, TeeKeyPair)>;
+    async fn get_token(&self) -> crate::Result<(Token, TeeKeyPair)>;
 }
 
 #[derive(Clone, Debug)]