redesign

src/main.rs

@@ -1,80 +1,206 @@
 use axum::{
     extract::{
         ws::{Message, WebSocket, WebSocketUpgrade},
-        State,
+        Query, State,
     },
-    response::{Html, IntoResponse},
+    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};
 
-// Define application state to hold the broadcast channel
+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() {
-    // Create a broadcast channel with a capacity of 100 messages
+    dotenv().ok();
+
+    // Create a broadcast channel
     let (tx, _rx) = broadcast::channel(100);
-    let app_state = Arc::new(AppState { tx });
+
+    // 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 using hyper, listening globally on port 3000
+    // 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();
 }
 
-// Handler to serve the HTML file
-async fn index() -> Html<&'static str> {
-    Html(include_str!("../index.html"))
+// 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()
 }
 
-// Handler to upgrade HTTP connection to WebSocket
 async fn websocket_handler(
     ws: WebSocketUpgrade,
     State(state): State<Arc<AppState>>,
+    session: Session,
 ) -> impl IntoResponse {
-    ws.on_upgrade(|socket| websocket(socket, state))
+    if let Some(user) = session.get::<User>("user").await.unwrap() {
+        return ws.on_upgrade(move |socket| websocket(socket, state, user));
+    }
+    StatusCode::UNAUTHORIZED.into_response()
 }
 
-// The actual WebSocket connection handler
-async fn websocket(stream: WebSocket, state: Arc<AppState>) {
-    // Split the stream into sender and receiver
+// WebSocket handler
+async fn websocket(stream: WebSocket, state: Arc<AppState>, user: User) {
     let (mut sender, mut receiver) = stream.split();
-
-    // Subscribe to the broadcast channel
     let mut rx = state.tx.subscribe();
 
-    // Spawn a task to send messages from the broadcast channel to this client
+    // 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 {
-            // In a real app, you might want to handle errors better
             if sender.send(Message::Text(msg.into())).await.is_err() {
                 break;
             }
         }
     });
 
-    // Spawn a task to receive messages from this client and broadcast them
+    // 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 {
-            // Broadcast the message to all subscribers
-            let _ = tx.send(text.to_string());
+            // Prepend username
+            let msg = format!("{}: {}", username, text);
+            let _ = tx.send(msg);
         }
     });
 
-    // If either task finishes (connection closed or error), abort the other
     tokio::select! {
         _ = (&mut send_task) => recv_task.abort(),
         _ = (&mut recv_task) => send_task.abort(),