1use crate::LOG_MODULE;
2use crate::SharedError;
3use crate::epochs::{epoch_to_slot_range, slot_to_epoch};
4use crate::firehose::FirehoseError;
5use crate::index::{SlotOffsetIndexError, slot_to_range};
6use crate::node::{Node, NodeWithCid, NodesWithCids, parse_any_from_cbordata};
7use crate::utils;
8use cid::Cid;
9use once_cell::sync::Lazy;
10use reqwest::RequestBuilder;
11use rseek::Seekable;
12use std::io;
13use std::io::SeekFrom;
14use std::sync::atomic::{AtomicU64, Ordering};
15use std::time::Instant;
16use std::vec::Vec;
17use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt};
18use tokio::time::sleep;
19
20const MAX_VARINT_LEN_64: usize = 10;
21const MIN_SEEK_SPACING_MS: u64 = 15;
22static SEEK_START_INSTANT: Lazy<Instant> = Lazy::new(Instant::now);
23static LAST_SEEK_HIT_TIME: AtomicU64 = AtomicU64::new(0);
24
25pub static TOTAL_BYTES_READ: AtomicU64 = AtomicU64::new(0);
28
29pub async fn read_uvarint<R: AsyncRead + Unpin>(reader: &mut R) -> io::Result<u64> {
31 let mut x = 0u64;
32 let mut s = 0u32;
33 let mut buffer = [0u8; 1];
34
35 for i in 0..MAX_VARINT_LEN_64 {
36 reader.read_exact(&mut buffer).await?;
37 let b = buffer[0];
38 if b < 0x80 {
39 if i == MAX_VARINT_LEN_64 - 1 && b > 1 {
40 return Err(io::Error::new(
41 io::ErrorKind::InvalidData,
42 "uvarint overflow",
43 ));
44 }
45 return Ok(x | ((b as u64) << s));
46 }
47 x |= ((b & 0x7f) as u64) << s;
48 s += 7;
49
50 if s > 63 {
51 return Err(io::Error::new(
52 io::ErrorKind::InvalidData,
53 "uvarint too long",
54 ));
55 }
56 }
57 Err(io::Error::new(
58 io::ErrorKind::InvalidData,
59 "uvarint overflow",
60 ))
61}
62
63#[derive(Clone, PartialEq, Eq, Hash)]
65pub struct RawNode {
66 pub cid: Cid,
68 pub data: Vec<u8>,
70}
71
72impl core::fmt::Debug for RawNode {
74 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
75 f.debug_struct("RawNode")
76 .field("cid", &self.cid)
77 .field("data", &self.data)
78 .finish()
79 }
80}
81
82impl RawNode {
83 pub const fn new(cid: Cid, data: Vec<u8>) -> RawNode {
85 RawNode { cid, data }
86 }
87
88 pub fn parse(&self) -> Result<Node, SharedError> {
90 match parse_any_from_cbordata(self.data.clone()) {
91 Ok(node) => Ok(node),
92 Err(err) => {
93 println!("Error: {:?}", err);
94 Err(Box::new(std::io::Error::other("Unknown type".to_owned())))
95 }
96 }
97 }
98
99 pub async fn from_cursor(cursor: &mut io::Cursor<Vec<u8>>) -> Result<RawNode, SharedError> {
101 let cid_version = read_uvarint(cursor).await?;
102 let multicodec = read_uvarint(cursor).await?;
105 let hash_function = read_uvarint(cursor).await?;
109 let digest_length = read_uvarint(cursor).await?;
113 if digest_length > 64 {
116 return Err(Box::new(std::io::Error::other(format!(
117 "Digest length too long, position={}",
118 cursor.position()
119 ))));
120 }
121
122 let mut digest = vec![0u8; digest_length as usize];
124 cursor.read_exact(&mut digest).await?;
125
126 let mut data = vec![];
128 cursor.read_to_end(&mut data).await?;
129
130 let ha = multihash::Multihash::wrap(hash_function, digest.as_slice())?;
133
134 match cid_version {
135 0 => {
136 let cid = Cid::new_v0(ha)?;
137 let raw_node = RawNode::new(cid, data);
138 Ok(raw_node)
139 }
140 1 => {
141 let cid = Cid::new_v1(multicodec, ha);
142 let raw_node = RawNode::new(cid, data);
143 Ok(raw_node)
144 }
145 _ => Err(Box::new(std::io::Error::other(
146 "Unknown CID version".to_owned(),
147 ))),
148 }
149 }
150}
151
152pub trait Len {
154 fn len(&self) -> u64;
156 fn is_empty(&self) -> bool {
158 self.len() == 0
159 }
160}
161
162impl<F> Len for Seekable<F>
163where
164 F: Fn() -> RequestBuilder + Send + Sync + 'static,
165{
166 fn len(&self) -> u64 {
167 self.file_size.unwrap_or(0)
168 }
169}
170
171pub struct NodeReader<R: AsyncRead + AsyncSeek + Len> {
173 pub reader: R,
175 pub header: Vec<u8>,
177 pub item_index: u64,
179 seek_permit_primed: bool,
181}
182
183impl<R: AsyncRead + Unpin + AsyncSeek + Len> NodeReader<R> {
184 pub const fn new(reader: R) -> NodeReader<R> {
186 NodeReader {
187 reader,
188 header: vec![],
189 item_index: 0,
190 seek_permit_primed: false,
191 }
192 }
193
194 pub async fn read_raw_header(&mut self) -> Result<Vec<u8>, SharedError> {
196 if !self.header.is_empty() {
197 return Ok(self.header.clone());
198 };
199 let header_length = read_uvarint(&mut self.reader).await?;
200 if header_length > 1024 {
201 return Err(Box::new(std::io::Error::other(
202 "Header length too long".to_owned(),
203 )));
204 }
205 let mut header = vec![0u8; header_length as usize];
206 self.reader.read_exact(&mut header).await?;
207
208 self.header.clone_from(&header);
209
210 let clone = header.clone();
211 Ok(clone.as_slice().to_owned())
212 }
213
214 pub async fn seek_to_slot(&mut self, slot: u64) -> Result<(), FirehoseError> {
219 self.seek_to_slot_inner(slot).await
220 }
221
222 pub async fn prime_seek_permit(&mut self) {
227 wait_for_seek_hit_slot().await;
228 self.seek_permit_primed = true;
229 }
230
231 async fn seek_to_slot_inner(&mut self, slot: u64) -> Result<(), FirehoseError> {
232 if self.header.is_empty() {
233 self.read_raw_header()
234 .await
235 .map_err(FirehoseError::SeekToSlotError)?;
236 };
237
238 let epoch = slot_to_epoch(slot);
239 let (epoch_start, epoch_end_inclusive) = epoch_to_slot_range(epoch);
240 let mut current = slot;
241 loop {
242 match slot_to_range(current).await {
243 Ok((offset, _)) => {
244 log::info!(
245 target: LOG_MODULE,
246 "Seeking to slot {} in epoch {} @ offset {}",
247 current,
248 epoch,
249 offset
250 );
251 return self.seek_to_offset(offset).await;
252 }
253 Err(SlotOffsetIndexError::SlotNotFound(..)) => {
254 if current >= epoch_end_inclusive {
255 return self.seek_to_epoch_tail(slot, epoch, epoch_start).await;
261 }
262 log::warn!(
263 target: LOG_MODULE,
264 "Slot {} not found in index, seeking to next slot",
265 current
266 );
267 current += 1;
268 }
269 Err(err) => return Err(FirehoseError::SlotOffsetIndexError(err)),
272 }
273 }
274 }
275
276 async fn seek_to_epoch_tail(
279 &mut self,
280 slot: u64,
281 epoch: u64,
282 epoch_start: u64,
283 ) -> Result<(), FirehoseError> {
284 let mut candidate = slot;
285 while candidate > epoch_start {
286 candidate -= 1;
287 match slot_to_range(candidate).await {
288 Ok((offset, length)) => {
289 log::warn!(
290 target: LOG_MODULE,
291 "No slot at or after {} is present in epoch {}; seeking past the last \
292 present slot {}",
293 slot,
294 epoch,
295 candidate
296 );
297 return self.seek_to_offset(offset + length).await;
298 }
299 Err(SlotOffsetIndexError::SlotNotFound(..)) => continue,
300 Err(err) => return Err(FirehoseError::SlotOffsetIndexError(err)),
301 }
302 }
303 Err(FirehoseError::SlotOffsetIndexError(
307 SlotOffsetIndexError::EpochHasNoIndexedSlots(epoch),
308 ))
309 }
310
311 async fn seek_to_offset(&mut self, offset: u64) -> Result<(), FirehoseError> {
312 if !std::mem::take(&mut self.seek_permit_primed) {
313 wait_for_seek_hit_slot().await;
314 }
315 self.reader
316 .seek(SeekFrom::Start(offset))
317 .await
318 .map_err(|e| FirehoseError::SeekToSlotError(Box::new(e)))?;
319 Ok(())
320 }
321
322 #[allow(clippy::should_implement_trait)]
323 pub async fn next(&mut self) -> Result<RawNode, SharedError> {
325 if self.header.is_empty() {
326 self.read_raw_header().await?;
327 };
328
329 self.item_index += 1;
331
332 let section_size = read_uvarint(&mut self.reader).await?;
334 if section_size > utils::MAX_ALLOWED_SECTION_SIZE as u64 {
337 return Err(Box::new(std::io::Error::other(
338 "Section size too long".to_owned(),
339 )));
340 }
341
342 let mut item = vec![0u8; section_size as usize];
344 self.reader.read_exact(&mut item).await?;
345 TOTAL_BYTES_READ.fetch_add(section_size, Ordering::Relaxed);
346
347 let mut cursor = io::Cursor::new(item);
352
353 RawNode::from_cursor(&mut cursor).await
354 }
355
356 pub async fn next_parsed(&mut self) -> Result<NodeWithCid, SharedError> {
358 let raw_node = self.next().await?;
359 let cid = raw_node.cid;
360 Ok(NodeWithCid::new(cid, raw_node.parse()?))
361 }
362
363 pub async fn read_until_block(&mut self) -> Result<NodesWithCids, SharedError> {
365 let mut nodes = NodesWithCids::new();
366 loop {
367 let node = match self.next_parsed().await {
368 Ok(node) => node,
369 Err(e)
370 if e.downcast_ref::<io::Error>()
371 .is_some_and(|io_err| io_err.kind() == io::ErrorKind::UnexpectedEof) =>
372 {
373 break;
374 }
375 Err(e) => return Err(e),
376 };
377 if node.get_node().is_block() {
378 nodes.push(node);
379 break;
380 }
381 nodes.push(node);
382 }
383 Ok(nodes)
384 }
385
386 pub const fn get_item_index(&self) -> u64 {
388 self.item_index
389 }
390}
391
392fn seek_monotonic_millis() -> u64 {
393 let elapsed = SEEK_START_INSTANT.elapsed().as_millis();
394 if elapsed > u64::MAX as u128 {
395 u64::MAX
396 } else {
397 elapsed as u64
398 }
399}
400
401async fn wait_for_seek_hit_slot() {
402 loop {
403 let now_ms = seek_monotonic_millis();
404 let last_hit = LAST_SEEK_HIT_TIME.load(Ordering::Relaxed);
405 let since_last = now_ms.saturating_sub(last_hit);
406 if since_last < MIN_SEEK_SPACING_MS {
407 sleep(std::time::Duration::from_millis(
410 MIN_SEEK_SPACING_MS - since_last,
411 ))
412 .await;
413 continue;
414 }
415 if LAST_SEEK_HIT_TIME
416 .compare_exchange(last_hit, now_ms, Ordering::SeqCst, Ordering::SeqCst)
417 .is_ok()
418 {
419 return;
420 }
421 sleep(std::time::Duration::from_millis(MIN_SEEK_SPACING_MS)).await;
423 }
424}
425
426pub fn cid_from_cbor_link(val: &serde_cbor::Value) -> Result<cid::Cid, SharedError> {
428 if let serde_cbor::Value::Bytes(b) = val
429 && b.first() == Some(&0)
430 {
431 return Ok(cid::Cid::try_from(b[1..].to_vec())?);
432 }
433 Err("invalid DAG‑CBOR link encoding".into())
434}
435
436#[cfg(test)]
437mod tests {
438 use super::*;
439 use serde_cbor::Value::{self, Array, Bytes, Integer, Map, Null, Text};
440 use sha2::{Digest, Sha256};
441
442 impl Len for io::Cursor<Vec<u8>> {
443 fn len(&self) -> u64 {
444 self.get_ref().len() as u64
445 }
446 }
447
448 fn append_varint(mut value: usize, bytes: &mut Vec<u8>) {
449 while value >= 128 {
450 bytes.push((value as u8 & 0x7f) | 0x80);
451 value >>= 7;
452 }
453 bytes.push(value as u8);
454 }
455
456 fn archive(nodes: &[Value]) -> (Vec<u8>, Vec<Cid>, Vec<u8>) {
457 let mut sections = Vec::new();
458 let mut cids = Vec::new();
459 for node in nodes {
460 let data = serde_cbor::to_vec(node).unwrap();
461 let digest = Sha256::digest(&data);
462 let cid = Cid::new_v1(0x71, multihash::Multihash::wrap(0x12, &digest).unwrap());
463 let cid_bytes = cid.to_bytes();
464 append_varint(cid_bytes.len() + data.len(), &mut sections);
465 sections.extend(cid_bytes);
466 sections.extend(data);
467 cids.push(cid);
468 }
469 let mut root = vec![0];
470 root.extend(cids.last().unwrap().to_bytes());
471 let header = serde_cbor::to_vec(&Map([
472 (Text("roots".into()), Array(vec![Bytes(root)])),
473 (Text("version".into()), Integer(1)),
474 ]
475 .into()))
476 .unwrap();
477 let mut car = Vec::new();
478 append_varint(header.len(), &mut car);
479 car.extend_from_slice(&header);
480 car.extend(sections);
481 (car, cids, header)
482 }
483
484 #[tokio::test]
485 async fn test_async_node_reader() {
486 let payload = vec![42; 256];
488 let block = |slot| {
489 Array(vec![
490 Integer(2),
491 Integer(slot),
492 Array(vec![]),
493 Array(vec![]),
494 Array(vec![Integer(slot - 1), Integer(1_700_000_000), Integer(7)]),
495 Null,
496 ])
497 };
498 let (car, cids, header) = archive(&[
499 Array(vec![Integer(6), Null, Null, Null, Bytes(payload.clone())]),
500 block(42),
501 block(43),
502 ]);
503 let mut reader = NodeReader::new(io::Cursor::new(car));
504
505 assert_eq!(reader.read_raw_header().await.unwrap(), header);
506 let section_start = reader.reader.position();
507 assert_eq!(reader.read_raw_header().await.unwrap(), header);
508 assert_eq!(reader.reader.position(), section_start);
509
510 let first = reader.read_until_block().await.unwrap();
511 assert_eq!(first.get_cids(), cids[..2]);
512 assert_eq!(
513 first
514 .get(0)
515 .get_node()
516 .get_dataframe()
517 .unwrap()
518 .data
519 .as_slice(),
520 payload,
521 );
522 let first_block = first.get_block().unwrap();
523 assert_eq!(first_block.slot, 42);
524 assert_eq!(first_block.meta.parent_slot, 41);
525 assert_eq!(first_block.meta.blocktime, 1_700_000_000);
526
527 let second = reader.read_until_block().await.unwrap();
528 assert_eq!(second.get_cids(), cids[2..]);
529 assert_eq!(second.get_block().unwrap().slot, 43);
530 assert_eq!(reader.get_item_index(), 3);
531 assert!(reader.read_until_block().await.unwrap().is_empty());
532 }
533
534 #[tokio::test]
535 async fn read_until_block_propagates_invalid_node() {
536 let (car, _, _) = archive(&[Array(vec![Integer(999)])]);
537 let mut reader = NodeReader::new(io::Cursor::new(car));
538 let error = reader.read_until_block().await.err().unwrap();
539 assert_eq!(
540 error.downcast_ref::<io::Error>().unwrap().kind(),
541 io::ErrorKind::Other,
542 );
543 }
544}