improvise/crates/improvise-ws-server/tests/socket.rs

//! The one test that uses a real socket — the ticket's AC 5. Its subject is
//! the shell, not the decisions: `context/test-evidence/clients.md`.

use futures_util::{SinkExt, StreamExt};
use tokio::net::TcpListener;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::http::StatusCode;

use improvise_command::ui::app::App;
use improvise_core::model::cell::CellValue;
use improvise_core::view::Axis;
use improvise_core::workbook::Workbook;
use improvise_protocol::{
    BasisCache, CellId, Command, GridCache, GridUpdate, ModelCommand, Projection, WidgetCache,
};
use improvise_protocol::{ClientFrame, ServerFrame};
use improvise_ws_server::origin::OriginPolicy;
use improvise_ws_server::registry::{DocumentId, Registry};

/// One document, served as `quarterly`.
fn registry() -> Registry {
    let mut registry = Registry::new();
    registry
        .open(DocumentId::new("quarterly"), app())
        .expect("a fresh registry has nothing to collide with");
    registry
}

fn app() -> App {
    let mut wb = Workbook::new("Test");
    wb.add_category("Region").unwrap();
    wb.add_item("Region", "North").unwrap();
    wb.add_item("Region", "South").unwrap();
    wb.add_item("_Measure", "Revenue").unwrap();
    wb.add_item("_Measure", "Cost").unwrap();
    wb.active_view_mut().set_axis("Region", Axis::Row);
    wb.active_view_mut().set_axis("_Measure", Axis::Column);
    let mut app = App::new(wb, None);
    app.rebuild_layout();
    app
}

#[test]
fn a_client_on_a_socket_is_seeded_and_gets_its_edit_back() {
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let local = tokio::task::LocalSet::new();

    runtime.block_on(local.run_until(async {
        let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::task::spawn_local(async move {
            let _ = improvise_ws_server::transport::serve(
                registry(),
                listener,
                OriginPolicy::default(),
            )
            .await;
        });

        // No path: with one document served, `/` reaches it, exactly as it
        // did before the server could host several (improvise-j31x).
        let (mut ws, _) = tokio_tungstenite::connect_async(format!("ws://{addr}"))
            .await
            .expect("the server must accept a websocket connection");

        // The handshake: the extent, first, because the snapshot is built at
        // it.
        send(&mut ws, ClientFrame::Viewport { rows: 10, cols: 10 }).await;
        let ServerFrame::Snapshot(snapshot) = recv(&mut ws).await else {
            panic!("the first frame back was not a Snapshot");
        };
        let mut grid = GridCache::default();
        let mut basis = BasisCache::default();
        for projection in &snapshot.seed {
            match projection {
                Projection::Grid(update) => grid.apply(update),
                Projection::Basis(update) => basis.apply(update),
                _ => {}
            }
        }
        assert!(
            basis.current().is_some(),
            "the bootstrap must publish a basis, or this client can never \
             compose a write at all (ADR 001)"
        );
        assert_eq!(
            grid.row_frame.leaf_labels(),
            vec!["North".to_string(), "South".to_string()]
        );
        assert_eq!(grid.cells[0][0].text, "");

        // The client composes the write from what it was *told* — ADR 002. A
        // `CellKey` here would be the test knowing what a browser does not.
        let CellId::Record(target) = grid.cells[0][0]
            .id
            .clone()
            .expect("the server must name the cell it projected")
        else {
            panic!("a raw cell must be projected as settable, not as an aggregate");
        };

        send(
            &mut ws,
            ClientFrame::command(Command::Model {
                // Composed from what the server published and nothing else: a
                // self-built basis claims to have seen what was never sent.
                basis: basis
                    .claim([target.record])
                    .expect("the bootstrap published a basis"),
                cmd: ModelCommand::SetCell {
                    view: "Default".into(),
                    cell: target,
                    value: CellValue::Number(12.0),
                },
            }),
        )
        .await;

        // The grid update may arrive behind the projections for the other
        // widgets, so fold frames until the cell lands.
        let mut waited = 0;
        while grid.cells[0][0].text.is_empty() && waited < 8 {
            let ServerFrame::Commands(commands) = recv(&mut ws).await else {
                panic!("a second Snapshot arrived for a cell edit");
            };
            for command in &commands {
                if let Command::Project(Projection::Grid(update)) = command {
                    grid.apply(update);
                }
            }
            waited += 1;
        }
        assert_eq!(
            grid.cells[0][0].text, "12.00",
            "the edit never came back down the socket"
        );
        // The acknowledgement follows everything the frame produced, so a
        // client discards its undo knowing what it is discarding (nxah).
        assert_eq!(
            recv(&mut ws).await,
            ServerFrame::Accepted,
            "a mutating frame that landed must be acknowledged, and last"
        );

        // A frame the vocabulary does not cover is skipped, not fatal. The
        // resize that follows genuinely reframes a 2x2 fixture to 1x1.
        ws.send(Message::Text("{\"NotAFrame\":true}".into()))
            .await
            .unwrap();
        send(&mut ws, ClientFrame::Viewport { rows: 1, cols: 1 }).await;
        let ServerFrame::Commands(commands) = recv(&mut ws).await else {
            panic!("a resize produced a Snapshot");
        };
        assert!(
            commands
                .iter()
                .any(|c| matches!(c, Command::Project(Projection::Grid(_)))),
            "shrinking the window sent no grid: {commands:?}"
        );
        assert!(
            commands.iter().any(|c| matches!(
                c,
                Command::Project(Projection::Grid(GridUpdate::ColWidths(_)))
            )),
            "the reframed grid carried no widths: {commands:?}"
        );
    }));
}

