Reject oversized memtable entries

[xingbowang]

Summary

Memtable entries encode their internal-key and value lengths as varint32, while MemTable::Add receives Slice lengths as size_t. Previously, it narrowed the key and value lengths and calculated the complete allocation size in uint32_t before allocating storage. A key that no longer fits after adding the eight-byte internal-key suffix, or a key/value pair whose complete encoding exceeds UINT32_MAX, could therefore wrap the calculated allocation size. Subsequent encoding and copying could write beyond the undersized allocation or silently truncate data.

WriteBatch validates individual key and value lengths on common public write paths, but that is not sufficient: it does not validate the aggregate memtable encoding size and MemTable::Add is also reached by internal and recovery paths. MemTable::Add is the final common boundary before allocation and copying, so validate the sizes there before narrowing or mutating the memtable. Compute the complete encoded length in uint64_t and return Status::InvalidArgument when the entry cannot be represented by the uint32-based format.

Keep the existing post-encoding assertion unchanged. It verifies private encoder bookkeeping for inputs that already satisfy the format contract. It cannot replace the runtime validation because it executes after allocation and copying, compares already-narrowed unsigned values that can hide wraparound, and is compiled out of release builds. The runtime checks reject invalid data safely, while the assertion continues to catch internal implementation mistakes during development.

The regression test covers an individually representable value whose aggregate entry size overflows, and a key whose internal-key length overflows. It also verifies rejection happens before the memtable is modified.

Test plan

• Built db_memtable_test and ran DBMemTableTest.RejectsOversizedEncodedEntry in normal and ASSERT_STATUS_CHECKED=1 configurations.
• Compiled production memtable.o with DEBUG_LEVEL=0.

[github-actions]

✅ clang-tidy: No findings on changed lines

Completed in 103.6s.

[meta-codesync]

@xingbowang has imported this pull request. If you are a Meta employee, you can view this in D117893558.

[github-actions]

✅ Claude Code Review

Auto-triggered after CI passed — reviewing commit c88b7b3

---

Summary

Well-designed defensive fix that adds size validation to MemTable::Add() before narrowing casts, preventing integer overflow in encoded entry size computation. The validation is placed at the right layer (final common boundary before allocation), uses correct arithmetic, and returns before any memtable mutation. All callers already handle non-OK status properly.

High-severity findings (0):
No high-severity findings.

Full review (click to expand)

Findings

🔴 HIGH

No high-severity findings.

🟡 MEDIUM

M1. Implicit narrowing conversion from uint64_t to uint32_t — memtable.cc:1137

• Issue: const uint32_t internal_key_size = key_size + kNumInternalBytes; where kNumInternalBytes is constexpr uint64_t (db/dbformat.h:171). The addition promotes to uint64_t, then the result is implicitly narrowed to uint32_t. While the value is guaranteed to fit after validation (max UINT32_MAX), this implicit narrowing is the kind of construct CLAUDE.md flags: the old code used key_size + 8 (int literal, no promotion beyond uint32_t), while the new code introduces a 64-to-32-bit implicit narrowing. Some strict compiler configurations or static analyzers may warn.
• Root cause: kNumInternalBytes was defined as uint64_t rather than a type that matches the assignment target.
• Suggested fix: Add an explicit static_cast<uint32_t>() to document the intentional narrowing:

  const uint32_t internal_key_size = key_size + static_cast<uint32_t>(kNumInternalBytes);

  Or, per CLAUDE.md preference for static_cast_with_check, use:

  const uint32_t internal_key_size =
      static_cast_with_check<uint32_t>(uint64_t{key_size} + kNumInternalBytes);

M2. MemTable::Update retains unguarded static_cast<uint32_t>(value.size()) — memtable.cc:1956

