Skip to content

Commit

Permalink
Fix panic when inserting pair exceeding ~4GiB
Browse files Browse the repository at this point in the history
StorageError::ValueTooLarge will now be returned for key-value pairs
that exceed 3.75GiB
  • Loading branch information
cberner committed Jul 15, 2024
1 parent 54d5cf9 commit 6c07fd1
Show file tree
Hide file tree
Showing 7 changed files with 59 additions and 8 deletions.
11 changes: 7 additions & 4 deletions src/multimap_table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::tree_store::{
btree_stats, AllPageNumbersBtreeIter, BranchAccessor, Btree, BtreeHeader, BtreeMut,
BtreeRangeIter, BtreeStats, CachePriority, Checksum, LeafAccessor, LeafMutator, Page, PageHint,
PageNumber, RawBtree, RawLeafBuilder, TransactionalMemory, UntypedBtreeMut, BRANCH, LEAF,
MAX_VALUE_LENGTH,
MAX_PAIR_LENGTH, MAX_VALUE_LENGTH,
};
use crate::types::{Key, TypeName, Value};
use crate::{AccessGuard, MultimapTableHandle, Result, StorageError, WriteTransaction};
Expand Down Expand Up @@ -833,9 +833,12 @@ impl<'txn, K: Key + 'static, V: Key + 'static> MultimapTable<'txn, K, V> {
if value_bytes_ref.len() > MAX_VALUE_LENGTH {
return Err(StorageError::ValueTooLarge(value_bytes_ref.len()));
}
let key_bytes = K::as_bytes(key.borrow());
if key_bytes.as_ref().len() > MAX_VALUE_LENGTH {
return Err(StorageError::ValueTooLarge(key_bytes.as_ref().len()));
let key_len = K::as_bytes(key.borrow()).as_ref().len();
if key_len > MAX_VALUE_LENGTH {
return Err(StorageError::ValueTooLarge(key_len));
}
if value_bytes_ref.len() + key_len > MAX_PAIR_LENGTH {
return Err(StorageError::ValueTooLarge(value_bytes_ref.len() + key_len));
}
let get_result = self.tree.get(key.borrow())?;
let existed = if get_result.is_some() {
Expand Down
8 changes: 7 additions & 1 deletion src/table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use crate::db::TransactionGuard;
use crate::sealed::Sealed;
use crate::tree_store::{
AccessGuardMut, Btree, BtreeExtractIf, BtreeHeader, BtreeMut, BtreeRangeIter, PageHint,
PageNumber, RawBtree, TransactionalMemory, MAX_VALUE_LENGTH,
PageNumber, RawBtree, TransactionalMemory, MAX_PAIR_LENGTH, MAX_VALUE_LENGTH,
};
use crate::types::{Key, MutInPlaceValue, Value};
use crate::{AccessGuard, StorageError, WriteTransaction};
Expand Down Expand Up @@ -201,6 +201,9 @@ impl<'txn, K: Key + 'static, V: Value + 'static> Table<'txn, K, V> {
if key_len > MAX_VALUE_LENGTH {
return Err(StorageError::ValueTooLarge(key_len));
}
if value_len + key_len > MAX_PAIR_LENGTH {
return Err(StorageError::ValueTooLarge(value_len + key_len));
}
self.tree.insert(key.borrow(), value.borrow())
}

Expand Down Expand Up @@ -233,6 +236,9 @@ impl<'txn, K: Key + 'static, V: MutInPlaceValue + 'static> Table<'txn, K, V> {
if key_len > MAX_VALUE_LENGTH {
return Err(StorageError::ValueTooLarge(key_len));
}
if value_length as usize + key_len > MAX_PAIR_LENGTH {
return Err(StorageError::ValueTooLarge(value_length as usize + key_len));
}
self.tree.insert_reserve(key.borrow(), value_length)
}
}
Expand Down
5 changes: 4 additions & 1 deletion src/transactions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use crate::transaction_tracker::{SavepointId, TransactionId, TransactionTracker}
use crate::tree_store::{
Btree, BtreeHeader, BtreeMut, FreedPageList, FreedTableKey, InternalTableDefinition, PageHint,
PageNumber, SerializedSavepoint, TableTree, TableTreeMut, TableType, TransactionalMemory,
MAX_VALUE_LENGTH,
MAX_PAIR_LENGTH, MAX_VALUE_LENGTH,
};
use crate::types::{Key, Value};
use crate::{
Expand Down Expand Up @@ -244,6 +244,9 @@ impl<'db, 's, K: Key + 'static, V: Value + 'static> SystemTable<'db, 's, K, V> {
if key_len > MAX_VALUE_LENGTH {
return Err(StorageError::ValueTooLarge(key_len));
}
if value_len + key_len > MAX_PAIR_LENGTH {
return Err(StorageError::ValueTooLarge(value_len + key_len));
}
self.tree.insert(key.borrow(), value.borrow())
}

Expand Down
2 changes: 1 addition & 1 deletion src/tree_store/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ pub(crate) use btree_iters::{AllPageNumbersBtreeIter, BtreeExtractIf, BtreeRange
pub use page_store::{file_backend, InMemoryBackend, Savepoint};
pub(crate) use page_store::{
CachePriority, Page, PageHint, PageNumber, SerializedSavepoint, TransactionalMemory,
FILE_FORMAT_VERSION2, MAX_VALUE_LENGTH, PAGE_SIZE,
FILE_FORMAT_VERSION2, MAX_PAIR_LENGTH, MAX_VALUE_LENGTH, PAGE_SIZE,
};
pub(crate) use table_tree::{
FreedPageList, FreedTableKey, InternalTableDefinition, TableTree, TableTreeMut, TableType,
Expand Down
1 change: 1 addition & 0 deletions src/tree_store/page_store/base.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use std::sync::Arc;
use std::sync::Mutex;

pub(crate) const MAX_VALUE_LENGTH: usize = 3 * 1024 * 1024 * 1024;
pub(crate) const MAX_PAIR_LENGTH: usize = 3 * 1024 * 1024 * 1024 + 768 * 1024 * 1024;
pub(crate) const MAX_PAGE_INDEX: u32 = 0x000F_FFFF;

// On-disk format is:
Expand Down
2 changes: 1 addition & 1 deletion src/tree_store/page_store/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ mod savepoint;
#[allow(dead_code)]
mod xxh3;

pub(crate) use base::{Page, PageHint, PageNumber, MAX_VALUE_LENGTH};
pub(crate) use base::{Page, PageHint, PageNumber, MAX_PAIR_LENGTH, MAX_VALUE_LENGTH};
pub(crate) use header::PAGE_SIZE;
pub use in_memory_backend::InMemoryBackend;
pub(crate) use page_manager::{xxh3_checksum, TransactionalMemory, FILE_FORMAT_VERSION2};
Expand Down
38 changes: 38 additions & 0 deletions tests/integration_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,44 @@ fn large_values() {
txn.commit().unwrap();
}

// Note: this test requires > 3GiB of memory
#[test]
fn value_too_large() {
let tmpfile = create_tempfile();

let db = Database::create(tmpfile.path()).unwrap();
let txn = db.begin_write().unwrap();

let small_value = vec![0u8; 1024];
let too_big_value = vec![0u8; 3 * 1024 * 1024 * 1024 + 1];
{
let mut table = txn.open_table(SLICE_TABLE).unwrap();
assert!(matches!(
table.insert(small_value.as_slice(), too_big_value.as_slice()),
Err(StorageError::ValueTooLarge(_))
));
assert!(matches!(
table.insert(too_big_value.as_slice(), small_value.as_slice()),
Err(StorageError::ValueTooLarge(_))
));
assert!(matches!(
table.insert(too_big_value.as_slice(), too_big_value.as_slice()),
Err(StorageError::ValueTooLarge(_))
));
drop(too_big_value);
let almost_big_value = vec![0u8; 2 * 1024 * 1024 * 1024];
assert!(matches!(
table.insert(almost_big_value.as_slice(), almost_big_value.as_slice()),
Err(StorageError::ValueTooLarge(_))
));
}
txn.commit().unwrap();

let txn = db.begin_read().unwrap();
let table = txn.open_table(SLICE_TABLE).unwrap();
assert!(table.is_empty().unwrap());
}

#[test]
fn many_pairs() {
let tmpfile = create_tempfile();
Expand Down

0 comments on commit 6c07fd1

Please sign in to comment.