maudit / crates / maudit-cli / src / dev / server.rs

use axum::{
    Router,
    body::{Body, to_bytes},
    extract::{
        Request, State,
        ws::{Message, WebSocket, WebSocketUpgrade},
    },
    http::{HeaderValue, StatusCode, Uri, header::CONTENT_LENGTH},
    middleware::{self, Next},
    response::{IntoResponse, Response},
    routing::get,
};
use quanta::Instant;
use serde_json::json;
use tokio::{
    net::TcpSocket,
    signal,
    sync::{RwLock, broadcast},
};
use tracing::{Level, debug};

use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use tower_http::{
    services::ServeDir,
    trace::{DefaultMakeSpan, TraceLayer},
};

use axum::extract::connect_info::ConnectInfo;
use futures::{SinkExt, stream::StreamExt};

use crate::consts::PORT;
use crate::server_utils::{CustomOnResponse, find_open_port, log_server_start};
use axum::http::header;
use local_ip_address::local_ip;
use tokio::fs;

#[derive(Clone, Debug)]
pub struct WebSocketMessage {
    pub data: String,
}

#[derive(Clone, Debug)]
pub enum StatusType {
    Success,
    Info,
    Error,
}

impl std::fmt::Display for StatusType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            StatusType::Success => write!(f, "success"),
            StatusType::Info => write!(f, "info"),
            StatusType::Error => write!(f, "error"),
        }
    }
}

// Persistent state for new connections
#[derive(Clone, Debug)]
pub struct PersistentStatus {
    pub status_type: StatusType, // Only Success or Error
    pub message: String,
}

#[derive(Clone)]
struct AppState {
    tx: broadcast::Sender<WebSocketMessage>,
    current_status: Arc<RwLock<Option<PersistentStatus>>>,
}

fn inject_live_reload_script(html_content: &str, socket_addr: SocketAddr, host: bool) -> String {
    let mut content = html_content.to_string();

    // Run cargo xtask build-cli-js to build the client.js file if missing
    let script_content = include_str!("../../js/dist/client.js").replace(
        "{SERVER_ADDRESS}",
        &format!(
            "{}:{}",
            if !host {
                socket_addr.ip().to_string()
            } else {
                local_ip().unwrap().to_string()
            },
            &socket_addr.port().to_string()
        ),
    );

    content.push_str(&format!("\n\n<script>{script_content}</script>"));
    content
}

pub async fn start_dev_web_server(
    start_time: Instant,
    tx: broadcast::Sender<WebSocketMessage>,
    host: bool,
    port: Option<u16>,
    initial_error: Option<String>,
    current_status: Arc<RwLock<Option<PersistentStatus>>>,
) {
    // TODO: The dist dir should be configurable
    let dist_dir = "dist";

    // Send initial error if present
    if let Some(error) = initial_error {
        let _ = tx.send(WebSocketMessage {
            data: json!({
                "type": StatusType::Error.to_string(),
                "message": error
            })
            .to_string(),
        });
    }

    async fn handle_404(socket_addr: SocketAddr, host: bool, dist_dir: &str) -> impl IntoResponse {
        let content = match fs::read_to_string(format!("{}/404.html", dist_dir)).await {
            Ok(custom_content) => custom_content,
            Err(_) => include_str!("./404.html").to_string(),
        };

        (
            StatusCode::NOT_FOUND,
            [(header::CONTENT_TYPE, "text/html; charset=utf-8")],
            inject_live_reload_script(&content, socket_addr, host),
        )
            .into_response()
    }

    // run it with hyper, if --host 0.0.0.0 otherwise localhost
    let addr = if host {
        IpAddr::from([0, 0, 0, 0])
    } else {
        IpAddr::from([127, 0, 0, 1])
    };

    // Use provided port or default to the constant PORT
    let starting_port = port.unwrap_or(PORT);

    let port = find_open_port(&addr, starting_port).await;
    let socket = TcpSocket::new_v4().unwrap();
    let _ = socket.set_reuseaddr(true);

    let socket_addr = SocketAddr::new(addr, port);
    socket.bind(socket_addr).unwrap();

    let listener = socket.listen(1024).unwrap();

    debug!("listening on {}", listener.local_addr().unwrap());

    let serve_dir =
        ServeDir::new(dist_dir).not_found_service(axum::routing::any(move || async move {
            handle_404(socket_addr, host, dist_dir).await
        }));

    // TODO: Return a `.well-known/appspecific/com.chrome.devtools.json` for Chrome

    let router = Router::new()
        .route("/ws", get(ws_handler))
        .fallback_service(serve_dir)
        .layer(middleware::from_fn(add_cache_headers))
        .layer(middleware::from_fn(move |req, next| {
            add_dev_client_script(req, next, socket_addr, host)
        }))
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(DefaultMakeSpan::default().include_headers(true)),
        )
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(DefaultMakeSpan::new().level(Level::INFO))
                .on_response(CustomOnResponse),
        )
        .with_state(AppState {
            tx: tx.clone(),
            current_status: current_status.clone(),
        });

    log_server_start(
        start_time,
        host,
        listener.local_addr().unwrap(),
        "Development",
    );

    axum::serve(
        listener,
        router.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .with_graceful_shutdown(shutdown_signal())
    .await
    .unwrap();
}

