ngit-grasp/tests/common/setup_drop_relay.rs

//! Relay fixture that disconnects the WebSocket during its NIP-11 setup fetch.

use std::sync::Arc;
use std::time::Duration;

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{oneshot, watch};

/// Endpoint for reproducing a stale successful connection result.
pub struct SetupDropRelay {
    url: String,
    shutdown_tx: Option<oneshot::Sender<()>>,
    handle: Option<tokio::task::JoinHandle<()>>,
}

impl SetupDropRelay {
    pub async fn start() -> Self {
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("SetupDropRelay failed to bind");
        let port = listener.local_addr().expect("setup-drop local_addr").port();
        let (nip11_tx, nip11_rx) = watch::channel(false);
        let (dropped_tx, dropped_rx) = watch::channel(false);
        let nip11_tx = Arc::new(nip11_tx);
        let dropped_tx = Arc::new(dropped_tx);
        let (shutdown_tx, mut shutdown_rx) = oneshot::channel();

        let handle = tokio::spawn(async move {
            loop {
                tokio::select! {
                    accepted = listener.accept() => {
                        let Ok((stream, _)) = accepted else { break };
                        let nip11_tx = nip11_tx.clone();
                        let nip11_rx = nip11_rx.clone();
                        let dropped_tx = dropped_tx.clone();
                        let dropped_rx = dropped_rx.clone();
                        tokio::spawn(async move {
                            handle_connection(
                                stream,
                                nip11_tx,
                                nip11_rx,
                                dropped_tx,
                                dropped_rx,
                            )
                            .await;
                        });
                    }
                    _ = &mut shutdown_rx => break,
                }
            }
        });

        Self {
            url: format!("ws://127.0.0.1:{port}"),
            shutdown_tx: Some(shutdown_tx),
            handle: Some(handle),
        }
    }

    pub fn url(&self) -> &str {
        &self.url
    }

    pub async fn stop(mut self) {
        if let Some(tx) = self.shutdown_tx.take() {
            let _ = tx.send(());
        }
        if let Some(handle) = self.handle.take() {
            let _ = handle.await;
        }
    }
}

async fn handle_connection(
    mut stream: TcpStream,
    nip11_tx: Arc<watch::Sender<bool>>,
    mut nip11_rx: watch::Receiver<bool>,
    dropped_tx: Arc<watch::Sender<bool>>,
    mut dropped_rx: watch::Receiver<bool>,
) {
    let mut header = [0_u8; 4096];
    let header_len = loop {
        let Ok(length) = stream.peek(&mut header).await else {
            return;
        };
        if length == 0 {
            return;
        }
        if header[..length]
            .windows(4)
            .any(|bytes| bytes == b"\r\n\r\n")
        {
            break length;
        }
        tokio::task::yield_now().await;
    };
    let request = String::from_utf8_lossy(&header[..header_len]).to_ascii_lowercase();

    if request.contains("upgrade: websocket") {
        let Ok(_websocket) = tokio_tungstenite::accept_async(stream).await else {
            return;
        };
        while !*nip11_rx.borrow() && nip11_rx.changed().await.is_ok() {}
        // Dropping the live socket is the behavior under test. Give the SDK a
        // bounded propagation window before allowing the HTTP setup to finish.
        drop(_websocket);
        let _ = dropped_tx.send(true);
    } else {
        let mut request_bytes = vec![0_u8; header_len];
        if stream.read_exact(&mut request_bytes).await.is_err() {
            return;
        }
        let _ = nip11_tx.send(true);
        while !*dropped_rx.borrow() && dropped_rx.changed().await.is_ok() {}
        tokio::time::sleep(Duration::from_millis(100)).await;
        let body = r#"{"limitation":{"max_subscriptions":20}}"#;
        let response = format!(
            "HTTP/1.1 200 OK\r\nContent-Type: application/nostr+json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
            body.len(),
            body
        );
        let _ = stream.write_all(response.as_bytes()).await;
    }
}

impl Drop for SetupDropRelay {
    fn drop(&mut self) {
        if let Some(tx) = self.shutdown_tx.take() {
            let _ = tx.send(());
        }
    }
}