chat/src/main.rs

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Query, State,
    },
    http::StatusCode,
    response::{Html, IntoResponse, Redirect},
    routing::get,
    Router,
};
use dotenvy::dotenv;
use futures::{sink::SinkExt, stream::StreamExt};
use oauth2::{
    basic::BasicClient,AuthUrl, ClientId, ClientSecret, RedirectUrl, TokenResponse,
    TokenUrl,
};
use std::sync::Arc;
use tokio::sync::broadcast;
use tower_sessions::{cookie::SameSite, Expiry, MemoryStore, Session, SessionManagerLayer};

mod user;
use user::User;

// Define application state
struct AppState {
    tx: broadcast::Sender<String>,
    oauth_client: BasicClient,
}

const AFTER_LOGIN_URL: &str = "/";

#[tokio::main]
async fn main() {
    dotenv().ok();

    // Create a broadcast channel
    let (tx, _rx) = broadcast::channel(100);

    // OAuth configuration
    let client_id = dotenvy::var("CLIENT_ID").unwrap();
    let client_secret = dotenvy::var("CLIENT_SECRET").unwrap();
    // NOTE: In production, do not hardcode localhost
    let auth_url = AuthUrl::new("https://idm.flegr.me/application/o/authorize/".to_string())
        .expect("Invalid authorization endpoint URL");
    let token_url = TokenUrl::new("https://idm.flegr.me/application/o/token/".to_string())
        .expect("Invalid token endpoint URL");

    let client = BasicClient::new(
        ClientId::new(client_id),
        Some(ClientSecret::new(client_secret)),
        auth_url,
        Some(token_url),
    )
    .set_redirect_uri(
        RedirectUrl::new("http://localhost:3001/auth/callback".to_string())
            .expect("Invalid redirect URL"),
    );

    let app_state = Arc::new(AppState {
        tx,
        oauth_client: client,
    });

    // Session configuration
    let session_store = MemoryStore::default();
    let session_layer = SessionManagerLayer::new(session_store)
        .with_secure(false) // For localhost; set to true in production with HTTPS
        .with_same_site(SameSite::Lax) // Ensure we can receive cookies from OAuth redirect
        .with_expiry(Expiry::OnInactivity(time::Duration::minutes(30)));       

    // Build application with routes
    let app = Router::new()
        .route("/", get(index))
        .route("/login", get(login))
        .route("/logout", get(logout))
        .route("/auth/callback", get(auth_callback))
        .route("/ws", get(websocket_handler))
        .layer(session_layer)
        .with_state(app_state);

    // Run the app
    let listener = tokio::net::TcpListener::bind("0.0.0.0:3001").await.unwrap();
    println!("Chat server listening on http://0.0.0.0:3001");
    axum::serve(listener, app).await.unwrap();
}

// Handlers

async fn index(session: Session) -> impl IntoResponse {
    if let Some(_user) = session.get::<User>("user").await.unwrap() {
        return Html(include_str!("../index.html")).into_response();
    }
    Redirect::to("/login").into_response()
}

async fn login(State(state): State<Arc<AppState>>) -> impl IntoResponse {
    let (auth_url, _csrf_token) = state
        .oauth_client
        .authorize_url(oauth2::CsrfToken::new_random)
        .url();

    Redirect::to(auth_url.as_str())
}

async fn logout(session: Session) -> impl IntoResponse {
    session.delete().await.unwrap();
    Redirect::to("/login")
}

#[derive(serde::Deserialize)]
struct AuthRequest {
    code: String,
    state: String,
}

async fn auth_callback(
    Query(query): Query<AuthRequest>,
    State(state): State<Arc<AppState>>,
    session: Session,
) -> impl IntoResponse {
    let token = match state
        .oauth_client
        .exchange_code(oauth2::AuthorizationCode::new(query.code))
        .request_async(oauth2::reqwest::async_http_client)
        .await
    {
        Ok(t) => t,
        Err(e) => {
            return (
                StatusCode::BAD_REQUEST,
                format!("Failed to exchange authorization code: {}", e),
            )
            .into_response();
        }
    };

    let client = reqwest::Client::new();
    let user_info_resp = client
        .get("https://idm.flegr.me/application/o/userinfo/")
        .header("User-Agent", "axum-chat-app")
        .bearer_auth(token.access_token().secret())
        .send()
        .await;
    
    let user_data: User = match user_info_resp {
        Ok(resp) => match resp.json().await {
            Ok(u) => u,
            Err(e) => return (
                StatusCode::INTERNAL_SERVER_ERROR, 
                format!("Failed to parse user info: {}", e)
            ).into_response()
        },
        Err(e) => return (
            StatusCode::INTERNAL_SERVER_ERROR, 
            format!("Failed to fetch user info: {}", e)
        ).into_response()
    };

    session.insert("user", user_data).await.unwrap();

    Redirect::to(AFTER_LOGIN_URL).into_response()
}

async fn websocket_handler(
    ws: WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
    session: Session,
) -> impl IntoResponse {
    if let Some(user) = session.get::<User>("user").await.unwrap() {
        return ws.on_upgrade(move |socket| websocket(socket, state, user));
    }
    StatusCode::UNAUTHORIZED.into_response()
}

// WebSocket handler
async fn websocket(stream: WebSocket, state: Arc<AppState>, user: User) {
    let (mut sender, mut receiver) = stream.split();
    let mut rx = state.tx.subscribe();

    // Send welcome message
    let msg = format!("System: Welcome, {}!", user.login);
    let _ = sender.send(Message::Text(msg.into())).await;

    // Send task
    let mut send_task = tokio::spawn(async move {
        while let Ok(msg) = rx.recv().await {
            if sender.send(Message::Text(msg.into())).await.is_err() {
                break;
            }
        }
    });

    // Receive task
    let tx = state.tx.clone();
    let username = user.login.clone();
    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(Message::Text(text))) = receiver.next().await {
            // Prepend username
            let msg = format!("{}: {}", username, text);
            let _ = tx.send(msg);
        }
    });

    tokio::select! {
        _ = (&mut send_task) => recv_task.abort(),
        _ = (&mut recv_task) => send_task.abort(),
    };
}