feat(iroh)!: make blobs::read_at more flexible (#2756)

Introduces `ReadAtLen` which allows to specify the behaviour of
`read_at` more closely

- `All` - reads until the end (formerly `None`)
- `Exact(size)` - Reads exactly this many bytes (formerly `Some(size)`
- `AtMost(size)` - Reads at most this many bytes, but allows for a
shorter response if not enough data is available

Closes #2738

## Breaking Changes

- `iroh::client::Blobs::read_at` and `read_at_to_bytes` now take
`ReadAtLen` instead of `Option<usize>`

---------

Co-authored-by: Asmir Avdicevic <[email protected]>

Author: dignifiedquire

.github/workflows/ci.yml

@@ -362,4 +362,4 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - run: pip install --user codespell[toml]
-    - run: codespell --ignore-words-list=ans,crate,inout,ratatui,ser,stayin,swarmin,worl --skip=CHANGELOG.md
+    - run: codespell --ignore-words-list=ans,atmost,crate,inout,ratatui,ser,stayin,swarmin,worl --skip=CHANGELOG.md

iroh/src/client/blobs.rs

@@ -159,7 +159,7 @@ impl Client {
     /// Read offset + len from a single blob.
     ///
     /// If `len` is `None` it will read the full blob.
-    pub async fn read_at(&self, hash: Hash, offset: u64, len: Option<usize>) -> Result<Reader> {
+    pub async fn read_at(&self, hash: Hash, offset: u64, len: ReadAtLen) -> Result<Reader> {
         Reader::from_rpc_read_at(&self.rpc, hash, offset, len).await
     }
 
@@ -178,12 +178,7 @@ impl Client {
     /// Read all bytes of single blob at `offset` for length `len`.
     ///
     /// This allocates a buffer for the full length.
-    pub async fn read_at_to_bytes(
-        &self,
-        hash: Hash,
-        offset: u64,
-        len: Option<usize>,
-    ) -> Result<Bytes> {
+    pub async fn read_at_to_bytes(&self, hash: Hash, offset: u64, len: ReadAtLen) -> Result<Bytes> {
         Reader::from_rpc_read_at(&self.rpc, hash, offset, len)
             .await?
             .read_to_bytes()
@@ -484,6 +479,28 @@ impl SimpleStore for Client {
     }
 }
 
+/// Defines the way to read bytes.
+#[derive(Debug, Serialize, Deserialize, Default, Clone, Copy)]
+pub enum ReadAtLen {
+    /// Reads all available bytes.
+    #[default]
+    All,
+    /// Reads exactly this many bytes, erroring out on larger or smaller.
+    Exact(u64),
+    /// Reads at most this many bytes.
+    AtMost(u64),
+}
+
+impl ReadAtLen {
+    pub(crate) fn as_result_len(&self, size_remaining: u64) -> u64 {
+        match self {
+            ReadAtLen::All => size_remaining,
+            ReadAtLen::Exact(len) => *len,
+            ReadAtLen::AtMost(len) => std::cmp::min(*len, size_remaining),
+        }
+    }
+}
+
 /// Whether to wrap the added data in a collection.
 #[derive(Debug, Serialize, Deserialize, Default, Clone)]
 pub enum WrapOption {
@@ -872,14 +889,14 @@ impl Reader {
     }
 
     pub(crate) async fn from_rpc_read(rpc: &RpcClient, hash: Hash) -> anyhow::Result<Self> {
-        Self::from_rpc_read_at(rpc, hash, 0, None).await
+        Self::from_rpc_read_at(rpc, hash, 0, ReadAtLen::All).await
     }
 
     async fn from_rpc_read_at(
         rpc: &RpcClient,
         hash: Hash,
         offset: u64,
-        len: Option<usize>,
+        len: ReadAtLen,
     ) -> anyhow::Result<Self> {
         let stream = rpc
             .server_streaming(ReadAtRequest { hash, offset, len })
@@ -898,9 +915,7 @@ impl Reader {
             Ok(_) => Err(io::Error::new(io::ErrorKind::Other, "Expected data frame")),
             Err(err) => Err(io::Error::new(io::ErrorKind::Other, format!("{err}"))),
         });
-        let len = len
-            .map(|l| l as u64)
-            .unwrap_or_else(|| size.value() - offset);
+        let len = len.as_result_len(size.value() - offset);
         Ok(Self::new(size.value(), len, is_complete, Box::pin(stream)))
     }
 
@@ -1120,66 +1135,114 @@ mod tests {
         assert_eq!(&res, &buf[..]);
 
         // Read at smaller than blob_get_chunk_size
-        let res = client.blobs().read_at_to_bytes(hash, 0, Some(100)).await?;
+        let res = client
+            .blobs()
+            .read_at_to_bytes(hash, 0, ReadAtLen::Exact(100))
+            .await?;
         assert_eq!(res.len(), 100);
         assert_eq!(&res[..], &buf[0..100]);
 
-        let res = client.blobs().read_at_to_bytes(hash, 20, Some(120)).await?;
+        let res = client
+            .blobs()
+            .read_at_to_bytes(hash, 20, ReadAtLen::Exact(120))
+            .await?;
         assert_eq!(res.len(), 120);
         assert_eq!(&res[..], &buf[20..140]);
 
         // Read at equal to blob_get_chunk_size
         let res = client
             .blobs()
-            .read_at_to_bytes(hash, 0, Some(1024 * 64))
+            .read_at_to_bytes(hash, 0, ReadAtLen::Exact(1024 * 64))
             .await?;
         assert_eq!(res.len(), 1024 * 64);
         assert_eq!(&res[..], &buf[0..1024 * 64]);
 
         let res = client
             .blobs()
-            .read_at_to_bytes(hash, 20, Some(1024 * 64))
+            .read_at_to_bytes(hash, 20, ReadAtLen::Exact(1024 * 64))
             .await?;
         assert_eq!(res.len(), 1024 * 64);
         assert_eq!(&res[..], &buf[20..(20 + 1024 * 64)]);
 
         // Read at larger than blob_get_chunk_size
         let res = client
             .blobs()
-            .read_at_to_bytes(hash, 0, Some(10 + 1024 * 64))
+            .read_at_to_bytes(hash, 0, ReadAtLen::Exact(10 + 1024 * 64))
             .await?;
         assert_eq!(res.len(), 10 + 1024 * 64);
         assert_eq!(&res[..], &buf[0..(10 + 1024 * 64)]);
 
         let res = client
             .blobs()
-            .read_at_to_bytes(hash, 20, Some(10 + 1024 * 64))
+            .read_at_to_bytes(hash, 20, ReadAtLen::Exact(10 + 1024 * 64))
             .await?;
         assert_eq!(res.len(), 10 + 1024 * 64);
         assert_eq!(&res[..], &buf[20..(20 + 10 + 1024 * 64)]);
 
         // full length
-        let res = client.blobs().read_at_to_bytes(hash, 20, None).await?;
+        let res = client
+            .blobs()
+            .read_at_to_bytes(hash, 20, ReadAtLen::All)
+            .await?;
         assert_eq!(res.len(), 1024 * 128 - 20);
         assert_eq!(&res[..], &buf[20..]);
 
         // size should be total
-        let reader = client.blobs().read_at(hash, 0, Some(20)).await?;
+        let reader = client
+            .blobs()
+            .read_at(hash, 0, ReadAtLen::Exact(20))
+            .await?;
         assert_eq!(reader.size(), 1024 * 128);
         assert_eq!(reader.response_size, 20);
 
+        // last chunk - exact
+        let res = client
+            .blobs()
+            .read_at_to_bytes(hash, 1024 * 127, ReadAtLen::Exact(1024))
+            .await?;
+        assert_eq!(res.len(), 1024);
+        assert_eq!(res, &buf[1024 * 127..]);
+
+        // last chunk - open
+        let res = client
+            .blobs()
+            .read_at_to_bytes(hash, 1024 * 127, ReadAtLen::All)
+            .await?;
+        assert_eq!(res.len(), 1024);
+        assert_eq!(res, &buf[1024 * 127..]);
+
+        // last chunk - larger
+        let mut res = client
+            .blobs()
+            .read_at(hash, 1024 * 127, ReadAtLen::AtMost(2048))
+            .await?;
+        assert_eq!(res.size, 1024 * 128);
+        assert_eq!(res.response_size, 1024);
+        let res = res.read_to_bytes().await?;
+        assert_eq!(res.len(), 1024);
+        assert_eq!(res, &buf[1024 * 127..]);
+
         // out of bounds - too long
-        let res = client.blobs().read_at(hash, 0, Some(1024 * 128 + 1)).await;
+        let res = client
+            .blobs()
+            .read_at(hash, 0, ReadAtLen::Exact(1024 * 128 + 1))
+            .await;
         let err = res.unwrap_err();
         assert!(err.to_string().contains("out of bound"));
 
         // out of bounds - offset larger than blob
-        let res = client.blobs().read_at(hash, 1024 * 128 + 1, None).await;
+        let res = client
+            .blobs()
+            .read_at(hash, 1024 * 128 + 1, ReadAtLen::All)
+            .await;
         let err = res.unwrap_err();
         assert!(err.to_string().contains("out of range"));
 
         // out of bounds - offset + length too large
-        let res = client.blobs().read_at(hash, 1024 * 127, Some(1025)).await;
+        let res = client
+            .blobs()
+            .read_at(hash, 1024 * 127, ReadAtLen::Exact(1025))
+            .await;
         let err = res.unwrap_err();
         assert!(err.to_string().contains("out of bound"));
 

iroh/src/node/rpc.rs

@@ -1238,7 +1238,10 @@ impl<D: BaoStore> Handler<D> {
                 size
             );
 
-            let len = req.len.unwrap_or((size.value() - req.offset) as usize);
+            let len: usize = req
+                .len
+                .as_result_len(size.value() - req.offset)
+                .try_into()?;
 
             anyhow::ensure!(
                 req.offset + len as u64 <= size.value(),

iroh/src/rpc_protocol/blobs.rs

@@ -21,7 +21,9 @@ use nested_enum_utils::enum_conversions;
 use quic_rpc_derive::rpc_requests;
 use serde::{Deserialize, Serialize};
 
-use crate::client::blobs::{BlobInfo, BlobStatus, DownloadMode, IncompleteBlobInfo, WrapOption};
+use crate::client::blobs::{
+    BlobInfo, BlobStatus, DownloadMode, IncompleteBlobInfo, ReadAtLen, WrapOption,
+};
 
 use super::RpcService;
 
@@ -190,7 +192,7 @@ pub struct ReadAtRequest {
     /// Offset to start reading at
     pub offset: u64,
     /// Length of the data to get
-    pub len: Option<usize>,
+    pub len: ReadAtLen,
 }
 
 /// Response to [`ReadAtRequest`]