use ninox_core::events::{Engine, Event};
use axum::{
extract::{
ws::{Message, WebSocket},
Path, State, WebSocketUpgrade,
},
response::IntoResponse,
routing::get,
Router,
};
use futures::{SinkExt, StreamExt};
use std::sync::Arc;
pub fn terminal_router(engine: Arc<Engine>) -> Router {
Router::new()
.route("/:id/terminal", get(ws_handler))
.with_state(engine)
}
async fn ws_handler(
ws: WebSocketUpgrade,
Path(session_id): Path<String>,
State(engine): State<Arc<Engine>>,
) -> impl IntoResponse {
ws.on_upgrade(move |socket| handle(socket, session_id, engine))
}
async fn handle(socket: WebSocket, session_id: String, engine: Arc<Engine>) {
let (mut tx, mut rx) = socket.split();
let e2 = engine.clone();
let sid = session_id.clone();
tokio::spawn(async move {
while let Some(Ok(Message::Binary(b))) = rx.next().await {
if let Some(w) = e2.get_pty_writer(&sid).await {
let _ = w.send(b.to_vec());
}
}
});
let mut event_rx = engine.subscribe();
loop {
match event_rx.recv().await {
Ok(Event::TerminalOutput { session_id: sid, bytes }) if sid == session_id => {
if tx.send(Message::Binary(bytes)).await.is_err() {
break;
}
}
Err(_) => break,
_ => {}
}
}
}