238 lines
7.4 KiB
Rust
238 lines
7.4 KiB
Rust
//! End-to-end tests for the per-session attach WS endpoint:
|
|
//! - bad / good JWT
|
|
//! - PTY bytes round-trip through real tmux
|
|
//! - multi-attach: two clients see the same output
|
|
|
|
mod common;
|
|
|
|
use std::collections::HashMap;
|
|
use std::net::{Ipv4Addr, SocketAddr};
|
|
use std::time::Duration;
|
|
|
|
use chrono::{Duration as ChronoDuration, Utc};
|
|
use common::DaemonFixture;
|
|
use futures_util::{SinkExt, StreamExt};
|
|
use maestro_daemon::{attach, issue_jwt, JwtVerifier, SessionManager};
|
|
use maestro_protocol::{CreateSessionRequest, JwtClaims};
|
|
use tempfile::TempDir;
|
|
use tokio::net::TcpListener;
|
|
use tokio::sync::oneshot;
|
|
use tokio_tungstenite::tungstenite::Message as WsMessage;
|
|
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
|
|
use uuid::Uuid;
|
|
|
|
const TOKEN: &str = "shared-secret-for-tests";
|
|
|
|
type WsClient = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;
|
|
|
|
struct Harness {
|
|
addr: SocketAddr,
|
|
manager: SessionManager,
|
|
_shutdown: oneshot::Sender<()>,
|
|
_fx: DaemonFixture,
|
|
}
|
|
|
|
async fn spawn_daemon_ws(daemon_id: &str) -> Harness {
|
|
let fx = DaemonFixture::new(daemon_id);
|
|
let manager = fx.manager.clone();
|
|
let verifier = JwtVerifier::new(TOKEN, daemon_id);
|
|
let router = attach::router(manager.clone(), verifier);
|
|
let listener = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
|
|
.await
|
|
.unwrap();
|
|
let addr = listener.local_addr().unwrap();
|
|
let (tx, rx) = oneshot::channel::<()>();
|
|
tokio::spawn(async move {
|
|
let _ = axum::serve(listener, router)
|
|
.with_graceful_shutdown(async move {
|
|
let _ = rx.await;
|
|
})
|
|
.await;
|
|
});
|
|
Harness {
|
|
addr,
|
|
manager,
|
|
_shutdown: tx,
|
|
_fx: fx,
|
|
}
|
|
}
|
|
|
|
fn token_for(session_id: Uuid, daemon_id: &str, valid_seconds: i64) -> String {
|
|
issue_jwt(
|
|
TOKEN,
|
|
&JwtClaims {
|
|
session_id,
|
|
daemon_id: daemon_id.into(),
|
|
exp: (Utc::now() + ChronoDuration::seconds(valid_seconds)).timestamp(),
|
|
iat: Some(Utc::now().timestamp()),
|
|
},
|
|
)
|
|
.unwrap()
|
|
}
|
|
|
|
async fn connect_attach(
|
|
addr: SocketAddr,
|
|
session: Uuid,
|
|
token: &str,
|
|
) -> Result<WsClient, tokio_tungstenite::tungstenite::Error> {
|
|
let url = format!("ws://{addr}/sessions/{session}/stream?token={token}");
|
|
connect_async(url).await.map(|(ws, _)| ws)
|
|
}
|
|
|
|
async fn create_bash_session(h: &Harness) -> Uuid {
|
|
let cwd = TempDir::new().unwrap();
|
|
let cwd_path = cwd.path().to_path_buf();
|
|
let s = h
|
|
.manager
|
|
.create(CreateSessionRequest {
|
|
daemon_id: "laptop".into(),
|
|
cwd: cwd_path,
|
|
command: "bash --noprofile --norc".into(),
|
|
env: HashMap::new(),
|
|
metadata: HashMap::new(),
|
|
worktree: None,
|
|
})
|
|
.await
|
|
.unwrap();
|
|
std::mem::forget(cwd); // keep dir for the test's duration
|
|
s.id
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn rejects_bad_token() {
|
|
let h = spawn_daemon_ws("laptop").await;
|
|
let id = create_bash_session(&h).await;
|
|
let err = connect_attach(h.addr, id, "not-a-jwt").await.unwrap_err();
|
|
let msg = format!("{err}");
|
|
assert!(msg.contains("401") || msg.contains("403"), "{msg}");
|
|
let _ = h.manager.kill(id).await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn rejects_token_for_different_session() {
|
|
let h = spawn_daemon_ws("laptop").await;
|
|
let real_id = create_bash_session(&h).await;
|
|
let other_id = Uuid::new_v4();
|
|
let token = token_for(other_id, "laptop", 300);
|
|
let err = connect_attach(h.addr, real_id, &token).await.unwrap_err();
|
|
assert!(format!("{err}").contains("403"));
|
|
let _ = h.manager.kill(real_id).await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn rejects_token_for_different_daemon() {
|
|
let h = spawn_daemon_ws("laptop").await;
|
|
let id = create_bash_session(&h).await;
|
|
let token = token_for(id, "dev-vm", 300);
|
|
let err = connect_attach(h.addr, id, &token).await.unwrap_err();
|
|
assert!(format!("{err}").contains("401"));
|
|
let _ = h.manager.kill(id).await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn pty_round_trip_through_real_tmux() {
|
|
let h = spawn_daemon_ws("laptop").await;
|
|
let id = create_bash_session(&h).await;
|
|
let token = token_for(id, "laptop", 300);
|
|
let mut ws = connect_attach(h.addr, id, &token).await.unwrap();
|
|
|
|
ws.send(WsMessage::Binary(b"echo MAESTRO_RT_OK\n".to_vec()))
|
|
.await
|
|
.unwrap();
|
|
|
|
let want = b"MAESTRO_RT_OK";
|
|
let mut got: Vec<u8> = Vec::new();
|
|
let deadline = std::time::Instant::now() + Duration::from_secs(5);
|
|
while std::time::Instant::now() < deadline {
|
|
if let Ok(Some(Ok(WsMessage::Binary(bytes)))) =
|
|
tokio::time::timeout(Duration::from_millis(50), ws.next()).await
|
|
{
|
|
got.extend_from_slice(&bytes);
|
|
if got.windows(want.len()).any(|w| w == want) {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
assert!(
|
|
got.windows(want.len()).any(|w| w == want),
|
|
"marker not seen in {} bytes",
|
|
got.len()
|
|
);
|
|
let _ = h.manager.kill(id).await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn multi_attach_broadcasts_to_every_client() {
|
|
let h = spawn_daemon_ws("laptop").await;
|
|
let id = create_bash_session(&h).await;
|
|
let token = token_for(id, "laptop", 300);
|
|
|
|
let mut ws_a = connect_attach(h.addr, id, &token).await.unwrap();
|
|
let mut ws_b = connect_attach(h.addr, id, &token).await.unwrap();
|
|
|
|
// give the broadcast subscribers a moment to register
|
|
tokio::time::sleep(Duration::from_millis(100)).await;
|
|
|
|
ws_a.send(WsMessage::Binary(b"echo MAESTRO_MULTI_OK\n".to_vec()))
|
|
.await
|
|
.unwrap();
|
|
|
|
let want = b"MAESTRO_MULTI_OK";
|
|
let mut got_a: Vec<u8> = Vec::new();
|
|
let mut got_b: Vec<u8> = Vec::new();
|
|
let deadline = std::time::Instant::now() + Duration::from_secs(5);
|
|
while std::time::Instant::now() < deadline {
|
|
if let Ok(Some(Ok(WsMessage::Binary(bytes)))) =
|
|
tokio::time::timeout(Duration::from_millis(50), ws_a.next()).await
|
|
{
|
|
got_a.extend_from_slice(&bytes);
|
|
}
|
|
if let Ok(Some(Ok(WsMessage::Binary(bytes)))) =
|
|
tokio::time::timeout(Duration::from_millis(50), ws_b.next()).await
|
|
{
|
|
got_b.extend_from_slice(&bytes);
|
|
}
|
|
if got_a.windows(want.len()).any(|w| w == want)
|
|
&& got_b.windows(want.len()).any(|w| w == want)
|
|
{
|
|
break;
|
|
}
|
|
}
|
|
assert!(
|
|
got_a.windows(want.len()).any(|w| w == want),
|
|
"client A missed marker"
|
|
);
|
|
assert!(
|
|
got_b.windows(want.len()).any(|w| w == want),
|
|
"client B missed marker"
|
|
);
|
|
|
|
let _ = h.manager.kill(id).await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn ping_text_frame_gets_pong() {
|
|
let h = spawn_daemon_ws("laptop").await;
|
|
let id = create_bash_session(&h).await;
|
|
let token = token_for(id, "laptop", 300);
|
|
let mut ws = connect_attach(h.addr, id, &token).await.unwrap();
|
|
|
|
let ping = serde_json::to_string(&maestro_protocol::ClientControl::Ping).unwrap();
|
|
ws.send(WsMessage::Text(ping)).await.unwrap();
|
|
|
|
let mut saw_pong = false;
|
|
let deadline = std::time::Instant::now() + Duration::from_secs(2);
|
|
while std::time::Instant::now() < deadline {
|
|
if let Ok(Some(Ok(WsMessage::Text(t)))) =
|
|
tokio::time::timeout(Duration::from_millis(100), ws.next()).await
|
|
{
|
|
if t.contains("Pong") {
|
|
saw_pong = true;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
assert!(saw_pong, "no Pong frame received");
|
|
let _ = h.manager.kill(id).await;
|
|
}
|