-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
19 changed files
with
1,930 additions
and
29 deletions.
There are no files selected for viewing
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,147 @@ | ||
| use candid::{ | ||
| utils::{encode_args, ArgumentEncoder}, | ||
| CandidType, Decode, Principal, | ||
| }; | ||
| use ic_agent::Agent; | ||
| use ic_tee_cdk::{format_error, SignInResponse, SignedDelegation}; | ||
| use serde_bytes::ByteBuf; | ||
| use std::sync::Arc; | ||
| use tokio::sync::RwLock; | ||
|
|
||
| use crate::TEEIdentity; | ||
|
|
||
| #[derive(Clone)] | ||
| pub struct TEEAgent { | ||
| identity: Arc<RwLock<TEEIdentity>>, | ||
| agent: Arc<RwLock<Agent>>, | ||
| authentication_canister: Principal, | ||
| configuration_canister: Principal, | ||
| } | ||
|
|
||
| impl TEEAgent { | ||
| pub fn new( | ||
| host: &str, | ||
| authentication_canister: Principal, | ||
| configuration_canister: Principal, | ||
| ) -> Result<Self, String> { | ||
| let identity = TEEIdentity::new(); | ||
| let agent = Agent::builder() | ||
| .with_url(host) | ||
| .with_verify_query_signatures(true) | ||
| .with_identity(identity.clone()) | ||
| .build() | ||
| .map_err(format_error)?; | ||
| Ok(Self { | ||
| identity: Arc::new(RwLock::new(identity)), | ||
| agent: Arc::new(RwLock::new(agent)), | ||
| authentication_canister, | ||
| configuration_canister, | ||
| }) | ||
| } | ||
|
|
||
| pub async fn principal(&self) -> Principal { | ||
| self.identity.read().await.principal() | ||
| } | ||
|
|
||
| pub async fn session_key(&self) -> Vec<u8> { | ||
| self.identity.read().await.session_key() | ||
| } | ||
|
|
||
| pub async fn is_authenticated(&self) -> bool { | ||
| self.identity.read().await.is_authenticated() | ||
| } | ||
|
|
||
| pub async fn sign_in(&self, kind: String, attestation: ByteBuf) -> Result<(), String> { | ||
| let res: Result<SignInResponse, String> = self | ||
| .update_call( | ||
| &self.authentication_canister, | ||
| "sign_in", | ||
| (kind, attestation), | ||
| ) | ||
| .await?; | ||
| let res = res?; | ||
| let mut id = { | ||
| let id = self.identity.read().await; | ||
| id.clone() | ||
| // drop read lock | ||
| }; | ||
|
|
||
| id.with_user_key(res.user_key.to_vec()); | ||
| let res: Result<SignedDelegation, String> = self | ||
| .query_call( | ||
| &self.authentication_canister, | ||
| "get_delegation", | ||
| ( | ||
| id.principal(), | ||
| ByteBuf::from(id.session_key()), | ||
| res.expiration, | ||
| ), | ||
| ) | ||
| .await?; | ||
| let res = res?; | ||
|
|
||
| id.with_delegation(res)?; | ||
| self.agent.write().await.set_identity(id.clone()); | ||
| let mut w = self.identity.write().await; | ||
| *w = id; | ||
|
|
||
| Ok(()) | ||
| } | ||
|
|
||
| pub async fn sign_in_with( | ||
| &self, | ||
| f: impl FnOnce(Vec<u8>) -> Result<(String, ByteBuf), String>, | ||
| ) -> Result<(), String> { | ||
| let session_key = self.session_key().await; | ||
| let (kind, attestation) = f(session_key)?; | ||
| self.sign_in(kind, attestation).await | ||
| } | ||
|
|
||
| pub async fn update_call<In, Out>( | ||
| &self, | ||
| canister_id: &Principal, | ||
| method_name: &str, | ||
| args: In, | ||
| ) -> Result<Out, String> | ||
| where | ||
| In: ArgumentEncoder + Send, | ||
| Out: CandidType + for<'a> candid::Deserialize<'a>, | ||
| { | ||
| let input = encode_args(args).map_err(format_error)?; | ||
| let res = self | ||
| .agent | ||
| .read() | ||
| .await | ||
| .update(canister_id, method_name) | ||
| .with_arg(input) | ||
| .call_and_wait() | ||
| .await | ||
| .map_err(format_error)?; | ||
| let output = Decode!(res.as_slice(), Out).map_err(format_error)?; | ||
| Ok(output) | ||
| } | ||
|
|
||
| pub async fn query_call<In, Out>( | ||
| &self, | ||
| canister_id: &Principal, | ||
| method_name: &str, | ||
| args: In, | ||
| ) -> Result<Out, String> | ||
| where | ||
| In: ArgumentEncoder + Send, | ||
| Out: CandidType + for<'a> candid::Deserialize<'a>, | ||
| { | ||
| let input = encode_args(args).map_err(format_error)?; | ||
| let res = self | ||
| .agent | ||
| .read() | ||
| .await | ||
| .query(canister_id, method_name) | ||
| .with_arg(input) | ||
| .call() | ||
| .await | ||
| .map_err(format_error)?; | ||
| let output = Decode!(res.as_slice(), Out).map_err(format_error)?; | ||
| Ok(output) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,170 @@ | ||
| use axum::{ | ||
| async_trait, | ||
| body::{Body, Bytes}, | ||
| extract::{FromRequest, Request}, | ||
| http::{ | ||
| header::{self, HeaderMap, HeaderValue}, | ||
| StatusCode, | ||
| }, | ||
| response::{IntoResponse, Response}, | ||
| }; | ||
| use bytes::{BufMut, BytesMut}; | ||
| use serde::{de::DeserializeOwned, Serialize}; | ||
|
|
||
| pub static CONTENT_TYPE_CBOR: &str = "application/cbor"; | ||
| pub static CONTENT_TYPE_JSON: &str = "application/json"; | ||
| pub static CONTENT_TYPE_TEXT: &str = "text/plain"; | ||
| pub enum ContentType<T> { | ||
| JSON(T, Option<StatusCode>), | ||
| CBOR(T, Option<StatusCode>), | ||
| Text(String, Option<StatusCode>), | ||
| Ohter(String, Option<StatusCode>), | ||
| } | ||
|
|
||
| impl ContentType<()> { | ||
| pub fn from(headers: &HeaderMap<HeaderValue>) -> Self { | ||
| if let Some(accept) = headers.get(header::ACCEPT) { | ||
| if let Ok(accept) = accept.to_str() { | ||
| if accept.contains(CONTENT_TYPE_CBOR) { | ||
| return ContentType::CBOR((), None); | ||
| } | ||
| if accept.contains(CONTENT_TYPE_JSON) { | ||
| return ContentType::JSON((), None); | ||
| } | ||
| if accept.contains(CONTENT_TYPE_TEXT) { | ||
| return ContentType::Text("".to_string(), None); | ||
| } | ||
| return ContentType::Ohter(accept.to_string(), None); | ||
| } | ||
| } | ||
|
|
||
| ContentType::Ohter("unknown".to_string(), None) | ||
| } | ||
|
|
||
| pub fn from_content_type(headers: &HeaderMap) -> Self { | ||
| if let Some(content_type) = headers.get(header::CONTENT_TYPE) { | ||
| if let Ok(content_type) = content_type.to_str() { | ||
| if let Ok(mime) = content_type.parse::<mime::Mime>() { | ||
| if mime.type_() == "application" { | ||
| if mime.subtype() == "cbor" | ||
| || mime.suffix().map_or(false, |name| name == "cbor") | ||
| { | ||
| return ContentType::CBOR((), None); | ||
| } else if mime.subtype() == "json" | ||
| || mime.suffix().map_or(false, |name| name == "json") | ||
| { | ||
| return ContentType::JSON((), None); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| ContentType::Ohter("unknown".to_string(), None) | ||
| } | ||
| } | ||
|
|
||
| #[async_trait] | ||
| impl<T, S> FromRequest<S> for ContentType<T> | ||
| where | ||
| T: DeserializeOwned + Send + Sync, | ||
| Bytes: FromRequest<S>, | ||
| S: Send + Sync, | ||
| { | ||
| type Rejection = Response; | ||
|
|
||
| async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> { | ||
| match ContentType::from_content_type(req.headers()) { | ||
| ContentType::JSON(_, _) => { | ||
| let body = Bytes::from_request(req, state) | ||
| .await | ||
| .map_err(IntoResponse::into_response)?; | ||
| let value: T = serde_json::from_slice(&body).map_err(|err| { | ||
| ContentType::Text::<String>(err.to_string(), Some(StatusCode::BAD_REQUEST)) | ||
| .into_response() | ||
| })?; | ||
| Ok(Self::JSON(value, None)) | ||
| } | ||
| ContentType::CBOR(_, _) => { | ||
| let body = Bytes::from_request(req, state) | ||
| .await | ||
| .map_err(IntoResponse::into_response)?; | ||
| let value: T = ciborium::from_reader(&body[..]).map_err(|err| { | ||
| ContentType::Text::<String>(err.to_string(), Some(StatusCode::BAD_REQUEST)) | ||
| .into_response() | ||
| })?; | ||
| Ok(Self::CBOR(value, None)) | ||
| } | ||
| _ => Err(StatusCode::UNSUPPORTED_MEDIA_TYPE.into_response()), | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl<T> IntoResponse for ContentType<T> | ||
| where | ||
| T: Serialize, | ||
| { | ||
| fn into_response(self) -> Response { | ||
| let mut buf = BytesMut::with_capacity(128).writer(); | ||
| match self { | ||
| ContentType::JSON(v, c) => match serde_json::to_writer(&mut buf, &v) { | ||
| Ok(()) => ( | ||
| c.unwrap_or_default(), | ||
| [( | ||
| header::CONTENT_TYPE, | ||
| HeaderValue::from_static(CONTENT_TYPE_JSON), | ||
| )], | ||
| buf.into_inner().freeze(), | ||
| ) | ||
| .into_response(), | ||
| Err(err) => ( | ||
| StatusCode::INTERNAL_SERVER_ERROR, | ||
| [( | ||
| header::CONTENT_TYPE, | ||
| HeaderValue::from_static(CONTENT_TYPE_TEXT), | ||
| )], | ||
| err.to_string(), | ||
| ) | ||
| .into_response(), | ||
| }, | ||
| ContentType::CBOR(v, c) => match ciborium::into_writer(&v, &mut buf) { | ||
| Ok(()) => ( | ||
| c.unwrap_or_default(), | ||
| [( | ||
| header::CONTENT_TYPE, | ||
| HeaderValue::from_static(CONTENT_TYPE_CBOR), | ||
| )], | ||
| buf.into_inner().freeze(), | ||
| ) | ||
| .into_response(), | ||
| Err(err) => ( | ||
| StatusCode::INTERNAL_SERVER_ERROR, | ||
| [( | ||
| header::CONTENT_TYPE, | ||
| HeaderValue::from_static(CONTENT_TYPE_TEXT), | ||
| )], | ||
| err.to_string(), | ||
| ) | ||
| .into_response(), | ||
| }, | ||
| ContentType::Text(v, c) => ( | ||
| c.unwrap_or_default(), | ||
| [( | ||
| header::CONTENT_TYPE, | ||
| HeaderValue::from_static(CONTENT_TYPE_TEXT), | ||
| )], | ||
| v, | ||
| ) | ||
| .into_response(), | ||
| ContentType::Ohter(v, c) => ( | ||
| c.unwrap_or(StatusCode::UNSUPPORTED_MEDIA_TYPE), | ||
| [( | ||
| header::CONTENT_TYPE, | ||
| HeaderValue::from_static(CONTENT_TYPE_TEXT), | ||
| )], | ||
| format!("Unsupported MIME type: {}", v), | ||
| ) | ||
| .into_response(), | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.