clawsync/crates/clawsync-transport/tests/e2e_sync.rs

//! End-to-end TCP sync integration test.
//!
//! Scenario: a source node has N revisions; a destination node already has K
//! revisions. The test runs the full wire protocol over a real TCP loopback
//! socket and asserts that exactly N-K layer packets are transferred.
//!
//! Protocol flow (client = source, server = destination):
//! ```text
//! Client                          Server
//!   │── ManifestRequest ─────────▶│
//!   │◀─ ManifestResponse ─────────│
//!   │── LayerPacket (rev K) ────▶ │  (for each missing revision)
//!   │◀─ Ack { revision: K } ──────│
//!   │   …                         │
//!   │── SyncComplete ─────────────▶│
//! ```

use std::net::{IpAddr, Ipv4Addr, SocketAddr};

use clawhdf5_onion::writer::OnionFile;
use clawsync_onion::{differ::diff_revisions, manifest::ClawSyncManifest, merger::merge_packets};
use clawsync_transport::{
    protocol::SyncMessage,
    tcp::{TcpConnection, TcpServer},
};
use tempfile::NamedTempFile;

fn localhost_any() -> SocketAddr {
    SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)
}

const H5_BASE: &[u8] = b"\x89HDF\r\n\x1a\n";

/// Create an `OnionFile` with `n` committed revisions.
/// Each revision writes a single page (4 KiB) at offset 0, filled with `i`.
fn make_onion(n: u8) -> (NamedTempFile, OnionFile) {
    let tmp = NamedTempFile::new().unwrap();
    let h5_path = tmp.path().with_extension("h5");
    std::fs::write(&h5_path, H5_BASE).unwrap();
    let mut onion = OnionFile::create(&h5_path, 4096).unwrap();
    for i in 0..n {
        let mut s = onion.begin_session(None).unwrap();
        s.record_page(0, &vec![i; 4096]);
        onion.commit_session(s, Some(&format!("rev {i}"))).unwrap();
    }
    (tmp, onion)
}

// ─────────────────────────────────────────────────────────────────────────────
// Core helper: run a full push from `src` → `dst` over TCP and return the
// number of LayerPackets actually transferred.
// ─────────────────────────────────────────────────────────────────────────────

async fn run_push(src: OnionFile, mut dst: OnionFile) -> (u64, OnionFile) {
    // Bind the server; the OS picks a free port.
    let server = TcpServer::bind(localhost_any()).await.unwrap();
    let server_addr = server.local_addr;

    // Build the client-side manifest now (before moving `src` into the task).
    let src_manifest = ClawSyncManifest::from_onion("src", &src, H5_BASE);

    // ── Server task ──────────────────────────────────────────────────────────
    // Returns (packets_received, dst) so the caller can inspect the result.
    let server_handle = tokio::spawn(async move {
        let (mut conn, _) = server.accept().await.unwrap();
        let mut packets_received: u64 = 0;

        loop {
            let msg = conn.recv().await.unwrap();
            match msg {
                SyncMessage::ManifestRequest { head_revision, .. } => {
                    // Build our local manifest and send it back.
                    let dst_manifest = ClawSyncManifest::from_onion("dst", &dst, H5_BASE);
                    conn.send(&SyncMessage::ManifestResponse {
                        manifest: dst_manifest,
                    })
                    .await
                    .unwrap();
                    let _ = head_revision; // used for negotiation by client
                }

                SyncMessage::LayerPacket { packet } => {
                    let rev = packet.revision;
                    merge_packets(&mut dst, vec![packet], false).unwrap();
                    packets_received += 1;
                    conn.send(&SyncMessage::Ack { revision: rev })
                        .await
                        .unwrap();
                }

                SyncMessage::SyncComplete { .. } => {
                    // Flush and stop accepting.
                    dst.flush().unwrap();
                    break;
                }

                other => panic!("server: unexpected message: {other:?}"),
            }
        }

        (packets_received, dst)
    });

    // ── Client logic ─────────────────────────────────────────────────────────
    let mut client = TcpConnection::connect(server_addr).await.unwrap();

    // 1. Send manifest request.
    client
        .send(&SyncMessage::ManifestRequest {
            agent_id: "src".to_string(),
            head_revision: src_manifest.head_revision,
            revision_count: src_manifest.revision_count,
        })
        .await
        .unwrap();

    // 2. Receive remote manifest; determine what to push.
    let remote_head = match client.recv().await.unwrap() {
        SyncMessage::ManifestResponse { manifest } => {
            if manifest.revision_count == 0 {
                clawhdf5_onion::format::NO_PARENT
            } else {
                manifest.head_revision
            }
        }
        other => panic!("client: expected ManifestResponse, got {other:?}"),
    };

    // 3. Compute and send missing revisions.
    let packets = diff_revisions(&src, remote_head).unwrap();
    let total_sent = packets.len() as u64;

    for packet in packets {
        let rev = packet.revision;
        client
            .send(&SyncMessage::LayerPacket { packet })
            .await
            .unwrap();
        // Wait for ack before sending next packet (ordered, reliable).
        match client.recv().await.unwrap() {
            SyncMessage::Ack { revision } => assert_eq!(revision, rev),
            other => panic!("client: expected Ack, got {other:?}"),
        }
    }

    // 4. Signal completion.
    client
        .send(&SyncMessage::SyncComplete {
            revisions_transferred: total_sent,
            bytes_transferred: 0,
        })
        .await
        .unwrap();

    let (server_received, dst_final) = server_handle.await.unwrap();
    assert_eq!(
        server_received, total_sent,
        "server count must match client send count"
    );

    (total_sent, dst_final)
}

