iroh_db_blobs/
stream.rs

1use std::{
2    collections::BTreeMap,
3    future::Future,
4    io::{Error as IoError, ErrorKind, SeekFrom},
5    pin::Pin,
6    sync::Arc,
7    task::{Context, Poll},
8};
9
10use iroh_db_core::BlobHash;
11use iroh_db_security::DomainKey;
12use iroh_db_store::BlobStore;
13use tokio::io::{AsyncRead, AsyncSeek, ReadBuf};
14
15use crate::{
16    BlobError,
17    engine::{MAX_CHUNK_SIZE, MIN_CHUNK_SIZE, decrypt_chunk},
18    model::{ChunkDescriptor, ManifestBody, ManifestEnvelope},
19    scheduler::BlobPriority,
20};
21
22type ChunkFuture = Pin<Box<dyn Future<Output = Result<Vec<u8>, BlobError>> + Send + 'static>>;
23pub(crate) type ChunkFetcher = Arc<
24    dyn Fn(
25            BlobHash,
26            u64,
27            BlobPriority,
28            bool,
29        ) -> Pin<Box<dyn Future<Output = Result<(), BlobError>> + Send + 'static>>
30        + Send
31        + Sync,
32>;
33
34/// Lazy authenticated plaintext reader over independently encrypted chunks.
35pub struct BlobStream {
36    store: Arc<dyn BlobStore>,
37    key: DomainKey,
38    envelope: ManifestEnvelope,
39    manifest: ManifestBody,
40    position: u64,
41    cached: Option<(usize, Vec<u8>)>,
42    pending: Option<(usize, ChunkFuture)>,
43    fetcher: Option<ChunkFetcher>,
44    priority: BlobPriority,
45    prefetch_chunks: usize,
46    prefetch_anchor: Option<usize>,
47    prefetch_tasks: BTreeMap<usize, tokio::task::JoinHandle<()>>,
48}
49
50impl BlobStream {
51    pub(crate) fn new(
52        store: Arc<dyn BlobStore>,
53        key: DomainKey,
54        envelope: ManifestEnvelope,
55        manifest: ManifestBody,
56    ) -> Self {
57        Self {
58            store,
59            key,
60            envelope,
61            manifest,
62            position: 0,
63            cached: None,
64            pending: None,
65            fetcher: None,
66            priority: BlobPriority::Normal,
67            prefetch_chunks: 0,
68            prefetch_anchor: None,
69            prefetch_tasks: BTreeMap::new(),
70        }
71    }
72
73    pub(crate) fn with_fetcher(
74        mut self,
75        fetcher: ChunkFetcher,
76        priority: BlobPriority,
77        prefetch_chunks: u8,
78    ) -> Self {
79        self.fetcher = Some(fetcher);
80        self.priority = priority;
81        self.prefetch_chunks = usize::from(prefetch_chunks.min(16));
82        self
83    }
84
85    /// Returns the authenticated plaintext length.
86    pub const fn len(&self) -> u64 {
87        self.manifest.plaintext_len
88    }
89
90    /// Returns whether the authenticated blob is empty.
91    pub const fn is_empty(&self) -> bool {
92        self.manifest.plaintext_len == 0
93    }
94
95    /// Returns the authenticated media type supplied at import time.
96    pub fn media_type(&self) -> Option<&str> {
97        self.manifest.media_type.as_deref()
98    }
99
100    fn begin_chunk_load(&mut self, index: usize) {
101        let store = self.store.clone();
102        let key = self.key.clone();
103        let domain_id = self.envelope.domain_id;
104        let epoch = self.envelope.epoch;
105        let salt = self.envelope.salt;
106        let descriptor = self.manifest.chunks[index].clone();
107        let fetcher = self.fetcher.clone();
108        let priority = self.priority;
109        let future = Box::pin(async move {
110            if let Some(fetcher) = fetcher {
111                fetcher(
112                    descriptor.ciphertext_hash,
113                    u64::from(descriptor.plaintext_len) + 16,
114                    priority,
115                    priority == BlobPriority::High,
116                )
117                .await?;
118            }
119            let ciphertext = store
120                .get(descriptor.ciphertext_hash)
121                .await
122                .map_err(BlobError::from)?
123                .ok_or(BlobError::MissingBlob(descriptor.ciphertext_hash))?;
124            decrypt_chunk(&key, domain_id, epoch, &salt, &descriptor, &ciphertext)
125        });
126        self.pending = Some((index, future));
127    }
128
129    fn schedule_prefetch(&mut self, index: usize) {
130        if self.prefetch_chunks == 0 || self.prefetch_anchor == Some(index) {
131            return;
132        }
133        self.prefetch_anchor = Some(index);
134        let Some(fetcher) = self.fetcher.clone() else {
135            return;
136        };
137        let priority = self.priority.lower();
138        let window_end = index
139            .saturating_add(self.prefetch_chunks)
140            .min(self.manifest.chunks.len().saturating_sub(1));
141        self.prefetch_tasks.retain(|task_index, task| {
142            let keep = !task.is_finished() && *task_index >= index && *task_index <= window_end;
143            if !keep && !task.is_finished() {
144                task.abort();
145            }
146            keep
147        });
148        let descriptors = self
149            .manifest
150            .chunks
151            .iter()
152            .enumerate()
153            .skip(index.saturating_add(1))
154            .take(self.prefetch_chunks)
155            .map(|(task_index, descriptor)| (task_index, descriptor.clone()))
156            .collect::<Vec<_>>();
157        for (task_index, descriptor) in descriptors {
158            if self.prefetch_tasks.contains_key(&task_index) {
159                continue;
160            }
161            let fetcher = fetcher.clone();
162            let store = self.store.clone();
163            let task = tokio::spawn(async move {
164                if store
165                    .get(descriptor.ciphertext_hash)
166                    .await
167                    .ok()
168                    .flatten()
169                    .is_none()
170                {
171                    let _ = fetcher(
172                        descriptor.ciphertext_hash,
173                        u64::from(descriptor.plaintext_len) + 16,
174                        priority,
175                        false,
176                    )
177                    .await;
178                }
179            });
180            self.prefetch_tasks.insert(task_index, task);
181        }
182    }
183
184    fn cancel_prefetch(&mut self) {
185        for (_, task) in std::mem::take(&mut self.prefetch_tasks) {
186            task.abort();
187        }
188        self.prefetch_anchor = None;
189    }
190}
191
192impl AsyncRead for BlobStream {
193    fn poll_read(
194        mut self: Pin<&mut Self>,
195        context: &mut Context<'_>,
196        output: &mut ReadBuf<'_>,
197    ) -> Poll<std::io::Result<()>> {
198        loop {
199            if self.position >= self.manifest.plaintext_len || output.remaining() == 0 {
200                return Poll::Ready(Ok(()));
201            }
202            let chunk_size = u64::from(self.manifest.chunk_size);
203            let chunk_index_u64 = self.position / chunk_size;
204            let chunk_index = usize::try_from(chunk_index_u64).map_err(|error| {
205                IoError::new(ErrorKind::InvalidData, BlobError::Codec(error.to_string()))
206            })?;
207            self.schedule_prefetch(chunk_index);
208            if let Some((cached_index, bytes)) = &self.cached
209                && *cached_index == chunk_index
210            {
211                let offset = usize::try_from(self.position % chunk_size).map_err(|error| {
212                    IoError::new(ErrorKind::InvalidData, BlobError::Codec(error.to_string()))
213                })?;
214                let available = &bytes[offset..];
215                let count = available.len().min(output.remaining());
216                output.put_slice(&available[..count]);
217                self.position += u64::try_from(count).map_err(|error| {
218                    IoError::new(ErrorKind::InvalidData, BlobError::Codec(error.to_string()))
219                })?;
220                return Poll::Ready(Ok(()));
221            }
222
223            if self
224                .pending
225                .as_ref()
226                .is_none_or(|(pending_index, _)| *pending_index != chunk_index)
227            {
228                self.begin_chunk_load(chunk_index);
229            }
230            let (pending_index, future) = self
231                .pending
232                .as_mut()
233                .expect("a chunk future was just installed");
234            match future.as_mut().poll(context) {
235                Poll::Pending => return Poll::Pending,
236                Poll::Ready(Ok(bytes)) => {
237                    let index = *pending_index;
238                    self.pending = None;
239                    self.cached = Some((index, bytes));
240                }
241                Poll::Ready(Err(error)) => {
242                    self.pending = None;
243                    return Poll::Ready(Err(IoError::new(ErrorKind::InvalidData, error)));
244                }
245            }
246        }
247    }
248}
249
250impl AsyncSeek for BlobStream {
251    fn start_seek(mut self: Pin<&mut Self>, position: SeekFrom) -> std::io::Result<()> {
252        let base = match position {
253            SeekFrom::Start(position) => i128::from(position),
254            SeekFrom::End(offset) => i128::from(self.manifest.plaintext_len) + i128::from(offset),
255            SeekFrom::Current(offset) => i128::from(self.position) + i128::from(offset),
256        };
257        let position = u64::try_from(base)
258            .map_err(|_| IoError::new(ErrorKind::InvalidInput, "seek before start"))?;
259        self.position = position;
260        self.pending = None;
261        self.cancel_prefetch();
262        Ok(())
263    }
264
265    fn poll_complete(
266        self: Pin<&mut Self>,
267        _context: &mut Context<'_>,
268    ) -> Poll<std::io::Result<u64>> {
269        Poll::Ready(Ok(self.position))
270    }
271}
272
273impl Drop for BlobStream {
274    fn drop(&mut self) {
275        self.cancel_prefetch();
276    }
277}
278
279pub(crate) fn validate_manifest(manifest: &ManifestBody) -> Result<(), BlobError> {
280    let chunk_size = manifest.chunk_size as usize;
281    if !(MIN_CHUNK_SIZE..=MAX_CHUNK_SIZE).contains(&chunk_size) {
282        return Err(BlobError::InvalidManifest);
283    }
284    let mut total = 0_u64;
285    for (index, chunk) in manifest.chunks.iter().enumerate() {
286        let expected_index = u32::try_from(index).map_err(|_| BlobError::InvalidManifest)?;
287        if chunk.index != expected_index
288            || chunk.plaintext_len == 0
289            || chunk.plaintext_len as usize > chunk_size
290        {
291            return Err(BlobError::InvalidManifest);
292        }
293        if index + 1 != manifest.chunks.len() && chunk.plaintext_len as usize != chunk_size {
294            return Err(BlobError::InvalidManifest);
295        }
296        total = total
297            .checked_add(u64::from(chunk.plaintext_len))
298            .ok_or(BlobError::InvalidManifest)?;
299    }
300    if total != manifest.plaintext_len
301        || (manifest.plaintext_len == 0) != manifest.chunks.is_empty()
302    {
303        return Err(BlobError::InvalidManifest);
304    }
305    Ok(())
306}
307
308#[allow(dead_code)]
309fn _descriptor_is_send(_: &ChunkDescriptor) {}
310
311#[cfg(test)]
312mod tests {
313    use std::sync::atomic::{AtomicUsize, Ordering};
314
315    use iroh_db_core::DomainId;
316    use iroh_db_store::{BlobFuture, BlobStoreError};
317    use tokio::io::{AsyncReadExt as _, AsyncSeekExt as _};
318
319    use super::*;
320
321    struct EmptyStore;
322
323    impl BlobStore for EmptyStore {
324        fn put(&self, _bytes: Vec<u8>) -> BlobFuture<'_, BlobHash> {
325            Box::pin(async { Err(BlobStoreError::new("put is unused")) })
326        }
327
328        fn get(&self, _hash: BlobHash) -> BlobFuture<'_, Option<Vec<u8>>> {
329            Box::pin(async { Ok(None) })
330        }
331
332        fn list(&self) -> BlobFuture<'_, Vec<BlobHash>> {
333            Box::pin(async { Ok(Vec::new()) })
334        }
335    }
336
337    #[tokio::test]
338    async fn cached_chunk_read_restarts_prefetch_after_seek() {
339        const CHUNK_SIZE: usize = 64 * 1024;
340        let chunk_size = u32::try_from(CHUNK_SIZE).unwrap();
341        let chunks = (0_u8..3)
342            .map(|index| ChunkDescriptor {
343                index: u32::from(index),
344                plaintext_len: chunk_size,
345                ciphertext_hash: BlobHash::from_bytes([index + 1; 32]),
346                nonce: [index; 24],
347            })
348            .collect();
349        let envelope = ManifestEnvelope {
350            version: 1,
351            domain_id: DomainId::from_bytes([1; 32]),
352            epoch: 0,
353            salt: [2; 32],
354            nonce: [3; 24],
355            ciphertext: Vec::new(),
356        };
357        let body = ManifestBody {
358            plaintext_len: u64::try_from(CHUNK_SIZE * 3).unwrap(),
359            plaintext_hash: [4; 32],
360            chunk_size,
361            chunks,
362            media_type: None,
363        };
364        let calls = Arc::new(AtomicUsize::new(0));
365        let fetch_calls = calls.clone();
366        let fetcher: ChunkFetcher = Arc::new(move |_, _, _, _| {
367            let fetch_calls = fetch_calls.clone();
368            Box::pin(async move {
369                fetch_calls.fetch_add(1, Ordering::Relaxed);
370                Ok(())
371            })
372        });
373        let mut stream = BlobStream::new(
374            Arc::new(EmptyStore),
375            DomainKey::from_bytes([5; 32]),
376            envelope,
377            body,
378        )
379        .with_fetcher(fetcher, BlobPriority::High, 2);
380        stream.cached = Some((0, vec![7; CHUNK_SIZE]));
381
382        let mut byte = [0_u8; 1];
383        stream.read_exact(&mut byte).await.unwrap();
384        wait_for_calls(&calls, 2).await;
385        stream.seek(SeekFrom::Start(1024)).await.unwrap();
386        stream.read_exact(&mut byte).await.unwrap();
387        wait_for_calls(&calls, 4).await;
388    }
389
390    async fn wait_for_calls(calls: &AtomicUsize, expected: usize) {
391        tokio::time::timeout(std::time::Duration::from_secs(1), async {
392            while calls.load(Ordering::Relaxed) < expected {
393                tokio::task::yield_now().await;
394            }
395        })
396        .await
397        .unwrap();
398    }
399}
iroh_db_blobs/stream.rs

iroh_db_blobs/
stream.rs
