osirose-new/auth-service/src/grpc.rs

use crate::auth::auth_service_server::AuthService;
use crate::auth::{LoginRequest, LoginResponse, PasswordResetRequest, PasswordResetResponse, RegisterRequest, RegisterResponse, ResetPasswordRequest, ResetPasswordResponse, ValidateTokenRequest, ValidateTokenResponse};
use crate::database_client::{DatabaseClientTrait};
use crate::jwt::{generate_token, validate_token};
use crate::users::{hash_password, verify_user};
use chrono::{Duration, Utc};
use rand::Rng;
use tonic::{Request, Response, Status};
use tracing::{info, warn};

pub struct MyAuthService<T: DatabaseClientTrait + Clone> {
    pub db_client: T,
}

#[tonic::async_trait]
impl<T: DatabaseClientTrait + Send + Sync + Clone + 'static> AuthService for MyAuthService<T> {
    async fn login(
        &self,
        request: Request<LoginRequest>,
    ) -> Result<Response<LoginResponse>, Status> {
        let req = request.into_inner();

        info!("Login attempt for username: {}", req.username);

        if let Some(user_id) = verify_user(self.db_client.clone(), &req.username, &req.password).await {
            let token = generate_token(&user_id, vec!["user".to_string()])
                .map_err(|_| Status::internal("Token generation failed"))?;

            info!("Login successful for username: {}", req.username);
            Ok(Response::new(LoginResponse { token, user_id }))
        } else {
            warn!("Invalid login attempt for username: {}", req.username);
            Err(Status::unauthenticated("Invalid credentials"))
        }
    }

    async fn validate_token(
        &self,
        request: Request<ValidateTokenRequest>,
    ) -> Result<Response<ValidateTokenResponse>, Status> {
        let req = request.into_inner();

        match validate_token(&req.token) {
            Ok(user_id) => Ok(Response::new(ValidateTokenResponse {
                valid: true,
                user_id,
            })),
            Err(_) => Ok(Response::new(ValidateTokenResponse {
                valid: false,
                user_id: "".to_string(),
            })),
        }
    }

    async fn register(
        &self,
        request: Request<RegisterRequest>,
    ) -> Result<Response<RegisterResponse>, Status> {
        let req = request.into_inner();

        // Hash the password
        let hashed_password = hash_password(&req.password);

        // Create user in the database
        let user = self.db_client.clone().create_user(&req.username, &req.email, &hashed_password)
            .await
            .map_err(|e| Status::internal(format!("Database error: {}", e)))?;

        Ok(Response::new(RegisterResponse { user_id: user.user_id }))
    }

    async fn request_password_reset(
        &self,
        request: Request<PasswordResetRequest>,
    ) -> Result<Response<PasswordResetResponse>, Status> {
        let email = request.into_inner().email;

        let user = self.db_client.clone().get_user_by_email(&email).await;

        // Check if the email exists
        if user.ok().is_some() {
            // Generate a reset token
            let reset_token: String = rand::thread_rng()
                .sample_iter(&rand::distributions::Alphanumeric)
                .take(32)
                .map(char::from)
                .collect();
            
            // Set token expiration (e.g., 1 hour)
            let expires_at = Utc::now() + Duration::hours(1);
            
            // Store the reset token in the database
            self.db_client.clone()
                .store_password_reset(&email, &reset_token, expires_at)
                .await
                .map_err(|e| Status::internal(format!("Database error: {}", e)))?;

            // Send the reset email
            // send_email(&email, "Password Reset Request", &format!(
            //     "Click the link to reset your password: https://example.com/reset?token={}",
            //     reset_token
            // ))
            //     .map_err(|e| Status::internal(format!("Email error: {}", e)))?;

            Ok(Response::new(PasswordResetResponse {
                message: "Password reset email sent".to_string(),
            }))
        } else {
            // Respond with a generic message to avoid information leaks
            Ok(Response::new(PasswordResetResponse {
                message: "If the email exists, a reset link has been sent.".to_string(),
            }))
        }
    }

    async fn reset_password(
        &self,
        request: Request<ResetPasswordRequest>,
    ) -> Result<Response<ResetPasswordResponse>, Status> {
        let req = request.into_inner();

        // Validate the reset token
        if let Some(password_reset) = self.db_client.clone().get_password_reset(&req.reset_token).await
                .map_err(|e| Status::internal(format!("Database error: {}", e)))? {
            if password_reset.expires_at < Utc::now() {
                return Err(Status::unauthenticated("Token expired"));
            }
            
            // Hash the new password
            let hashed_password = hash_password(&req.new_password);
            
            // Update the user's password
            self.db_client
                .update_user_password(&password_reset.email, &hashed_password)
                .await
                .map_err(|e| Status::internal(format!("Database error: {}", e)))?;
            
            // Delete the reset token
            self.db_client
                .delete_password_reset(&req.reset_token)
                .await
                .map_err(|e| Status::internal(format!("Database error: {}", e)))?;
            
            Ok(Response::new(ResetPasswordResponse {
                message: "Password successfully reset".to_string(),
            }))
        } else {
            Err(Status::unauthenticated("Invalid reset token"))
        }
    }
}