-
-
+
+
+
+
+
+
+
+
Click on a contact to begin transmission...
+
+
+
+
+
diff --git a/src/main.rs b/src/main.rs
index 69fb1d7..fb9106f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,211 +1,243 @@
use axum::{
+ Router,
extract::{
- ws::{Message, WebSocket, WebSocketUpgrade},
Query, State,
+ ws::{Message, WebSocket, WebSocketUpgrade},
},
http::StatusCode,
- response::{Html, IntoResponse, Redirect},
+ response::{Html, IntoResponse},
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};
+use serde::{Deserialize, Serialize};
+use std::{collections::HashMap, sync::Arc};
+use tokio::sync::{RwLock, mpsc};
mod user;
use user::User;
-// Define application state
-struct AppState {
- tx: broadcast::Sender
,
- oauth_client: BasicClient,
+// WebSocket message types for client-server communication
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "type")]
+enum WsMessage {
+ #[serde(rename = "user_list")]
+ UserList { users: Vec },
+ #[serde(rename = "private_message")]
+ PrivateMessage {
+ from: String,
+ to: String,
+ content: String,
+ },
+ #[serde(rename = "user_joined")]
+ UserJoined { username: String },
+ #[serde(rename = "system")]
+ System { content: String },
}
-const AFTER_LOGIN_URL: &str = "/";
+// Per-user connection channel
+type UserTx = mpsc::UnboundedSender;
+
+// Shared state for tracking online users and their connections
+struct ChatState {
+ // Map of username -> sender channel
+ connections: RwLock>,
+}
+
+impl ChatState {
+ fn new() -> Self {
+ Self {
+ connections: RwLock::new(HashMap::new()),
+ }
+ }
+
+ async fn add_user(&self, username: String, tx: UserTx) {
+ let mut conns = self.connections.write().await;
+ conns.insert(username, tx);
+ }
+
+ async fn remove_user(&self, username: &str) {
+ let mut conns = self.connections.write().await;
+ conns.remove(username);
+ }
+
+ async fn get_online_users(&self) -> Vec {
+ let conns = self.connections.read().await;
+ conns.keys().cloned().collect()
+ }
+
+ async fn send_to_user(&self, username: &str, message: &str) -> bool {
+ let conns = self.connections.read().await;
+ if let Some(tx) = conns.get(username) {
+ tx.send(message.to_string()).is_ok()
+ } else {
+ false
+ }
+ }
+
+ async fn broadcast_except(&self, message: &str, exclude: &str) {
+ let conns = self.connections.read().await;
+ for (username, tx) in conns.iter() {
+ if username != exclude {
+ let _ = tx.send(message.to_string());
+ }
+ }
+ }
+}
+
+// Application state
+struct AppState {
+ chat: ChatState,
+}
#[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();
- let port = dotenvy::var("PORT").unwrap_or_else(|_| "3001".to_string());
- let host = dotenvy::var("HOST").unwrap_or_else(|_| "http://localhost:".to_string() + &port);
- // 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(host + "/auth/callback")
- .expect("Invalid redirect URL"),
- );
-
let app_state = Arc::new(AppState {
- tx,
- oauth_client: client,
+ chat: ChatState::new(),
});
- // 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 port = dotenvy::var("PORT").unwrap_or_else(|_| "3001".to_string());
- let listener = tokio::net::TcpListener::bind("0.0.0.0:".to_string() + &port).await.unwrap();
+ let listener = tokio::net::TcpListener::bind("0.0.0.0:".to_string() + &port)
+ .await
+ .unwrap();
println!("Chat server listening on http://0.0.0.0:{}", port);
axum::serve(listener, app).await.unwrap();
}
// Handlers
-async fn index(session: Session) -> impl IntoResponse {
- if let Some(_user) = session.get::("user").await.unwrap() {
- return Html(include_str!("../index.html")).into_response();
- }
- Redirect::to("/login").into_response()
+async fn index() -> impl IntoResponse {
+ Html(include_str!("../index.html"))
}
-async fn login(State(state): State>) -> impl IntoResponse {
- let (auth_url, _csrf_token) = state
- .oauth_client
- .authorize_url(oauth2::CsrfToken::new_random)
- .url();
-
- Redirect::to(auth_url.as_str())
+#[derive(Deserialize)]
+struct WsParams {
+ token: String,
}
-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,
- State(state): State>,
- 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()
+#[derive(Debug, Serialize, Deserialize)]
+struct Claims {
+ preferred_username: String,
+ // Add other fields as needed
}
async fn websocket_handler(
ws: WebSocketUpgrade,
+ Query(params): Query,
State(state): State>,
- session: Session,
) -> impl IntoResponse {
- if let Some(user) = session.get::("user").await.unwrap() {
- return ws.on_upgrade(move |socket| websocket(socket, state, user));
+ // For now, we'll verify the token by calling the userinfo endpoint.
+ // In a production app, you should verify the JWT signature locally using JWKS.
+ 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(¶ms.token)
+ .send()
+ .await;
+
+ match user_info_resp {
+ Ok(resp) if resp.status().is_success() => {
+ let user: User = match resp.json().await {
+ Ok(u) => u,
+ Err(_) => return StatusCode::INTERNAL_SERVER_ERROR.into_response(),
+ };
+ ws.on_upgrade(move |socket| handle_websocket(socket, state, user))
+ }
+ _ => StatusCode::UNAUTHORIZED.into_response(),
}
- StatusCode::UNAUTHORIZED.into_response()
}
-// WebSocket handler
-async fn websocket(stream: WebSocket, state: Arc, user: User) {
- let (mut sender, mut receiver) = stream.split();
- let mut rx = state.tx.subscribe();
+// WebSocket connection handler
+async fn handle_websocket(stream: WebSocket, state: Arc, user: User) {
+ let (mut ws_sender, mut ws_receiver) = stream.split();
+ let username = user.login.clone();
+
+ // Create channel for this user
+ let (tx, mut rx) = mpsc::unbounded_channel::();
+
+ // Add user to connections
+ state.chat.add_user(username.clone(), tx).await;
+ println!("User connected: {}", username);
+
+ // Send current user list to the newly connected user
+ let users = state.chat.get_online_users().await;
+ println!("Sending user list to {}: {:?}", username, users);
+ let user_list_msg = serde_json::to_string(&WsMessage::UserList { users }).unwrap();
+ let _ = ws_sender.send(Message::Text(user_list_msg.into())).await;
// Send welcome message
- let msg = format!("System: Welcome, {}!", user.login);
- let _ = sender.send(Message::Text(msg.into())).await;
+ let welcome = serde_json::to_string(&WsMessage::System {
+ content: format!("Welcome, {}!", username),
+ })
+ .unwrap();
+ let _ = ws_sender.send(Message::Text(welcome.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() {
+ // Task to forward messages from channel to WebSocket
+ let send_task = tokio::spawn(async move {
+ while let Some(msg) = rx.recv().await {
+ if ws_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);
+ // Broadcast user joined to everyone else (after send_task is spawned so others receive it)
+ let join_msg = serde_json::to_string(&WsMessage::UserJoined {
+ username: username.clone(),
+ })
+ .unwrap();
+ state.chat.broadcast_except(&join_msg, &username).await;
+
+ // Task to receive messages from WebSocket and route them
+ let state_clone = state.clone();
+ let username_clone = username.clone();
+ let recv_task = tokio::spawn(async move {
+ while let Some(Ok(Message::Text(text))) = ws_receiver.next().await {
+ // Parse incoming message
+ if let Ok(msg) = serde_json::from_str::(&text) {
+ // Create private message
+ let private_msg = serde_json::to_string(&WsMessage::PrivateMessage {
+ from: username_clone.clone(),
+ to: msg.to.clone(),
+ content: msg.content.clone(),
+ })
+ .unwrap();
+
+ // Send to recipient
+ state_clone.chat.send_to_user(&msg.to, &private_msg).await;
+
+ // Also send back to sender (for their own chat view)
+ state_clone
+ .chat
+ .send_to_user(&username_clone, &private_msg)
+ .await;
+ }
}
});
+ // Wait for either task to complete
tokio::select! {
- _ = (&mut send_task) => recv_task.abort(),
- _ = (&mut recv_task) => send_task.abort(),
- };
+ _ = send_task => {},
+ _ = recv_task => {},
+ }
+
+ // User disconnected - clean up
+ state.chat.remove_user(&username).await;
+}
+
+// Message structure from client
+#[derive(Debug, Deserialize)]
+struct ClientMessage {
+ to: String,
+ content: String,
}