hyperhive/hive-agent/src/web_ui/screen.rs

//! VNC screen websocket relay + agent icon.

use axum::{
    extract::State,
    http::StatusCode,
    response::{IntoResponse, Response},
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

use super::AppState;

/// This agent's icon. Serves the operator-configured SVG from
/// `/etc/hyperhive/icon.svg` (set via the `services.hyperhive.agent.icon` agent.nix
/// option) when present, otherwise **404** — there is no bundled
/// server-side default any more (that was `hive_sh4re::assets::
/// branding_svg`, now removed). Consumers fall back client-side: the
/// dashboard's `agents.js` and this agent's own `app.js` both swap an
/// `/icon` load failure to the frontend-bundled `/favicon.svg` rather
/// than probing first, so a 404 here is the expected "unconfigured"
/// signal, not an error case to work around.
pub(super) async fn serve_icon() -> Response {
    match std::fs::read_to_string("/etc/hyperhive/icon.svg") {
        Ok(body) => ([("content-type", "image/svg+xml")], body).into_response(),
        Err(_) => (StatusCode::NOT_FOUND, "no icon configured").into_response(),
    }
}

/// WebSocket handler: upgrade then pump bytes between the WS client and
/// the VNC server on `127.0.0.1:<vnc_port>`. Returns 404 when gui is not
/// enabled for this agent.
pub(super) async fn screen_ws(
    ws: axum::extract::ws::WebSocketUpgrade,
    State(state): State<AppState>,
) -> Response {
    let Some(vnc_port) = state.gui_vnc_port else {
        return (StatusCode::NOT_FOUND, "gui not enabled for this agent").into_response();
    };
    ws.on_upgrade(move |socket| relay_ws_vnc(socket, vnc_port))
}

/// Pure byte pump: forwards raw bytes between the WebSocket client and
/// the VNC TCP stream. Transparent to any RFB variant (plain, `VeNCrypt`).
async fn relay_ws_vnc(socket: axum::extract::ws::WebSocket, vnc_port: u16) {
    // Import futures traits locally so they don't conflict with
    // tokio_stream::StreamExt used at module scope.
    use axum::extract::ws::Message;
    use futures_util::{SinkExt, StreamExt as _};

    let addr = format!("127.0.0.1:{vnc_port}");
    let Ok(tcp) = tokio::net::TcpStream::connect(&addr).await else {
        tracing::warn!(%addr, "screen/ws: could not connect to VNC server");
        return;
    };
    let (mut tcp_rx, mut tcp_tx) = tcp.into_split();
    let (mut ws_tx, mut ws_rx) = socket.split();

    // WS → TCP
    let ws_to_tcp = tokio::spawn(async move {
        while let Some(Ok(msg)) = futures_util::StreamExt::next(&mut ws_rx).await {
            match msg {
                Message::Binary(data) if tcp_tx.write_all(&data).await.is_err() => {
                    break;
                }
                Message::Close(_) => break,
                _ => {} // ping/pong/text: ignore
            }
        }
    });

    // TCP → WS
    let tcp_to_ws = tokio::spawn(async move {
        let mut buf = vec![0u8; 8192];
        loop {
            match tcp_rx.read(&mut buf).await {
                Ok(0) | Err(_) => break,
                Ok(n) => {
                    if ws_tx
                        .send(Message::Binary(buf[..n].to_vec().into()))
                        .await
                        .is_err()
                    {
                        break;
                    }
                }
            }
        }
    });

    // Wait for either direction to close, then let both tasks drop.
    tokio::select! {
        _ = ws_to_tcp => {}
        _ = tcp_to_ws => {}
    }
}