diff --git a/src/car.rs b/src/car.rs index 410d1ee..97b1140 100644 --- a/src/car.rs +++ b/src/car.rs @@ -5,6 +5,7 @@ use cbor4ii::core::dec::Read; use pyo3::prelude::*; use pyo3::types::*; +use crate::cid::parse_cid_prefix; use crate::dag_cbor::de::to_pyobject; use crate::error::value_error; use crate::ffi::recursion::current_recursion_limit; @@ -68,25 +69,20 @@ pub fn decode_car<'py>(py: Python<'py>, data: &[u8]) -> PyResult<(Py, Bou } let cid_bytes_before = buf.buf; - // `&[u8]` is itself an `io::Read`, so we hand it to `Cid::read_bytes` - // directly and recover the consumed length from the slice shrink. - let mut slice: &[u8] = cid_bytes_before; - let cid_result = ::cid::Cid::read_bytes(&mut slice); - let Ok(cid) = cid_result else { + let Some((consumed, codec)) = parse_cid_prefix(cid_bytes_before) else { return Err(value_error( "Failed to read CID of block", - cid_result.unwrap_err().to_string(), + "Invalid CID".to_string(), )); }; - if cid.codec() != 0x71 { + if codec != 0x71 { return Err(value_error( "Failed to read CAR block", "Unsupported codec. For now we support only DAG-CBOR (0x71)".to_string(), )); } - let consumed = cid_bytes_before.len() - slice.len(); buf.advance(consumed); let cid_raw = &cid_bytes_before[..consumed]; diff --git a/src/cid.rs b/src/cid.rs index 834d6db..c7b010e 100644 --- a/src/cid.rs +++ b/src/cid.rs @@ -27,6 +27,53 @@ pub(crate) fn looks_like_cid(bytes: &[u8]) -> bool { bytes.len() == 34 && bytes[0] == 0x12 && bytes[1] == 0x20 } +// Minimal-encoding unsigned varint, as `unsigned_varint::decode::u64`: +// ≤10 bytes, last byte of a multi-byte varint must be non-zero. +#[inline] +fn read_varint(bytes: &[u8], pos: &mut usize) -> Option { + let mut n: u64 = 0; + for i in 0..10 { + let &b = bytes.get(*pos + i)?; + n |= ((b & 0x7f) as u64) << (i * 7); + if b & 0x80 == 0 { + if b == 0 && i > 0 { + return None; + } + *pos += i + 1; + return Some(n); + } + } + None +} + +/// Structural check that `bytes` starts with a valid binary CID; returns +/// `(consumed, codec)`. Accepts exactly what `::cid::Cid::try_from` accepts +/// (trailing bytes are the caller's concern) but skips the `Multihash` +/// construction and its 64-byte digest copy. +#[inline] +pub(crate) fn parse_cid_prefix(bytes: &[u8]) -> Option<(usize, u64)> { + let mut pos = 0; + let version = read_varint(bytes, &mut pos)?; + let codec = read_varint(bytes, &mut pos)?; + + // CIDv0: `0x12 0x20` + 32-byte sha2-256 digest; codec is implicitly dag-pb + if (version, codec) == (0x12, 0x20) { + let end = pos + 32; + return (bytes.len() >= end).then_some((end, 0x70)); + } + if version != 1 { + return None; + } + + read_varint(bytes, &mut pos)?; // multihash code + let hash_size = read_varint(bytes, &mut pos)?; + if hash_size > 64 { + return None; + } + let end = pos + hash_size as usize; + (bytes.len() >= end).then_some((end, codec)) +} + pub(crate) fn extract_cid(data: &Bound) -> PyResult<::cid::Cid> { let cid = if let Ok(s) = data.cast::() { ::cid::Cid::try_from(s.to_str()?) diff --git a/src/dag_cbor/de.rs b/src/dag_cbor/de.rs index 07e9e0a..a13d6b9 100644 --- a/src/dag_cbor/de.rs +++ b/src/dag_cbor/de.rs @@ -5,6 +5,7 @@ use cbor4ii::core::{ }; use pyo3::{ffi, prelude::*, types::*, BoundObject}; +use crate::cid::parse_cid_prefix; use crate::error::value_error; use crate::ffi::dict::new_presized; use crate::ffi::key_cache::intern; @@ -158,7 +159,7 @@ where } let cid_without_prefix = &cid[1..]; - if ::cid::Cid::try_from(cid_without_prefix).is_err() { + if parse_cid_prefix(cid_without_prefix).is_none() { return Err(anyhow!("Invalid CID")); } diff --git a/src/dag_cbor/ser.rs b/src/dag_cbor/ser.rs index 280408d..5b2b725 100644 --- a/src/dag_cbor/ser.rs +++ b/src/dag_cbor/ser.rs @@ -6,7 +6,7 @@ use cbor4ii::core::{ use pyo3::pybacked::PyBackedStr; use pyo3::{ffi, prelude::*, types::*}; -use crate::cid::looks_like_cid; +use crate::cid::{looks_like_cid, parse_cid_prefix}; use crate::error::value_error; use crate::io::VecWriter; @@ -133,7 +133,7 @@ where if tp == &raw mut ffi::PyBytes_Type { let b = obj.cast_unchecked::(); let bytes = b.as_bytes(); - if looks_like_cid(bytes) && ::cid::Cid::try_from(bytes).is_ok() { + if looks_like_cid(bytes) && parse_cid_prefix(bytes).is_some() { // by providing custom encoding we avoid extra allocation types::Tag(42, PrefixedCidBytes(bytes)).encode(w)?; } else { @@ -187,7 +187,7 @@ where Ok(()) } else if let Ok(b) = obj.cast::() { let bytes = b.as_bytes(); - if looks_like_cid(bytes) && ::cid::Cid::try_from(bytes).is_ok() { + if looks_like_cid(bytes) && parse_cid_prefix(bytes).is_some() { types::Tag(42, PrefixedCidBytes(bytes)).encode(w)?; } else { types::Bytes(bytes).encode(w)?; diff --git a/src/io/writer.rs b/src/io/writer.rs index ea79670..e500c1b 100644 --- a/src/io/writer.rs +++ b/src/io/writer.rs @@ -1,15 +1,27 @@ +use std::cell::Cell; use std::convert::Infallible; use cbor4ii::core::enc; -// `enc::Write` over a raw `Vec`: no syscalls behind it, so a `BufWriter` -// wrapper would just add a memcpy per push for no benefit. +// Retaining bigger buffers would pin worst-case memory per thread forever. +const MAX_POOLED_CAPACITY: usize = 1 << 20; + +thread_local! { + static POOL: Cell> = const { Cell::new(Vec::new()) }; +} + +// `enc::Write` over a `Vec` recycled through a thread-local pool: encoding +// many small records reuses one grown allocation instead of re-growing from +// zero on every call. `Cell::take` leaves an empty `Vec` behind, so a +// re-entrant encode just falls back to a fresh buffer. pub(crate) struct VecWriter(Vec); impl VecWriter { #[inline] pub(crate) fn new() -> Self { - VecWriter(Vec::new()) + let mut buf = POOL.take(); + buf.clear(); + VecWriter(buf) } #[inline] @@ -18,6 +30,15 @@ impl VecWriter { } } +impl Drop for VecWriter { + fn drop(&mut self) { + let buf = std::mem::take(&mut self.0); + if buf.capacity() <= MAX_POOLED_CAPACITY { + POOL.set(buf); + } + } +} + impl enc::Write for VecWriter { type Error = Infallible;