add dedicated docs crate and refactor for improved stability

add tests

.gitignore

@@ -2,4 +2,5 @@
 /node_modules/
 /.wrangler/
 /.idea/
-/build/
+/build/
+/project

README.md

@@ -1,200 +1,41 @@
-# zitadel-session-worker
+# axum-tower-sessions-edge
+[![Rust](https://github.com/seemueller-io/axum-tower-sessions-edge/actions/workflows/test.yaml/badge.svg)](https://github.com/seemueller-io/axum-tower-sessions-edge/actions/workflows/test.yaml)
+[![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT)
 
-> ⚠️ **WARNING**: This project is currently in development and **NOT** production-ready. Use at your own risk. It may
-> contain bugs, security vulnerabilities, or incomplete features. This should
-> serve as a starting point for anyone building similar technology. All feedback is welcome.
-
-A Rust Cloudflare Worker that provides authentication and session management for web applications using ZITADEL as the identity provider. It adopts the implementation for oauth2 token introspection from [smartive/zitadel-rs](https://github.com/smartive/zitadel-rust). 
-
-## Overview
-
-This project is a Rust-based Cloudflare Worker that acts as an authentication proxy for web applications. It handles:
-
-- oauth2/oidc w/PKCE via Zitadel
-- Session management using Cloudflare KV storage
-- Token introspection and validation
-- Proxying authenticated requests to backend services
-
-When deployed, the worker sits between your users and your application services. It:
-1. Intercepts incoming requests
-2. Verifies if the user has a valid session
-3. If not, redirects to ZITADEL for authentication
-4. After successful authentication, creates a session and proxies the request to your service
-5. For subsequent requests, validates the session and proxies authenticated requests
-
-
-> **Note**: Caches are used by the introspection and session modules. They prevent excessive r/w.
- 
-## Prerequisites
-
-- [Rust](https://www.rust-lang.org/tools/install) (latest stable version)
-- LLVM and clang
-- [Bun](https://bun.sh/) JavaScript runtime
-- [Wrangler CLI](https://developers.cloudflare.com/workers/wrangler/) for Cloudflare Workers development
-- ZITADEL Administrator Access
-
-## Installation
-
-1. Clone the repository:
-   ```bash
-   git clone <repository-url>
-   cd zitadel-session-worker
-   ```
-
-2. Install dependencies:
-   ```bash
-   # Install JavaScript dependencies
-   bun install
-   ```
-
-## Configuration
-
-> **Note**: There is a docker compose file with Zitadel in this repository that can be used for testing.
-
-### Environment Variables
-
-Create a `.dev.vars` file in the project root with the following variables:
-
-```
-CLIENT_ID="your-client-id"
-CLIENT_SECRET="your-client-secret"
-AUTH_SERVER_URL="your-zitadel-instance-url"
-ZITADEL_ORG_ID="your-organization-id"
-ZITADEL_PROJECT_ID="your-project-id"
-APP_URL="http://localhost:3000"
-DEV_MODE="true"
-```
-
-### Wrangler Configuration
-
-- `wrangler.jsonc` - Base configuration
-
-## Development
-
-### Running Locally
+> **Warning**: This API may be unstable.
 
 ```bash
-# Start the development server
-bun run dev
+git clone https://github.com/seemueller-io/axum-tower-sessions-edge.git
+cd axum-tower-sessions-edge
+bun install
+# Create a `.dev.vars` file in the project root with the following variables:
+#CLIENT_ID="your-client-id"
+#CLIENT_SECRET="your-client-secret"
+#AUTH_SERVER_URL="https://your-zitadel-instance-url"
+#ZITADEL_ORG_ID="your-organization-id"
+#ZITADEL_PROJECT_ID="your-project-id"
+#APP_URL="http://localhost:3000"
+
+# Update the wrangler.jsonc and replace the value of PROXY_TARGET with a worker script name. 
+
+npx wrangler dev 
+# Open `http://localhost:3000` in your browser. If everything is configured correctly, you should be taken to a Zitadel login page.
 ```
-
-This will start the worker on `localhost:3000`.
-
-### Building
-
-```bash
-# Build the project
-cargo clean && cargo install -q worker-build && worker-build --release
-```
-
-## Deployment
-
-### Deploying to Cloudflare
-
-```bash
-# Deploy to development environment
-bun run deploy:dev
-
-# Deploy with updated secrets
-bun run deploy:dev:secrets
-```
-
-### Viewing Logs
-
-```bash
-# View logs from the deployed worker
-bun run tail:dev
-```
-
-## Integration with Your Application
-
-To integrate this worker with your existing application:
-
-1. **Configure Cloudflare**:
-   - Set up a Cloudflare Worker route that points to your application domain
-   - Deploy this worker to that route
-
-2. **Configure ZITADEL**:
-   - Create an application in ZITADEL
-   - Configure the redirect URI to `https://your-worker-domain/login/callback`
-   - Get the client ID and client secret
-
-3. **Configure this Worker**:
-   - Update the environment variables with your ZITADEL credentials
-   - Set the `APP_URL` to your application's URL
-   - Set an http route in `wrangler.jsonc`
-
-4. **Access Control**:
-   - The worker will automatically handle authentication
-   - Your application will receive authenticated requests with user information
-   - You can access user information via the `/api/whoami` endpoint
-
-## Testing
-
-The project uses Rust's built-in testing framework with tokio for async tests.
-
-```bash
-# Run all tests
-cargo test
-```
-
-### Adding New Tests
-
-1. For unit tests, add them to the `tests` module in the relevant source file
-2. For async tests, use the `#[tokio::test]` attribute
-3. Follow the existing pattern of testing both success and error cases
-4. Mock external dependencies when necessary
-
-## Debugging
-
-1. For local development, use `console_log!` macros to output debug information
-2. View logs in the wrangler development console
-3. For deployed workers, use `bun run tail:dev` to stream logs
-4. Check the `/api/whoami` endpoint to verify user authentication and session data
-
-## Project Structure
-
-- `src/` - Rust source code
-  - `api/` - API endpoints and routing
-  - `axum_introspector/` - Axum framework integration for token introspection
-  - `credentials/` - Credential management
-  - `oidc/` - OpenID Connect implementation
-  - `session_storage/` - Session storage implementations
-  - `utilities.rs` - Utility functions
-  - `lib.rs` - Main entry point and worker setup
-
-## Contributing
-
-Contributions to this project are welcome! Here are some guidelines:
-
-1. **Fork the repository** and create your branch from `main`
-2. **Install dependencies** and ensure you can build the project
-3. **Make your changes** and add or update tests as necessary
-4. **Ensure tests pass** by running `cargo test`
-5. **Format your code** with `cargo fmt`
-6. **Submit a pull request** with a clear description of your changes
-
-### Code Style
-
-- Follow Rust's standard code style and idioms
-- Use `cargo fmt` to format code
-- Use `cargo clippy` for linting
-
 ## Acknowledgements
 
 This project is made possible thanks to:
 
-- **ZITADEL**: For providing the robust identity management platform that powers this authentication proxy
-- **Smartive**: For [zitadel-rs](https://github.com/smartive/zitadel-rust) 
-- **Cloudflare**: For their Workers platform and KV storage solution
-- **Open Source Community**: For the various dependencies and tools that make this project possible:
-    - The Rust ecosystem and its crates
-    - The Axum web framework
-    - The Tower middleware ecosystem
-    - Various other open-source projects listed in our dependencies
-
-I appreciate the hard work and dedication of all the developers and organizations that contribute to the open-source
-ecosystem.
+- **Open Source Community**: For the various dependencies and tools that make this project possible.
+- [The Rust ecosystem](https://www.rust-lang.org/ecosystem) and its crates
+- [ZITADEL](https://zitadel.com/): For providing the robust identity management platform that powers this authentication
+  proxy 
+- [Smartive](https://github.com/smartive): For [zitadel-rs](https://github.com/smartive/zitadel-rust)
+- [Cloudflare](https://github.com/cloudflare): For their [Workers](https://workers.cloudflare.com/) platform and KV storage
+  solution
+- [Fermyon/Spin](https://www.fermyon.com/spin): [http-auth-middleware](https://github.com/fermyon/http-auth-middleware) (Reference implementation)
+- [The Axum web framework](https://github.com/tokio-rs/axum)
+- [The Tower middleware ecosystem](https://github.com/tower-rs)
+- Various other open-source projects listed in [Cargo.toml](./Cargo.toml)
 
 
 ## License

src/api/mod.rs

@@ -1,2 +1,3 @@
 pub mod public;
-pub mod authenticated;
+pub mod authenticated;
+pub mod router;

src/api/router.rs

@@ -0,0 +1,188 @@
+use crate::api::authenticated::AuthenticatedApi;
+use crate::api::public::PublicApi;
+use crate::axum_introspector::introspection::{IntrospectionState, IntrospectionStateBuilder};
+use crate::oidc::introspection::cache::in_memory::InMemoryIntrospectionCache;
+use crate::session_storage::in_memory::MemoryStore;
+use axum::response::{IntoResponse, Redirect};
+use axum::routing::{any, get};
+use axum::{Router, ServiceExt};
+use http::HeaderName;
+use std::iter::once;
+use std::sync::Arc;
+use tower_cookies::CookieManagerLayer;
+use tower_http::cors::CorsLayer;
+use tower_http::propagate_header::PropagateHeaderLayer;
+use tower_http::sensitive_headers::SetSensitiveRequestHeadersLayer;
+use tower_sessions::cookie::{Key, SameSite};
+use tower_sessions::SessionManagerLayer;
+use tower_sessions_core::Expiry;
+
+// Test configuration struct
+#[derive(Clone)]
+pub struct TestConfig {
+    pub auth_server_url: String,
+    pub client_id: String,
+    pub client_secret: String,
+    pub app_url: String,
+    pub dev_mode: bool,
+}
+
+impl Default for TestConfig {
+    fn default() -> Self {
+        Self {
+            auth_server_url: "https://test-auth-server.example.com".to_string(),
+            client_id: "test-client-id".to_string(),
+            client_secret: "test-client-secret".to_string(),
+            app_url: "http://localhost:3000".to_string(),
+            dev_mode: true,
+        }
+    }
+}
+
+// App state for testing
+#[derive(Clone)]
+pub struct TestAppState {
+    pub introspection_state: IntrospectionState,
+    pub session_store: MemoryStore,
+}
+
+impl From<TestAppState> for IntrospectionState {
+    fn from(state: TestAppState) -> Self {
+        state.introspection_state
+    }
+}
+
+// Create a router for testing
+pub async fn create_router(config: TestConfig) -> Router<TestAppState> {
+    // Create a memory-based introspection cache for testing
+    let cache = InMemoryIntrospectionCache::new();
+
+    // Create introspection state
+    let introspection_state = IntrospectionStateBuilder::new(&config.auth_server_url)
+        .with_basic_auth(&config.client_id, &config.client_secret)
+        .with_introspection_cache(cache)
+        .build()
+        .await
+        .unwrap();
+
+    // Create a memory-based session store for testing
+    let session_store = MemoryStore::default();
+
+    // Create app state
+    let state = TestAppState {
+        introspection_state,
+        session_store: session_store.clone(),
+    };
+
+    // Generate keys for session encryption and signing
+    let signing_key = Key::generate();
+    let encryption_key = Key::generate();
+
+    // Parse the app URL to get the host for cookies
+    let cookie_host_uri = config.app_url.parse::<http::Uri>().unwrap();
+    let mut cookie_host = cookie_host_uri.authority().unwrap().to_string();
+
+    if cookie_host.starts_with("localhost:") {
+        cookie_host = "localhost".to_string();
+    }
+
+    // Create session layer
+    let session_layer = SessionManagerLayer::new(session_store)
+        .with_name("session")
+        .with_expiry(Expiry::OnSessionEnd)
+        .with_domain(cookie_host)
+        .with_same_site(SameSite::Lax)
+        .with_signed(signing_key)
+        .with_private(encryption_key)
+        .with_path("/")
+        .with_secure(!config.dev_mode)
+        .with_always_save(false);
+
+    // Error handling middleware
+    async fn handle_introspection_errors(
+        mut response: axum_core::response::Response,
+    ) -> axum_core::response::Response {
+        let x_error_header_value = response
+            .headers()
+            .get("x-introspection-error")
+            .and_then(|header_value| header_value.to_str().ok());
+
+        match response.status() {
+            http::StatusCode::UNAUTHORIZED => {
+                if let Some(x_error) = x_error_header_value {
+                    if x_error == "unauthorized" {
+                        return Redirect::to("/login").into_response();
+                    }
+                }
+                response
+            }
+            http::StatusCode::BAD_REQUEST => {
+                if let Some(x_error) = x_error_header_value {
+                    if x_error == "invalid schema"
+                        || x_error == "invalid header"
+                        || x_error == "introspection error"
+                    {
+                        return Redirect::to("/login").into_response();
+                    }
+                }
+                response
+            }
+            http::StatusCode::FORBIDDEN => {
+                if let Some(x_error) = x_error_header_value {
+                    if x_error == "user is inactive" {
+                        return Redirect::to("/login").into_response();
+                    }
+                }
+                response
+            }
+            http::StatusCode::NOT_FOUND => {
+                if let Some(x_error) = x_error_header_value {
+                    if x_error == "user was not found" {
+                        return Redirect::to("/login").into_response();
+                    }
+                }
+                response
+            }
+            http::StatusCode::INTERNAL_SERVER_ERROR => {
+                if let Some(x_error) = x_error_header_value {
+                    if x_error == "missing config" {
+                        return Redirect::to("/login").into_response();
+                    }
+                }
+                response
+            }
+            _ => response,
+        }
+    }
+
+    // Create the router with test-specific routes
+    Router::new()
+        .route("/api/whoami", get(whoami))
+        .route("/public", get(public_test_route))
+        .route("/protected", get(protected_test_route))
+        .layer(PropagateHeaderLayer::new(HeaderName::from_static(
+            "x-request-id",
+        )))
+        .layer(axum::middleware::map_response(handle_introspection_errors))
+        .with_state(state)
+        .layer(session_layer)
+        .layer(CookieManagerLayer::new())
+        .layer(CorsLayer::very_permissive())
+        .layer(SetSensitiveRequestHeadersLayer::new(once(
+            http::header::AUTHORIZATION,
+        )))
+}
+
+// Test routes
+async fn whoami() -> impl IntoResponse {
+    "test user"
+}
+
+async fn public_test_route() -> impl IntoResponse {
+    "public route"
+}
+
+async fn protected_test_route() -> impl IntoResponse {
+    "protected route"
+}
+

src/api/tests/middleware.rs

@@ -0,0 +1,96 @@
+use super::*;
+use axum::http::Method;
+
+#[tokio::test]
+async fn test_auth_middleware_rejects_invalid_token() {
+    let app = test_app().await;
+    
+    let (status, _) = make_request(
+        app,
+        Method::GET,
+        "/protected",
+        None,
+        Some(vec![("Authorization".to_string(), "Bearer invalid-token".to_string())]),
+    ).await;
+    
+    // Should redirect to login or return unauthorized
+    assert!(status == StatusCode::UNAUTHORIZED || status == StatusCode::FOUND);
+}
+
+#[tokio::test]
+async fn test_auth_middleware_accepts_valid_token() {
+    let app = test_app().await;
+    
+    // Create a valid token for testing
+    let token = create_test_token();
+    
+    let (status, _) = make_request(
+        app,
+        Method::GET,
+        "/protected",
+        None,
+        Some(vec![("Authorization".to_string(), format!("Bearer {}", token))]),
+    ).await;
+    
+    assert_eq!(status, StatusCode::OK);
+}
+
+#[tokio::test]
+async fn test_session_middleware_creates_session() {
+    let app = test_app().await;
+    
+    let (status, headers) = make_request_with_response_headers(
+        app,
+        Method::GET,
+        "/login",
+        None,
+        None,
+    ).await;
+    
+    assert_eq!(status, StatusCode::OK);
+    
+    // Check that a session cookie was set
+    let has_session_cookie = headers.iter()
+        .any(|(name, value)| name.to_lowercase() == "set-cookie" && value.contains("session="));
+    
+    assert!(has_session_cookie);
+}
+
+#[tokio::test]
+async fn test_error_handling_middleware_redirects_to_login() {
+    let app = test_app().await;
+    
+    // Make a request that will trigger an unauthorized error with the specific header
+    let (status, _) = make_request(
+        app,
+        Method::GET,
+        "/protected",
+        None,
+        Some(vec![
+            ("Authorization".to_string(), "Bearer invalid-token".to_string()),
+            ("X-Introspection-Error".to_string(), "unauthorized".to_string()),
+        ]),
+    ).await;
+    
+    // Should redirect to login
+    assert_eq!(status, StatusCode::FOUND);
+}
+
+#[tokio::test]
+async fn test_cors_middleware() {
+    let app = test_app().await;
+    
+    let (_, headers) = make_request_with_response_headers(
+        app,
+        Method::GET,
+        "/public",
+        None,
+        Some(vec![("Origin".to_string(), "http://example.com".to_string())]),
+    ).await;
+    
+    // Check that CORS headers were set
+    let has_cors_headers = headers.iter()
+        .any(|(name, _)| name.to_lowercase() == "access-control-allow-origin");
+    
+    assert!(has_cors_headers);
+}

src/api/tests/mod.rs

@@ -0,0 +1,126 @@
+use axum::{
+    body::Body,
+    http::{Request, StatusCode},
+    Router,
+};
+use tower::ServiceExt;
+
+// Import your API router
+use crate::api::router;
+
+// Helper function to create a test app
+async fn test_app() -> Router {
+    // Create a test configuration
+    let config = TestConfig::default();
+    
+    // Create the router with test configuration
+    router::create_router(config).await
+}
+
+// Helper function to make a test request
+async fn make_request(
+    app: Router,
+    method: http::Method,
+    uri: &str,
+    body: Option<String>,
+    headers: Option<Vec<(String, String)>>,
+) -> (StatusCode, String) {
+    let mut req_builder = Request::builder()
+        .method(method)
+        .uri(uri);
+    
+    // Add headers if provided
+    if let Some(headers) = headers {
+        for (name, value) in headers {
+            req_builder = req_builder.header(name, value);
+        }
+    }
+    
+    // Add body if provided
+    let body = match body {
+        Some(b) => Body::from(b),
+        None => Body::empty(),
+    };
+    
+    let req = req_builder.body(Body::from(body)).unwrap();
+    
+    // Process the request
+    let response = app.oneshot(req).await.unwrap();
+    
+    // Extract status code
+    let status = response.status();
+    
+    // Extract body
+    let body = hyper::body::to_bytes(response.into_body())
+        .await
+        .unwrap();
+    let body = String::from_utf8(body.to_vec()).unwrap();
+    
+    (status, body)
+}
+
+// Helper function to make a request and return headers
+async fn make_request_with_response_headers(
+    app: Router,
+    method: http::Method,
+    uri: &str,
+    body: Option<String>,
+    headers: Option<Vec<(String, String)>>,
+) -> (StatusCode, Vec<(String, String)>) {
+    let mut req_builder = Request::builder()
+        .method(method)
+        .uri(uri);
+    
+    // Add headers if provided
+    if let Some(headers) = headers {
+        for (name, value) in headers {
+            req_builder = req_builder.header(name, value);
+        }
+    }
+    
+    // Add body if provided
+    let body = match body {
+        Some(b) => Body::from(b),
+        None => Body::empty(),
+    };
+    
+    let req = req_builder.body(Body::from(body)).unwrap();
+    
+    // Process the request
+    let response = app.oneshot(req).await.unwrap();
+    
+    // Extract status code
+    let status = response.status();
+    
+    // Extract headers
+    let headers = response.headers().iter()
+        .map(|(name, value)| (name.to_string(), value.to_str().unwrap_or("").to_string()))
+        .collect();
+    
+    (status, headers)
+}
+
+// Helper function to create a test token
+fn create_test_token() -> String {
+    // In a real implementation, this would create a valid JWT token
+    // For testing purposes, we can use a placeholder
+    "test-token".to_string()
+}
+
+// Helper struct for test configuration
+#[derive(Clone)]
+struct TestConfig {
+    // Add fields as needed for your tests
+}
+
+impl Default for TestConfig {
+    fn default() -> Self {
+        Self {
+            // Initialize with default values
+        }
+    }
+}
+
+// Export the test modules
+pub mod routes;
+pub mod middleware;

src/api/tests/routes.rs

@@ -0,0 +1,87 @@
+use super::*;
+use axum::http::Method;
+
+#[tokio::test]
+async fn test_public_route_accessible() {
+    let app = test_app().await;
+    
+    let (status, body) = make_request(
+        app,
+        Method::GET,
+        "/public",
+        None,
+        None,
+    ).await;
+    
+    assert_eq!(status, StatusCode::OK);
+    assert_eq!(body, "public route");
+}
+
+#[tokio::test]
+async fn test_protected_route_requires_auth() {
+    let app = test_app().await;
+    
+    let (status, _) = make_request(
+        app,
+        Method::GET,
+        "/protected",
+        None,
+        None,
+    ).await;
+    
+    // Should redirect to login or return unauthorized
+    assert!(status == StatusCode::UNAUTHORIZED || status == StatusCode::FOUND);
+}
+
+#[tokio::test]
+async fn test_protected_route_with_valid_token() {
+    let app = test_app().await;
+    
+    // Create a valid token for testing
+    let token = create_test_token();
+    
+    let (status, body) = make_request(
+        app,
+        Method::GET,
+        "/protected",
+        None,
+        Some(vec![("Authorization".to_string(), format!("Bearer {}", token))]),
+    ).await;
+    
+    assert_eq!(status, StatusCode::OK);
+    assert_eq!(body, "protected route");
+}
+
+#[tokio::test]
+async fn test_login_page_accessible() {
+    let app = test_app().await;
+    
+    let (status, _) = make_request(
+        app,
+        Method::GET,
+        "/login",
+        None,
+        None,
+    ).await;
+    
+    assert_eq!(status, StatusCode::OK);
+}
+
+#[tokio::test]
+async fn test_whoami_endpoint() {
+    let app = test_app().await;
+    
+    // Create a valid token for testing
+    let token = create_test_token();
+    
+    let (status, body) = make_request(
+        app,
+        Method::GET,
+        "/api/whoami",
+        None,
+        Some(vec![("Authorization".to_string(), format!("Bearer {}", token))]),
+    ).await;
+    
+    assert_eq!(status, StatusCode::OK);
+    assert_eq!(body, "test user");
+}

src/config.rs

@@ -0,0 +1,81 @@
+//! Configuration management for the application.
+//!
+//! This module centralizes all configuration settings and provides validation
+//! for required configuration at startup.
+
+use std::fmt::Debug;
+use worker::Env;
+
+/// Constants for KV storage keys
+pub const KV_STORAGE_BINDING: &str = "KV_STORAGE";
+pub const SIGNING_KEY: &str = "keystore::sig";
+pub const ENCRYPTION_KEY: &str = "keystore::enc";
+
+/// Application configuration
+#[derive(Clone, Debug)]
+pub struct Config {
+    /// The URL of the authentication server
+    pub auth_server_url: String,
+    /// The client ID for OAuth authentication
+    pub client_id: String,
+    /// The client secret for OAuth authentication
+    pub client_secret: String,
+    /// The application URL
+    pub app_url: String,
+    /// Whether the application is running in development mode
+    pub dev_mode: bool,
+}
+
+impl Config {
+    /// Create a new configuration from environment variables
+    ///
+    /// # Arguments
+    ///
+    /// * `env` - The environment containing configuration values
+    ///
+    /// # Returns
+    ///
+    /// A Result containing the configuration or an error if required values are missing
+    pub fn from_env(env: &Env) -> Result<Self, ConfigError> {
+        let auth_server_url = env
+            .secret("AUTH_SERVER_URL")
+            .map_err(|_| ConfigError::MissingValue("AUTH_SERVER_URL"))?
+            .to_string();
+
+        let client_id = env
+            .secret("CLIENT_ID")
+            .map_err(|_| ConfigError::MissingValue("CLIENT_ID"))?
+            .to_string();
+
+        let client_secret = env
+            .secret("CLIENT_SECRET")
+            .map_err(|_| ConfigError::MissingValue("CLIENT_SECRET"))?
+            .to_string();
+
+        let app_url = env
+            .secret("APP_URL")
+            .map_err(|_| ConfigError::MissingValue("APP_URL"))?
+            .to_string();
+
+        let dev_mode = env
+            .var("DEV_MODE")
+            .map(|var| var.to_string() == "true")
+            .unwrap_or(false);
+
+        Ok(Config {
+            auth_server_url,
+            client_id,
+            client_secret,
+            app_url,
+            dev_mode,
+        })
+    }
+}
+
+/// Errors that can occur when loading configuration
+#[derive(Debug, thiserror::Error)]
+pub enum ConfigError {
+    /// A required configuration value is missing
+    #[error("Missing required configuration value: {0}")]
+    MissingValue(&'static str),
+}

src/docs.rs

@@ -0,0 +1,89 @@
+//! # axum-tower-sessions-edge Documentation
+//!
+//! This module provides comprehensive documentation for the axum-tower-sessions-edge project.
+//! It serves as a central place for understanding the project's architecture, components,
+//! and usage patterns.
+//!
+//! ## Overview
+//!
+//! axum-tower-sessions-edge is a Rust library that validates incoming requests for defined routes
+//! and forwards traffic to the service defined as `PROXY_TARGET`. It's designed to work with
+//! Cloudflare Workers and targets the `wasm32-unknown-unknown` platform.
+//!
+//! ## Features
+//!
+//! - **OAuth 2.0**: Implementation of the OAuth 2.0 authorization framework
+//! - **PKCE (Proof Key for Code Exchange)**: Enhanced security for OAuth 2.0
+//! - **Token Introspection**: Validation of OAuth 2.0 tokens
+//!
+//! ## Architecture
+//!
+//! The project is organized into several modules:
+//!
+//! - **api**: Contains the API endpoints for both authenticated and public routes
+//! - **axum_introspector**: Handles token introspection with Axum
+//! - **credentials**: Manages authentication credentials
+//! - **oidc**: Implements OpenID Connect functionality
+//! - **session_storage**: Handles session management
+//! - **utilities**: Provides utility functions
+//! - **zitadel_http**: HTTP client for Zitadel
+//!
+//! ## Usage
+//!
+//! ### Basic Setup
+//!
+//! To use this library, you need to configure it with your OAuth 2.0 provider details:
+//!
+//! ```rust
+//! // Example configuration (not actual code)
+//! let introspection_state = IntrospectionStateBuilder::new("https://your-auth-server-url")
+//!     .with_basic_auth("your-client-id", "your-client-secret")
+//!     .with_introspection_cache(cache)
+//!     .build()
+//!     .await
+//!     .unwrap();
+//! ```
+//!
+//! ### Authentication Flow
+//!
+//! The library implements a standard OAuth 2.0 flow:
+//!
+//! 1. User accesses a protected route
+//! 2. If not authenticated, they are redirected to the login page
+//! 3. User authenticates with the OAuth provider
+//! 4. Provider redirects back with an authorization code
+//! 5. The code is exchanged for tokens
+//! 6. User session is established
+//! 7. User is granted access to protected resources
+//!
+//! ## Components
+//!
+//! ### IntrospectionState
+//!
+//! Central component for token introspection and validation:
+//!
+//! ```rust
+//! // Example usage (not actual code)
+//! let introspection_state = IntrospectionStateBuilder::new(auth_server_url)
+//!     .with_basic_auth(client_id, client_secret)
+//!     .with_introspection_cache(cache)
+//!     .build()
+//!     .await?;
+//! ```
+//!
+//! ### Session Management
+//!
+//! The library uses tower-sessions for session management:
+//!
+//! ```rust
+//! // Example session setup (not actual code)
+//! let session_layer = SessionManagerLayer::new(session_store)
+//!     .with_name("session")
+//!     .with_expiry(Expiry::OnSessionEnd)
+//!     .with_secure(!is_dev);
+//! ```
+//!
+//! ## Deployment
+//!
+//! This library is designed to be deployed as a Cloudflare Worker. See the README.md for
+//! detailed deployment instructions.

src/error.rs

@@ -0,0 +1,71 @@
+//! Error handling for the application.
+//!
+//! This module provides centralized error handling functionality,
+//! including middleware for handling introspection errors.
+
+use axum::response::{IntoResponse, Redirect, Response};
+use http::StatusCode;
+
+/// Middleware for handling introspection errors.
+///
+/// This middleware checks for specific error headers and redirects
+/// to the login page when appropriate.
+pub async fn handle_introspection_errors(mut response: Response) -> Response {
+    let x_error_header_value = response
+        .headers()
+        .get("x-introspection-error")
+        .and_then(|header_value| header_value.to_str().ok());
+
+    // not used but is available
+    let x_session_header_value = response
+        .headers()
+        .get("x-session")
+        .and_then(|header_value| header_value.to_str().ok());
+
+    match response.status() {
+        StatusCode::UNAUTHORIZED => {
+            if let Some(x_error) = x_error_header_value {
+                if x_error == "unauthorized" {
+                    return Redirect::to("/login").into_response();
+                }
+            }
+            response
+        }
+        StatusCode::BAD_REQUEST => {
+            if let Some(x_error) = x_error_header_value {
+                if x_error == "invalid schema"
+                    || x_error == "invalid header"
+                    || x_error == "introspection error"
+                {
+                    return Redirect::to("/login").into_response();
+                }
+            }
+            response
+        }
+        StatusCode::FORBIDDEN => {
+            if let Some(x_error) = x_error_header_value {
+                if x_error == "user is inactive" {
+                    return Redirect::to("/login").into_response();
+                }
+            }
+            response
+        }
+        StatusCode::NOT_FOUND => {
+            if let Some(x_error) = x_error_header_value {
+                if x_error == "user was not found" {
+                    return Redirect::to("/login").into_response();
+                }
+            }
+            response
+        }
+        StatusCode::INTERNAL_SERVER_ERROR => {
+            if let Some(x_error) = x_error_header_value {
+                if x_error == "missing config" {
+                    return Redirect::to("/login").into_response();
+                }
+            }
+            response
+        }
+        _ => response,
+    }
+}

src/lib.rs

@@ -1,38 +1,43 @@
+//! # axum-tower-sessions-edge
+//!
+//! A Rust library that validates incoming requests for defined routes and forwards traffic 
+//! to the service defined as `PROXY_TARGET`. It's designed to work with Cloudflare Workers 
+//! and targets the `wasm32-unknown-unknown` platform.
+//!
+//! ## Features
+//!
+//! - OAuth 2.0 authentication flow
+//! - Proof Key for Code Exchange (PKCE) for enhanced security
+//! - OAuth 2.0 Token Introspection for token validation
+//! - Session management with tower-sessions
+//! - Cloudflare Workers integration
+//!
+//! See the [docs](crate::docs) module for comprehensive documentation.
+
 mod api;
 mod axum_introspector;
+mod config;
 mod credentials;
+mod docs;
+mod error;
 mod oidc;
+mod router;
+mod session;
 mod session_storage;
 mod utilities;
 mod zitadel_http;
 
-use crate::api::authenticated::AuthenticatedApi;
-use crate::api::public::PublicApi;
-use crate::axum_introspector::introspection::{
-    IntrospectedUser, IntrospectionState, IntrospectionStateBuilder,
-};
+use axum::handler::Handler;
+use crate::axum_introspector::introspection::IntrospectionStateBuilder;
+use crate::config::{Config, KV_STORAGE_BINDING};
 use crate::oidc::introspection::cache::cloudflare::CloudflareIntrospectionCache;
+use crate::router::{create_router, AppState};
+use crate::session::{create_session_layer, SessionConfig};
 use crate::session_storage::cloudflare::CloudflareKvStore;
-use axum::extract::FromRef;
-use axum::response::{IntoResponse, Redirect};
-use axum::routing::{any, get};
-use axum::{Router, ServiceExt};
-use bytes::Bytes;
-use http::HeaderName;
 use serde::{Deserialize, Serialize};
-use serde_json::to_string;
-use std::fmt::Debug;
-use std::iter::once;
-use std::ops::Deref;
-use tower::ServiceExt as TowerServiceExt;
+use tower::ServiceExt;
 use tower_cookies::cookie::SameSite;
-use tower_cookies::CookieManagerLayer;
-use tower_http::cors::CorsLayer;
-use tower_http::propagate_header::PropagateHeaderLayer;
-use tower_http::sensitive_headers::SetSensitiveRequestHeadersLayer;
 use tower_service::Service;
-use tower_sessions::cookie::Key;
-use tower_sessions::SessionManagerLayer;
 use tower_sessions_core::Expiry;
 use tracing::instrument::WithSubscriber;
 use tracing_subscriber::prelude::*;
@@ -54,227 +59,105 @@ fn start() {
         .init()
 }
 
-const SIGNING_KEY: &str = "keystore::sig";
-const ENCRYPTION_KEY: &str = "keystore::enc";
-
-// main entrypoint
-
-#[event(fetch)]
-async fn fetch(
-    req: HttpRequest,
-    _env: Env,
-    _ctx: Context,
-) -> Result<axum::http::Response<axum::body::Body>> {
-    console_error_panic_hook::set_once();
-
-    Ok(route(req, _env).await)
-}
-
 #[derive(Serialize, Deserialize, Clone, Debug)]
 struct Callback {
     code: String,
     state: String,
 }
 
-#[derive(Clone)]
-struct AppState {
-    introspection_state: IntrospectionState,
+#[event(fetch)]
+async fn fetch(
+    req: HttpRequest,
     env: Env,
-    session_store: CloudflareKvStore,
-}
-impl FromRef<AppState> for IntrospectionState {
-    fn from_ref(input: &AppState) -> Self {
-        input.introspection_state.clone()
-    }
+    _ctx: Context,
+) -> Result<axum::http::Response<axum::body::Body>> {
+    console_error_panic_hook::set_once();
+
+    Ok(route(req, env).await)
 }
 
-async fn route(req: HttpRequest, _env: Env) -> axum_core::response::Response {
-    let kv = _env.kv("KV_STORAGE").unwrap();
+async fn route(req: HttpRequest, env: Env) -> axum::http::Response<axum::body::Body> {
+    // Load configuration from environment
+    let config = match Config::from_env(&env) {
+        Ok(config) => config,
+        Err(err) => {
+            console_error!("Configuration error: {}", err);
+            return axum::http::Response::builder()
+                .status(500)
+                .body(axum::body::Body::from("Internal Server Error: Configuration error"))
+                .unwrap();
+        }
+    };
+
+    // Initialize KV store
+    let kv = match env.kv(KV_STORAGE_BINDING) {
+        Ok(kv) => kv,
+        Err(err) => {
+            console_error!("KV store error: {}", err);
+            return axum::http::Response::builder()
+                .status(500)
+                .body(axum::body::Body::from("Internal Server Error: KV store error"))
+                .unwrap();
+        }
+    };
+
+    // Initialize introspection cache
     let cache = CloudflareIntrospectionCache::new(kv.clone());
 
-    let introspection_state = IntrospectionStateBuilder::new(
-        _env.secret("AUTH_SERVER_URL")
-            .unwrap()
-            .to_string()
-            .as_str(),
-    )
-    .with_basic_auth(
-        _env.secret("CLIENT_ID")
-            .unwrap()
-            .to_string()
-            .as_str(),
-        _env.secret("CLIENT_SECRET")
-            .unwrap()
-            .to_string()
-            .as_str(),
-    )
-    .with_introspection_cache(cache)
-    .build()
-    .await
-    .unwrap();
+    // Build introspection state
+    let introspection_state = match IntrospectionStateBuilder::new(&config.auth_server_url)
+        .with_basic_auth(&config.client_id, &config.client_secret)
+        .with_introspection_cache(cache)
+        .build()
+        .await
+    {
+        Ok(state) => state,
+        Err(err) => {
+            console_error!("Introspection state error: {}", err);
+            return axum::http::Response::builder()
+                .status(500)
+                .body(axum::body::Body::from("Internal Server Error: Introspection state error"))
+                .unwrap();
+        }
+    };
 
+    // Initialize session store
     let session_store = CloudflareKvStore::new(kv.clone());
 
+    // Create application state
     let state = AppState {
         introspection_state,
         session_store: session_store.clone(),
-        env: _env.clone(),
+        env: env.clone(),
     };
 
-    let dev_mode = _env.var("DEV_MODE").unwrap().to_string(); // Example check
-
-    let is_dev = dev_mode == "true";
-
-    let keystore = _env.kv("KV_STORAGE").unwrap();
-
-    let signing = if let Some(bytes) = keystore.get(SIGNING_KEY).bytes().await.unwrap() {
-        Key::derive_from(bytes.as_slice())
-    } else {
-        let key = Key::generate();
-        keystore
-            .put_bytes(SIGNING_KEY, key.master())
-            .unwrap()
-            .execute()
-            .await
-            .unwrap();
-        key
+    // Create session configuration
+    let session_config = SessionConfig {
+        cookie_name: "session".to_string(),
+        expiry: Expiry::OnSessionEnd,
+        domain: "localhost".to_string(), // Will be overridden in create_session_layer
+        path: "/".to_string(),
+        secure: !config.dev_mode,
+        same_site: SameSite::Lax,
     };
 
-    let encryption = if let Some(bytes) = keystore.get(ENCRYPTION_KEY).bytes().await.unwrap() {
-        Key::derive_from(bytes.as_slice())
-    } else {
-        let key = Key::generate();
-        keystore
-            .put_bytes(ENCRYPTION_KEY, key.master())
-            .unwrap()
-            .execute()
-            .await
-            .unwrap();
-        key
-    };
+    // Create session layer
+    let session_layer = create_session_layer(
+        &config,
+        Some(session_config),
+        session_store,
+        kv,
+    ).await;
 
-    let host_string = _env.secret("APP_URL").unwrap().to_string().as_str().to_owned();
+    // Create router
+    let router = create_router(state, session_layer);
 
-    let cookie_host_uri = host_string.parse::<http::Uri>().unwrap();
+    // Handle request
+    // Convert the worker request to an axum request
+    let axum_req = axum::extract::Request::try_from(req).unwrap();
 
-    let mut cookie_host = cookie_host_uri.authority().unwrap().to_string();
-    
-    if cookie_host.starts_with("localhost:") {
-        cookie_host = "localhost".to_string();
-    }
-
-    let session_layer = SessionManagerLayer::new(state.session_store.clone())
-        .with_name("session")
-        .with_expiry(Expiry::OnSessionEnd)
-        .with_domain(cookie_host)
-        .with_same_site(SameSite::Lax)
-        .with_signed(signing)
-        .with_private(encryption)
-        .with_path("/")
-        .with_secure(!is_dev)
-        .with_always_save(false);
-
-    async fn handle_introspection_errors(
-        mut response: axum_core::response::Response,
-    ) -> axum_core::response::Response {
-        let x_error_header_value = response
-            .headers()
-            .get("x-introspection-error")
-            .and_then(|header_value| header_value.to_str().ok());
-
-        // not used but is available
-        let x_session_header_value = response
-            .headers()
-            .get("x-session")
-            .and_then(|header_value| header_value.to_str().ok());
-
-        match response.status() {
-            http::StatusCode::UNAUTHORIZED => {
-                if let Some(x_error) = x_error_header_value {
-                    if x_error == "unauthorized" {
-                        return Redirect::to("/login").into_response();
-                    }
-                }
-                response
-            }
-            http::StatusCode::BAD_REQUEST => {
-                if let Some(x_error) = x_error_header_value {
-                    if x_error == "invalid schema"
-                        || x_error == "invalid header"
-                        || x_error == "introspection error"
-                    {
-                        return Redirect::to("/login").into_response();
-                    }
-                }
-                response
-            }
-            http::StatusCode::FORBIDDEN => {
-                if let Some(x_error) = x_error_header_value {
-                    if x_error == "user is inactive" {
-                        return Redirect::to("/login").into_response();
-                    }
-                }
-                response
-            }
-            http::StatusCode::NOT_FOUND => {
-                if let Some(x_error) = x_error_header_value {
-                    if x_error == "user was not found" {
-                        return Redirect::to("/login").into_response();
-                    }
-                }
-                response
-            }
-            http::StatusCode::INTERNAL_SERVER_ERROR => {
-                if let Some(x_error) = x_error_header_value {
-                    if x_error == "missing config" {
-                        return Redirect::to("/login").into_response();
-                    }
-                }
-                response
-            }
-            _ => response,
-        }
-    }
-
-    let mut router = Router::new()
-        .route("/", any(AuthenticatedApi::proxy))
-        .route("/login", get(PublicApi::login_page)) // Add the login page route
-        .route("/login/callback", get(PublicApi::callback))
-        .route("/login/authorize", get(PublicApi::authorize))
-        .route("/api/whoami", get(whoami))
-        .route("/*path", any(AuthenticatedApi::proxy))
-        .layer(PropagateHeaderLayer::new(HeaderName::from_static(
-            "x-request-id",
-        )))
-        .layer(axum::middleware::map_response(handle_introspection_errors))
-        .with_state(state)
-        .layer(session_layer)
-        .layer(CookieManagerLayer::new())
-        .layer(CorsLayer::very_permissive())
-        .layer(SetSensitiveRequestHeadersLayer::new(once(
-            http::header::AUTHORIZATION,
-        )));
-
-    router
-        .as_service()
-        .ready()
-        .await
-        .unwrap()
-        .oneshot(req)
-        .await
-        .unwrap()
-}
-
-async fn whoami(
-    session: tower_sessions::Session,
-    introspected_user: IntrospectedUser,
-) -> impl IntoResponse {
-    console_log!("calling whoami");
-    to_string(&introspected_user).unwrap()
-}
-
-impl FromRef<AppState> for CloudflareKvStore {
-    fn from_ref(input: &AppState) -> Self {
-        input.session_store.clone()
-    }
+    // Use the router to handle the request
+    // Since we've modified create_router to return a Router with empty state,
+    // we can now use the oneshot method directly
+    router.oneshot(axum_req).await.unwrap()
 }

src/router.rs

@@ -0,0 +1,99 @@
+//! Routing configuration for the application.
+//!
+//! This module provides centralized routing functionality,
+//! including router configuration and middleware setup.
+
+use crate::api::authenticated::AuthenticatedApi;
+use crate::api::public::PublicApi;
+use crate::error::handle_introspection_errors;
+use worker::console_log;
+use axum::extract::FromRef;
+use axum::response::IntoResponse;
+use axum::routing::{any, get};
+use axum::{Router, ServiceExt};
+use http::HeaderName;
+use serde_json::to_string;
+use std::iter::once;
+use tower_cookies::CookieManagerLayer;
+use tower_http::cors::CorsLayer;
+use tower_http::propagate_header::PropagateHeaderLayer;
+use tower_http::sensitive_headers::SetSensitiveRequestHeadersLayer;
+use tower_sessions::SessionManagerLayer;
+
+use crate::axum_introspector::introspection::{IntrospectedUser, IntrospectionState};
+use crate::session_storage::cloudflare::CloudflareKvStore;
+
+/// Application state shared across handlers
+#[derive(Clone)]
+pub struct AppState {
+    /// State for token introspection
+    pub introspection_state: IntrospectionState,
+    /// Cloudflare environment
+    pub env: worker::Env,
+    /// Session store
+    pub session_store: CloudflareKvStore,
+}
+
+impl FromRef<AppState> for IntrospectionState {
+    fn from_ref(input: &AppState) -> Self {
+        input.introspection_state.clone()
+    }
+}
+
+impl FromRef<AppState> for CloudflareKvStore {
+    fn from_ref(input: &AppState) -> Self {
+        input.session_store.clone()
+    }
+}
+
+/// Create a router with the given state and session layer
+///
+/// # Arguments
+///
+/// * `state` - The application state
+/// * `session_layer` - The session manager layer
+///
+/// # Returns
+///
+/// A configured router
+pub fn create_router(
+    state: AppState,
+    session_layer: SessionManagerLayer<CloudflareKvStore, tower_sessions::service::PrivateCookie>,
+) -> Router {
+    Router::new()
+        .route("/", any(AuthenticatedApi::proxy))
+        .route("/login", get(PublicApi::login_page))
+        .route("/login/callback", get(PublicApi::callback))
+        .route("/login/authorize", get(PublicApi::authorize))
+        .route("/api/whoami", get(whoami))
+        .route("/*path", any(AuthenticatedApi::proxy))
+        .layer(PropagateHeaderLayer::new(HeaderName::from_static(
+            "x-request-id",
+        )))
+        .layer(axum::middleware::map_response(handle_introspection_errors))
+        .with_state(state)
+        .layer(session_layer)
+        .layer(CookieManagerLayer::new())
+        .layer(CorsLayer::very_permissive())
+        .layer(SetSensitiveRequestHeadersLayer::new(once(
+            http::header::AUTHORIZATION,
+        )))
+}
+
+/// Handler for the whoami endpoint
+///
+/// # Arguments
+///
+/// * `session` - The user's session
+/// * `introspected_user` - The introspected user information
+///
+/// # Returns
+///
+/// The user information as JSON
+pub async fn whoami(
+    session: tower_sessions::Session,
+    introspected_user: IntrospectedUser,
+) -> impl IntoResponse {
+    console_log!("calling whoami");
+    to_string(&introspected_user).unwrap()
+}

src/session.rs

@@ -0,0 +1,128 @@
+//! Session management for the application.
+//!
+//! This module provides centralized session management functionality,
+//! including session configuration and key management.
+
+use crate::config::{Config, ENCRYPTION_KEY, SIGNING_KEY};
+use crate::session_storage::cloudflare::CloudflareKvStore;
+use tower_cookies::cookie::SameSite;
+use tower_sessions::cookie::Key;
+use tower_sessions::service::PrivateCookie;
+use tower_sessions::SessionManagerLayer;
+use tower_sessions_core::Expiry;
+use worker::kv::KvStore as Kv;
+
+/// Session configuration options
+#[derive(Clone, Debug)]
+pub struct SessionConfig {
+    /// The name of the session cookie
+    pub cookie_name: String,
+    /// The expiry policy for the session
+    pub expiry: Expiry,
+    /// The domain for the session cookie
+    pub domain: String,
+    /// The path for the session cookie
+    pub path: String,
+    /// Whether the session cookie should be secure
+    pub secure: bool,
+    /// The same-site policy for the session cookie
+    pub same_site: SameSite,
+}
+
+impl Default for SessionConfig {
+    fn default() -> Self {
+        Self {
+            cookie_name: "session".to_string(),
+            expiry: Expiry::OnSessionEnd,
+            domain: "localhost".to_string(),
+            path: "/".to_string(),
+            secure: true,
+            same_site: SameSite::Lax,
+        }
+    }
+}
+
+/// Create a session manager layer with the given configuration
+///
+/// # Arguments
+///
+/// * `config` - The application configuration
+/// * `session_config` - The session configuration
+/// * `session_store` - The session store
+/// * `keystore` - The KV store for key management
+///
+/// # Returns
+///
+/// A session manager layer
+pub async fn create_session_layer(
+    config: &Config,
+    session_config: Option<SessionConfig>,
+    session_store: CloudflareKvStore,
+    keystore: Kv,
+) -> SessionManagerLayer<CloudflareKvStore, PrivateCookie> {
+    let session_config = session_config.unwrap_or_default();
+
+    let (signing, encryption) = get_or_create_keys(keystore).await;
+
+    let mut domain = session_config.domain;
+
+    // Handle localhost special case
+    if let Ok(uri) = config.app_url.parse::<http::Uri>() {
+        if let Some(authority) = uri.authority() {
+            domain = authority.to_string();
+            if domain.starts_with("localhost:") {
+                domain = "localhost".to_string();
+            }
+        }
+    }
+
+    SessionManagerLayer::new(session_store)
+        .with_name(session_config.cookie_name)
+        .with_expiry(session_config.expiry)
+        .with_domain(domain)
+        .with_same_site(session_config.same_site)
+        .with_signed(signing)
+        .with_private(encryption)
+        .with_path(session_config.path)
+        .with_secure(!config.dev_mode)
+        .with_always_save(false)
+}
+
+/// Get or create signing and encryption keys
+///
+/// # Arguments
+///
+/// * `keystore` - The KV store for key management
+///
+/// # Returns
+///
+/// A tuple of (signing_key, encryption_key)
+async fn get_or_create_keys(keystore: Kv) -> (Key, Key) {
+    let signing = if let Some(bytes) = keystore.get(SIGNING_KEY).bytes().await.unwrap() {
+        Key::derive_from(bytes.as_slice())
+    } else {
+        let key = Key::generate();
+        keystore
+            .put_bytes(SIGNING_KEY, key.master())
+            .unwrap()
+            .execute()
+            .await
+            .unwrap();
+        key
+    };
+
+    let encryption = if let Some(bytes) = keystore.get(ENCRYPTION_KEY).bytes().await.unwrap() {
+        Key::derive_from(bytes.as_slice())
+    } else {
+        let key = Key::generate();
+        keystore
+            .put_bytes(ENCRYPTION_KEY, key.master())
+            .unwrap()
+            .execute()
+            .await
+            .unwrap();
+        key
+    };
+
+    (signing, encryption)
+}