individual messages
This commit is contained in:
parent
b098214a9d
commit
7fdf18ff1b
4 changed files with 617 additions and 312 deletions
328
src/main.rs
328
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<String>,
|
||||
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<String> },
|
||||
#[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<String>;
|
||||
|
||||
// Shared state for tracking online users and their connections
|
||||
struct ChatState {
|
||||
// Map of username -> sender channel
|
||||
connections: RwLock<HashMap<String, UserTx>>,
|
||||
}
|
||||
|
||||
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<String> {
|
||||
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>("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<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())
|
||||
#[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<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()
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct Claims {
|
||||
preferred_username: String,
|
||||
// Add other fields as needed
|
||||
}
|
||||
|
||||
async fn websocket_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
Query(params): Query<WsParams>,
|
||||
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));
|
||||
// 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<AppState>, 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<AppState>, 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::<String>();
|
||||
|
||||
// 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::<ClientMessage>(&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,
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue