[runtime] Change send and recv signatures (#929)

Author: danlaine

runtime/src/lib.rs

@@ -17,6 +17,7 @@
 //! `commonware-runtime` is **ALPHA** software and is not yet recommended for production use. Developers should
 //! expect breaking changes and occasional instability.
 
+use commonware_utils::{StableBuf, StableBufMut};
 use prometheus_client::registry::Metric;
 use std::io::Error as IoError;
 use std::{
@@ -272,15 +273,15 @@ pub trait Listener: Sync + Send + 'static {
 /// messages over a network connection.
 pub trait Sink: Sync + Send + 'static {
     /// Send a message to the sink.
-    fn send(&mut self, msg: &[u8]) -> impl Future<Output = Result<(), Error>> + Send;
+    fn send<B: StableBuf>(&mut self, msg: B) -> impl Future<Output = Result<(), Error>> + Send;
 }
 
 /// Interface that any runtime must implement to receive
 /// messages over a network connection.
 pub trait Stream: Sync + Send + 'static {
     /// Receive a message from the stream, storing it in the given buffer.
     /// Reads exactly the number of bytes that fit in the buffer.
-    fn recv(&mut self, buf: &mut [u8]) -> impl Future<Output = Result<(), Error>> + Send;
+    fn recv<B: StableBufMut>(&mut self, buf: B) -> impl Future<Output = Result<B, Error>> + Send;
 }
 
 /// Interface to interact with storage.
@@ -356,6 +357,7 @@ pub trait Blob: Clone + Send + Sync + 'static {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use bytes::Bytes;
     use commonware_macros::select;
     use futures::channel::oneshot;
     use futures::{channel::mpsc, future::ready, join, SinkExt, StreamExt};
@@ -1325,8 +1327,7 @@ mod tests {
             async fn read_line<St: Stream>(stream: &mut St) -> Result<String, Error> {
                 let mut line = Vec::new();
                 loop {
-                    let mut byte = [0; 1];
-                    stream.recv(&mut byte).await?;
+                    let byte = stream.recv(vec![0; 1]).await?;
                     if byte[0] == b'\n' {
                         if line.last() == Some(&b'\r') {
                             line.pop(); // Remove trailing \r
@@ -1359,9 +1360,8 @@ mod tests {
                 stream: &mut St,
                 content_length: usize,
             ) -> Result<String, Error> {
-                let mut body = vec![0; content_length];
-                stream.recv(&mut body).await?;
-                String::from_utf8(body).map_err(|_| Error::ReadFailed)
+                let read = stream.recv(vec![0; content_length]).await?;
+                String::from_utf8(read).map_err(|_| Error::ReadFailed)
             }
 
             // Simulate a client connecting to the server
@@ -1384,7 +1384,7 @@ mod tests {
                         "GET /metrics HTTP/1.1\r\nHost: {}\r\nConnection: close\r\n\r\n",
                         address
                     );
-                    sink.send(request.as_bytes()).await.unwrap();
+                    sink.send(Bytes::from(request)).await.unwrap();
 
                     // Read and verify the HTTP status line
                     let status_line = read_line(&mut stream).await.unwrap();

runtime/src/mocks.rs

@@ -2,6 +2,7 @@
 
 use crate::{Error, Sink as SinkTrait, Stream as StreamTrait};
 use bytes::Bytes;
+use commonware_utils::{StableBuf, StableBufMut};
 use futures::channel::oneshot;
 use std::{
     collections::VecDeque,
@@ -41,10 +42,10 @@ pub struct Sink {
 }
 
 impl SinkTrait for Sink {
-    async fn send(&mut self, msg: &[u8]) -> Result<(), Error> {
+    async fn send<B: StableBuf>(&mut self, msg: B) -> Result<(), Error> {
         let (os_send, data) = {
             let mut channel = self.channel.lock().unwrap();
-            channel.buffer.extend(msg);
+            channel.buffer.extend(msg.as_ref());
 
             // If there is a waiter and the buffer is large enough,
             // return the waiter (while clearing the waiter field).
@@ -74,16 +75,16 @@ pub struct Stream {
 }
 
 impl StreamTrait for Stream {
-    async fn recv(&mut self, buf: &mut [u8]) -> Result<(), Error> {
+    async fn recv<B: StableBufMut>(&mut self, mut buf: B) -> Result<B, Error> {
         let os_recv = {
             let mut channel = self.channel.lock().unwrap();
 
             // If the message is fully available in the buffer,
             // drain the value into buf and return.
             if channel.buffer.len() >= buf.len() {
                 let b: Vec<u8> = channel.buffer.drain(0..buf.len()).collect();
-                buf.copy_from_slice(&b);
-                return Ok(());
+                buf.put_slice(&b);
+                return Ok(buf);
             }
 
             // Otherwise, populate the waiter.
@@ -95,8 +96,9 @@ impl StreamTrait for Stream {
 
         // Wait for the waiter to be resolved.
         let data = os_recv.await.map_err(|_| Error::RecvFailed)?;
-        buf.copy_from_slice(&data);
-        Ok(())
+        assert_eq!(data.len(), buf.len());
+        buf.put_slice(&data);
+        Ok(buf)
     }
 }
 
@@ -111,52 +113,45 @@ mod tests {
     #[test]
     fn test_send_recv() {
         let (mut sink, mut stream) = Channel::init();
-
-        let data = b"hello world";
-        let mut buf = vec![0; data.len()];
+        let data = b"hello world".to_vec();
 
         block_on(async {
-            sink.send(data).await.unwrap();
-            stream.recv(&mut buf).await.unwrap();
+            sink.send(data.clone()).await.unwrap();
+            let buf = stream.recv(vec![0; data.len()]).await.unwrap();
+            assert_eq!(buf, data);
         });
-
-        assert_eq!(buf, data);
     }
 
     #[test]
     fn test_send_recv_partial_multiple() {
         let (mut sink, mut stream) = Channel::init();
-
-        let data1 = b"hello";
-        let data2 = b"world";
-        let mut buf1 = vec![0; data1.len()];
-        let mut buf2 = vec![0; data2.len()];
+        let data = b"hello".to_vec();
+        let data2 = b" world".to_vec();
 
         block_on(async {
-            sink.send(data1).await.unwrap();
+            sink.send(data).await.unwrap();
             sink.send(data2).await.unwrap();
-            stream.recv(&mut buf1[0..3]).await.unwrap();
-            stream.recv(&mut buf1[3..]).await.unwrap();
-            stream.recv(&mut buf2).await.unwrap();
+            let buf = stream.recv(vec![0; 5]).await.unwrap();
+            assert_eq!(buf, b"hello");
+            let buf = stream.recv(buf).await.unwrap();
+            assert_eq!(buf, b" worl");
+            let buf = stream.recv(vec![0; 1]).await.unwrap();
+            assert_eq!(buf, b"d");
         });
-
-        assert_eq!(buf1, data1);
-        assert_eq!(buf2, data2);
     }
 
     #[test]
     fn test_send_recv_async() {
         let (mut sink, mut stream) = Channel::init();
 
         let data = b"hello world";
-        let mut buf = vec![0; data.len()];
-
-        block_on(async {
-            futures::try_join!(stream.recv(&mut buf), async {
+        let buf = block_on(async {
+            futures::try_join!(stream.recv(vec![0; data.len()]), async {
                 sleep(Duration::from_millis(10_000));
-                sink.send(data).await
+                sink.send(data.to_vec()).await
             },)
-            .unwrap();
+            .unwrap()
+            .0
         });
 
         assert_eq!(buf, data);
@@ -170,8 +165,7 @@ mod tests {
         // If the oneshot sender is dropped before the oneshot receiver is resolved,
         // the recv function should return an error.
         executor.start(|_| async move {
-            let mut buf = vec![0; 5];
-            let (v, _) = join!(stream.recv(&mut buf), async {
+            let (v, _) = join!(stream.recv(vec![0; 5]), async {
                 // Take the waiter and drop it.
                 sink.channel.lock().unwrap().waiter.take();
             },);
@@ -187,12 +181,10 @@ mod tests {
         // If the waiter value has a min, but the oneshot receiver is dropped,
         // the send function should return an error when attempting to send the data.
         executor.start(|context| async move {
-            let mut buf = vec![0; 5];
-
             // Create a waiter using a recv call.
             // But then drop the receiver.
             select! {
-                v = stream.recv(&mut buf) => {
+                v = stream.recv( vec![0;5]) => {
                     panic!("unexpected value: {:?}", v);
                 },
                 _ = context.sleep(Duration::from_millis(100)) => {
@@ -202,7 +194,7 @@ mod tests {
             drop(stream);
 
             // Try to send a message (longer than the requested amount), but the receiver is dropped.
-            let result = sink.send(b"hello world").await;
+            let result = sink.send(b"hello world".to_vec()).await;
             assert!(matches!(result, Err(Error::SendFailed)));
         });
     }
@@ -214,9 +206,8 @@ mod tests {
 
         // If there is no data to read, test that the recv function just blocks. A timeout should return first.
         executor.start(|context| async move {
-            let mut buf = vec![0; 5];
             select! {
-                v = stream.recv(&mut buf) => {
+                v = stream.recv(vec![0;5]) => {
                     panic!("unexpected value: {:?}", v);
                 },
                 _ = context.sleep(Duration::from_millis(100)) => {

runtime/src/network/audited.rs

@@ -1,4 +1,5 @@
 use crate::{deterministic::Auditor, Error, SinkOf, StreamOf};
+use commonware_utils::{StableBuf, StableBufMut};
 use sha2::Digest;
 use std::{net::SocketAddr, sync::Arc};
 
@@ -10,10 +11,10 @@ pub struct Sink<S: crate::Sink> {
 }
 
 impl<S: crate::Sink> crate::Sink for Sink<S> {
-    async fn send(&mut self, data: &[u8]) -> Result<(), Error> {
+    async fn send<B: StableBuf>(&mut self, data: B) -> Result<(), Error> {
         self.auditor.event(b"send_attempt", |hasher| {
             hasher.update(self.remote_addr.to_string().as_bytes());
-            hasher.update(data);
+            hasher.update(data.as_ref());
         });
 
         self.inner.send(data).await.inspect_err(|e| {
@@ -25,7 +26,6 @@ impl<S: crate::Sink> crate::Sink for Sink<S> {
 
         self.auditor.event(b"send_success", |hasher| {
             hasher.update(self.remote_addr.to_string().as_bytes());
-            hasher.update(data);
         });
         Ok(())
     }
@@ -39,12 +39,12 @@ pub struct Stream<S: crate::Stream> {
 }
 
 impl<S: crate::Stream> crate::Stream for Stream<S> {
-    async fn recv(&mut self, buf: &mut [u8]) -> Result<(), Error> {
+    async fn recv<B: StableBufMut>(&mut self, buf: B) -> Result<B, Error> {
         self.auditor.event(b"recv_attempt", |hasher| {
             hasher.update(self.remote_addr.to_string().as_bytes());
         });
 
-        self.inner.recv(buf).await.inspect_err(|e| {
+        let buf = self.inner.recv(buf).await.inspect_err(|e| {
             self.auditor.event(b"recv_failure", |hasher| {
                 hasher.update(self.remote_addr.to_string().as_bytes());
                 hasher.update(e.to_string().as_bytes());
@@ -53,9 +53,9 @@ impl<S: crate::Stream> crate::Stream for Stream<S> {
 
         self.auditor.event(b"recv_success", |hasher| {
             hasher.update(self.remote_addr.to_string().as_bytes());
-            hasher.update(buf);
+            hasher.update(buf.as_ref());
         });
-        Ok(())
+        Ok(buf)
     }
 }
 
@@ -242,12 +242,11 @@ mod tests {
                 let (_, mut sink, mut stream) = listener.accept().await.unwrap();
 
                 // Receive data from client
-                let mut buf = [0u8; CLIENT_MSG.len()];
-                stream.recv(&mut buf).await.unwrap();
+                let buf = stream.recv(vec![0; CLIENT_MSG.len()]).await.unwrap();
                 assert_eq!(&buf, CLIENT_MSG.as_bytes());
 
                 // Send response
-                sink.send(SERVER_MSG.as_bytes()).await.unwrap();
+                sink.send(Vec::from(SERVER_MSG)).await.unwrap();
             });
             server_handles.push(handle);
         }
@@ -261,11 +260,10 @@ mod tests {
                 let (mut sink, mut stream) = network.dial(listener_addr).await.unwrap();
 
                 // Send data to server
-                sink.send(CLIENT_MSG.as_bytes()).await.unwrap();
+                sink.send(Vec::from(CLIENT_MSG)).await.unwrap();
 
                 // Receive response
-                let mut buf = [0u8; SERVER_MSG.len()];
-                stream.recv(&mut buf).await.unwrap();
+                let buf = stream.recv(vec![0; SERVER_MSG.len()]).await.unwrap();
                 assert_eq!(&buf, SERVER_MSG.as_bytes());
             });
             client_handles.push(handle);

runtime/src/network/deterministic.rs

@@ -1,4 +1,4 @@
-use crate::{mocks, Error};
+use crate::{mocks, Error, StableBuf, StableBufMut};
 use futures::{channel::mpsc, SinkExt as _, StreamExt as _};
 use std::{
     collections::HashMap,
@@ -16,7 +16,7 @@ pub struct Sink {
 }
 
 impl crate::Sink for Sink {
-    async fn send(&mut self, msg: &[u8]) -> Result<(), Error> {
+    async fn send<B: StableBuf>(&mut self, msg: B) -> Result<(), Error> {
         self.sender.send(msg).await.map_err(|_| Error::SendFailed)
     }
 }
@@ -27,7 +27,7 @@ pub struct Stream {
 }
 
 impl crate::Stream for Stream {
-    async fn recv(&mut self, buf: &mut [u8]) -> Result<(), Error> {
+    async fn recv<B: StableBufMut>(&mut self, buf: B) -> Result<B, Error> {
         self.receiver.recv(buf).await.map_err(|_| Error::RecvFailed)
     }
 }