+
+
+
+
+
+
+
diff --git a/src/api.rs b/src/api.rs
index 88048cd..4e5a82f 100644
--- a/src/api.rs
+++ b/src/api.rs
@@ -1,102 +1,115 @@
-use futures_util::{SinkExt, StreamExt, TryFutureExt};
-use std::{
- collections::HashMap,
- sync::{
- atomic::{AtomicUsize, Ordering},
- Arc,
+use axum::{
+ extract::{
+ ws::{Message, WebSocket, WebSocketUpgrade},
+ State,
},
+ response::IntoResponse,
};
-use tokio::sync::{mpsc, RwLock};
-use tokio_stream::wrappers::UnboundedReceiverStream;
-use warp::ws::{Message, WebSocket};
+use futures::{sink::SinkExt, stream::StreamExt};
+use std::{
+ collections::HashSet,
+ sync::{Arc, Mutex},
+};
+use tokio::sync::broadcast;
-/// Our global unique user id counter.
-static NEXT_USER_ID: AtomicUsize = AtomicUsize::new(1);
+// Our shared state
+pub struct AppState {
+ // We require unique usernames. This tracks which usernames have been taken.
+ pub user_set: Mutex>,
+ // Channel used to send messages to all connected clients.
+ pub tx: broadcast::Sender,
+}
-/// Our state of currently connected users.
-///
-/// - Key is their id
-/// - Value is a sender of `warp::ws::Message`
-pub type Users = Arc>>>;
+pub async fn websocket_handler(
+ ws: WebSocketUpgrade,
+ State(state): State>,
+) -> impl IntoResponse {
+ ws.on_upgrade(|socket| websocket(socket, state))
+}
-pub async fn on_user_connected(ws: WebSocket, users: Users) {
- // Use a counter to assign a new unique ID for this user.
- let my_id = NEXT_USER_ID.fetch_add(1, Ordering::Relaxed);
+// This function deals with a single websocket connection, i.e., a single
+// connected client / user, for which we will spawn two independent tasks (for
+// receiving / sending chat messages).
+pub async fn websocket(stream: WebSocket, state: Arc) {
+ // By splitting, we can send and receive at the same time.
+ let (mut sender, mut receiver) = stream.split();
- eprintln!("User {} connected!", my_id);
+ // Username gets set in the receive loop, if it's valid.
+ let mut username = String::new();
+ // Loop until a text message is found.
+ while let Some(Ok(message)) = receiver.next().await {
+ if let Message::Text(name) = message {
+ // If username that is sent by client is not taken, fill username string.
+ check_username(&state, &mut username, &name);
- // Split the socket into a sender and receive of messages.
- let (mut user_ws_tx, mut user_ws_rx) = ws.split();
+ // If not empty we want to quit the loop else we want to quit function.
+ if !username.is_empty() {
+ break;
+ } else {
+ // Only send our client that username is taken.
+ let _ = sender
+ .send(Message::Text(String::from("Username already taken.")))
+ .await;
- // Use an unbounded channel to handle buffering and flushing of messages
- // to the websocket...
- let (tx, rx) = mpsc::unbounded_channel();
- let mut rx = UnboundedReceiverStream::new(rx);
+ return;
+ }
+ }
+ }
- let _ = user_ws_tx
- .send(Message::text(format!("Server: Welcome User {}", my_id)))
- .await;
+ // We subscribe *before* sending the "joined" message, so that we will also
+ // display it to our client.
+ let mut rx = state.tx.subscribe();
- tokio::task::spawn(async move {
- while let Some(message) = rx.next().await {
- user_ws_tx
- .send(message)
- .unwrap_or_else(|e| {
- eprintln!("websocket send error: {}", e);
- })
- .await;
+ // Now send the "joined" message to all subscribers.
+ let msg = format!("{username} joined.");
+ tracing::debug!("{msg}");
+ let _ = state.tx.send(msg);
+
+ // Spawn the first task that will receive broadcast messages and send text
+ // messages over the websocket to our client.
+ let mut send_task = tokio::spawn(async move {
+ while let Ok(msg) = rx.recv().await {
+ // In any websocket error, break loop.
+ if sender.send(Message::Text(msg)).await.is_err() {
+ break;
+ }
}
});
- // Save the sender in our list of connected users.
- users.write().await.insert(my_id, tx);
+ // Clone things we want to pass (move) to the receiving task.
+ let tx = state.tx.clone();
+ let name = username.clone();
- // Return a `Future` that is basically a state machine managing
- // this specific user's connection.
+ // Spawn a task that takes messages from the websocket, prepends the user
+ // name, and sends them to all broadcast subscribers.
+ let mut recv_task = tokio::spawn(async move {
+ while let Some(Ok(Message::Text(text))) = receiver.next().await {
+ // Add username before message.
+ let _ = tx.send(format!("{name}: {text}"));
+ }
+ });
- // Every time the user sends a message, broadcast it to
- // all other users...
- while let Some(result) = user_ws_rx.next().await {
- let msg = match result {
- Ok(msg) => msg,
- Err(e) => {
- eprintln!("websocket error(uid={}): {}", my_id, e);
- break;
- }
- };
- handle_incoming_message(my_id, msg, &users).await;
- }
-
- // user_ws_rx stream will keep processing as long as the user stays
- // connected. Once they disconnect, then...
- on_user_disconnected(my_id, &users).await;
-}
-
-pub async fn handle_incoming_message(my_id: usize, msg: Message, users: &Users) {
- // Skip any non-Text messages...
- let msg = if let Ok(s) = msg.to_str() {
- s
- } else {
- return;
+ // If any one of the tasks run to completion, we abort the other.
+ tokio::select! {
+ _ = (&mut send_task) => recv_task.abort(),
+ _ = (&mut recv_task) => send_task.abort(),
};
- let new_msg = format!("User {}: {}", my_id, msg);
+ // Send "user left" message (similar to "joined" above).
+ let msg = format!("{username} left.");
+ tracing::debug!("{msg}");
+ let _ = state.tx.send(msg);
- // New message from this user, send it to everyone else (except same uid)...
- for (&uid, tx) in users.read().await.iter() {
- if my_id != uid {
- if let Err(_disconnected) = tx.send(Message::text(new_msg.clone())) {
- // The tx is disconnected, our `on_user_disconnected` code
- // should be happening in another task, nothing more to
- // do here.
- }
- }
+ // Remove username from map so new clients can take it again.
+ state.user_set.lock().unwrap().remove(&username);
+}
+
+pub fn check_username(state: &AppState, string: &mut String, name: &str) {
+ let mut user_set = state.user_set.lock().unwrap();
+
+ if !user_set.contains(name) {
+ user_set.insert(name.to_owned());
+
+ string.push_str(name);
}
}
-
-pub async fn on_user_disconnected(my_id: usize, users: &Users) {
- eprintln!("User {} left.", my_id);
-
- // Stream closed up, so remove from the user list
- users.write().await.remove(&my_id);
-}
diff --git a/src/main.rs b/src/main.rs
index 84d2202..8b480b2 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,15 @@
-use std::fs::{read, read_to_string};
+use axum::{
+ response::{Html},
+ routing::get,
+ Router,
+};
+use std::{
+ collections::HashSet,
+ sync::{Arc, Mutex},
+};
use std::{error::Error, fs, result::Result};
-use warp::Filter;
+use tokio::sync::broadcast;
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
#[allow(non_snake_case)]
pub mod CAHd_game;
@@ -67,34 +76,37 @@ fn test() -> Result<(), Box> {
}
#[tokio::main]
-async fn main() -> Result<(), Box> {
- pretty_env_logger::init();
- test()?;
+async fn main() {
+ tracing_subscriber::registry()
+ .with(
+ tracing_subscriber::EnvFilter::try_from_default_env()
+ .unwrap_or_else(|_| "example_chat=trace".into()),
+ )
+ .with(tracing_subscriber::fmt::layer())
+ .init();
- // Keep track of all connected users, key is usize, value
- // is a websocket sender.
- let users = api::Users::default();
- // Turn our "state" into a new Filter...
- let users = warp::any().map(move || users.clone());
+ let _ = test();
- // GET /chat -> websocket upgrade
- let chat = warp::path("chat")
- // The `ws()` filter will prepare Websocket handshake...
- .and(warp::ws())
- .and(users)
- .map(|ws: warp::ws::Ws, users| {
- // This will call our function if the handshake succeeds.
- ws.on_upgrade(move |socket| on_user_connected(socket, users))
- });
+ // Set up application state for use with with_state().
+ let user_set = Mutex::new(HashSet::new());
+ let (tx, _rx) = broadcast::channel(100);
- // GET / -> index html
- let index = warp::get()
- .and(warp::path::end())
- .and(warp::fs::file("./test_client.html"));
+ let app_state = Arc::new(AppState { user_set, tx });
- let routes = index.or(chat);
+ let app = Router::new()
+ .route("/", get(index))
+ .route("/websocket", get(websocket_handler))
+ .with_state(app_state);
- warp::serve(routes).run(([0, 0, 0, 0], 3030)).await;
-
- Ok(())
+ let listener = tokio::net::TcpListener::bind("127.0.0.1:3000")
+ .await
+ .unwrap();
+ tracing::debug!("listening on {}", listener.local_addr().unwrap());
+ axum::serve(listener, app).await.unwrap();
+}
+
+
+// Include utf-8 file at **compile** time.
+async fn index() -> Html<&'static str> {
+ Html(std::include_str!("../chat.html"))
}
diff --git a/test_client.html b/test_client.html
deleted file mode 100644
index 53508c8..0000000
--- a/test_client.html
+++ /dev/null
@@ -1,51 +0,0 @@
-
-
-
- Cards For Humanity Test Client
-
-
-