// ─────────────────────────────────────────────────────────────────────────────
// Tests
// ─────────────────────────────────────────────────────────────────────────────

/// Push 5 revisions to an empty destination — all 5 must be transferred.
#[tokio::test]
async fn push_all_revisions_to_empty_destination() {
    let (_src_tmp, src) = make_onion(5);
    let (_dst_tmp, dst) = make_onion(0);

    let (transferred, dst_final) = run_push(src, dst).await;

    assert_eq!(transferred, 5, "all 5 revisions should be transferred");
    assert_eq!(dst_final.revision_count(), 5);
}

/// Push when destination is already at the same HEAD — nothing to transfer.
#[tokio::test]
async fn push_when_already_in_sync() {
    let (_src_tmp, src) = make_onion(3);
    let (_dst_tmp, dst) = make_onion(3);

    let (transferred, dst_final) = run_push(src, dst).await;

    assert_eq!(transferred, 0, "nothing should be transferred when in sync");
    assert_eq!(dst_final.revision_count(), 3);
}

/// Destination has K revisions, source has N; exactly N-K should transfer.
#[tokio::test]
async fn push_partial_delta_transfers_exactly_n_minus_k() {
    const K: u8 = 3;
    const N: u8 = 8;

    let (_src_tmp, src) = make_onion(N);
    let (_dst_tmp, dst) = make_onion(K);

    let (transferred, dst_final) = run_push(src, dst).await;

    let expected = (N - K) as u64;
    assert_eq!(
        transferred, expected,
        "exactly N-K={expected} revisions should be transferred"
    );
    assert_eq!(dst_final.revision_count(), N as u64);
}

/// After sync the destination can reconstruct every revision that the source has.
#[tokio::test]
async fn destination_revisions_match_source_after_sync() {
    let (_src_tmp, src) = make_onion(4);
    let (_dst_tmp, dst) = make_onion(0);

    // Capture source revision metadata before the move.
    let src_revisions = src.list_revisions();

    let (_transferred, dst_final) = run_push(src, dst).await;

    let dst_revisions = dst_final.list_revisions();
    assert_eq!(src_revisions.len(), dst_revisions.len());

    for (s, d) in src_revisions.iter().zip(dst_revisions.iter()) {
        assert_eq!(s.revision, d.revision, "revision numbers must match");
        assert_eq!(
            s.annotation, d.annotation,
            "annotations must survive the wire"
        );
        assert_eq!(
            s.blake3_hex, d.blake3_hex,
            "BLAKE3 hashes must be identical — page data is intact"
        );
    }
}