type Socket =
    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;

async fn send(ws: &mut Socket, msg: ClientFrame) {
    ws.send(Message::Text(serde_json::to_string(&msg).unwrap().into()))
        .await
        .expect("the socket must accept a frame");
}

async fn recv(ws: &mut Socket) -> ServerFrame {
    loop {
        let message = tokio::time::timeout(std::time::Duration::from_secs(3), ws.next())
            .await
            .expect("the server sent nothing within three seconds")
            .expect("the server closed the connection")
            .expect("the socket errored");
        if let Message::Text(text) = message {
            return serde_json::from_str(&text).expect("the server sent something unparseable");
        }
    }
}

/// The other thing only a socket can check: the document is named in the URL
/// path, and this covers that the path reaches `Registry::resolve` at all.
#[test]
fn the_url_path_picks_the_document() {
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let local = tokio::task::LocalSet::new();

    runtime.block_on(local.run_until(async {
        let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::task::spawn_local(async move {
            let mut registry = Registry::new();
            registry
                .open(DocumentId::new("alpha"), named_app("Alpha", "Aardvark"))
                .unwrap();
            registry
                .open(DocumentId::new("beta"), named_app("Beta", "Bison"))
                .unwrap();
            let _ =
                improvise_ws_server::transport::serve(registry, listener, OriginPolicy::default())
                    .await;
        });

        for (path, row) in [("alpha", "Aardvark"), ("beta", "Bison")] {
            let (mut ws, _) = tokio_tungstenite::connect_async(format!("ws://{addr}/{path}"))
                .await
                .expect("the server must accept a named document");
            send(&mut ws, ClientFrame::Viewport { rows: 10, cols: 10 }).await;
            let ServerFrame::Snapshot(snapshot) = recv(&mut ws).await else {
                panic!("the first frame back was not a Snapshot");
            };
            let mut grid = GridCache::default();
            for projection in &snapshot.seed {
                if let Projection::Grid(update) = projection {
                    grid.apply(update);
                }
            }
            assert_eq!(
                grid.row_frame.leaf_labels(),
                vec![row.to_string()],
                "/{path} served the wrong document"
            );
        }

        // A name nothing serves: the handshake succeeds — a websocket fact,
        // not a document one — and the connection is closed with the reason.
        let (mut ws, _) = tokio_tungstenite::connect_async(format!("ws://{addr}/gamma"))
            .await
            .expect("the handshake itself is not where a document is checked");
        send(&mut ws, ClientFrame::Viewport { rows: 10, cols: 10 }).await;
        let closed = tokio::time::timeout(std::time::Duration::from_secs(3), ws.next())
            .await
            .expect("the server neither answered nor closed")
            .expect("the stream ended without a close frame")
            .expect("the socket errored");
        let Message::Close(Some(frame)) = closed else {
            panic!("an unserved document was answered with {closed:?}");
        };
        assert!(
            frame.reason.contains("gamma"),
            "the close frame did not say what was wrong: {frame:?}"
        );

        // And the server is still serving: one refused connection is not the
        // registry task falling over.
        let (mut ws, _) = tokio_tungstenite::connect_async(format!("ws://{addr}/alpha"))
            .await
            .expect("the server must still be accepting");
        send(&mut ws, ClientFrame::Viewport { rows: 10, cols: 10 }).await;
        assert!(matches!(recv(&mut ws).await, ServerFrame::Snapshot(_)));
    }));
}

