Phase 2c: streaming Blob RPC #8
@@ -36,7 +36,7 @@
|
|||||||
use anyhow::{bail, Context, Result};
|
use anyhow::{bail, Context, Result};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use tokio::io::AsyncWriteExt;
|
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||||||
|
|
||||||
/// Physical chunk size. 4 MB is the sweet spot for our workloads:
|
/// Physical chunk size. 4 MB is the sweet spot for our workloads:
|
||||||
/// small enough that dedup catches typical file-tree overlaps between
|
/// small enough that dedup catches typical file-tree overlaps between
|
||||||
@@ -260,6 +260,99 @@ impl BlobStore {
|
|||||||
Ok(Some(out))
|
Ok(Some(out))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Store bytes from an async reader, returning the deterministic
|
||||||
|
/// content hash. Streams input in 4 MiB frames — memory ceiling is
|
||||||
|
/// one chunk buffer regardless of blob size. Same idempotence +
|
||||||
|
/// dedup guarantees as [`put_bytes`]: identical content across two
|
||||||
|
/// calls produces the same BlobId and stores each unique chunk
|
||||||
|
/// exactly once on disk.
|
||||||
|
///
|
||||||
|
/// EOF at a chunk boundary is honoured: an empty reader produces
|
||||||
|
/// the empty-blob BlobId with zero chunks, matching `put_bytes(&[])`.
|
||||||
|
pub async fn put_stream<R>(&self, mut reader: R) -> Result<BlobId>
|
||||||
|
where
|
||||||
|
R: AsyncRead + Unpin,
|
||||||
|
{
|
||||||
|
let mut blob_hasher = blake3::Hasher::new();
|
||||||
|
let mut chunk_hashes: Vec<ChunkHash> = Vec::new();
|
||||||
|
let mut total_size: u64 = 0;
|
||||||
|
let mut buf = vec![0u8; CHUNK_SIZE];
|
||||||
|
|
||||||
|
loop {
|
||||||
|
// Fill the chunk buffer or hit EOF. `read` may return fewer
|
||||||
|
// bytes than requested; loop until we have CHUNK_SIZE or the
|
||||||
|
// reader is drained.
|
||||||
|
let mut have = 0usize;
|
||||||
|
while have < CHUNK_SIZE {
|
||||||
|
let n = reader
|
||||||
|
.read(&mut buf[have..])
|
||||||
|
.await
|
||||||
|
.context("reading from streaming source")?;
|
||||||
|
if n == 0 {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
have += n;
|
||||||
|
}
|
||||||
|
if have == 0 {
|
||||||
|
break; // Reader drained on a clean chunk boundary.
|
||||||
|
}
|
||||||
|
let bytes = &buf[..have];
|
||||||
|
blob_hasher.update(bytes);
|
||||||
|
let chunk_hash = ChunkHash(blake3::hash(bytes).into());
|
||||||
|
self.write_chunk_if_absent(&chunk_hash, bytes).await?;
|
||||||
|
chunk_hashes.push(chunk_hash);
|
||||||
|
total_size += have as u64;
|
||||||
|
if have < CHUNK_SIZE {
|
||||||
|
break; // Short final chunk; no more data possible.
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let blob_id = BlobId(blob_hasher.finalize().into());
|
||||||
|
let manifest = BlobManifest {
|
||||||
|
blob_id,
|
||||||
|
total_size,
|
||||||
|
chunks: chunk_hashes,
|
||||||
|
};
|
||||||
|
self.write_manifest_if_absent(&manifest).await?;
|
||||||
|
Ok(blob_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Stream a stored blob into `writer` chunk-by-chunk. Returns
|
||||||
|
/// `Ok(false)` when no manifest exists for `id` (writer is left
|
||||||
|
/// untouched); `Ok(true)` on success. Each chunk is hash-verified
|
||||||
|
/// before it's emitted, so store corruption surfaces as `Err` mid
|
||||||
|
/// stream — callers who need to abort a partial write should
|
||||||
|
/// invalidate the destination buffer on error.
|
||||||
|
pub async fn stream_to<W>(&self, id: &BlobId, writer: &mut W) -> Result<bool>
|
||||||
|
where
|
||||||
|
W: AsyncWrite + Unpin,
|
||||||
|
{
|
||||||
|
let manifest = match self.load_manifest(id).await? {
|
||||||
|
Some(m) => m,
|
||||||
|
None => return Ok(false),
|
||||||
|
};
|
||||||
|
for chunk_hash in &manifest.chunks {
|
||||||
|
let path = self.chunk_path(chunk_hash);
|
||||||
|
let bytes = tokio::fs::read(&path)
|
||||||
|
.await
|
||||||
|
.with_context(|| format!("reading chunk {}", path.display()))?;
|
||||||
|
let recomputed = ChunkHash(blake3::hash(&bytes).into());
|
||||||
|
if recomputed != *chunk_hash {
|
||||||
|
bail!(
|
||||||
|
"chunk hash mismatch at {}: manifest says {}, disk hashes to {}",
|
||||||
|
path.display(),
|
||||||
|
chunk_hash.to_hex(),
|
||||||
|
recomputed.to_hex()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
writer
|
||||||
|
.write_all(&bytes)
|
||||||
|
.await
|
||||||
|
.context("writing chunk bytes to destination")?;
|
||||||
|
}
|
||||||
|
Ok(true)
|
||||||
|
}
|
||||||
|
|
||||||
/// Whether a manifest exists for `id`. Cheap — no chunk reads.
|
/// Whether a manifest exists for `id`. Cheap — no chunk reads.
|
||||||
pub async fn contains(&self, id: &BlobId) -> Result<bool> {
|
pub async fn contains(&self, id: &BlobId) -> Result<bool> {
|
||||||
match tokio::fs::metadata(self.manifest_path(id)).await {
|
match tokio::fs::metadata(self.manifest_path(id)).await {
|
||||||
@@ -780,6 +873,113 @@ mod tests {
|
|||||||
assert_eq!(m.chunks.len(), 2);
|
assert_eq!(m.chunks.len(), 2);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ── Phase 2c: streaming APIs ────────────────────────────────────
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn put_stream_produces_same_hash_as_put_bytes() {
|
||||||
|
let (_tmp, store) = open_store();
|
||||||
|
let data = vec![0x77u8; CHUNK_SIZE * 2 + 500]; // 3 chunks
|
||||||
|
// Put via the bounded API for the baseline hash.
|
||||||
|
let expected = store.put_bytes(&data).await.unwrap();
|
||||||
|
// Put via the streaming API from a Cursor. Must produce the same
|
||||||
|
// BlobId — streaming and bounded are the same content addressing.
|
||||||
|
let cursor = std::io::Cursor::new(data.clone());
|
||||||
|
let via_stream = store.put_stream(cursor).await.unwrap();
|
||||||
|
assert_eq!(via_stream, expected);
|
||||||
|
let manifest = store.load_manifest(&via_stream).await.unwrap().unwrap();
|
||||||
|
assert_eq!(manifest.chunks.len(), 3);
|
||||||
|
assert_eq!(manifest.total_size, data.len() as u64);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn put_stream_handles_empty_reader() {
|
||||||
|
let (_tmp, store) = open_store();
|
||||||
|
let cursor = std::io::Cursor::new(Vec::<u8>::new());
|
||||||
|
let id = store.put_stream(cursor).await.unwrap();
|
||||||
|
let expected = store.put_bytes(&[]).await.unwrap();
|
||||||
|
assert_eq!(id, expected);
|
||||||
|
let stat = store.stat(&id).await.unwrap().unwrap();
|
||||||
|
assert_eq!(stat.total_size, 0);
|
||||||
|
assert_eq!(stat.chunk_count, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn put_stream_handles_short_reads() {
|
||||||
|
// A reader that only serves 100 bytes per `read` call must still
|
||||||
|
// produce full CHUNK_SIZE chunks — the outer loop refills.
|
||||||
|
struct Trickle {
|
||||||
|
buf: Vec<u8>,
|
||||||
|
pos: usize,
|
||||||
|
}
|
||||||
|
impl AsyncRead for Trickle {
|
||||||
|
fn poll_read(
|
||||||
|
mut self: std::pin::Pin<&mut Self>,
|
||||||
|
_cx: &mut std::task::Context<'_>,
|
||||||
|
dst: &mut tokio::io::ReadBuf<'_>,
|
||||||
|
) -> std::task::Poll<std::io::Result<()>> {
|
||||||
|
if self.pos >= self.buf.len() {
|
||||||
|
return std::task::Poll::Ready(Ok(()));
|
||||||
|
}
|
||||||
|
let take = 100.min(self.buf.len() - self.pos).min(dst.remaining());
|
||||||
|
dst.put_slice(&self.buf[self.pos..self.pos + take]);
|
||||||
|
self.pos += take;
|
||||||
|
std::task::Poll::Ready(Ok(()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let (_tmp, store) = open_store();
|
||||||
|
let data = vec![0xa5u8; CHUNK_SIZE + 300]; // exactly one full + one partial
|
||||||
|
let reader = Trickle {
|
||||||
|
buf: data.clone(),
|
||||||
|
pos: 0,
|
||||||
|
};
|
||||||
|
let via_stream = store.put_stream(reader).await.unwrap();
|
||||||
|
let via_bytes = store.put_bytes(&data).await.unwrap();
|
||||||
|
assert_eq!(via_stream, via_bytes);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn stream_to_writes_full_blob() {
|
||||||
|
let (_tmp, store) = open_store();
|
||||||
|
let data = vec![0x33u8; CHUNK_SIZE * 2]; // 2 full chunks
|
||||||
|
let id = store.put_bytes(&data).await.unwrap();
|
||||||
|
let mut sink: Vec<u8> = Vec::new();
|
||||||
|
let ok = store.stream_to(&id, &mut sink).await.unwrap();
|
||||||
|
assert!(ok);
|
||||||
|
assert_eq!(sink, data);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn stream_to_returns_false_when_missing() {
|
||||||
|
let (_tmp, store) = open_store();
|
||||||
|
let ghost = BlobId::from_bytes([0u8; 32]);
|
||||||
|
let mut sink: Vec<u8> = Vec::new();
|
||||||
|
let ok = store.stream_to(&ghost, &mut sink).await.unwrap();
|
||||||
|
assert!(!ok);
|
||||||
|
assert!(sink.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn stream_to_detects_chunk_corruption() {
|
||||||
|
let (_tmp, store) = open_store();
|
||||||
|
let id = store.put_bytes(b"corrupt me").await.unwrap();
|
||||||
|
let hex = id.to_hex();
|
||||||
|
let bucket = store.root().join("chunks").join(&hex[..2]);
|
||||||
|
let entries: Vec<_> = std::fs::read_dir(&bucket)
|
||||||
|
.unwrap()
|
||||||
|
.filter_map(|r| r.ok())
|
||||||
|
.collect();
|
||||||
|
assert_eq!(entries.len(), 1);
|
||||||
|
std::fs::write(entries[0].path(), b"tampered").unwrap();
|
||||||
|
|
||||||
|
let mut sink: Vec<u8> = Vec::new();
|
||||||
|
let err = store.stream_to(&id, &mut sink).await.unwrap_err();
|
||||||
|
assert!(
|
||||||
|
err.to_string().contains("chunk hash mismatch"),
|
||||||
|
"unexpected error: {err}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
/// Recursive count of regular files under `root`. Test helper.
|
/// Recursive count of regular files under `root`. Test helper.
|
||||||
fn count_files_under(root: &Path) -> usize {
|
fn count_files_under(root: &Path) -> usize {
|
||||||
if !root.exists() {
|
if !root.exists() {
|
||||||
|
|||||||
+248
-506
@@ -31,6 +31,12 @@ use anyhow::{bail, Context, Result};
|
|||||||
use quinn::{Connection, ConnectionError};
|
use quinn::{Connection, ConnectionError};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||||||
|
|
||||||
|
/// First byte of a streaming-method reply that indicates "success — the
|
||||||
|
/// content follows". Distinct from every [`ErrorCode`] value so the
|
||||||
|
/// client can trivially route on this single byte.
|
||||||
|
pub const STREAM_STATUS_OK: u8 = 0x00;
|
||||||
|
|
||||||
/// Cap on a single request or response, including the method tag.
|
/// Cap on a single request or response, including the method tag.
|
||||||
/// 16 MiB is generous enough to hold one 4 MiB blob chunk with plenty
|
/// 16 MiB is generous enough to hold one 4 MiB blob chunk with plenty
|
||||||
@@ -57,6 +63,19 @@ pub enum Method {
|
|||||||
/// `payload`: 32-byte `BlobId`. Reply: JSON `BlobManifest` or
|
/// `payload`: 32-byte `BlobId`. Reply: JSON `BlobManifest` or
|
||||||
/// single-byte [`ErrorCode::NotFound`].
|
/// single-byte [`ErrorCode::NotFound`].
|
||||||
BlobLoadManifest = 0x06,
|
BlobLoadManifest = 0x06,
|
||||||
|
/// Streaming variant of [`Method::BlobPut`] (Phase 2c). Wire:
|
||||||
|
/// method_byte followed by arbitrarily-many bytes of blob content
|
||||||
|
/// until the client half-closes the send stream. Reply:
|
||||||
|
/// `STREAM_STATUS_OK` (1 byte) followed by 32-byte `BlobId`, OR a
|
||||||
|
/// single-byte error code. Unlike `BlobPut` this is not bounded by
|
||||||
|
/// [`MAX_MESSAGE_BYTES`]; multi-GB blobs are the target workload.
|
||||||
|
BlobPutStream = 0x07,
|
||||||
|
/// Streaming variant of [`Method::BlobGet`] (Phase 2c). Wire:
|
||||||
|
/// method_byte followed by 32-byte `BlobId`. Reply:
|
||||||
|
/// `STREAM_STATUS_OK` (1 byte) followed by the blob content
|
||||||
|
/// streamed until the server half-closes, OR a single-byte error
|
||||||
|
/// code (typically [`ErrorCode::NotFound`]).
|
||||||
|
BlobGetStream = 0x08,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Method {
|
impl Method {
|
||||||
@@ -70,10 +89,19 @@ impl Method {
|
|||||||
0x04 => Some(Method::BlobGet),
|
0x04 => Some(Method::BlobGet),
|
||||||
0x05 => Some(Method::BlobPut),
|
0x05 => Some(Method::BlobPut),
|
||||||
0x06 => Some(Method::BlobLoadManifest),
|
0x06 => Some(Method::BlobLoadManifest),
|
||||||
|
0x07 => Some(Method::BlobPutStream),
|
||||||
|
0x08 => Some(Method::BlobGetStream),
|
||||||
_ => None,
|
_ => None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Whether this method uses the streaming wire format (status byte
|
||||||
|
/// + arbitrary-length content on the reply). Non-streaming methods
|
||||||
|
/// use the bounded `payload | reply` shape with `read_to_end`.
|
||||||
|
pub fn is_streaming(self) -> bool {
|
||||||
|
matches!(self, Method::BlobPutStream | Method::BlobGetStream)
|
||||||
|
}
|
||||||
|
|
||||||
/// Byte tag as an owned u8. `as u8` also works; this exists for symmetry.
|
/// Byte tag as an owned u8. `as u8` also works; this exists for symmetry.
|
||||||
pub fn as_byte(self) -> u8 {
|
pub fn as_byte(self) -> u8 {
|
||||||
self as u8
|
self as u8
|
||||||
@@ -278,6 +306,16 @@ impl RpcRouter {
|
|||||||
None => Ok(HandlerOutcome::Error(ErrorCode::NotFound)),
|
None => Ok(HandlerOutcome::Error(ErrorCode::NotFound)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
Method::BlobPutStream | Method::BlobGetStream => {
|
||||||
|
// Streaming methods go through a different wire path
|
||||||
|
// (`serve_connection` peeks at the method tag and hands
|
||||||
|
// the raw send/recv streams to the streaming handler).
|
||||||
|
// Reaching this arm would mean the caller tried to
|
||||||
|
// dispatch a streaming method through the bounded
|
||||||
|
// request path — reject with UnknownMethod-shaped
|
||||||
|
// error so it's obvious what happened.
|
||||||
|
Ok(HandlerOutcome::Error(ErrorCode::InvalidRequest))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -306,16 +344,218 @@ pub async fn serve_connection(conn: Connection, router: Arc<RpcRouter>) -> Resul
|
|||||||
| Err(ConnectionError::TimedOut) => return Ok(()),
|
| Err(ConnectionError::TimedOut) => return Ok(()),
|
||||||
Err(e) => return Err(anyhow::Error::from(e)),
|
Err(e) => return Err(anyhow::Error::from(e)),
|
||||||
};
|
};
|
||||||
let request = recv
|
|
||||||
.read_to_end(MAX_MESSAGE_BYTES)
|
// Peek at the method tag byte to decide whether to hand the
|
||||||
|
// stream off to a streaming handler or drain it into a bounded
|
||||||
|
// request buffer.
|
||||||
|
let mut tag = [0u8; 1];
|
||||||
|
match recv.read_exact(&mut tag).await {
|
||||||
|
Ok(()) => {}
|
||||||
|
Err(_) => {
|
||||||
|
let _ = send.write_all(&[ErrorCode::EmptyRequest.as_byte()]).await;
|
||||||
|
let _ = send.finish();
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
match Method::from_byte(tag[0]) {
|
||||||
|
Some(Method::BlobPutStream) => {
|
||||||
|
if let Err(e) = handle_blob_put_stream(&router, recv, send).await {
|
||||||
|
tracing::warn!(error = %e, "BlobPutStream handler failed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(Method::BlobGetStream) => {
|
||||||
|
if let Err(e) = handle_blob_get_stream(&router, recv, send).await {
|
||||||
|
tracing::warn!(error = %e, "BlobGetStream handler failed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(_) => {
|
||||||
|
// Bounded methods: drain the remainder of the request
|
||||||
|
// into memory and dispatch as before.
|
||||||
|
let rest = recv
|
||||||
|
.read_to_end(MAX_MESSAGE_BYTES.saturating_sub(1))
|
||||||
.await
|
.await
|
||||||
.context("reading RPC request")?;
|
.context("reading bounded RPC request")?;
|
||||||
|
let mut request = Vec::with_capacity(1 + rest.len());
|
||||||
|
request.push(tag[0]);
|
||||||
|
request.extend_from_slice(&rest);
|
||||||
let reply = dispatch(&router, &request).await;
|
let reply = dispatch(&router, &request).await;
|
||||||
send.write_all(&reply)
|
send.write_all(&reply).await.context("writing RPC reply")?;
|
||||||
.await
|
|
||||||
.context("writing RPC reply")?;
|
|
||||||
send.finish().context("finishing RPC send stream")?;
|
send.finish().context("finishing RPC send stream")?;
|
||||||
}
|
}
|
||||||
|
None => {
|
||||||
|
let _ = send.write_all(&[ErrorCode::UnknownMethod.as_byte()]).await;
|
||||||
|
let _ = send.finish();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Streaming handler for [`Method::BlobPutStream`]. Feeds the incoming
|
||||||
|
/// bytes straight into `BlobStore::put_stream` — memory ceiling is one
|
||||||
|
/// chunk buffer regardless of blob size.
|
||||||
|
async fn handle_blob_put_stream(
|
||||||
|
router: &RpcRouter,
|
||||||
|
recv: quinn::RecvStream,
|
||||||
|
mut send: quinn::SendStream,
|
||||||
|
) -> Result<()> {
|
||||||
|
let store = match router.blob_store() {
|
||||||
|
Some(s) => s.clone(),
|
||||||
|
None => {
|
||||||
|
send.write_all(&[ErrorCode::NotConfigured.as_byte()]).await?;
|
||||||
|
send.finish()?;
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
match store.put_stream(recv).await {
|
||||||
|
Ok(id) => {
|
||||||
|
let mut reply = Vec::with_capacity(33);
|
||||||
|
reply.push(STREAM_STATUS_OK);
|
||||||
|
reply.extend_from_slice(id.as_bytes());
|
||||||
|
send.write_all(&reply).await?;
|
||||||
|
send.finish()?;
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!(error = %e, "BlobPutStream put_stream failed");
|
||||||
|
let _ = send.write_all(&[ErrorCode::HandlerFailure.as_byte()]).await;
|
||||||
|
let _ = send.finish();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Streaming handler for [`Method::BlobGetStream`]. Verifies the blob
|
||||||
|
/// exists (writes `NotFound` on absence), then pipes each chunk from
|
||||||
|
/// disk straight into the send stream. Callers see:
|
||||||
|
/// `STREAM_STATUS_OK` (1 byte) followed by the blob bytes.
|
||||||
|
async fn handle_blob_get_stream(
|
||||||
|
router: &RpcRouter,
|
||||||
|
mut recv: quinn::RecvStream,
|
||||||
|
mut send: quinn::SendStream,
|
||||||
|
) -> Result<()> {
|
||||||
|
let store = match router.blob_store() {
|
||||||
|
Some(s) => s.clone(),
|
||||||
|
None => {
|
||||||
|
send.write_all(&[ErrorCode::NotConfigured.as_byte()]).await?;
|
||||||
|
send.finish()?;
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let mut id_bytes = [0u8; 32];
|
||||||
|
if recv.read_exact(&mut id_bytes).await.is_err() {
|
||||||
|
send.write_all(&[ErrorCode::InvalidRequest.as_byte()]).await?;
|
||||||
|
send.finish()?;
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
let id = BlobId::from_bytes(id_bytes);
|
||||||
|
match store.load_manifest(&id).await? {
|
||||||
|
None => {
|
||||||
|
send.write_all(&[ErrorCode::NotFound.as_byte()]).await?;
|
||||||
|
}
|
||||||
|
Some(_) => {
|
||||||
|
send.write_all(&[STREAM_STATUS_OK]).await?;
|
||||||
|
store.stream_to(&id, &mut send).await?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
send.finish()?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Streaming client helpers ─────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Upload a blob by streaming from `reader` into the peer's store.
|
||||||
|
/// Never buffers the whole blob — memory usage is dominated by
|
||||||
|
/// tokio's default copy buffer. Returns the peer-assigned BlobId,
|
||||||
|
/// which is NOT locally re-verified (the peer doesn't get the whole
|
||||||
|
/// content in memory at once so we can't cheaply hash from our side);
|
||||||
|
/// callers who need cryptographic proof of end-to-end integrity
|
||||||
|
/// should follow up with `call_blob_stat` and compare hashes.
|
||||||
|
pub async fn call_blob_put_stream<R>(
|
||||||
|
conn: &Connection,
|
||||||
|
mut reader: R,
|
||||||
|
) -> Result<BlobId>
|
||||||
|
where
|
||||||
|
R: AsyncRead + Unpin,
|
||||||
|
{
|
||||||
|
let (mut send, mut recv) = conn
|
||||||
|
.open_bi()
|
||||||
|
.await
|
||||||
|
.context("opening bidi stream for BlobPutStream")?;
|
||||||
|
send.write_all(&[Method::BlobPutStream.as_byte()])
|
||||||
|
.await
|
||||||
|
.context("writing method tag")?;
|
||||||
|
tokio::io::copy(&mut reader, &mut send)
|
||||||
|
.await
|
||||||
|
.context("streaming blob content to peer")?;
|
||||||
|
send.finish().context("finishing BlobPutStream send")?;
|
||||||
|
let mut status = [0u8; 1];
|
||||||
|
recv.read_exact(&mut status)
|
||||||
|
.await
|
||||||
|
.context("reading BlobPutStream status byte")?;
|
||||||
|
if let Some(code) = decode_error(status[0]) {
|
||||||
|
bail!("peer replied with error: {}", code.describe());
|
||||||
|
}
|
||||||
|
if status[0] != STREAM_STATUS_OK {
|
||||||
|
bail!(
|
||||||
|
"peer replied with unknown status byte 0x{:02x}",
|
||||||
|
status[0]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let mut id_bytes = [0u8; 32];
|
||||||
|
recv.read_exact(&mut id_bytes)
|
||||||
|
.await
|
||||||
|
.context("reading BlobId from BlobPutStream reply")?;
|
||||||
|
Ok(BlobId::from_bytes(id_bytes))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fetch a blob by streaming its bytes into `writer`. Returns
|
||||||
|
/// `Ok(false)` when the peer replies [`ErrorCode::NotFound`]; other
|
||||||
|
/// error codes surface as `Err`. `writer` receives exactly the blob
|
||||||
|
/// content — no leading status byte, no framing.
|
||||||
|
pub async fn call_blob_get_stream<W>(
|
||||||
|
conn: &Connection,
|
||||||
|
id: &BlobId,
|
||||||
|
writer: &mut W,
|
||||||
|
) -> Result<bool>
|
||||||
|
where
|
||||||
|
W: AsyncWrite + Unpin,
|
||||||
|
{
|
||||||
|
let (mut send, mut recv) = conn
|
||||||
|
.open_bi()
|
||||||
|
.await
|
||||||
|
.context("opening bidi stream for BlobGetStream")?;
|
||||||
|
let mut req = Vec::with_capacity(33);
|
||||||
|
req.push(Method::BlobGetStream.as_byte());
|
||||||
|
req.extend_from_slice(id.as_bytes());
|
||||||
|
send.write_all(&req)
|
||||||
|
.await
|
||||||
|
.context("writing BlobGetStream request")?;
|
||||||
|
send.finish().context("finishing BlobGetStream send")?;
|
||||||
|
|
||||||
|
let mut status = [0u8; 1];
|
||||||
|
recv.read_exact(&mut status)
|
||||||
|
.await
|
||||||
|
.context("reading BlobGetStream status byte")?;
|
||||||
|
if let Some(code) = decode_error(status[0]) {
|
||||||
|
if code == ErrorCode::NotFound {
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
bail!("peer replied with error: {}", code.describe());
|
||||||
|
}
|
||||||
|
if status[0] != STREAM_STATUS_OK {
|
||||||
|
bail!(
|
||||||
|
"peer replied with unknown status byte 0x{:02x}",
|
||||||
|
status[0]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
tokio::io::copy(&mut recv, writer)
|
||||||
|
.await
|
||||||
|
.context("streaming blob content from peer")?;
|
||||||
|
writer
|
||||||
|
.flush()
|
||||||
|
.await
|
||||||
|
.context("flushing destination after BlobGetStream")?;
|
||||||
|
Ok(true)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Turn a raw wire-format request into a response — either the
|
/// Turn a raw wire-format request into a response — either the
|
||||||
@@ -504,503 +744,5 @@ pub async fn call_blob_load_manifest(
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
#[path = "rpc/tests.rs"]
|
||||||
use super::*;
|
mod tests;
|
||||||
use crate::cluster::transport::{NodeIdentity, QuicClient, QuicServer};
|
|
||||||
use crate::config::{ClusterConfig, PeerEntry};
|
|
||||||
use std::net::SocketAddr;
|
|
||||||
use std::sync::atomic::{AtomicU16, Ordering};
|
|
||||||
use std::time::Duration;
|
|
||||||
|
|
||||||
/// Dedicated port range for RPC tests, distinct from gossip (41000+)
|
|
||||||
/// and transport (42000+) so parallel tests never conflict.
|
|
||||||
static NEXT_PORT: AtomicU16 = AtomicU16::new(43001);
|
|
||||||
fn next_port() -> u16 {
|
|
||||||
NEXT_PORT.fetch_add(1, Ordering::Relaxed)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn loopback(port: u16) -> SocketAddr {
|
|
||||||
format!("127.0.0.1:{port}").parse().unwrap()
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn bootstrap_gossip(name: &str, port: u16) -> Arc<ClusterGossip> {
|
|
||||||
let cfg = ClusterConfig {
|
|
||||||
zone: "fabric-10g".into(),
|
|
||||||
bind_lan: Some(loopback(port)),
|
|
||||||
..Default::default()
|
|
||||||
};
|
|
||||||
Arc::new(ClusterGossip::bootstrap(&cfg, name).await.unwrap())
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn method_round_trips_byte_encoding() {
|
|
||||||
assert_eq!(Method::Ping.as_byte(), 0x01);
|
|
||||||
assert_eq!(Method::PeerStatus.as_byte(), 0x02);
|
|
||||||
assert_eq!(Method::BlobStat.as_byte(), 0x03);
|
|
||||||
assert_eq!(Method::BlobGet.as_byte(), 0x04);
|
|
||||||
assert_eq!(Method::BlobPut.as_byte(), 0x05);
|
|
||||||
assert_eq!(Method::BlobLoadManifest.as_byte(), 0x06);
|
|
||||||
for m in [
|
|
||||||
Method::Ping,
|
|
||||||
Method::PeerStatus,
|
|
||||||
Method::BlobStat,
|
|
||||||
Method::BlobGet,
|
|
||||||
Method::BlobPut,
|
|
||||||
Method::BlobLoadManifest,
|
|
||||||
] {
|
|
||||||
assert_eq!(Method::from_byte(m.as_byte()), Some(m));
|
|
||||||
}
|
|
||||||
assert_eq!(Method::from_byte(0x00), None);
|
|
||||||
assert_eq!(Method::from_byte(0xff), None);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn open_blob_store() -> (tempfile::TempDir, Arc<BlobStore>) {
|
|
||||||
let tmp = tempfile::TempDir::new().unwrap();
|
|
||||||
let store = Arc::new(BlobStore::open(tmp.path().to_path_buf()).unwrap());
|
|
||||||
(tmp, store)
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn router_with_blobs(name: &str, port: u16) -> (tempfile::TempDir, Arc<RpcRouter>) {
|
|
||||||
let gossip = bootstrap_gossip(name, port).await;
|
|
||||||
let (tmp, store) = open_blob_store();
|
|
||||||
let router = Arc::new(
|
|
||||||
RpcRouter::new(gossip, name.into(), "fabric-10g".into())
|
|
||||||
.with_blob_store(store),
|
|
||||||
);
|
|
||||||
(tmp, router)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn dispatch_returns_pong_for_ping() {
|
|
||||||
let gossip = bootstrap_gossip("solo", next_port()).await;
|
|
||||||
let router = RpcRouter::new(gossip.clone(), "solo".into(), "test-zone".into());
|
|
||||||
|
|
||||||
// Direct dispatch — no network involved.
|
|
||||||
let reply = dispatch(&router, &[Method::Ping.as_byte(), b'h', b'i']).await;
|
|
||||||
assert_eq!(reply, b"pong:hi");
|
|
||||||
gossip.peer("nobody").await; // touch to keep gossip alive
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn dispatch_returns_json_for_peer_status() {
|
|
||||||
let gossip = bootstrap_gossip("architect", next_port()).await;
|
|
||||||
let router = RpcRouter::new(
|
|
||||||
gossip.clone(),
|
|
||||||
"architect".into(),
|
|
||||||
"fabric-10g".into(),
|
|
||||||
);
|
|
||||||
let reply = dispatch(&router, &[Method::PeerStatus.as_byte()]).await;
|
|
||||||
let decoded: PeerStatusReply = serde_json::from_slice(&reply).unwrap();
|
|
||||||
assert_eq!(decoded.local_name, "architect");
|
|
||||||
assert_eq!(decoded.local_zone, "fabric-10g");
|
|
||||||
// Solo node — no peers yet.
|
|
||||||
assert!(decoded.peers.is_empty(), "solo node has no peers");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn dispatch_returns_empty_request_error() {
|
|
||||||
let gossip = bootstrap_gossip("solo", next_port()).await;
|
|
||||||
let router = RpcRouter::new(gossip.clone(), "solo".into(), "z".into());
|
|
||||||
let reply = dispatch(&router, &[]).await;
|
|
||||||
assert_eq!(reply, vec![ErrorCode::EmptyRequest.as_byte()]);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn dispatch_returns_unknown_method_error() {
|
|
||||||
let gossip = bootstrap_gossip("solo", next_port()).await;
|
|
||||||
let router = RpcRouter::new(gossip.clone(), "solo".into(), "z".into());
|
|
||||||
let reply = dispatch(&router, &[0xab, 0x01, 0x02]).await;
|
|
||||||
assert_eq!(reply, vec![ErrorCode::UnknownMethod.as_byte()]);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn rpc_call_rejects_oversize_payload() {
|
|
||||||
// No network involved — the size check runs client-side before we
|
|
||||||
// even try to open the bidi stream. Use a dummy connection built
|
|
||||||
// via generate_test_pair; we won't actually connect.
|
|
||||||
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
|
||||||
let server = QuicServer::bind(loopback(0), id_b).unwrap();
|
|
||||||
let server_addr = server.local_addr().unwrap();
|
|
||||||
let _accept = tokio::spawn(async move {
|
|
||||||
let _ = server.accept().await;
|
|
||||||
});
|
|
||||||
let client = QuicClient::new(loopback(0), id_a).unwrap();
|
|
||||||
let conn = client.connect(server_addr, "b").await.unwrap();
|
|
||||||
let huge = vec![0u8; MAX_MESSAGE_BYTES];
|
|
||||||
let err = rpc_call(&conn, Method::Ping, &huge)
|
|
||||||
.await
|
|
||||||
.err()
|
|
||||||
.expect("must reject oversize");
|
|
||||||
assert!(err.to_string().contains("exceeds cap"));
|
|
||||||
conn.close(quinn::VarInt::from_u32(0), b"done");
|
|
||||||
client.shutdown().await;
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn end_to_end_ping_and_peer_status_over_real_quic() {
|
|
||||||
// Two full nodes: A runs gossip + a QuicServer serving RpcRouter.
|
|
||||||
// B is a client that dials A and calls both RPCs.
|
|
||||||
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
|
||||||
|
|
||||||
let port_a = next_port();
|
|
||||||
let gossip_a = bootstrap_gossip("a", port_a).await;
|
|
||||||
let router = Arc::new(RpcRouter::new(
|
|
||||||
gossip_a.clone(),
|
|
||||||
"a".into(),
|
|
||||||
"fabric-10g".into(),
|
|
||||||
));
|
|
||||||
|
|
||||||
// A's advertised state so the PeerStatus reply exercises the
|
|
||||||
// gossip → PeerView path even though B isn't in the peer table.
|
|
||||||
gossip_a.set_hot_used(4096).await;
|
|
||||||
gossip_a.set_hot_max(1_000_000).await;
|
|
||||||
gossip_a
|
|
||||||
.set_warm_projects(&["osobh/clawverse", "osobh/clawmates"])
|
|
||||||
.await;
|
|
||||||
|
|
||||||
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
|
||||||
let server_addr = server.local_addr().unwrap();
|
|
||||||
let accept_task = tokio::spawn(async move {
|
|
||||||
if let Some(Ok(conn)) = server.accept().await {
|
|
||||||
let _ = serve_connection(conn, router).await;
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
|
||||||
let conn = client.connect(server_addr, "a").await.unwrap();
|
|
||||||
|
|
||||||
// Ping.
|
|
||||||
let pong = call_ping(&conn, b"hello").await.unwrap();
|
|
||||||
assert_eq!(pong, b"hello");
|
|
||||||
|
|
||||||
// PeerStatus.
|
|
||||||
let status = call_peer_status(&conn).await.unwrap();
|
|
||||||
assert_eq!(status.local_name, "a");
|
|
||||||
assert_eq!(status.local_zone, "fabric-10g");
|
|
||||||
assert!(status.peers.is_empty(), "A has no peers configured");
|
|
||||||
|
|
||||||
conn.close(quinn::VarInt::from_u32(0), b"done");
|
|
||||||
client.shutdown().await;
|
|
||||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
|
||||||
accept_task.abort();
|
|
||||||
// Keep gossip_a alive to the end so its background task doesn't
|
|
||||||
// drop mid-serve.
|
|
||||||
let _ = gossip_a.peers().await;
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn peer_status_reflects_peer_gossip_state() {
|
|
||||||
// A and B both run gossip; A serves RPC. When A's PeerStatus is
|
|
||||||
// called by a third-party client, the reply's `peers` field
|
|
||||||
// contains B (as long as gossip has converged).
|
|
||||||
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
|
||||||
let port_a = next_port();
|
|
||||||
let port_b = next_port();
|
|
||||||
|
|
||||||
let cfg_a = ClusterConfig {
|
|
||||||
zone: "fabric-10g".into(),
|
|
||||||
bind_lan: Some(loopback(port_a)),
|
|
||||||
..Default::default()
|
|
||||||
};
|
|
||||||
let cfg_b = ClusterConfig {
|
|
||||||
zone: "lan-1g".into(),
|
|
||||||
bind_lan: Some(loopback(port_b)),
|
|
||||||
peers: vec![PeerEntry {
|
|
||||||
name: "a".into(),
|
|
||||||
zone: "fabric-10g".into(),
|
|
||||||
lan_addr: Some(loopback(port_a)),
|
|
||||||
tailscale_addr: None,
|
|
||||||
}],
|
|
||||||
..Default::default()
|
|
||||||
};
|
|
||||||
let gossip_a = Arc::new(ClusterGossip::bootstrap(&cfg_a, "a").await.unwrap());
|
|
||||||
let gossip_b = Arc::new(ClusterGossip::bootstrap(&cfg_b, "b").await.unwrap());
|
|
||||||
|
|
||||||
let router = Arc::new(RpcRouter::new(
|
|
||||||
gossip_a.clone(),
|
|
||||||
"a".into(),
|
|
||||||
"fabric-10g".into(),
|
|
||||||
));
|
|
||||||
|
|
||||||
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
|
||||||
let server_addr = server.local_addr().unwrap();
|
|
||||||
let accept_task = tokio::spawn(async move {
|
|
||||||
if let Some(Ok(conn)) = server.accept().await {
|
|
||||||
let _ = serve_connection(conn, router).await;
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
// Wait for gossip convergence: A must see B.
|
|
||||||
let deadline = std::time::Instant::now() + Duration::from_secs(10);
|
|
||||||
loop {
|
|
||||||
if let Some(v) = gossip_a.peer("b").await {
|
|
||||||
if v.alive {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if std::time::Instant::now() >= deadline {
|
|
||||||
panic!("A never saw B alive within 10s");
|
|
||||||
}
|
|
||||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
|
||||||
let conn = client.connect(server_addr, "a").await.unwrap();
|
|
||||||
let status = call_peer_status(&conn).await.unwrap();
|
|
||||||
assert_eq!(status.local_name, "a");
|
|
||||||
assert_eq!(status.peers.len(), 1, "A should report exactly B");
|
|
||||||
assert_eq!(status.peers[0].name, "b");
|
|
||||||
assert_eq!(status.peers[0].zone, "lan-1g");
|
|
||||||
assert!(status.peers[0].alive);
|
|
||||||
|
|
||||||
conn.close(quinn::VarInt::from_u32(0), b"done");
|
|
||||||
client.shutdown().await;
|
|
||||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
|
||||||
accept_task.abort();
|
|
||||||
// Keep gossip services alive until end.
|
|
||||||
drop(gossip_b);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn error_code_describe_covers_all_variants() {
|
|
||||||
assert_eq!(ErrorCode::EmptyRequest.describe(), "empty request");
|
|
||||||
assert_eq!(ErrorCode::UnknownMethod.describe(), "unknown method");
|
|
||||||
assert_eq!(ErrorCode::HandlerFailure.describe(), "handler failure");
|
|
||||||
assert_eq!(ErrorCode::NotFound.describe(), "not found");
|
|
||||||
assert_eq!(ErrorCode::InvalidRequest.describe(), "invalid request");
|
|
||||||
assert_eq!(
|
|
||||||
ErrorCode::NotConfigured.describe(),
|
|
||||||
"server subsystem not configured"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn decode_error_covers_all_known_codes() {
|
|
||||||
for code in [
|
|
||||||
ErrorCode::EmptyRequest,
|
|
||||||
ErrorCode::UnknownMethod,
|
|
||||||
ErrorCode::HandlerFailure,
|
|
||||||
ErrorCode::NotFound,
|
|
||||||
ErrorCode::InvalidRequest,
|
|
||||||
ErrorCode::NotConfigured,
|
|
||||||
] {
|
|
||||||
assert_eq!(decode_error(code.as_byte()), Some(code));
|
|
||||||
}
|
|
||||||
assert_eq!(decode_error(0x00), None);
|
|
||||||
assert_eq!(decode_error(0xff), None);
|
|
||||||
}
|
|
||||||
|
|
||||||
// ── Phase 2b: Blob RPC ─────────────────────────────────────────────
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn blob_rpcs_return_not_configured_without_store() {
|
|
||||||
// Router built via `new` alone (no `.with_blob_store`) must
|
|
||||||
// refuse Blob* methods with a well-known error code.
|
|
||||||
let gossip = bootstrap_gossip("solo", next_port()).await;
|
|
||||||
let router = RpcRouter::new(gossip, "solo".into(), "z".into());
|
|
||||||
for method in [
|
|
||||||
Method::BlobStat,
|
|
||||||
Method::BlobGet,
|
|
||||||
Method::BlobPut,
|
|
||||||
Method::BlobLoadManifest,
|
|
||||||
] {
|
|
||||||
let mut req = vec![method.as_byte()];
|
|
||||||
req.extend_from_slice(&[0u8; 32]);
|
|
||||||
let reply = dispatch(&router, &req).await;
|
|
||||||
assert_eq!(
|
|
||||||
reply,
|
|
||||||
vec![ErrorCode::NotConfigured.as_byte()],
|
|
||||||
"method {method:?} should be NotConfigured without a store"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn blob_stat_returns_not_found_for_missing() {
|
|
||||||
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
|
||||||
let missing = BlobId::from_bytes([0u8; 32]);
|
|
||||||
let mut req = vec![Method::BlobStat.as_byte()];
|
|
||||||
req.extend_from_slice(missing.as_bytes());
|
|
||||||
let reply = dispatch(&router, &req).await;
|
|
||||||
assert_eq!(reply, vec![ErrorCode::NotFound.as_byte()]);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn blob_stat_returns_json_for_existing() {
|
|
||||||
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
|
||||||
let id = router
|
|
||||||
.blob_store()
|
|
||||||
.unwrap()
|
|
||||||
.put_bytes(b"tiny content")
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
let mut req = vec![Method::BlobStat.as_byte()];
|
|
||||||
req.extend_from_slice(id.as_bytes());
|
|
||||||
let reply = dispatch(&router, &req).await;
|
|
||||||
let stat: BlobStat = serde_json::from_slice(&reply).unwrap();
|
|
||||||
assert_eq!(stat.total_size, b"tiny content".len() as u64);
|
|
||||||
assert_eq!(stat.chunk_count, 1);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn blob_stat_returns_invalid_request_for_bad_length() {
|
|
||||||
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
|
||||||
// Payload is only 5 bytes; a valid BlobId is 32.
|
|
||||||
let req = vec![Method::BlobStat.as_byte(), 1, 2, 3, 4, 5];
|
|
||||||
let reply = dispatch(&router, &req).await;
|
|
||||||
assert_eq!(reply, vec![ErrorCode::InvalidRequest.as_byte()]);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn blob_get_returns_content_bytes() {
|
|
||||||
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
|
||||||
let payload: &[u8] = b"contents to retrieve";
|
|
||||||
let id = router.blob_store().unwrap().put_bytes(payload).await.unwrap();
|
|
||||||
let mut req = vec![Method::BlobGet.as_byte()];
|
|
||||||
req.extend_from_slice(id.as_bytes());
|
|
||||||
let reply = dispatch(&router, &req).await;
|
|
||||||
assert_eq!(reply, payload);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn blob_put_stores_bytes_and_returns_hash() {
|
|
||||||
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
|
||||||
let payload = b"put via rpc";
|
|
||||||
let mut req = vec![Method::BlobPut.as_byte()];
|
|
||||||
req.extend_from_slice(payload);
|
|
||||||
let reply = dispatch(&router, &req).await;
|
|
||||||
assert_eq!(reply.len(), 32);
|
|
||||||
let mut id_bytes = [0u8; 32];
|
|
||||||
id_bytes.copy_from_slice(&reply);
|
|
||||||
let assigned = BlobId::from_bytes(id_bytes);
|
|
||||||
let expected = BlobId::from_bytes(blake3::hash(payload).into());
|
|
||||||
assert_eq!(assigned, expected);
|
|
||||||
// Round-trip: the bytes are now readable via the store.
|
|
||||||
let round = router
|
|
||||||
.blob_store()
|
|
||||||
.unwrap()
|
|
||||||
.get_bytes(&assigned)
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
assert_eq!(round.as_deref(), Some(payload.as_slice()));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn blob_load_manifest_returns_json_for_existing() {
|
|
||||||
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
|
||||||
let data = vec![0x77u8; 4 * 1024 * 1024 + 1]; // 2 chunks
|
|
||||||
let id = router.blob_store().unwrap().put_bytes(&data).await.unwrap();
|
|
||||||
let mut req = vec![Method::BlobLoadManifest.as_byte()];
|
|
||||||
req.extend_from_slice(id.as_bytes());
|
|
||||||
let reply = dispatch(&router, &req).await;
|
|
||||||
let manifest: BlobManifest = serde_json::from_slice(&reply).unwrap();
|
|
||||||
assert_eq!(manifest.blob_id, id);
|
|
||||||
assert_eq!(manifest.total_size, data.len() as u64);
|
|
||||||
assert_eq!(manifest.chunks.len(), 2);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn blob_load_manifest_returns_not_found_for_missing() {
|
|
||||||
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
|
||||||
let missing = BlobId::from_bytes([0u8; 32]);
|
|
||||||
let mut req = vec![Method::BlobLoadManifest.as_byte()];
|
|
||||||
req.extend_from_slice(missing.as_bytes());
|
|
||||||
let reply = dispatch(&router, &req).await;
|
|
||||||
assert_eq!(reply, vec![ErrorCode::NotFound.as_byte()]);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn end_to_end_blob_put_stat_get_over_real_quic() {
|
|
||||||
// The full loop: B → A over real QUIC + mTLS.
|
|
||||||
// 1. B puts a blob on A (BlobPut).
|
|
||||||
// 2. B queries stat + fetches it back (BlobStat + BlobGet).
|
|
||||||
// 3. B asks for the manifest (BlobLoadManifest).
|
|
||||||
// Every step goes through the actual wire, no shortcuts.
|
|
||||||
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
|
||||||
let (_tmp, router) = router_with_blobs("a", next_port()).await;
|
|
||||||
|
|
||||||
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
|
||||||
let server_addr = server.local_addr().unwrap();
|
|
||||||
let router_srv = router.clone();
|
|
||||||
let accept_task = tokio::spawn(async move {
|
|
||||||
if let Some(Ok(conn)) = server.accept().await {
|
|
||||||
let _ = serve_connection(conn, router_srv).await;
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
|
||||||
let conn = client.connect(server_addr, "a").await.unwrap();
|
|
||||||
|
|
||||||
let payload: &[u8] = b"cross-node payload";
|
|
||||||
let assigned = call_blob_put(&conn, payload).await.unwrap();
|
|
||||||
let expected = BlobId::from_bytes(blake3::hash(payload).into());
|
|
||||||
assert_eq!(assigned, expected);
|
|
||||||
|
|
||||||
let stat = call_blob_stat(&conn, &assigned).await.unwrap().unwrap();
|
|
||||||
assert_eq!(stat.total_size, payload.len() as u64);
|
|
||||||
assert_eq!(stat.chunk_count, 1);
|
|
||||||
|
|
||||||
let round = call_blob_get(&conn, &assigned).await.unwrap().unwrap();
|
|
||||||
assert_eq!(round, payload);
|
|
||||||
|
|
||||||
let manifest = call_blob_load_manifest(&conn, &assigned)
|
|
||||||
.await
|
|
||||||
.unwrap()
|
|
||||||
.unwrap();
|
|
||||||
assert_eq!(manifest.blob_id, assigned);
|
|
||||||
assert_eq!(manifest.chunks.len(), 1);
|
|
||||||
|
|
||||||
// NotFound path also works over the wire.
|
|
||||||
let ghost = BlobId::from_bytes([0u8; 32]);
|
|
||||||
assert!(call_blob_stat(&conn, &ghost).await.unwrap().is_none());
|
|
||||||
assert!(call_blob_get(&conn, &ghost).await.unwrap().is_none());
|
|
||||||
assert!(call_blob_load_manifest(&conn, &ghost)
|
|
||||||
.await
|
|
||||||
.unwrap()
|
|
||||||
.is_none());
|
|
||||||
|
|
||||||
conn.close(quinn::VarInt::from_u32(0), b"done");
|
|
||||||
client.shutdown().await;
|
|
||||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
|
||||||
accept_task.abort();
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn end_to_end_multi_chunk_blob_over_real_quic() {
|
|
||||||
// 6 MB blob → 2 chunks. Round-trips whole via BlobPut/Get.
|
|
||||||
// Also confirms the RPC layer's MAX_MESSAGE_BYTES bump from
|
|
||||||
// 16 KiB to 16 MiB actually took effect end-to-end.
|
|
||||||
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
|
||||||
let (_tmp, router) = router_with_blobs("a", next_port()).await;
|
|
||||||
|
|
||||||
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
|
||||||
let server_addr = server.local_addr().unwrap();
|
|
||||||
let router_srv = router.clone();
|
|
||||||
let accept_task = tokio::spawn(async move {
|
|
||||||
if let Some(Ok(conn)) = server.accept().await {
|
|
||||||
let _ = serve_connection(conn, router_srv).await;
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
|
||||||
let conn = client.connect(server_addr, "a").await.unwrap();
|
|
||||||
|
|
||||||
let payload: Vec<u8> = (0..6 * 1024 * 1024)
|
|
||||||
.map(|i| (i % 251) as u8)
|
|
||||||
.collect();
|
|
||||||
let assigned = call_blob_put(&conn, &payload).await.unwrap();
|
|
||||||
let manifest = call_blob_load_manifest(&conn, &assigned)
|
|
||||||
.await
|
|
||||||
.unwrap()
|
|
||||||
.unwrap();
|
|
||||||
assert_eq!(manifest.chunks.len(), 2, "6 MB should split into 2 chunks");
|
|
||||||
let round = call_blob_get(&conn, &assigned).await.unwrap().unwrap();
|
|
||||||
assert_eq!(round.len(), payload.len());
|
|
||||||
assert_eq!(round, payload);
|
|
||||||
|
|
||||||
conn.close(quinn::VarInt::from_u32(0), b"done");
|
|
||||||
client.shutdown().await;
|
|
||||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
|
||||||
accept_task.abort();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -0,0 +1,694 @@
|
|||||||
|
//! Tests split out of rpc.rs to keep the parent under the 1300-line ceiling.
|
||||||
|
//! Same visibility as `#[cfg(test)] mod tests { ... }` — everything reaches
|
||||||
|
//! parent private items via `use super::*`.
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
use crate::cluster::transport::{NodeIdentity, QuicClient, QuicServer};
|
||||||
|
use crate::config::{ClusterConfig, PeerEntry};
|
||||||
|
use std::net::SocketAddr;
|
||||||
|
use std::sync::atomic::{AtomicU16, Ordering};
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
/// Dedicated port range for RPC tests, distinct from gossip (41000+)
|
||||||
|
/// and transport (42000+) so parallel tests never conflict.
|
||||||
|
static NEXT_PORT: AtomicU16 = AtomicU16::new(43001);
|
||||||
|
fn next_port() -> u16 {
|
||||||
|
NEXT_PORT.fetch_add(1, Ordering::Relaxed)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn loopback(port: u16) -> SocketAddr {
|
||||||
|
format!("127.0.0.1:{port}").parse().unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn bootstrap_gossip(name: &str, port: u16) -> Arc<ClusterGossip> {
|
||||||
|
let cfg = ClusterConfig {
|
||||||
|
zone: "fabric-10g".into(),
|
||||||
|
bind_lan: Some(loopback(port)),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
Arc::new(ClusterGossip::bootstrap(&cfg, name).await.unwrap())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn method_round_trips_byte_encoding() {
|
||||||
|
assert_eq!(Method::Ping.as_byte(), 0x01);
|
||||||
|
assert_eq!(Method::PeerStatus.as_byte(), 0x02);
|
||||||
|
assert_eq!(Method::BlobStat.as_byte(), 0x03);
|
||||||
|
assert_eq!(Method::BlobGet.as_byte(), 0x04);
|
||||||
|
assert_eq!(Method::BlobPut.as_byte(), 0x05);
|
||||||
|
assert_eq!(Method::BlobLoadManifest.as_byte(), 0x06);
|
||||||
|
for m in [
|
||||||
|
Method::Ping,
|
||||||
|
Method::PeerStatus,
|
||||||
|
Method::BlobStat,
|
||||||
|
Method::BlobGet,
|
||||||
|
Method::BlobPut,
|
||||||
|
Method::BlobLoadManifest,
|
||||||
|
] {
|
||||||
|
assert_eq!(Method::from_byte(m.as_byte()), Some(m));
|
||||||
|
}
|
||||||
|
assert_eq!(Method::from_byte(0x00), None);
|
||||||
|
assert_eq!(Method::from_byte(0xff), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn open_blob_store() -> (tempfile::TempDir, Arc<BlobStore>) {
|
||||||
|
let tmp = tempfile::TempDir::new().unwrap();
|
||||||
|
let store = Arc::new(BlobStore::open(tmp.path().to_path_buf()).unwrap());
|
||||||
|
(tmp, store)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn router_with_blobs(name: &str, port: u16) -> (tempfile::TempDir, Arc<RpcRouter>) {
|
||||||
|
let gossip = bootstrap_gossip(name, port).await;
|
||||||
|
let (tmp, store) = open_blob_store();
|
||||||
|
let router = Arc::new(
|
||||||
|
RpcRouter::new(gossip, name.into(), "fabric-10g".into())
|
||||||
|
.with_blob_store(store),
|
||||||
|
);
|
||||||
|
(tmp, router)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn dispatch_returns_pong_for_ping() {
|
||||||
|
let gossip = bootstrap_gossip("solo", next_port()).await;
|
||||||
|
let router = RpcRouter::new(gossip.clone(), "solo".into(), "test-zone".into());
|
||||||
|
|
||||||
|
// Direct dispatch — no network involved.
|
||||||
|
let reply = dispatch(&router, &[Method::Ping.as_byte(), b'h', b'i']).await;
|
||||||
|
assert_eq!(reply, b"pong:hi");
|
||||||
|
gossip.peer("nobody").await; // touch to keep gossip alive
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn dispatch_returns_json_for_peer_status() {
|
||||||
|
let gossip = bootstrap_gossip("architect", next_port()).await;
|
||||||
|
let router = RpcRouter::new(
|
||||||
|
gossip.clone(),
|
||||||
|
"architect".into(),
|
||||||
|
"fabric-10g".into(),
|
||||||
|
);
|
||||||
|
let reply = dispatch(&router, &[Method::PeerStatus.as_byte()]).await;
|
||||||
|
let decoded: PeerStatusReply = serde_json::from_slice(&reply).unwrap();
|
||||||
|
assert_eq!(decoded.local_name, "architect");
|
||||||
|
assert_eq!(decoded.local_zone, "fabric-10g");
|
||||||
|
// Solo node — no peers yet.
|
||||||
|
assert!(decoded.peers.is_empty(), "solo node has no peers");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn dispatch_returns_empty_request_error() {
|
||||||
|
let gossip = bootstrap_gossip("solo", next_port()).await;
|
||||||
|
let router = RpcRouter::new(gossip.clone(), "solo".into(), "z".into());
|
||||||
|
let reply = dispatch(&router, &[]).await;
|
||||||
|
assert_eq!(reply, vec![ErrorCode::EmptyRequest.as_byte()]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn dispatch_returns_unknown_method_error() {
|
||||||
|
let gossip = bootstrap_gossip("solo", next_port()).await;
|
||||||
|
let router = RpcRouter::new(gossip.clone(), "solo".into(), "z".into());
|
||||||
|
let reply = dispatch(&router, &[0xab, 0x01, 0x02]).await;
|
||||||
|
assert_eq!(reply, vec![ErrorCode::UnknownMethod.as_byte()]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn rpc_call_rejects_oversize_payload() {
|
||||||
|
// No network involved — the size check runs client-side before we
|
||||||
|
// even try to open the bidi stream. Use a dummy connection built
|
||||||
|
// via generate_test_pair; we won't actually connect.
|
||||||
|
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
||||||
|
let server = QuicServer::bind(loopback(0), id_b).unwrap();
|
||||||
|
let server_addr = server.local_addr().unwrap();
|
||||||
|
let _accept = tokio::spawn(async move {
|
||||||
|
let _ = server.accept().await;
|
||||||
|
});
|
||||||
|
let client = QuicClient::new(loopback(0), id_a).unwrap();
|
||||||
|
let conn = client.connect(server_addr, "b").await.unwrap();
|
||||||
|
let huge = vec![0u8; MAX_MESSAGE_BYTES];
|
||||||
|
let err = rpc_call(&conn, Method::Ping, &huge)
|
||||||
|
.await
|
||||||
|
.err()
|
||||||
|
.expect("must reject oversize");
|
||||||
|
assert!(err.to_string().contains("exceeds cap"));
|
||||||
|
conn.close(quinn::VarInt::from_u32(0), b"done");
|
||||||
|
client.shutdown().await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn end_to_end_ping_and_peer_status_over_real_quic() {
|
||||||
|
// Two full nodes: A runs gossip + a QuicServer serving RpcRouter.
|
||||||
|
// B is a client that dials A and calls both RPCs.
|
||||||
|
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
||||||
|
|
||||||
|
let port_a = next_port();
|
||||||
|
let gossip_a = bootstrap_gossip("a", port_a).await;
|
||||||
|
let router = Arc::new(RpcRouter::new(
|
||||||
|
gossip_a.clone(),
|
||||||
|
"a".into(),
|
||||||
|
"fabric-10g".into(),
|
||||||
|
));
|
||||||
|
|
||||||
|
// A's advertised state so the PeerStatus reply exercises the
|
||||||
|
// gossip → PeerView path even though B isn't in the peer table.
|
||||||
|
gossip_a.set_hot_used(4096).await;
|
||||||
|
gossip_a.set_hot_max(1_000_000).await;
|
||||||
|
gossip_a
|
||||||
|
.set_warm_projects(&["osobh/clawverse", "osobh/clawmates"])
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
||||||
|
let server_addr = server.local_addr().unwrap();
|
||||||
|
let accept_task = tokio::spawn(async move {
|
||||||
|
if let Some(Ok(conn)) = server.accept().await {
|
||||||
|
let _ = serve_connection(conn, router).await;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
||||||
|
let conn = client.connect(server_addr, "a").await.unwrap();
|
||||||
|
|
||||||
|
// Ping.
|
||||||
|
let pong = call_ping(&conn, b"hello").await.unwrap();
|
||||||
|
assert_eq!(pong, b"hello");
|
||||||
|
|
||||||
|
// PeerStatus.
|
||||||
|
let status = call_peer_status(&conn).await.unwrap();
|
||||||
|
assert_eq!(status.local_name, "a");
|
||||||
|
assert_eq!(status.local_zone, "fabric-10g");
|
||||||
|
assert!(status.peers.is_empty(), "A has no peers configured");
|
||||||
|
|
||||||
|
conn.close(quinn::VarInt::from_u32(0), b"done");
|
||||||
|
client.shutdown().await;
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
accept_task.abort();
|
||||||
|
// Keep gossip_a alive to the end so its background task doesn't
|
||||||
|
// drop mid-serve.
|
||||||
|
let _ = gossip_a.peers().await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn peer_status_reflects_peer_gossip_state() {
|
||||||
|
// A and B both run gossip; A serves RPC. When A's PeerStatus is
|
||||||
|
// called by a third-party client, the reply's `peers` field
|
||||||
|
// contains B (as long as gossip has converged).
|
||||||
|
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
||||||
|
let port_a = next_port();
|
||||||
|
let port_b = next_port();
|
||||||
|
|
||||||
|
let cfg_a = ClusterConfig {
|
||||||
|
zone: "fabric-10g".into(),
|
||||||
|
bind_lan: Some(loopback(port_a)),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let cfg_b = ClusterConfig {
|
||||||
|
zone: "lan-1g".into(),
|
||||||
|
bind_lan: Some(loopback(port_b)),
|
||||||
|
peers: vec![PeerEntry {
|
||||||
|
name: "a".into(),
|
||||||
|
zone: "fabric-10g".into(),
|
||||||
|
lan_addr: Some(loopback(port_a)),
|
||||||
|
tailscale_addr: None,
|
||||||
|
}],
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
let gossip_a = Arc::new(ClusterGossip::bootstrap(&cfg_a, "a").await.unwrap());
|
||||||
|
let gossip_b = Arc::new(ClusterGossip::bootstrap(&cfg_b, "b").await.unwrap());
|
||||||
|
|
||||||
|
let router = Arc::new(RpcRouter::new(
|
||||||
|
gossip_a.clone(),
|
||||||
|
"a".into(),
|
||||||
|
"fabric-10g".into(),
|
||||||
|
));
|
||||||
|
|
||||||
|
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
||||||
|
let server_addr = server.local_addr().unwrap();
|
||||||
|
let accept_task = tokio::spawn(async move {
|
||||||
|
if let Some(Ok(conn)) = server.accept().await {
|
||||||
|
let _ = serve_connection(conn, router).await;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// Wait for gossip convergence: A must see B.
|
||||||
|
let deadline = std::time::Instant::now() + Duration::from_secs(10);
|
||||||
|
loop {
|
||||||
|
if let Some(v) = gossip_a.peer("b").await {
|
||||||
|
if v.alive {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if std::time::Instant::now() >= deadline {
|
||||||
|
panic!("A never saw B alive within 10s");
|
||||||
|
}
|
||||||
|
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
||||||
|
let conn = client.connect(server_addr, "a").await.unwrap();
|
||||||
|
let status = call_peer_status(&conn).await.unwrap();
|
||||||
|
assert_eq!(status.local_name, "a");
|
||||||
|
assert_eq!(status.peers.len(), 1, "A should report exactly B");
|
||||||
|
assert_eq!(status.peers[0].name, "b");
|
||||||
|
assert_eq!(status.peers[0].zone, "lan-1g");
|
||||||
|
assert!(status.peers[0].alive);
|
||||||
|
|
||||||
|
conn.close(quinn::VarInt::from_u32(0), b"done");
|
||||||
|
client.shutdown().await;
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
accept_task.abort();
|
||||||
|
// Keep gossip services alive until end.
|
||||||
|
drop(gossip_b);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn error_code_describe_covers_all_variants() {
|
||||||
|
assert_eq!(ErrorCode::EmptyRequest.describe(), "empty request");
|
||||||
|
assert_eq!(ErrorCode::UnknownMethod.describe(), "unknown method");
|
||||||
|
assert_eq!(ErrorCode::HandlerFailure.describe(), "handler failure");
|
||||||
|
assert_eq!(ErrorCode::NotFound.describe(), "not found");
|
||||||
|
assert_eq!(ErrorCode::InvalidRequest.describe(), "invalid request");
|
||||||
|
assert_eq!(
|
||||||
|
ErrorCode::NotConfigured.describe(),
|
||||||
|
"server subsystem not configured"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn decode_error_covers_all_known_codes() {
|
||||||
|
for code in [
|
||||||
|
ErrorCode::EmptyRequest,
|
||||||
|
ErrorCode::UnknownMethod,
|
||||||
|
ErrorCode::HandlerFailure,
|
||||||
|
ErrorCode::NotFound,
|
||||||
|
ErrorCode::InvalidRequest,
|
||||||
|
ErrorCode::NotConfigured,
|
||||||
|
] {
|
||||||
|
assert_eq!(decode_error(code.as_byte()), Some(code));
|
||||||
|
}
|
||||||
|
assert_eq!(decode_error(0x00), None);
|
||||||
|
assert_eq!(decode_error(0xff), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Phase 2b: Blob RPC ─────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn blob_rpcs_return_not_configured_without_store() {
|
||||||
|
// Router built via `new` alone (no `.with_blob_store`) must
|
||||||
|
// refuse Blob* methods with a well-known error code.
|
||||||
|
let gossip = bootstrap_gossip("solo", next_port()).await;
|
||||||
|
let router = RpcRouter::new(gossip, "solo".into(), "z".into());
|
||||||
|
for method in [
|
||||||
|
Method::BlobStat,
|
||||||
|
Method::BlobGet,
|
||||||
|
Method::BlobPut,
|
||||||
|
Method::BlobLoadManifest,
|
||||||
|
] {
|
||||||
|
let mut req = vec![method.as_byte()];
|
||||||
|
req.extend_from_slice(&[0u8; 32]);
|
||||||
|
let reply = dispatch(&router, &req).await;
|
||||||
|
assert_eq!(
|
||||||
|
reply,
|
||||||
|
vec![ErrorCode::NotConfigured.as_byte()],
|
||||||
|
"method {method:?} should be NotConfigured without a store"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn blob_stat_returns_not_found_for_missing() {
|
||||||
|
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
||||||
|
let missing = BlobId::from_bytes([0u8; 32]);
|
||||||
|
let mut req = vec![Method::BlobStat.as_byte()];
|
||||||
|
req.extend_from_slice(missing.as_bytes());
|
||||||
|
let reply = dispatch(&router, &req).await;
|
||||||
|
assert_eq!(reply, vec![ErrorCode::NotFound.as_byte()]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn blob_stat_returns_json_for_existing() {
|
||||||
|
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
||||||
|
let id = router
|
||||||
|
.blob_store()
|
||||||
|
.unwrap()
|
||||||
|
.put_bytes(b"tiny content")
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let mut req = vec![Method::BlobStat.as_byte()];
|
||||||
|
req.extend_from_slice(id.as_bytes());
|
||||||
|
let reply = dispatch(&router, &req).await;
|
||||||
|
let stat: BlobStat = serde_json::from_slice(&reply).unwrap();
|
||||||
|
assert_eq!(stat.total_size, b"tiny content".len() as u64);
|
||||||
|
assert_eq!(stat.chunk_count, 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn blob_stat_returns_invalid_request_for_bad_length() {
|
||||||
|
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
||||||
|
// Payload is only 5 bytes; a valid BlobId is 32.
|
||||||
|
let req = vec![Method::BlobStat.as_byte(), 1, 2, 3, 4, 5];
|
||||||
|
let reply = dispatch(&router, &req).await;
|
||||||
|
assert_eq!(reply, vec![ErrorCode::InvalidRequest.as_byte()]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn blob_get_returns_content_bytes() {
|
||||||
|
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
||||||
|
let payload: &[u8] = b"contents to retrieve";
|
||||||
|
let id = router.blob_store().unwrap().put_bytes(payload).await.unwrap();
|
||||||
|
let mut req = vec![Method::BlobGet.as_byte()];
|
||||||
|
req.extend_from_slice(id.as_bytes());
|
||||||
|
let reply = dispatch(&router, &req).await;
|
||||||
|
assert_eq!(reply, payload);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn blob_put_stores_bytes_and_returns_hash() {
|
||||||
|
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
||||||
|
let payload = b"put via rpc";
|
||||||
|
let mut req = vec![Method::BlobPut.as_byte()];
|
||||||
|
req.extend_from_slice(payload);
|
||||||
|
let reply = dispatch(&router, &req).await;
|
||||||
|
assert_eq!(reply.len(), 32);
|
||||||
|
let mut id_bytes = [0u8; 32];
|
||||||
|
id_bytes.copy_from_slice(&reply);
|
||||||
|
let assigned = BlobId::from_bytes(id_bytes);
|
||||||
|
let expected = BlobId::from_bytes(blake3::hash(payload).into());
|
||||||
|
assert_eq!(assigned, expected);
|
||||||
|
// Round-trip: the bytes are now readable via the store.
|
||||||
|
let round = router
|
||||||
|
.blob_store()
|
||||||
|
.unwrap()
|
||||||
|
.get_bytes(&assigned)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(round.as_deref(), Some(payload.as_slice()));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn blob_load_manifest_returns_json_for_existing() {
|
||||||
|
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
||||||
|
let data = vec![0x77u8; 4 * 1024 * 1024 + 1]; // 2 chunks
|
||||||
|
let id = router.blob_store().unwrap().put_bytes(&data).await.unwrap();
|
||||||
|
let mut req = vec![Method::BlobLoadManifest.as_byte()];
|
||||||
|
req.extend_from_slice(id.as_bytes());
|
||||||
|
let reply = dispatch(&router, &req).await;
|
||||||
|
let manifest: BlobManifest = serde_json::from_slice(&reply).unwrap();
|
||||||
|
assert_eq!(manifest.blob_id, id);
|
||||||
|
assert_eq!(manifest.total_size, data.len() as u64);
|
||||||
|
assert_eq!(manifest.chunks.len(), 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn blob_load_manifest_returns_not_found_for_missing() {
|
||||||
|
let (_tmp, router) = router_with_blobs("solo", next_port()).await;
|
||||||
|
let missing = BlobId::from_bytes([0u8; 32]);
|
||||||
|
let mut req = vec![Method::BlobLoadManifest.as_byte()];
|
||||||
|
req.extend_from_slice(missing.as_bytes());
|
||||||
|
let reply = dispatch(&router, &req).await;
|
||||||
|
assert_eq!(reply, vec![ErrorCode::NotFound.as_byte()]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn end_to_end_blob_put_stat_get_over_real_quic() {
|
||||||
|
// The full loop: B → A over real QUIC + mTLS.
|
||||||
|
// 1. B puts a blob on A (BlobPut).
|
||||||
|
// 2. B queries stat + fetches it back (BlobStat + BlobGet).
|
||||||
|
// 3. B asks for the manifest (BlobLoadManifest).
|
||||||
|
// Every step goes through the actual wire, no shortcuts.
|
||||||
|
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
||||||
|
let (_tmp, router) = router_with_blobs("a", next_port()).await;
|
||||||
|
|
||||||
|
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
||||||
|
let server_addr = server.local_addr().unwrap();
|
||||||
|
let router_srv = router.clone();
|
||||||
|
let accept_task = tokio::spawn(async move {
|
||||||
|
if let Some(Ok(conn)) = server.accept().await {
|
||||||
|
let _ = serve_connection(conn, router_srv).await;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
||||||
|
let conn = client.connect(server_addr, "a").await.unwrap();
|
||||||
|
|
||||||
|
let payload: &[u8] = b"cross-node payload";
|
||||||
|
let assigned = call_blob_put(&conn, payload).await.unwrap();
|
||||||
|
let expected = BlobId::from_bytes(blake3::hash(payload).into());
|
||||||
|
assert_eq!(assigned, expected);
|
||||||
|
|
||||||
|
let stat = call_blob_stat(&conn, &assigned).await.unwrap().unwrap();
|
||||||
|
assert_eq!(stat.total_size, payload.len() as u64);
|
||||||
|
assert_eq!(stat.chunk_count, 1);
|
||||||
|
|
||||||
|
let round = call_blob_get(&conn, &assigned).await.unwrap().unwrap();
|
||||||
|
assert_eq!(round, payload);
|
||||||
|
|
||||||
|
let manifest = call_blob_load_manifest(&conn, &assigned)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(manifest.blob_id, assigned);
|
||||||
|
assert_eq!(manifest.chunks.len(), 1);
|
||||||
|
|
||||||
|
// NotFound path also works over the wire.
|
||||||
|
let ghost = BlobId::from_bytes([0u8; 32]);
|
||||||
|
assert!(call_blob_stat(&conn, &ghost).await.unwrap().is_none());
|
||||||
|
assert!(call_blob_get(&conn, &ghost).await.unwrap().is_none());
|
||||||
|
assert!(call_blob_load_manifest(&conn, &ghost)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.is_none());
|
||||||
|
|
||||||
|
conn.close(quinn::VarInt::from_u32(0), b"done");
|
||||||
|
client.shutdown().await;
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
accept_task.abort();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn end_to_end_multi_chunk_blob_over_real_quic() {
|
||||||
|
// 6 MB blob → 2 chunks. Round-trips whole via BlobPut/Get.
|
||||||
|
// Also confirms the RPC layer's MAX_MESSAGE_BYTES bump from
|
||||||
|
// 16 KiB to 16 MiB actually took effect end-to-end.
|
||||||
|
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
||||||
|
let (_tmp, router) = router_with_blobs("a", next_port()).await;
|
||||||
|
|
||||||
|
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
||||||
|
let server_addr = server.local_addr().unwrap();
|
||||||
|
let router_srv = router.clone();
|
||||||
|
let accept_task = tokio::spawn(async move {
|
||||||
|
if let Some(Ok(conn)) = server.accept().await {
|
||||||
|
let _ = serve_connection(conn, router_srv).await;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
||||||
|
let conn = client.connect(server_addr, "a").await.unwrap();
|
||||||
|
|
||||||
|
let payload: Vec<u8> = (0..6 * 1024 * 1024)
|
||||||
|
.map(|i| (i % 251) as u8)
|
||||||
|
.collect();
|
||||||
|
let assigned = call_blob_put(&conn, &payload).await.unwrap();
|
||||||
|
let manifest = call_blob_load_manifest(&conn, &assigned)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(manifest.chunks.len(), 2, "6 MB should split into 2 chunks");
|
||||||
|
let round = call_blob_get(&conn, &assigned).await.unwrap().unwrap();
|
||||||
|
assert_eq!(round.len(), payload.len());
|
||||||
|
assert_eq!(round, payload);
|
||||||
|
|
||||||
|
conn.close(quinn::VarInt::from_u32(0), b"done");
|
||||||
|
client.shutdown().await;
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
accept_task.abort();
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Phase 2c: streaming Blob RPC ─────────────────────────────────
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn method_reports_streaming_variants() {
|
||||||
|
assert!(!Method::Ping.is_streaming());
|
||||||
|
assert!(!Method::PeerStatus.is_streaming());
|
||||||
|
assert!(!Method::BlobStat.is_streaming());
|
||||||
|
assert!(!Method::BlobGet.is_streaming());
|
||||||
|
assert!(!Method::BlobPut.is_streaming());
|
||||||
|
assert!(!Method::BlobLoadManifest.is_streaming());
|
||||||
|
assert!(Method::BlobPutStream.is_streaming());
|
||||||
|
assert!(Method::BlobGetStream.is_streaming());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn end_to_end_stream_put_and_get_over_real_quic() {
|
||||||
|
// The whole point of Phase 2c: put a big blob without holding
|
||||||
|
// it in memory on either side. This test drives a 12 MiB
|
||||||
|
// payload (3 chunks) through BlobPutStream and BlobGetStream.
|
||||||
|
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
||||||
|
let (_tmp, router) = router_with_blobs("a", next_port()).await;
|
||||||
|
|
||||||
|
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
||||||
|
let server_addr = server.local_addr().unwrap();
|
||||||
|
let router_srv = router.clone();
|
||||||
|
let accept_task = tokio::spawn(async move {
|
||||||
|
if let Some(Ok(conn)) = server.accept().await {
|
||||||
|
let _ = serve_connection(conn, router_srv).await;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
||||||
|
let conn = client.connect(server_addr, "a").await.unwrap();
|
||||||
|
|
||||||
|
// 12 MiB payload; deliberately not a chunk multiple so the last
|
||||||
|
// chunk is short.
|
||||||
|
let mut payload: Vec<u8> = Vec::with_capacity(12 * 1024 * 1024 + 777);
|
||||||
|
for i in 0..12 * 1024 * 1024 + 777 {
|
||||||
|
payload.push((i % 251) as u8);
|
||||||
|
}
|
||||||
|
let reader = std::io::Cursor::new(payload.clone());
|
||||||
|
let id = call_blob_put_stream(&conn, reader).await.unwrap();
|
||||||
|
let expected = BlobId::from_bytes(blake3::hash(&payload).into());
|
||||||
|
assert_eq!(id, expected);
|
||||||
|
|
||||||
|
// Manifest verifies the chunk split.
|
||||||
|
let manifest = call_blob_load_manifest(&conn, &id).await.unwrap().unwrap();
|
||||||
|
assert_eq!(manifest.chunks.len(), 4, "12 MiB + 777 → 4 chunks");
|
||||||
|
assert_eq!(manifest.total_size, payload.len() as u64);
|
||||||
|
|
||||||
|
// Stream it back.
|
||||||
|
let mut sink: Vec<u8> = Vec::new();
|
||||||
|
let ok = call_blob_get_stream(&conn, &id, &mut sink).await.unwrap();
|
||||||
|
assert!(ok);
|
||||||
|
assert_eq!(sink.len(), payload.len());
|
||||||
|
assert_eq!(sink, payload);
|
||||||
|
|
||||||
|
conn.close(quinn::VarInt::from_u32(0), b"done");
|
||||||
|
client.shutdown().await;
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
accept_task.abort();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn stream_get_returns_false_for_missing_blob() {
|
||||||
|
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
||||||
|
let (_tmp, router) = router_with_blobs("a", next_port()).await;
|
||||||
|
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
||||||
|
let server_addr = server.local_addr().unwrap();
|
||||||
|
let router_srv = router.clone();
|
||||||
|
let accept_task = tokio::spawn(async move {
|
||||||
|
if let Some(Ok(conn)) = server.accept().await {
|
||||||
|
let _ = serve_connection(conn, router_srv).await;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
||||||
|
let conn = client.connect(server_addr, "a").await.unwrap();
|
||||||
|
|
||||||
|
let ghost = BlobId::from_bytes([0u8; 32]);
|
||||||
|
let mut sink: Vec<u8> = Vec::new();
|
||||||
|
let ok = call_blob_get_stream(&conn, &ghost, &mut sink)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert!(!ok);
|
||||||
|
assert!(sink.is_empty());
|
||||||
|
|
||||||
|
conn.close(quinn::VarInt::from_u32(0), b"done");
|
||||||
|
client.shutdown().await;
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
accept_task.abort();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn stream_methods_return_not_configured_without_store() {
|
||||||
|
// Router built without a store: streaming methods must reply
|
||||||
|
// with NotConfigured as their first status byte.
|
||||||
|
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
||||||
|
let gossip = bootstrap_gossip("a", next_port()).await;
|
||||||
|
// NOTE: no `.with_blob_store(...)` — Blob* methods should error.
|
||||||
|
let router = Arc::new(RpcRouter::new(gossip, "a".into(), "z".into()));
|
||||||
|
|
||||||
|
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
||||||
|
let server_addr = server.local_addr().unwrap();
|
||||||
|
let router_srv = router.clone();
|
||||||
|
let accept_task = tokio::spawn(async move {
|
||||||
|
if let Some(Ok(conn)) = server.accept().await {
|
||||||
|
let _ = serve_connection(conn, router_srv).await;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
||||||
|
let conn = client.connect(server_addr, "a").await.unwrap();
|
||||||
|
|
||||||
|
// BlobPutStream on a store-less router → NotConfigured.
|
||||||
|
let reader = std::io::Cursor::new(b"never lands".to_vec());
|
||||||
|
let put_err = call_blob_put_stream(&conn, reader)
|
||||||
|
.await
|
||||||
|
.err()
|
||||||
|
.expect("no store → error");
|
||||||
|
assert!(
|
||||||
|
put_err.to_string().contains("not configured"),
|
||||||
|
"put_stream err: {put_err}"
|
||||||
|
);
|
||||||
|
|
||||||
|
// BlobGetStream: same error surface.
|
||||||
|
let ghost = BlobId::from_bytes([0u8; 32]);
|
||||||
|
let mut sink: Vec<u8> = Vec::new();
|
||||||
|
let get_err = call_blob_get_stream(&conn, &ghost, &mut sink)
|
||||||
|
.await
|
||||||
|
.err()
|
||||||
|
.expect("no store → error");
|
||||||
|
assert!(
|
||||||
|
get_err.to_string().contains("not configured"),
|
||||||
|
"get_stream err: {get_err}"
|
||||||
|
);
|
||||||
|
|
||||||
|
conn.close(quinn::VarInt::from_u32(0), b"done");
|
||||||
|
client.shutdown().await;
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
accept_task.abort();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn stream_put_deduplicates_with_prior_put_bytes() {
|
||||||
|
// Uploading the same content twice — once bounded, once
|
||||||
|
// streaming — yields the same BlobId AND doesn't double-store
|
||||||
|
// chunks. Proves stream + bounded are consistent addresses.
|
||||||
|
let (id_a, id_b) = NodeIdentity::generate_test_pair("a", "b").unwrap();
|
||||||
|
let (_tmp, router) = router_with_blobs("a", next_port()).await;
|
||||||
|
|
||||||
|
// Pre-populate via the local store's bounded API.
|
||||||
|
let payload = vec![0xdeu8; 4 * 1024 * 1024 + 100];
|
||||||
|
let pre_id = router
|
||||||
|
.blob_store()
|
||||||
|
.unwrap()
|
||||||
|
.put_bytes(&payload)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let server = QuicServer::bind(loopback(0), id_a).unwrap();
|
||||||
|
let server_addr = server.local_addr().unwrap();
|
||||||
|
let router_srv = router.clone();
|
||||||
|
let accept_task = tokio::spawn(async move {
|
||||||
|
if let Some(Ok(conn)) = server.accept().await {
|
||||||
|
let _ = serve_connection(conn, router_srv).await;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let client = QuicClient::new(loopback(0), id_b).unwrap();
|
||||||
|
let conn = client.connect(server_addr, "a").await.unwrap();
|
||||||
|
|
||||||
|
let reader = std::io::Cursor::new(payload.clone());
|
||||||
|
let stream_id = call_blob_put_stream(&conn, reader).await.unwrap();
|
||||||
|
assert_eq!(pre_id, stream_id);
|
||||||
|
|
||||||
|
// The manifest is still there and its chunk count matches the
|
||||||
|
// pre-existing one — no fork.
|
||||||
|
let manifest = router
|
||||||
|
.blob_store()
|
||||||
|
.unwrap()
|
||||||
|
.load_manifest(&stream_id)
|
||||||
|
.await
|
||||||
|
.unwrap()
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(manifest.chunks.len(), 2);
|
||||||
|
|
||||||
|
conn.close(quinn::VarInt::from_u32(0), b"done");
|
||||||
|
client.shutdown().await;
|
||||||
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
|
accept_task.abort();
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user