/// Verify idempotency: pushing the same revisions twice leaves the destination
/// with exactly the same number of revisions (no duplicates).
#[tokio::test]
async fn push_twice_is_idempotent() {
    let (_src_tmp1, src1) = make_onion(3);
    let (_src_tmp2, _src2) = make_onion(3);
    let (_dst_tmp, dst) = make_onion(0);

    // First push.
    let (_t1, dst_after_first) = run_push(src1, dst).await;
    assert_eq!(dst_after_first.revision_count(), 3);

    // Second push of the identical 3 revisions.
    // `run_push` moves `dst_after_first`, so we need to rebuild one.
    // Re-open the same .onion file from disk.
    let (_src_tmp3, src3) = make_onion(3);
    let (_dst_tmp2, dst2) = make_onion(3); // simulates "already has all"
    let (t2, dst_final) = run_push(src3, dst2).await;

    assert_eq!(t2, 0, "second push should transfer nothing");
    assert_eq!(dst_final.revision_count(), 3, "no duplicate revisions");
    let _ = dst_after_first; // keep alive until here
}

/// Large push: 20 revisions, destination has 7 — assert 13 packets.
#[tokio::test]
async fn large_push_correct_packet_count() {
    const K: u8 = 7;
    const N: u8 = 20;

    let (_src_tmp, src) = make_onion(N);
    let (_dst_tmp, dst) = make_onion(K);

    let (transferred, dst_final) = run_push(src, dst).await;

    assert_eq!(transferred, (N - K) as u64);
    assert_eq!(dst_final.revision_count(), N as u64);
}

// ─────────────────────────────────────────────────────────────────────────────
// Pull helper: server has `src`; client starts with `dst` and pulls from it.
// Returns (packets_transferred, final_client_onion).
// ─────────────────────────────────────────────────────────────────────────────

