126 lines
3.6 KiB
Rust
126 lines
3.6 KiB
Rust
use std::{net::SocketAddr, process::exit, sync::Arc, time::Duration};
|
|
|
|
use axum::{
|
|
extract::{ws::Message, ConnectInfo, Query, State, WebSocketUpgrade},
|
|
http::{self, HeaderMap, HeaderValue},
|
|
response::{IntoResponse, Response},
|
|
routing::get,
|
|
Router,
|
|
};
|
|
use axum_extra::TypedHeader;
|
|
use axum_streams::StreamBodyAs;
|
|
use futures::{never::Never, stream::repeat_with, Stream};
|
|
use rust_embed::RustEmbed;
|
|
use serde::Deserialize;
|
|
use tokio::{net::TcpListener, time::interval};
|
|
use tower_http::trace::{DefaultMakeSpan, TraceLayer};
|
|
|
|
use crate::{
|
|
config::{WEB_HOST, WEB_UPDATE_INTERVAL},
|
|
grid,
|
|
stream::Multipart,
|
|
AsyncResult, CLIENTS, COUNTER,
|
|
};
|
|
|
|
#[derive(RustEmbed, Clone)]
|
|
#[folder = "assets/"]
|
|
struct Assets;
|
|
|
|
#[derive(Clone)]
|
|
pub struct WebApiContext {
|
|
pub grids: Arc<[grid::Flut<u32>]>,
|
|
}
|
|
|
|
pub async fn serve(ctx: WebApiContext) -> AsyncResult<Never> {
|
|
let assets = axum_embed::ServeEmbed::<Assets>::with_parameters(
|
|
Some("404.html".to_string()),
|
|
axum_embed::FallbackBehavior::NotFound,
|
|
Some("index.html".to_string()),
|
|
);
|
|
let app = Router::new()
|
|
.route("/imgstream", get(image_stream))
|
|
.route("/stats", get(stats_stream))
|
|
.fallback_service(assets)
|
|
.with_state(ctx)
|
|
// logging middleware
|
|
.layer(
|
|
TraceLayer::new_for_http()
|
|
.make_span_with(DefaultMakeSpan::default().include_headers(true)),
|
|
);
|
|
|
|
// run it with hyper
|
|
let Ok(listener) = TcpListener::bind(WEB_HOST).await else {
|
|
tracing::error!(
|
|
"Was unable to bind to {WEB_HOST}, please check if a different process is bound"
|
|
);
|
|
exit(1);
|
|
};
|
|
|
|
tracing::debug!("listening on {}", listener.local_addr()?);
|
|
|
|
axum::serve(
|
|
listener,
|
|
app.into_make_service_with_connect_info::<SocketAddr>(),
|
|
)
|
|
.await?;
|
|
|
|
Err("Web api exited".into())
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct CanvasQuery {
|
|
canvas: u8,
|
|
}
|
|
|
|
fn make_image_stream(
|
|
ctx: WebApiContext,
|
|
canvas: u8,
|
|
) -> impl Stream<Item = Result<Vec<u8>, axum::Error>> {
|
|
use tokio_stream::StreamExt;
|
|
let mut buf = Vec::new();
|
|
repeat_with(move || {
|
|
buf.clear();
|
|
buf.extend_from_slice(&ctx.grids[canvas as usize].read_jpg_buffer());
|
|
Ok(buf.clone())
|
|
})
|
|
.throttle(WEB_UPDATE_INTERVAL)
|
|
}
|
|
|
|
fn make_stats() -> Message {
|
|
let pixels: u64 = COUNTER.load(std::sync::atomic::Ordering::Relaxed);
|
|
let clients: u64 = CLIENTS.load(std::sync::atomic::Ordering::Relaxed);
|
|
format!("{{\"c\":{clients}, \"p\":{pixels}}}").into()
|
|
}
|
|
|
|
async fn stats_stream(ws: WebSocketUpgrade) -> Response {
|
|
ws.on_upgrade(|mut c| async move {
|
|
let mut interval = interval(Duration::from_millis(100));
|
|
loop {
|
|
interval.tick().await;
|
|
if let Err(e) = c.send(make_stats()).await {
|
|
tracing::warn!("websocket disconnected with {e:?}")
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
async fn image_stream(
|
|
user_agent: Option<TypedHeader<headers::UserAgent>>,
|
|
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
|
State(ctx): State<WebApiContext>,
|
|
Query(CanvasQuery { canvas }): Query<CanvasQuery>,
|
|
) -> impl IntoResponse {
|
|
let user_agent = if let Some(TypedHeader(user_agent)) = user_agent {
|
|
user_agent.to_string()
|
|
} else {
|
|
String::from("Unknown browser")
|
|
};
|
|
tracing::info!("`{user_agent}` at {addr} connected.");
|
|
let mut headers = HeaderMap::new();
|
|
headers.insert(
|
|
http::header::CONTENT_TYPE,
|
|
HeaderValue::from_static("image/jpeg"),
|
|
);
|
|
|
|
StreamBodyAs::new(Multipart::new(10, headers), make_image_stream(ctx, canvas))
|
|
}
|