1 files changed, 68 insertions, 18 deletions

diff --git a/src/http/mod.rs b/src/http/mod.rs
index 1bc7951..778a3de 100644
--- a/src/http/mod.rs
+++ b/src/http/mod.rs
@@ -1,21 +1,38 @@
 use crate::config::Config;
-use crate::oauth::OAuthServer;
+use crate::container::ServiceContainer;
+use crate::oauth::{OAuthServer, OAuthService};
 use std::collections::HashMap;
 use std::fs;
 use std::io::prelude::*;
 use std::net::{TcpListener, TcpStream};
+use std::sync::Arc;
 use url::Url;
 
 pub struct Server {
     config: Config,
-    pub oauth_server: OAuthServer,
+    oauth_server: Option<OAuthServer>,
+    oauth_service: Option<Arc<ServiceContainer>>,
 }
 
 impl Server {
     pub fn new(config: Config) -> Result<Server, Box<dyn std::error::Error>> {
         Ok(Server {
-            oauth_server: OAuthServer::new(&config)
-                .map_err(|e| format!("Failed to create OAuth server: {}", e))?,
+            oauth_server: Some(
+                OAuthServer::new(&config)
+                    .map_err(|e| format!("Failed to create OAuth server: {}", e))?,
+            ),
+            oauth_service: None,
+            config,
+        })
+    }
+
+    pub fn new_with_container(
+        config: Config,
+        container: Arc<ServiceContainer>,
+    ) -> Result<Server, Box<dyn std::error::Error>> {
+        Ok(Server {
+            oauth_server: None,
+            oauth_service: Some(container),
             config,
         })
     }
@@ -194,7 +211,13 @@ impl Server {
     }
 
     fn handle_jwks(&self, stream: &mut TcpStream) {
-        let jwks = self.oauth_server.get_jwks();
+        let jwks = if let Some(ref oauth_server) = self.oauth_server {
+            oauth_server.get_jwks()
+        } else if let Some(ref container) = self.oauth_service {
+            container.get_jwks()
+        } else {
+            "{\"keys\":[]}".to_string()
+        };
         self.send_json_response(stream, 200, "OK", &jwks);
     }
 
@@ -204,7 +227,16 @@ impl Server {
         params: &HashMap<String, String>,
         ip_address: Option<String>,
     ) {
-        match self.oauth_server.handle_authorize(params, ip_address) {
+        let result = if let Some(ref oauth_server) = self.oauth_server {
+            oauth_server.handle_authorize(params, ip_address)
+        } else if let Some(ref container) = self.oauth_service {
+            let oauth_service = OAuthService::new(container.clone());
+            oauth_service.handle_authorize(params, ip_address)
+        } else {
+            Err("{\"error\": \"server_error\", \"error_description\": \"No OAuth service available\"}".to_string())
+        };
+
+        match result {
             Ok(redirect_url) => {
                 let security_headers = self.get_security_headers();
                 let response = format!(
@@ -227,10 +259,16 @@ impl Server {
         // Extract Authorization header from request
         let auth_header = self.extract_auth_header(request);
 
-        match self
-            .oauth_server
-            .handle_token(&form_params, auth_header.as_deref(), ip_address)
-        {
+        let result = if let Some(ref oauth_server) = self.oauth_server {
+            oauth_server.handle_token(&form_params, auth_header.as_deref(), ip_address)
+        } else if let Some(ref container) = self.oauth_service {
+            let oauth_service = OAuthService::new(container.clone());
+            oauth_service.handle_token(&form_params, auth_header.as_deref(), ip_address)
+        } else {
+            Err("{\"error\": \"server_error\", \"error_description\": \"No OAuth service available\"}".to_string())
+        };
+
+        match result {
             Ok(token_response) => {
                 self.send_json_response(stream, 200, "OK", &token_response);
             }
@@ -245,10 +283,16 @@ impl Server {
         let form_params = self.parse_form_data(&body);
         let auth_header = self.extract_auth_header(request);
 
-        match self
-            .oauth_server
-            .handle_token_introspection(&form_params, auth_header.as_deref())
-        {
+        let result = if let Some(ref oauth_server) = self.oauth_server {
+            oauth_server.handle_token_introspection(&form_params, auth_header.as_deref())
+        } else if let Some(ref container) = self.oauth_service {
+            let oauth_service = OAuthService::new(container.clone());
+            oauth_service.handle_token_introspection(&form_params, auth_header.as_deref())
+        } else {
+            Err("{\"error\": \"server_error\", \"error_description\": \"No OAuth service available\"}".to_string())
+        };
+
+        match result {
             Ok(introspection_response) => {
                 self.send_json_response(stream, 200, "OK", &introspection_response);
             }
@@ -263,10 +307,16 @@ impl Server {
         let form_params = self.parse_form_data(&body);
         let auth_header = self.extract_auth_header(request);
 
-        match self
-            .oauth_server
-            .handle_token_revocation(&form_params, auth_header.as_deref())
-        {
+        let result = if let Some(ref oauth_server) = self.oauth_server {
+            oauth_server.handle_token_revocation(&form_params, auth_header.as_deref())
+        } else if let Some(ref container) = self.oauth_service {
+            let oauth_service = OAuthService::new(container.clone());
+            oauth_service.handle_token_revocation(&form_params, auth_header.as_deref())
+        } else {
+            Err("{\"error\": \"server_error\", \"error_description\": \"No OAuth service available\"}".to_string())
+        };
+
+        match result {
             Ok(_) => {
                 self.send_empty_response(stream, 200, "OK");
             }