fn named_app(name: &str, row: &str) -> App {
    let mut wb = Workbook::new(name);
    wb.add_category("Region").unwrap();
    wb.add_item("Region", row).unwrap();
    wb.add_item("_Measure", "Revenue").unwrap();
    wb.active_view_mut().set_axis("Region", Axis::Row);
    wb.active_view_mut().set_axis("_Measure", Axis::Column);
    let mut app = App::new(wb, None);
    app.rebuild_layout();
    app
}

/// A hostile page an allowed user visits opens a socket from that user's own
/// allowed address. Only `Origin` separates it from the real page.
#[test]
fn a_handshake_from_a_foreign_origin_is_refused() {
    with_server(OriginPolicy::default(), |addr| async move {
        let Err(refused) =
            tokio_tungstenite::connect_async(request(addr, Some("https://evil.example"))).await
        else {
            panic!("a foreign origin must not reach the document");
        };
        // The kind, not the wording: a refused upgrade and a dropped socket
        // are different outcomes and only one of them is this fix.
        let tokio_tungstenite::tungstenite::Error::Http(answer) = refused else {
            panic!("the handshake failed for some other reason: {refused}");
        };
        assert_eq!(answer.status(), StatusCode::FORBIDDEN);
    });
}

/// The shipped local recipe: the page is served by `python3 -m http.server`
/// on loopback and opens a socket to the server on another loopback port.
#[test]
fn a_handshake_from_a_loopback_page_is_accepted() {
    with_server(OriginPolicy::default(), |addr| async move {
        let (mut ws, _) =
            tokio_tungstenite::connect_async(request(addr, Some("http://localhost:8080")))
                .await
                .expect("the local pairing recipe must keep working");
        send(&mut ws, ClientFrame::Viewport { rows: 10, cols: 10 }).await;
        assert!(matches!(recv(&mut ws).await, ServerFrame::Snapshot(_)));
    });
}

/// Non-browser clients — curl, a native client, this file's other tests —
/// send no `Origin` at all, and absent is decided rather than defaulted.
#[test]
fn a_handshake_with_no_origin_is_accepted() {
    with_server(OriginPolicy::default(), |addr| async move {
        let (mut ws, _) = tokio_tungstenite::connect_async(request(addr, None))
            .await
            .expect("a client that sends no Origin is not a browser page");
        send(&mut ws, ClientFrame::Viewport { rows: 10, cols: 10 }).await;
        assert!(matches!(recv(&mut ws).await, ServerFrame::Snapshot(_)));
    });
}

/// The deployed page reaches the server through a TLS proxy, and `Origin` is
/// the one thing that survives the proxy to say which page it was.
#[test]
fn a_named_origin_reaches_the_document() {
    let named = OriginPolicy::new(["https://fiddlerwoaroof.github.io".to_string()]);
    with_server(named, |addr| async move {
        let (mut ws, _) = tokio_tungstenite::connect_async(request(
            addr,
            Some("https://fiddlerwoaroof.github.io"),
        ))
        .await
        .expect("--allow-origin must actually widen the policy");
        send(&mut ws, ClientFrame::Viewport { rows: 10, cols: 10 }).await;
        assert!(matches!(recv(&mut ws).await, ServerFrame::Snapshot(_)));

        // Widening for one page does not widen for its neighbours.
        assert!(
            tokio_tungstenite::connect_async(request(addr, Some("https://evil.example")))
                .await
                .is_err()
        );
    });
}

/// A handshake request for `addr`, carrying `origin` when there is one.
fn request(
    addr: std::net::SocketAddr,
    origin: Option<&str>,
) -> tokio_tungstenite::tungstenite::http::Request<()> {
    use tokio_tungstenite::tungstenite::client::IntoClientRequest;
    let mut request = format!("ws://{addr}/")
        .into_client_request()
        .expect("a loopback address is a valid websocket url");
    if let Some(origin) = origin {
        request
            .headers_mut()
            .insert("Origin", origin.parse().expect("a header value"));
    }
    request
}

/// Serve the one-document fixture on an ephemeral port and run `body`.
fn with_server<F, Fut>(origins: OriginPolicy, body: F)
where
    F: FnOnce(std::net::SocketAddr) -> Fut,
    Fut: std::future::Future<Output = ()>,
{
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let local = tokio::task::LocalSet::new();
    runtime.block_on(local.run_until(async move {
        let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::task::spawn_local(async move {
            let _ = improvise_ws_server::transport::serve(registry(), listener, origins).await;
        });
        body(addr).await;
    }));
}