6 files changed, 47 insertions, 26 deletions

diff --git a/src/authorization/cedar_authorizer.rs b/src/authorization/cedar_authorizer.rs
index 658de7a6..a877cf87 100644
--- a/src/authorization/cedar_authorizer.rs
+++ b/src/authorization/cedar_authorizer.rs
@@ -51,6 +51,7 @@ impl CedarAuthorizer {
         Ok(policies)
     }
 }
+
 impl Default for CedarAuthorizer {
     fn default() -> Self {
         Self::new_from(std::path::Path::new("/etc/authzd"))
diff --git a/src/authorization/mod.rs b/src/authorization/mod.rs
index 7d3856a5..d664815b 100644
--- a/src/authorization/mod.rs
+++ b/src/authorization/mod.rs
@@ -1,7 +1,9 @@
 pub mod authorizer;
 pub mod cedar_authorizer;
 pub mod check_service;
+pub mod server;
 
 pub use authorizer::Authorizer;
 pub use cedar_authorizer::CedarAuthorizer;
 pub use check_service::CheckService;
+pub use server::Server;
diff --git a/src/authorization/server.rs b/src/authorization/server.rs
new file mode 100644
index 00000000..f11d0465
--- /dev/null
+++ b/src/authorization/server.rs
@@ -0,0 +1,41 @@
+use super::cedar_authorizer::CedarAuthorizer;
+use super::check_service::CheckService;
+use envoy_types::ext_authz::v3::pb::AuthorizationServer;
+use std::sync::Arc;
+
+pub fn create_router() -> Result<tonic::transport::server::Router, Box<dyn std::error::Error>> {
+    let (_health_reporter, health_service) = tonic_health::server::health_reporter();
+    let authorizer = Arc::new(CedarAuthorizer::default());
+    let check_service = CheckService::new(authorizer);
+    let server = tonic::transport::Server::builder()
+        .add_service(AuthorizationServer::new(check_service))
+        .add_service(health_service)
+        .add_service(
+            tonic_reflection::server::Builder::configure()
+                .register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
+                .build_v1()
+                .unwrap(),
+        );
+    Ok(server)
+}
+
+pub struct Server {
+    router: tonic::transport::server::Router,
+}
+
+impl Server {
+    pub fn new() -> Result<Server, Box<dyn std::error::Error>> {
+        let router = create_router()?;
+        Ok(Server { router: router })
+    }
+
+    pub async fn serve(self, addr: std::net::SocketAddr) -> Result<(), tonic::transport::Error> {
+        self.router.serve(addr).await
+    }
+}
+
+impl Default for Server {
+    fn default() -> Self {
+        Self::new().unwrap()
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index a82c2ace..3bd8fbd1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,22 +1,2 @@
 pub mod authorization;
-pub use authorization::{Authorizer, CedarAuthorizer, CheckService};
-
-use envoy_types::ext_authz::v3::pb::AuthorizationServer;
-use std::sync::Arc;
-use tonic::transport::Server;
-
-pub fn create_server() -> Result<tonic::transport::server::Router, Box<dyn std::error::Error>> {
-    let (_health_reporter, health_service) = tonic_health::server::health_reporter();
-    let authorizer = Arc::new(authorization::CedarAuthorizer::default());
-    let check_service = authorization::CheckService::new(authorizer);
-    let server = Server::builder()
-        .add_service(AuthorizationServer::new(check_service))
-        .add_service(health_service)
-        .add_service(
-            tonic_reflection::server::Builder::configure()
-                .register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
-                .build_v1()
-                .unwrap(),
-        );
-    Ok(server)
-}
+pub use authorization::{Authorizer, CedarAuthorizer, CheckService, Server};
diff --git a/src/main.rs b/src/main.rs
index 13d313d7..8638e14b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,5 +1,3 @@
-use authzd::create_server;
-
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     tracing_subscriber::fmt().json().init();
@@ -8,9 +6,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .unwrap_or_else(|_| "[::1]:50051".to_string())
         .parse()?;
 
-    let server = create_server()?;
-
     log::info!("Listening on... {addr}");
+    let server = authzd::authorization::Server::new()?;
     server.serve(addr).await?;
 
     Ok(())
diff --git a/tests/authorization/server_test.rs b/tests/authorization/server_test.rs
index 55645dd4..fe8c8a73 100644
--- a/tests/authorization/server_test.rs
+++ b/tests/authorization/server_test.rs
@@ -15,7 +15,7 @@ mod tests {
 
     async fn start_server() -> (SocketAddr, tokio::task::JoinHandle<()>) {
         let addr = available_port().await;
-        let server = authzd::create_server().expect("Failed to create server");
+        let server = authzd::authorization::Server::default();
 
         let handle = tokio::spawn(async move {
             server.serve(addr).await.expect("Failed to start server");