Finish AuthService

ZenGo-X · Apr 5, 2022 · d230efb · d230efb
1 parent 374037c
commit d230efb
Show file tree

Hide file tree

Showing 5 changed files with 107 additions and 7 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/delivery/trusted-delivery/server/Cargo.toml b/delivery/trusted-delivery/server/Cargo.toml
@@ -14,6 +14,8 @@ trusted-delivery-core = { path = "../core" }
 tonic = "0.6"
 futures = "0.3"
 async-stream = "0.3"
+tower-service = "0.3"
+http = "0.2"
 
 # Crypto deps
 jsonwebtoken = "8"

diff --git a/delivery/trusted-delivery/server/src/auth/inner.rs b/delivery/trusted-delivery/server/src/auth/inner.rs
@@ -93,13 +93,13 @@ impl<C: CryptoSuite> TokenVerifier<C> {
     ) -> Result<(), InvalidToken> {
         let validation = jwt::Validation::new(JWT_ALGORITHM);
         let claims = jwt::decode::<Claims<C>>(&token.0, &self.secret_key, &validation)
-            .map_err(InvalidToken::ValidateToken)?
+            .map_err(InvalidTokenReason::ValidateToken)?
             .claims;
         if *identity != claims.identity {
-            return Err(InvalidToken::MismatchedIdentity);
+            return Err(InvalidTokenReason::MismatchedIdentity.into());
         }
         if *room_id != claims.room_id {
-            return Err(InvalidToken::MismatchedRoomId);
+            return Err(InvalidTokenReason::MismatchedRoomId.into());
         }
         Ok(())
     }
@@ -128,7 +128,11 @@ pub enum AuthError {
 }
 
 #[derive(Debug, Error)]
-pub enum InvalidToken {
+#[error(transparent)]
+pub struct InvalidToken(#[from] InvalidTokenReason);
+
+#[derive(Debug, Error)]
+enum InvalidTokenReason {
     #[error(transparent)]
     ValidateToken(jwt::errors::Error),
     #[error("mismatched room id")]

diff --git a/delivery/trusted-delivery/server/src/auth/mod.rs b/delivery/trusted-delivery/server/src/auth/mod.rs
@@ -13,10 +13,19 @@ use tonic::{Code, Request, Response, Status, Streaming};
 
 use crate::auth::inner::AuthError;
 
+pub use self::inner::{InvalidToken, TokenVerifier};
+
 mod inner;
 
 pub struct Auth<C: CryptoSuite>(Arc<inner::AuthInner<C>>);
 
+impl<C: CryptoSuite> Auth<C> {
+    pub fn new() -> (Self, TokenVerifier<C>) {
+        let (inner, token_verifier) = inner::AuthInner::new();
+        (Self(Arc::new(inner)), token_verifier)
+    }
+}
+
 #[tonic::async_trait]
 impl<C: CryptoSuite> trusted_delivery_api::auth_server::Auth for Auth<C> {
     type AuthStream = Pin<Box<dyn Stream<Item = Result<AuthResponse, Status>> + Send>>;
@@ -86,17 +95,19 @@ impl<C: CryptoSuite> trusted_delivery_api::auth_server::Auth for Auth<C> {
                             token: String::from(token).into(),
                         })),
                     };
+                    yield msg4;
+                    Ok(())
                 }
                 Err(AuthError::InvalidSignature) => {
-                    Err(Status::new(Code::Aborted, "invalid signature"))?
+                    Err(Status::new(Code::Aborted, "invalid signature"))
                 }
                 Err(AuthError::GenerateToken(err)) => {
                     Err(Status::new(
                         Code::Internal,
                         format!("generate jwt token: {}", err),
-                    ))?
+                    ))
                 }
-            }
+            }?
         })))
     }
 }
diff --git a/delivery/trusted-delivery/server/src/lib.rs b/delivery/trusted-delivery/server/src/lib.rs
@@ -1,6 +1,87 @@
+use std::task::{Context, Poll};
+
 use trusted_delivery_api::auth_server::AuthServer;
 use trusted_delivery_core::crypto::CryptoSuite;
 
+pub use self::auth::{InvalidToken, TokenVerifier};
+
 mod auth;
 
 pub struct AuthService<C: CryptoSuite>(AuthServer<auth::Auth<C>>);
+
+impl<C: CryptoSuite> AuthService<C> {
+    pub fn new() -> (Self, TokenVerifier<C>) {
+        let (auth, token_verifier) = auth::Auth::new();
+        (AuthService(AuthServer::new(auth)), token_verifier)
+    }
+}
+
+macro_rules! derive_service_for_wrapper {
+    ($($service:ident as $underlying:ty),+) => {
+        use http::{Request, Response};
+        use tonic::body::BoxBody;
+        use tonic::transport::Body;
+        use tonic::transport::NamedService;
+        use tower_service::Service;
+
+        $(
+        impl<C: CryptoSuite> Service<Request<Body>> for $service<C> {
+            type Response = Response<BoxBody>;
+            type Error = <$underlying as tower_service::Service<Request<Body>>>::Error;
+            type Future = <$underlying as tower_service::Service<Request<Body>>>::Future;
+
+            fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+                <$underlying as tower_service::Service<Request<Body>>>::poll_ready(
+                    &mut self.0,
+                    cx,
+                )
+            }
+
+            fn call(&mut self, req: Request<Body>) -> Self::Future {
+                <$underlying as tower_service::Service<Request<Body>>>::call(&mut self.0, req)
+            }
+        }
+        impl<C: CryptoSuite> NamedService for $service<C> {
+            const NAME: &'static str = <$underlying as tonic::transport::NamedService>::NAME;
+        }
+        impl<C: CryptoSuite> Clone for $service<C> {
+            fn clone(&self) -> Self {
+                Self(self.0.clone())
+            }
+        }
+        )+
+    };
+}
+
+derive_service_for_wrapper!(AuthService as AuthServer<auth::Auth<C>>);
+
+#[cfg(test)]
+mod tests {
+    use std::error::Error;
+
+    use http::{Request, Response};
+    use tonic::body::BoxBody;
+    use tonic::transport::Body;
+    use tonic::transport::NamedService;
+    use tower_service::Service;
+    use trusted_delivery_core::crypto::CryptoSuite;
+
+    fn _ensure_service_is_valid<S>()
+    where
+        S: Service<Request<Body>, Response = Response<BoxBody>>
+            + NamedService
+            + Clone
+            + Send
+            + 'static,
+        S::Future: Send + 'static,
+        S::Error: Into<Box<dyn Error + Send + Sync>> + Send,
+    {
+    }
+
+    /// Compile-time check that services are compliant with [Server::add_service function]
+    ///
+    /// [Server::add_service function]: tonic::transport::Server::add_service
+    fn _ensure_services_are_valid<C: CryptoSuite>() {
+        _ensure_service_is_valid::<crate::AuthService<C>>();
+    }
+}