path: root/src/oauth/server.rs

use crate::clients::{ClientManager, parse_basic_auth};
use crate::config::Config;
use crate::keys::KeyManager;
use crate::oauth::types::{AuthCode, Claims, ErrorResponse, TokenResponse};
use jsonwebtoken::{Algorithm, Header, encode};
use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};
use url::Url;
use uuid::Uuid;

pub struct OAuthServer {
    config: Config,
    key_manager: std::sync::Mutex<KeyManager>,
    auth_codes: std::sync::Mutex<HashMap<String, AuthCode>>,
    client_manager: std::sync::Mutex<ClientManager>,
}

impl OAuthServer {
    pub fn new(config: &Config) -> Result<Self, Box<dyn std::error::Error>> {
        let key_manager = KeyManager::new()?;
        let client_manager = ClientManager::new();

        Ok(Self {
            key_manager: std::sync::Mutex::new(key_manager),
            auth_codes: std::sync::Mutex::new(HashMap::new()),
            client_manager: std::sync::Mutex::new(client_manager),
            config: config.clone(),
        })
    }

    pub fn get_jwks(&self) -> String {
        let key_manager = self.key_manager.lock().unwrap();
        match key_manager.get_jwks() {
            Ok(jwks) => serde_json::to_string(&jwks).unwrap_or_else(|_| "{}".to_string()),
            Err(_) => serde_json::json!({"keys": []}).to_string(),
        }
    }

    pub fn handle_authorize(&self, params: &HashMap<String, String>) -> Result<String, String> {
        let client_id = params
            .get("client_id")
            .ok_or_else(|| self.error_response("invalid_request", "Missing client_id"))?;

        let redirect_uri = params
            .get("redirect_uri")
            .ok_or_else(|| self.error_response("invalid_request", "Missing redirect_uri"))?;

        let response_type = params
            .get("response_type")
            .ok_or_else(|| self.error_response("invalid_request", "Missing response_type"))?;

        // Validate client exists
        let client_manager = self.client_manager.lock().unwrap();
        let _client = client_manager.get_client(client_id)
            .ok_or_else(|| self.error_response("invalid_client", "Invalid client_id"))?;

        // Validate redirect URI is registered for this client
        if !client_manager.is_redirect_uri_valid(client_id, redirect_uri) {
            return Err(self.error_response("invalid_request", "Invalid redirect_uri"));
        }

        // Validate requested scopes
        let scope = params.get("scope").cloned();
        if !client_manager.is_scope_valid(client_id, &scope) {
            return Err(self.error_response("invalid_scope", "Invalid scope"));
        }

        if response_type != "code" {
            return Err(self.error_response(
                "unsupported_response_type",
                "Only code response type supported",
            ));
        }

        let code = Uuid::new_v4().to_string();
        let expires_at = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs()
            + 600;

        let auth_code = AuthCode {
            client_id: client_id.clone(),
            redirect_uri: redirect_uri.clone(),
            scope: scope,
            expires_at,
            user_id: "test_user".to_string(),
        };

        {
            let mut codes = self.auth_codes.lock().unwrap();
            codes.insert(code.clone(), auth_code);
        }

        let mut redirect_url = Url::parse(redirect_uri)
            .map_err(|_| self.error_response("invalid_request", "Invalid redirect_uri"))?;

        redirect_url.query_pairs_mut().append_pair("code", &code);

        if let Some(state) = params.get("state") {
            redirect_url.query_pairs_mut().append_pair("state", state);
        }

        Ok(redirect_url.to_string())
    }

    pub fn handle_token(&self, params: &HashMap<String, String>, auth_header: Option<&str>) -> Result<String, String> {
        let grant_type = params
            .get("grant_type")
            .ok_or_else(|| self.error_response("invalid_request", "Missing grant_type"))?;

        if grant_type != "authorization_code" {
            return Err(self.error_response(
                "unsupported_grant_type",
                "Only authorization_code grant type supported",
            ));
        }

        let code = params
            .get("code")
            .ok_or_else(|| self.error_response("invalid_request", "Missing code"))?;

        // Client authentication - RFC 6749 Section 3.2.1
        // Clients can authenticate via HTTP Basic Auth or form parameters
        let (client_id, client_secret) = if let Some(auth_header) = auth_header {
            // HTTP Basic Authentication (preferred method)
            parse_basic_auth(auth_header)
                .ok_or_else(|| self.error_response("invalid_client", "Invalid Authorization header"))?
        } else {
            // Form-based authentication (fallback)
            let client_id = params
                .get("client_id")
                .ok_or_else(|| self.error_response("invalid_request", "Missing client_id"))?;
            let client_secret = params
                .get("client_secret")
                .ok_or_else(|| self.error_response("invalid_request", "Missing client_secret"))?;
            (client_id.clone(), client_secret.clone())
        };

        // Authenticate the client
        let client_manager = self.client_manager.lock().unwrap();
        if !client_manager.authenticate_client(&client_id, &client_secret) {
            return Err(self.error_response("invalid_client", "Client authentication failed"));
        }

        let auth_code = {
            let mut codes = self.auth_codes.lock().unwrap();
            codes.remove(code).ok_or_else(|| {
                self.error_response("invalid_grant", "Invalid or expired authorization code")
            })?
        };

        if auth_code.client_id != client_id {
            return Err(self.error_response("invalid_grant", "Client ID mismatch"));
        }

        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        if now > auth_code.expires_at {
            return Err(self.error_response("invalid_grant", "Authorization code expired"));
        }

        let access_token =
            self.generate_access_token(&auth_code.user_id, &client_id, &auth_code.scope)?;

        let token_response = TokenResponse {
            access_token,
            token_type: "Bearer".to_string(),
            expires_in: 3600,
            refresh_token: None,
            scope: auth_code.scope,
        };

        serde_json::to_string(&token_response)
            .map_err(|_| self.error_response("server_error", "Failed to serialize token response"))
    }

    fn generate_access_token(
        &self,
        user_id: &str,
        client_id: &str,
        scope: &Option<String>,
    ) -> Result<String, String> {
        let mut key_manager = self.key_manager.lock().unwrap();

        // Check if we need to rotate keys
        if key_manager.should_rotate() {
            if let Err(_) = key_manager.rotate_keys() {
                return Err(self.error_response("server_error", "Key rotation failed"));
            }
        }

        let current_key = key_manager
            .get_current_key()
            .ok_or_else(|| self.error_response("server_error", "No signing key available"))?;

        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        let claims = Claims {
            sub: user_id.to_string(),
            iss: self.config.issuer_url.clone(),
            aud: client_id.to_string(),
            exp: now + 3600,
            iat: now,
            scope: scope.clone(),
        };

        let mut header = Header::new(Algorithm::RS256);
        header.kid = Some(current_key.kid.clone());

        encode(&header, &claims, &current_key.encoding_key)
            .map_err(|_| self.error_response("server_error", "Failed to generate token"))
    }

    fn error_response(&self, error: &str, description: &str) -> String {
        let error_resp = ErrorResponse {
            error: error.to_string(),
            error_description: Some(description.to_string()),
        };
        serde_json::to_string(&error_resp).unwrap_or_else(|_| "{}".to_string())
    }
}