96 lines
3.5 KiB
Rust
96 lines
3.5 KiB
Rust
//! 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 `hyperhive.icon` agent.nix
|
|
/// option) when present, otherwise the bundled default hyperhive logo.
|
|
/// Always returns an image, so consumers (dashboard, favicon) can hit
|
|
/// `/icon` unconditionally without probing whether one is configured.
|
|
pub(super) async fn serve_icon() -> impl IntoResponse {
|
|
// Per-agent icon overrides go through `/etc/hyperhive/icon.svg`
|
|
// (set via the `hyperhive.icon` agent.nix option); the bundled
|
|
// default is resolved at runtime from
|
|
// `$HIVE_ASSETS_DIR/branding/hyperhive.svg`. If neither file can
|
|
// be read we serve an empty body — keeps the response a valid SVG
|
|
// content-type without a panic on a misconfigured container.
|
|
let body = std::fs::read_to_string("/etc/hyperhive/icon.svg").unwrap_or_else(|_| {
|
|
std::fs::read_to_string(hive_sh4re::assets::branding_svg()).unwrap_or_default()
|
|
});
|
|
([("content-type", "image/svg+xml")], body)
|
|
}
|
|
|
|
/// 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 => {}
|
|
}
|
|
}
|