• Issue: MemTable::Update() at line 1956 does uint32_t new_size = static_cast<uint32_t>(value.size()) without the new validation. While the in-place update path only triggers when new_size <= prev_size (so the existing entry already passed validation), and the fallback at line 1987 calls Add() (which now validates), the static_cast is still technically an unguarded narrowing of a size_t. If value.size() > UINT32_MAX, new_size wraps, and the truncated new_size could be <= prev_size, causing the in-place path to execute with wrong data length.
• Root cause: The Update path's in-place branch silently narrows value.size() before the new_size <= prev_size comparison.
• Suggested fix: This is defense-in-depth since WriteBatch::Put already validates value.size() <= UINT32_MAX on public API paths. However, for consistency with the new Add() validation, consider adding a similar early check in Update(), or at minimum using static_cast_with_check.

🟢 LOW / NIT

L1. Test does not exercise the aggregate overflow check (second if) independently — db_memtable_test.cc:139

• Issue: The first test case (value.size() == UINT32_MAX with key.size() == 1) passes the first if (since value.size() == kMaxEncodedEntrySize, not >), but is caught by the second if (encoded_len64 > kMaxEncodedEntrySize). This is correct but means the test implicitly covers the aggregate check. A case where both key and value are individually representable but their combined encoding overflows would more clearly exercise the second check. For example: key.size() == UINT32_MAX/2 and value.size() == UINT32_MAX/2.
• Suggested fix: Consider adding a case that explicitly tests the aggregate overflow with two moderate-sized slices whose combined encoding exceeds UINT32_MAX.

L2. Missing positive boundary test — db_memtable_test.cc

• Issue: The test only verifies rejection. A complementary test showing that a just-below-limit entry (e.g., small key + small value) succeeds would confirm the validation doesn't over-reject. The existing tests elsewhere implicitly cover this, but a boundary positive case in this test would be more complete.
• Suggested fix: Optional — the existing test suite exercises normal Add() calls extensively.

L3. const qualification improvement is good but undocumented — memtable.cc:1134-1137

• Issue: key_size, val_size, and internal_key_size are now const. This is a minor improvement for readability and optimizer hints. No action needed; just noting the improvement.

Cross-Component Analysis

Context	Impact	Assessment
WriteBatch (public API)	Already validates key and value sizes before reaching Add()	Defense-in-depth; Add() validation catches edge cases like aggregate overflow
WAL Recovery	Replays previously-valid entries that already fit in uint32_t	Safe — no regression
MemTable::Update	Falls back to Add() for non-in-place updates	Safe via fallback; in-place path has pre-existing trust assumption (M2)
AddLogicallyRedundantRangeTombstone	Calls Add() with iterator-derived keys (bounded by existing SST sizes)	Safe; Add() returns non-OK, caller returns false (line 1019-1023)
Concurrent writes	Validation is read-only on inputs, before any shared state mutation	Thread-safe
WritePreparedTxnDB	Uses same WriteBatch → MemTableInserter → Add() path	Safe

Positive Observations

• Correct placement: Validation at MemTable::Add() is the right layer — it's the final common boundary before allocation and encoding, catching all paths including internal and recovery.
• No mutation before validation: All counter updates (num_entries_, data_size_), bloom filter updates, and memory allocation happen strictly after the validation checks.
• Consistent with existing patterns: Uses UNLIKELY for error paths, Status::InvalidArgument matching WriteBatch conventions, and kNumInternalBytes instead of magic number 8.
• Clean uint64_t arithmetic: Computing encoded_len64 in uint64_t with explicit uint64_t{} widening prevents any intermediate overflow.
• Negligible performance impact: Two branch-predicted-away comparisons on the write hot path; 64-bit arithmetic is native on x86_64/ARM64.
• Good test technique: Using Slice with fake lengths to exercise validation without allocating gigabytes.
• Documentation updated: Header comment in memtable.h documents the new InvalidArgument return.

---

ℹ️ About this response

Generated by Claude Code.
Review methodology: claude_md/ci_review_prompt.md

Limitations:

• Claude may miss context from files not in the diff
• Large PRs may be truncated
• Always apply human judgment to AI suggestions

Commands:

• /claude-review [context] — Request a code review
• /claude-query <question> — Ask about the PR or codebase