Skip to content

Commit

Permalink
kbs_protocol: use thiserror for error handling
Browse files Browse the repository at this point in the history
Signed-off-by: Xynnn007 <xynnn@linux.alibaba.com>
  • Loading branch information
Xynnn007 authored and jialez0 committed Aug 7, 2023
1 parent 63335e1 commit c338b44
Show file tree
Hide file tree
Showing 11 changed files with 153 additions and 45 deletions.
1 change: 1 addition & 0 deletions attestation-agent/kbs_protocol/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion attestation-agent/kbs_protocol/src/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
65 changes: 47 additions & 18 deletions attestation-agent/kbs_protocol/src/client/rcar_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -19,6 +19,7 @@ use crate::{
evidence_provider::EvidenceProvider,
keypair::TeeKeyPair,
token_provider::Token,
Error, Result,
};

#[derive(Deserialize, Debug, Clone)]
Expand All @@ -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());
Expand All @@ -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"),
Expand Down Expand Up @@ -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)
}
Expand All @@ -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)
}
}

Expand Down Expand Up @@ -266,4 +295,4 @@ mod test {
println!("Get token : {token:?}");
println!("Get key: {key:?}");
}
}
}
45 changes: 32 additions & 13 deletions attestation-agent/kbs_protocol/src/client/token_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -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(())
Expand All @@ -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)
}
}
47 changes: 47 additions & 0 deletions attestation-agent/kbs_protocol/src/error.rs
Original file line number Diff line number Diff line change
@@ -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,
}
3 changes: 2 additions & 1 deletion attestation-agent/kbs_protocol/src/evidence_provider/mock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
21 changes: 15 additions & 6 deletions attestation-agent/kbs_protocol/src/evidence_provider/native.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()))
}
}
2 changes: 2 additions & 0 deletions attestation-agent/kbs_protocol/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
2 changes: 1 addition & 1 deletion attestation-agent/kbs_protocol/src/token_provider/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
Loading

0 comments on commit c338b44

Please sign in to comment.