atproto-devtool 0.1.1

//! Ephemeral `did:web` document server for self-mint conformance checks.
//!
//! Binds `127.0.0.1:0`, accepts the first inbound TCP connection, and
//! serves a single hand-crafted HTTP/1.1 response carrying the DID
//! document JSON at `/.well-known/did.json`. Any other path returns
//! 404. Other requests to `/.well-known/did.json` are honored too —
//! the labeler may retry.
//!
//! Shuts down on drop: the RAII handle aborts the background task and
//! closes the listener.

use std::net::SocketAddr;
use std::sync::Arc;

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::task::JoinHandle;

/// A running DID document server.
///
/// The server runs on `127.0.0.1:{os-assigned-port}` in a background
/// task; the listening address is exposed via `local_addr()`. When the
/// `DidDocServer` is dropped, the background task is aborted.
pub struct DidDocServer {
    local_addr: SocketAddr,
    task: JoinHandle<()>,
}

impl DidDocServer {
    /// Bind `127.0.0.1:0` and start serving DID-document JSON bytes at
    /// `/.well-known/did.json`. The body is built **after** the listener
    /// has bound, by invoking `build_body` with the known `SocketAddr`.
    /// This lets callers embed the bound port into the DID document
    /// atomically — there is no probe phase, no possibility of port drift
    /// between binding and serving.
    pub async fn spawn<F>(build_body: F) -> std::io::Result<Self>
    where
        F: FnOnce(SocketAddr) -> Vec<u8>,
    {
        let listener = TcpListener::bind("127.0.0.1:0").await?;
        let local_addr = listener.local_addr()?;
        let did_doc_json = build_body(local_addr);
        let body: Arc<[u8]> = did_doc_json.into();

        let task = tokio::spawn(async move {
            loop {
                let accept = listener.accept().await;
                let (mut stream, _peer) = match accept {
                    Ok(sp) => sp,
                    Err(_) => return,
                };
                let body = body.clone();
                tokio::spawn(async move {
                    let _ = Self::handle_connection(&mut stream, &body).await;
                });
            }
        });

        Ok(Self { local_addr, task })
    }

    /// The listening address (always `127.0.0.1:{port}`).
    pub fn local_addr(&self) -> SocketAddr {
        self.local_addr
    }

    /// Minimal HTTP/1.1 handler: reads the request line, routes on path.
    async fn handle_connection(
        stream: &mut tokio::net::TcpStream,
        did_doc: &[u8],
    ) -> std::io::Result<()> {
        // Read up to 8 KiB of request headers; servers don't send large
        // GET requests but we cap to avoid unbounded reads.
        let mut buf = [0u8; 8192];
        let mut total = 0usize;
        while total < buf.len() {
            let n = stream.read(&mut buf[total..]).await?;
            if n == 0 {
                break;
            }
            total += n;
            // Headers end at CRLFCRLF.
            if buf[..total].windows(4).any(|w| w == b"\r\n\r\n") {
                break;
            }
        }
        let request = &buf[..total];

        // Parse just the request line: `GET /path HTTP/1.1\r\n`.
        let first_line_end = request
            .iter()
            .position(|&b| b == b'\r' || b == b'\n')
            .unwrap_or(request.len());
        let first_line = std::str::from_utf8(&request[..first_line_end]).unwrap_or("");
        let mut parts = first_line.split_whitespace();
        let method = parts.next().unwrap_or("");
        let path = parts.next().unwrap_or("");

        let (status_line, body, content_type) =
            if method == "GET" && path == "/.well-known/did.json" {
                ("HTTP/1.1 200 OK", did_doc, "application/json")
            } else {
                (
                    "HTTP/1.1 404 Not Found",
                    b"not found" as &[u8],
                    "text/plain",
                )
            };

        let response_head = format!(
            "{status_line}\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
            body.len()
        );
        stream.write_all(response_head.as_bytes()).await?;
        stream.write_all(body).await?;
        stream.flush().await?;
        let _ = stream.shutdown().await;
        Ok(())
    }
}

impl Drop for DidDocServer {
    fn drop(&mut self) {
        self.task.abort();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn serves_did_doc_on_well_known_path() {
        let body = br#"{"id":"did:web:127.0.0.1%3A0"}"#.to_vec();
        let body_for_assert = body.clone();
        let server = DidDocServer::spawn(move |_addr| body).await.expect("spawn");
        let url = format!("http://{}/.well-known/did.json", server.local_addr());
        let resp = reqwest::Client::new().get(&url).send().await.expect("http");
        assert_eq!(resp.status(), 200);
        assert_eq!(resp.headers()["content-type"], "application/json");
        let got = resp.bytes().await.expect("bytes");
        assert_eq!(got.as_ref(), body_for_assert.as_slice());
    }

    #[tokio::test]
    async fn returns_404_for_other_paths() {
        let server = DidDocServer::spawn(|_addr| b"{}".to_vec())
            .await
            .expect("spawn");
        let url = format!("http://{}/nope", server.local_addr());
        let resp = reqwest::Client::new().get(&url).send().await.expect("http");
        assert_eq!(resp.status(), 404);
    }

    #[tokio::test]
    async fn body_builder_receives_bound_addr() {
        let captured = std::sync::Arc::new(std::sync::Mutex::new(None));
        let captured_clone = captured.clone();
        let server = DidDocServer::spawn(move |addr| {
            *captured_clone.lock().unwrap() = Some(addr);
            format!(r#"{{"port":{}}}"#, addr.port()).into_bytes()
        })
        .await
        .expect("spawn");
        let addr = *captured.lock().unwrap();
        assert_eq!(addr, Some(server.local_addr()));
    }
}