pub async fn update_status(
    tx: &broadcast::Sender<WebSocketMessage>,
    current_status: Arc<RwLock<Option<PersistentStatus>>>,
    status_type: StatusType,
    message: &str,
) {
    // Only store persistent states (Success clears errors, Error stores the error)
    let persistent_status = match status_type {
        StatusType::Success => None, // Clear any error state
        StatusType::Error => Some(PersistentStatus {
            status_type: StatusType::Error,
            message: message.to_string(),
        }),
        // Everything else just keeps the current state
        _ => {
            let status = current_status.read().await;
            status.clone() // Keep existing persistent state
        }
    };

    // Update the stored status
    {
        let mut status = current_status.write().await;
        *status = persistent_status;
    }

    // Send the message to all connected clients
    let _ = tx.send(WebSocketMessage {
        data: json!({
            "type": status_type.to_string(),
            "message": message
        })
        .to_string(),
    });
}

async fn add_dev_client_script(
    req: Request,
    next: Next,
    socket_addr: SocketAddr,
    host: bool,
) -> Response {
    let uri = req.uri().clone();
    let mut res: axum::http::Response<Body> = next.run(req).await;

    res.extensions_mut().insert(uri.clone());

    if res.headers().get(axum::http::header::CONTENT_TYPE)
        == Some(&HeaderValue::from_static("text/html"))
    {
        let original_headers = res.headers().clone();
        let body = res.into_body();
        let bytes = to_bytes(body, usize::MAX).await.unwrap();

        let body = String::from_utf8_lossy(&bytes).into_owned();

        let body_with_script = inject_live_reload_script(&body, socket_addr, host);
        let new_body_length = body_with_script.len();

        // Copy the headers from the original response
        let mut res = Response::new(body_with_script.into());
        *res.headers_mut() = original_headers;

        // Update Content-Length header to match new body size
        res.headers_mut().insert(
            CONTENT_LENGTH,
            HeaderValue::from_str(&new_body_length.to_string()).unwrap(),
        );

        res.extensions_mut().insert(uri);

        return res;
    }

    res
}

async fn add_cache_headers(req: Request, next: Next) -> Response {
    let uri = req.uri().clone();
    let mut res = next.run(req).await;

    if let Some(content_type) = res.headers().get(axum::http::header::CONTENT_TYPE) {
        let cache_header = cache_header_by_content(&uri, content_type);
        if let Some(cache_header) = cache_header {
            res.headers_mut()
                .insert(header::CACHE_CONTROL, cache_header);
        }
    }

    res
}

fn cache_header_by_content(uri: &Uri, content_type: &HeaderValue) -> Option<HeaderValue> {
    if content_type == HeaderValue::from_static("text/html") {
        // No cache for HTML files
        Some(HeaderValue::from_static(
            "no-cache, no-store, must-revalidate",
        ))
    }
    // If something comes from the assets path, assume that it's fingerprinted and can be cached for a long time
    // TODO: Same as dist, shouldn't be hardcoded
    else if uri.path().starts_with("/_maudit/") {
        Some(HeaderValue::from_static(
            "public, max-age=31536000, immutable",
        ))
    } else {
        // Don't try to cache anything else, the browser will decide based on the last-modified header
        None
    }
}

async fn ws_handler(
    ws: WebSocketUpgrade,
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    debug!("`{addr} connected.");
    // finalize the upgrade process by returning upgrade callback.
    // we can customize the callback by sending additional info such as address.
    ws.on_upgrade(move |socket| handle_socket(socket, addr, state.tx, state.current_status))
}

async fn handle_socket(
    socket: WebSocket,
    who: SocketAddr,
    tx: broadcast::Sender<WebSocketMessage>,
    current_status: Arc<RwLock<Option<PersistentStatus>>>,
) {
    let (mut sender, mut receiver) = socket.split();

    // Send current persistent status to new connection if there is one
    {
        let status = current_status.read().await;
        if let Some(persistent_status) = status.as_ref() {
            let _ = sender
                .send(Message::Text(
                    json!({
                        "type": persistent_status.status_type.to_string(),
                        "message": persistent_status.message
                    })
                    .to_string()
                    .into(),
                ))
                .await;
        }
    }

    let mut rx = tx.subscribe();

    tokio::select! {
        _ = async {
            while let Some(Ok(msg)) = receiver.next().await {
                match msg {
                    Message::Text(_) => {}
                    Message::Binary(_) => {
                    }
                    _ => {}
                }
            }
        } => {},
        _ = async {
            while let Ok(msg) = rx.recv().await {
                debug!(">>> got messages from higher level: {0}", msg.data);
                let _ = sender.send(Message::Text(msg.data.into())).await;
            }
        } => {},
    }

    // returning from the handler closes the websocket connection
    debug!("Websocket context {who} destroyed");
}

async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }
}