fix(io/time_index): cap declared track length to prevent OOM on malformed file#223
fix(io/time_index): cap declared track length to prevent OOM on malformed file#223itsuzef wants to merge 2 commits into
Conversation
…rmed file read_track allocates Vec<TimeIndexEntry> with capacity derived from the on-disk count, gated only by `count * 16 == payload_bytes`. A crafted .mv2 with a self-consistent but absurd `bytes_length` in TimeIndexManifest passes that check and triggers an unbounded allocation before any read_exact would fail. Mirror the safety ceiling the sibling temporal_index reader already enforces (16 GiB) and reject oversized lengths up front.
There was a problem hiding this comment.
I like this patch — the cap is a strict improvement over main. The reason I went request-changes rather than approve is that I think the fix as written closes a narrower window than the PR description implies, and I'd like to talk through that before it lands.
What I saw when I tried to reproduce
Built a small harness using a custom GlobalAlloc that intercepts allocations ≥ 1 GiB and returns a fake aligned pointer (so the test process doesn't actually commit memory), then called memvid_core::io::time_index::read_track with a synthesized 12-byte header and a declared length = 12 + count * 16.
| Branch | declared_length |
count |
Result | large_allocations |
max_alloc (bytes) |
|---|---|---|---|---|---|
main (178e277) |
17,179,869,212 | 1,073,741,825 | Io { UnexpectedEof } |
1 | 17,179,869,200 |
main |
17,179,869,180 | 1,073,741,823 | Io { UnexpectedEof } |
1 | 17,179,869,168 |
| this PR (2cc7ea1) | 17,179,869,212 | — | InvalidTimeIndex { "length exceeds supported limit" } |
0 | 0 |
| this PR | 17,179,869,180 | 1,073,741,823 | Io { UnexpectedEof } |
1 | 17,179,869,168 |
Results were stable across runs.
Where I think the gap is
MAX_TIME_INDEX_BYTES = 1 << 34 = 17,179,869,184. The reader still does, at src/io/time_index.rs#L112:
let mut entries = Vec::with_capacity(count as usize);after only checking payload_bytes == count * 16. With length = 17,179,869,180 (≤ cap, payload aligned), count = 1,073,741,823, and sizeof(TimeIndexEntry) = 16, that with_capacity asks the allocator for ~16 GiB before any payload byte has been read. So the property the PR claims to establish — that count-proportional allocation can no longer be driven by file metadata — still holds only for declared lengths over the ceiling. Anything aligned and just under the ceiling exercises the same path, just bounded at 16 GiB instead of unbounded.
The new test (read_rejects_oversized_declared_length) covers MAX_TIME_INDEX_BYTES + 1, which is the easy side. The interesting boundary for a malformed-file invariant is the largest aligned value strictly under the cap.
Things worth considering
A few options, ordered roughly by how invasive they are:
- Bound the allocation rather than the declared length — either stream entries into the
Vecwithoutwith_capacity(count), or clamp the initial capacity (count.min(INITIAL_CAPACITY) as usize) and letVecgrow asread_exactsucceeds. That way boguscountonly allocates what the file actually delivers. - Add a boundary test at
MAX_TIME_INDEX_BYTES - 16(or whatever the largest aligned, accepted value works out to) that asserts either an early error or bounded allocation. The current test would still pass if the underlying path were wide open. - If the temporal-index reader is the model, it's worth double-checking — src/io/temporal_index.rs does a
vec![0u8; (length - HEADER_SIZE) as usize]body read before parsing, so the same near-cap amplifier exists there. Whichever shape we pick, it'd be good to apply it to both readers in one pass. - Minor: the
Cargo.lock2.0.135 → 2.0.139 bump looks incidental to this change. Keeping it out of this diff makes the security-relevant change easier to review on its own and avoids merge churn.
Totally possible the intent here is "raise the ceiling, accept up to 16 GiB of pre-read allocation as the cost of avoiding streaming" — that's a defensible call and I'm happy to revisit on that framing. If so, I think it's worth saying so explicitly in the PR description and dropping a comment near the with_capacity call, because today's framing reads as "prevents OOM" and the data above shows it narrows rather than prevents.
Repro, for reference
// custom GlobalAlloc that intercepts allocations >= 1 GiB and returns a fake aligned
// pointer + AtomicUsize counters for n_calls and max requested size.
let count: u64 = std::env::var("COUNT").ok()
.and_then(|s| s.parse().ok())
.unwrap_or((1u64 << 30) + 1);
let declared_length = 12 + count * 16;
let mut buf = Vec::with_capacity(12);
buf.extend_from_slice(&memvid_core::io::time_index::TIME_INDEX_MAGIC);
buf.extend_from_slice(&count.to_le_bytes());
let mut cursor = std::io::Cursor::new(buf);
let err = memvid_core::io::time_index::read_track(&mut cursor, 0, declared_length).unwrap_err();COUNT=1073741823 against this branch is the case that exercises the gap.
The existing length cap rejects declared lengths above 16 GiB but, for any
self-consistent header whose declared length sits just under the cap, the
readers still reserved count-proportional memory before validating a single
byte of payload:
- io::time_index::read_track used Vec::with_capacity(count) — up to ~16 GiB
of pre-read capacity for a hostile count aligned under the ceiling.
- io::temporal_index::read_track slurped the entire body into a single
Vec<u8> sized by (length - header) and then allocated two more vectors
sized by entry_count / anchor_count.
Both readers now bound the eager allocation:
- time_index: initial Vec capacity clamped to INITIAL_ENTRIES_CAPACITY (1024);
the vector grows naturally as read_exact succeeds.
- temporal_index: body is streamed one record at a time into fixed-size
stack buffers, hashed incrementally, and decoded directly into vectors
whose capacity is clamped to INITIAL_RECORDS_CAPACITY (1024). The body
Vec<u8> is gone entirely.
New near-cap tests construct the largest aligned declared length that passes
both the ceiling and the count-consistency checks, provide only the header
bytes, and assert that the reader fails fast at the first body read with
UnexpectedEof. If the eager-allocation path had not been bounded, the test
process would have requested ~16 GiB up front.
Tests:
- cargo test --lib → 340/340 pass
- cargo test --lib --features temporal_track → 346/346 pass
- cargo clippy --lib --no-deps → clean
- cargo fmt -- --check → clean
|
Thanks for the careful repro and the table — that was exactly the right way to push on this. New commit (
PR description updated to reframe the change as "bound allocation to actual payload read" rather than "prevent OOM" — your data showed the original framing was overclaiming, and the new framing matches what the code actually does. Test results:
|
Great PR though! This is a great catch! |
Summary
io::time_index::read_trackandio::temporal_index::read_trackaccept an on-disklengthandcount/entry_countand use them to size up-front allocations before validating any payload. A self-consistent but adversarially constructed.mv2can declare values that satisfy the structural checks but drive ~16 GiB of allocation per call.This PR closes that gap on both readers. The 16 GiB ceiling (
MAX_TIME_INDEX_BYTES/MAX_TEMPORAL_TRACK_BYTES) is preserved as a defence-in-depth upper bound — the actual fix is to bound the eager allocation to what the file has actually delivered.Changes
src/io/time_index.rsMAX_TIME_INDEX_BYTES = 1 << 34(16 GiB) and reject declared lengths above it (mirrors the temporal-index reader).Vec::with_capacity(count)toINITIAL_ENTRIES_CAPACITY = 1024. The vector still grows naturally asread_exactsucceeds, so the committed memory tracks the payload actually read.src/io/temporal_index.rsvec![0u8; length - HEADER_SIZE]) — records are now streamed one at a time into a fixed-size stack buffer ([u8; MENTION_RECORD_SIZE],[u8; ANCHOR_RECORD_SIZE]), hashed incrementally, and decoded directly into the output vectors.Vec::with_capacity(entry_count)/Vec::with_capacity(anchor_count)toINITIAL_RECORDS_CAPACITY = 1024. Same growth semantics as above.Why this matters (response to review)
The previous revision of this PR only enforced the ceiling. A near-cap aligned input (
declared_length = 17_179_869_180,count = 1_073_741_823) passed both the ceiling check and thecount * 16 == payload_bytescheck, soVec::with_capacity(count)still asked the allocator for ~16 GiB before any payload byte was read. With this revision, the same input commits onlyINITIAL_ENTRIES_CAPACITYslots up front (~16 KiB) and then fails at the firstread_exactwithUnexpectedEof.The same shape existed — and is now fixed — in the temporal-index reader.
Tests
Two new near-cap boundary tests, one per reader. Each constructs the largest aligned, accepted declared length, provides only the header bytes, and asserts that
read_trackreturnsIo { UnexpectedEof, .. }rather than completing or OOMing. If the eager-allocation path were still wide open, the test process would have requested ~16 GiB up front and the test would be unrunnable on a normal machine.cargo test --lib— 340/340 passcargo test --lib --features temporal_track— 346/346 passcargo clippy --lib --no-deps— cleancargo fmt -- --check— cleanNotes
Cargo.lockbump (2.0.135→2.0.139) is a lockfile sync with workspace version onmain(released ina79dfdd), not an incidental dep upgrade.