async fn run_pull(
    src: OnionFile,     // server-side onion (richer)
    mut dst: OnionFile, // client-side onion (missing revisions)
) -> (u64, OnionFile) {
    let server = TcpServer::bind(localhost_any()).await.unwrap();
    let server_addr = server.local_addr;

    // ── Server task ──────────────────────────────────────────────────────────
    let server_handle = tokio::spawn(async move {
        let (mut conn, _) = server.accept().await.unwrap();

        // 1. Receive client's ManifestRequest.
        let client_rev_count = match conn.recv().await.unwrap() {
            SyncMessage::ManifestRequest { revision_count, .. } => revision_count,
            other => panic!("server: expected ManifestRequest, got {other:?}"),
        };

        // 2. Send back the server manifest.
        let srv_manifest = ClawSyncManifest::from_onion("srv", &src, H5_BASE);
        conn.send(&SyncMessage::ManifestResponse {
            manifest: srv_manifest,
        })
        .await
        .unwrap();

        // 3. Server has more revisions — push the delta to the client.
        if src.revision_count() > client_rev_count {
            let remote_head = if client_rev_count == 0 {
                clawhdf5_onion::format::NO_PARENT
            } else {
                client_rev_count - 1
            };
            let packets = diff_revisions(&src, remote_head).unwrap();
            let total = packets.len() as u64;
            let mut bytes_sent = 0u64;
            for packet in packets {
                bytes_sent += packet.page_data_size() as u64;
                conn.send(&SyncMessage::LayerPacket { packet })
                    .await
                    .unwrap();
                match conn.recv().await.unwrap() {
                    SyncMessage::Ack { .. } => {}
                    other => panic!("server: expected Ack, got {other:?}"),
                }
            }
            conn.send(&SyncMessage::SyncComplete {
                revisions_transferred: total,
                bytes_transferred: bytes_sent,
            })
            .await
            .unwrap();
        } else {
            // Already in sync: wait for client's SyncComplete.
            match conn.recv().await.unwrap() {
                SyncMessage::SyncComplete { .. } => {}
                other => panic!("server: expected SyncComplete, got {other:?}"),
            }
        }
    });

    // ── Client logic ─────────────────────────────────────────────────────────
    let mut client = TcpConnection::connect(server_addr).await.unwrap();
    let local_rev_count = dst.revision_count();

    // 1. Send manifest request.
    client
        .send(&SyncMessage::ManifestRequest {
            agent_id: "client".to_string(),
            head_revision: local_rev_count.saturating_sub(1),
            revision_count: local_rev_count,
        })
        .await
        .unwrap();

    // 2. Receive server manifest; check if there is anything to pull.
    let server_rev_count = match client.recv().await.unwrap() {
        SyncMessage::ManifestResponse { manifest } => manifest.revision_count,
        other => panic!("client: expected ManifestResponse, got {other:?}"),
    };

    if server_rev_count <= local_rev_count {
        // Nothing to pull — tell server we're done.
        client
            .send(&SyncMessage::SyncComplete {
                revisions_transferred: 0,
                bytes_transferred: 0,
            })
            .await
            .unwrap();
        server_handle.await.unwrap();
        return (0, dst);
    }

    // 3. Receive packets from server.
    let mut packets = Vec::new();
    loop {
        match client.recv().await.unwrap() {
            SyncMessage::LayerPacket { packet } => {
                let rev = packet.revision;
                client
                    .send(&SyncMessage::Ack { revision: rev })
                    .await
                    .unwrap();
                packets.push(packet);
            }
            SyncMessage::SyncComplete {
                revisions_transferred,
                ..
            } => {
                merge_packets(&mut dst, packets, true).unwrap();
                server_handle.await.unwrap();
                return (revisions_transferred, dst);
            }
            other => panic!("client: unexpected {other:?}"),
        }
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Pull tests
// ─────────────────────────────────────────────────────────────────────────────

/// Pull all revisions from a full server to an empty client.
#[tokio::test]
async fn pull_all_from_server() {
    let (_srv_tmp, srv) = make_onion(5);
    let (_cli_tmp, cli) = make_onion(0);

    let (transferred, cli_final) = run_pull(srv, cli).await;

    assert_eq!(transferred, 5);
    assert_eq!(cli_final.revision_count(), 5);
}

/// Pull when client and server are already in sync — zero packets.
#[tokio::test]
async fn pull_noop_when_already_in_sync() {
    let (_srv_tmp, srv) = make_onion(4);
    let (_cli_tmp, cli) = make_onion(4);

    let (transferred, cli_final) = run_pull(srv, cli).await;

    assert_eq!(transferred, 0, "nothing to pull when already in sync");
    assert_eq!(cli_final.revision_count(), 4);
}

/// Pull partial delta: server has N, client has K — client receives N-K packets.
#[tokio::test]
async fn pull_partial_delta() {
    const K: u8 = 2;
    const N: u8 = 7;

    let (_srv_tmp, srv) = make_onion(N);
    let (_cli_tmp, cli) = make_onion(K);

    let (transferred, cli_final) = run_pull(srv, cli).await;

    assert_eq!(transferred, (N - K) as u64);
    assert_eq!(cli_final.revision_count(), N as u64);
}

/// BLAKE3 hashes must survive the pull wire path intact.
#[tokio::test]
async fn pull_revision_hashes_match() {
    let (_srv_tmp, srv) = make_onion(4);
    let (_cli_tmp, cli) = make_onion(0);

    let srv_revisions = srv.list_revisions();
    let (_, cli_final) = run_pull(srv, cli).await;
    let cli_revisions = cli_final.list_revisions();

    assert_eq!(srv_revisions.len(), cli_revisions.len());
    for (s, c) in srv_revisions.iter().zip(cli_revisions.iter()) {
        assert_eq!(
            s.blake3_hex, c.blake3_hex,
            "rev {} hash mismatch",
            s.revision
        );
        assert_eq!(s.annotation, c.annotation);
    }
}