lockwatch/lockwatch/src/main.rs
finga afa1d19dd8 Exchange data via websockets
Use websockets to send binary messages to the server. The server logs
the time taken by the client, as well as the time calculated with the
received messages sent by the client.
2022-10-02 21:14:32 +02:00

123 lines
3.8 KiB
Rust

use anyhow::Result;
use axum::{
extract::{
ws::{Message, WebSocket},
WebSocketUpgrade,
},
response::IntoResponse,
routing::get,
Router, TypedHeader,
};
use axum_extra::routing::SpaRouter;
use clap::{ArgAction, Parser};
use headers::UserAgent;
use log::{debug, error, info, trace, warn};
use std::{
env,
net::{IpAddr, SocketAddr},
time::Instant,
};
use tower::ServiceBuilder;
use tower_http::trace::TraceLayer;
#[derive(Parser)]
#[clap(about, author, version)]
struct Cli {
/// The address to listen on
#[clap(short, long, default_value = "::1")]
address: IpAddr,
/// The port to listen on
#[clap(short, long, default_value = "3000")]
port: u16,
/// Path to the directory containing the statically served files
#[clap(short, long, default_value = "dist")]
static_dir: String,
/// Set the log level. Multiple -v options increase the verbosity
#[clap(short, action = ArgAction::Count)]
verbosity: u8,
}
async fn handle_websocket(mut socket: WebSocket) {
let mut start: Option<Instant> = None;
loop {
if let Some(msg) = socket.recv().await {
if let Ok(msg) = msg {
match msg {
Message::Binary(b) => {
if let Ok(message) = bincode::deserialize::<client::Message>(&b) {
trace!("received message on websocket: {:?}", message);
match message {
client::Message::Start => start = Some(Instant::now()),
client::Message::Stop => {
if let Some(start) = start {
info!("server measured {:?}", start.elapsed());
} else {
error!("received stop message without running clock");
}
}
client::Message::Duration(d) => info!("client measured {:?}", d),
}
} else {
error!("received invalid message: {:?}", b);
}
}
Message::Ping(p) => {
debug!("websocket ping: {:?}", p);
}
Message::Pong(p) => {
debug!("websocket pong: {:?}", p);
}
Message::Close(m) => {
info!("client disconnected: {:?}", m);
return;
}
Message::Text(_) => warn!("websocket text messages are not implemented"),
}
} else {
info!("client disconnected");
return;
}
}
}
}
#[allow(clippy::unused_async)]
async fn websocket_handler(
ws: WebSocketUpgrade,
user_agent: Option<TypedHeader<UserAgent>>,
) -> impl IntoResponse {
if let Some(TypedHeader(user_agent)) = user_agent {
info!("`{}` connected", user_agent.as_str());
}
ws.on_upgrade(handle_websocket)
}
#[tokio::main]
async fn main() -> Result<()> {
let cli = Cli::parse();
if env::var("RUST_LOG").is_err() {
env::set_var("RUST_LOG", format!("{},hyper=info,mio=info", cli.verbosity));
}
tracing_subscriber::fmt::init();
let app = Router::new()
.route("/api/ws", get(websocket_handler))
.merge(SpaRouter::new("/client", cli.static_dir))
.layer(ServiceBuilder::new().layer(TraceLayer::new_for_http()));
let addr = SocketAddr::from((cli.address, cli.port));
info!("listening on {}", addr);
axum::Server::bind(&addr)
.serve(app.into_make_service())
.await?;
Ok(())
}