Added bytestream ALPN

crates/federation-net/tests/integration.rs

@@ -359,3 +359,112 @@ async fn rendezvous_discovers_peers_without_tickets() {
     .await
     .expect("test timed out");
 }
+
+#[tokio::test]
+async fn byte_streams_between_peers() {
+    let _net = NET_LOCK.lock().await;
+    tokio::time::timeout(TEST_TIMEOUT, async {
+        const ECHO_ALPN: &[u8] = b"test/echo/1";
+
+        let dir_a = tempfile::tempdir().expect("tempdir");
+        let dir_b = tempfile::tempdir().expect("tempdir");
+
+        // Peer A declares a stream protocol and echoes one message per stream.
+        let config_a = NetworkConfig::builder()
+            .data_dir(dir_a.path())
+            .network_id(NetworkId::from_name("test-net"))
+            .schema_id(SchemaId::from_name("test-schema-v1"))
+            .request_timeout(Duration::from_secs(10))
+            .stream_protocol(ECHO_ALPN)
+            .build()
+            .expect("valid config");
+        let (engine_a, _events_a) = Engine::start(config_a).await.expect("engine a starts");
+        let mut acceptor = engine_a.stream_acceptor(ECHO_ALPN).expect("acceptor");
+        // The acceptor is exclusive and unknown protocols are rejected.
+        assert!(matches!(
+            engine_a.stream_acceptor(ECHO_ALPN),
+            Err(NetworkError::StreamAcceptorTaken(_))
+        ));
+        assert!(matches!(
+            engine_a.stream_acceptor(b"test/unknown/1"),
+            Err(NetworkError::UnknownStreamProtocol(_))
+        ));
+        let echo_task = tokio::spawn(async move {
+            let mut stream = acceptor.accept().await.expect("incoming stream");
+            let mut buf = Vec::new();
+            let mut chunk = [0u8; 1024];
+            while let Some(n) = stream.recv.read(&mut chunk).await.expect("read") {
+                buf.extend_from_slice(&chunk[..n]);
+            }
+            stream.send.write_all(&buf).await.expect("write echo");
+            stream.send.finish().expect("finish");
+            // Keep the stream alive until the peer read everything.
+            let _ = stream.send.stopped().await;
+            (stream.peer_id, buf.len())
+        });
+
+        let (engine_b, _events_b) = start(dir_b.path(), "test-net", "test-schema-v1").await;
+        let ticket = engine_a.ticket().await.expect("ticket");
+
+        // B opens a byte stream to A using the address from the ticket and
+        // sends a payload larger than one network frame.
+        let payload = vec![0xAB_u8; 512 * 1024];
+        let mut stream = engine_b
+            .open_stream(ticket.endpoint_addr.clone(), ECHO_ALPN)
+            .await
+            .expect("open stream");
+        assert_eq!(stream.peer_id, engine_a.endpoint_id());
+        stream.send.write_all(&payload).await.expect("write");
+        stream.send.finish().expect("finish");
+        let mut echoed = Vec::new();
+        let mut chunk = [0u8; 1024];
+        while let Some(n) = stream.recv.read(&mut chunk).await.expect("read echo") {
+            echoed.extend_from_slice(&chunk[..n]);
+        }
+        assert_eq!(echoed, payload);
+
+        let (peer_seen_by_a, len_seen_by_a) = echo_task.await.expect("echo task");
+        assert_eq!(peer_seen_by_a, engine_b.endpoint_id());
+        assert_eq!(len_seen_by_a, payload.len());
+
+        engine_a.shutdown().await.expect("shutdown a");
+        engine_b.shutdown().await.expect("shutdown b");
+    })
+    .await
+    .expect("test timed out");
+}
+
+#[tokio::test]
+async fn byte_stream_to_wrong_network_is_rejected() {
+    let _net = NET_LOCK.lock().await;
+    tokio::time::timeout(TEST_TIMEOUT, async {
+        const BLOB_ALPN: &[u8] = b"test/blob/1";
+
+        let dir_a = tempfile::tempdir().expect("tempdir");
+        let dir_b = tempfile::tempdir().expect("tempdir");
+
+        let config_a = NetworkConfig::builder()
+            .data_dir(dir_a.path())
+            .network_id(NetworkId::from_name("net-one"))
+            .schema_id(SchemaId::from_name("test-schema-v1"))
+            .request_timeout(Duration::from_secs(10))
+            .stream_protocol(BLOB_ALPN)
+            .build()
+            .expect("valid config");
+        let (engine_a, _events_a) = Engine::start(config_a).await.expect("engine a starts");
+        let _acceptor = engine_a.stream_acceptor(BLOB_ALPN).expect("acceptor");
+
+        // B lives in a different network; the stream handshake must fail.
+        let (engine_b, _events_b) = start(dir_b.path(), "net-two", "test-schema-v1").await;
+        let ticket = engine_a.ticket().await.expect("ticket");
+        let result = engine_b
+            .open_stream(ticket.endpoint_addr.clone(), BLOB_ALPN)
+            .await;
+        assert!(matches!(result, Err(NetworkError::NetworkMismatch)));
+
+        engine_a.shutdown().await.expect("shutdown a");
+        engine_b.shutdown().await.expect("shutdown b");
+    })
+    .await
+    .expect("test timed out");
+}