use ankurah_core::storage::StorageEngine;
use ankurah_proto as proto;
use anyhow::Result;
use axum::{
extract::{
ws::{WebSocket, WebSocketUpgrade},
State,
},
response::{IntoResponse, Response},
routing::get,
Router,
};
use axum_extra::{headers, TypedHeader};
use bincode::deserialize;
use futures_util::StreamExt;
use std::net::IpAddr;
use std::{future::Future, net::SocketAddr, ops::ControlFlow, pin::Pin};
use tower::ServiceBuilder;
use tower_http::trace::{DefaultMakeSpan, DefaultOnRequest, DefaultOnResponse, TraceLayer};
#[cfg(feature = "instrument")]
use tracing::instrument;
use tracing::{debug, error, info, warn, Level};
use ankurah_core::{node::Node, policy::PolicyAgent};
use crate::{client_ip::SmartClientIp, OptionalUserAgent};
use super::state::Connection;
pub struct WebsocketServer<SE, PA>
where
SE: StorageEngine + Send + Sync + 'static,
PA: PolicyAgent + Send + Sync + 'static,
{
node: Option<Node<SE, PA>>,
}
impl<SE, PA> WebsocketServer<SE, PA>
where
SE: StorageEngine + Send + Sync + 'static,
PA: PolicyAgent + Send + Sync + 'static,
{
pub fn new(node: Node<SE, PA>) -> Self { Self { node: Some(node) } }
pub fn route_handler(
&self,
) -> impl Clone + Send + 'static + Fn(WebSocketUpgrade, SmartClientIp, OptionalUserAgent) -> Pin<Box<dyn Future<Output = Response> + Send>>
{
let node = self.node.as_ref().expect("websocket server cannot produce a route after being run").clone();
move |ws: WebSocketUpgrade, SmartClientIp(client_ip): SmartClientIp, OptionalUserAgent(user_agent)| {
let node = node.clone();
Box::pin(async move { upgrade_connection(ws, client_ip, user_agent, node) })
}
}
pub async fn run(&mut self, bind_address: &str) -> Result<()> {
let Some(node) = self.node.take() else {
return Err(anyhow::anyhow!("Already been run"));
};
let app = Router::new().route("/ws", get(ws_handler)).with_state(node).layer(
ServiceBuilder::new()
.layer(
TraceLayer::new_for_http()
.make_span_with(DefaultMakeSpan::new().level(Level::INFO))
.on_request(DefaultOnRequest::new().level(Level::INFO))
.on_response(DefaultOnResponse::new().level(Level::INFO)),
)
.into_inner(),
);
let listener = tokio::net::TcpListener::bind(bind_address).await?;
info!("Websocket server listening on {}", listener.local_addr()?);
axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>()).await?;
Ok(())
}
}
async fn ws_handler<SE, PA>(
ws: WebSocketUpgrade,
SmartClientIp(client_ip): SmartClientIp,
user_agent: Option<TypedHeader<headers::UserAgent>>,
State(node): State<Node<SE, PA>>,
) -> impl IntoResponse
where
SE: StorageEngine + Send + Sync + 'static,
PA: PolicyAgent + Send + Sync + 'static,
{
let user_agent = user_agent.map(|TypedHeader(user_agent)| user_agent.to_string());
upgrade_connection(ws, client_ip, user_agent, node)
}
fn upgrade_connection<SE, PA>(ws: WebSocketUpgrade, client_ip: IpAddr, user_agent: Option<String>, node: Node<SE, PA>) -> Response
where
SE: StorageEngine + Send + Sync + 'static,
PA: PolicyAgent + Send + Sync + 'static,
{
debug!("Websocket server upgrading connection");
let user_agent = format_user_agent(user_agent);
debug!("`{user_agent}` at {client_ip} connected.");
ws.on_upgrade(move |socket| handle_websocket(socket, client_ip, node))
}
fn format_user_agent(user_agent: Option<String>) -> String { user_agent.unwrap_or_else(|| String::from("Unknown browser")) }
#[cfg_attr(feature = "instrument", instrument(level = "debug", skip_all, fields(client_ip = %client_ip)))]
async fn handle_websocket<SE, PA>(socket: WebSocket, client_ip: IpAddr, node: Node<SE, PA>)
where
SE: StorageEngine + Send + Sync + 'static,
PA: PolicyAgent + Send + Sync + 'static,
{
info!("Websocket server connected to {}", client_ip);
let (sender, mut receiver) = socket.split();
let mut conn = Connection::Initial(Some(sender));
if let Err(e) = conn
.send(proto::Message::Presence(proto::Presence { node_id: node.id, durable: node.durable, system_root: node.system.root() }))
.await
{
debug!("Error sending presence to {client_ip}: {:?}", e);
return;
}
while let Some(msg) = receiver.next().await {
if let Ok(msg) = msg {
if process_message(msg, client_ip, &mut conn, node.clone()).await.is_break() {
break;
}
} else {
debug!("client {client_ip} abruptly disconnected");
break;
}
}
if let Connection::Established(peer_sender) = conn {
use ankurah_core::connector::PeerSender;
node.deregister_peer(peer_sender.recipient_node_id());
}
debug!("Websocket context {client_ip} destroyed");
}
#[cfg_attr(feature = "instrument", instrument(level = "debug", skip_all, fields(client_ip = %client_ip)))]
async fn process_message<SE, PA>(
msg: axum::extract::ws::Message,
client_ip: IpAddr,
state: &mut Connection,
node: Node<SE, PA>,
) -> ControlFlow<(), ()>
where
SE: StorageEngine + Send + Sync + 'static,
PA: PolicyAgent + Send + Sync + 'static,
{
match msg {
axum::extract::ws::Message::Binary(d) => {
debug!(">>> {} sent {} bytes", client_ip, d.len());
if let Ok(message) = deserialize::<proto::Message>(&d) {
match message {
proto::Message::Presence(presence) => {
match state {
Connection::Initial(sender) => {
if let Some(sender) = sender.take() {
debug!("Received client presence from {}", client_ip);
use super::sender::WebSocketClientSender;
let sender = WebSocketClientSender::new(presence.node_id, sender);
node.register_peer(presence, Box::new(sender.clone()));
*state = Connection::Established(sender);
}
}
_ => warn!("Received presence from {} but already have a peer sender - ignoring", client_ip),
}
}
proto::Message::PeerMessage(msg) => {
if let Connection::Established(_) = state {
tokio::spawn(async move {
if let Err(e) = node.handle_message(msg).await {
error!("Error handling message from {}: {:?}", client_ip, e);
}
});
} else {
warn!("Received peer message from {} but not connected as a peer", client_ip);
}
}
}
} else {
error!("Failed to deserialize message from {}", client_ip);
}
}
axum::extract::ws::Message::Text(t) => {
debug!(">>> {client_ip} sent str: {t:?}");
}
axum::extract::ws::Message::Close(c) => {
if let Some(cf) = c {
debug!(">>> {} sent close with code {} and reason `{}`", client_ip, cf.code, cf.reason);
} else {
debug!(">>> {client_ip} somehow sent close message without CloseFrame");
}
return ControlFlow::Break(());
}
axum::extract::ws::Message::Pong(v) => {
debug!(">>> {client_ip} sent pong with {v:?}");
}
axum::extract::ws::Message::Ping(v) => {
debug!(">>> {client_ip} sent ping with {v:?}");
}
}
ControlFlow::Continue(())
}