Skip to content

Commit a407990

Browse files
authored
fix: out-of-bounds read in ArrowBytesMap on a short/long hash collision (#25218)
## Which issue does this PR close? - No existing issue. Found while running the extended CI `force_hash_collisions` job against unrelated work; happy to file one first if you prefer. ## Rationale for this change `ArrowBytesMap` keeps a value's bytes inline in an entry's `offset_or_inline` field when the value is at most `SHORT_VALUE_LEN` (8) bytes, and keeps an offset into the value buffer there otherwise. The lookup for a **long** value checked only that the hashes matched, then built `header.range()` from that field and read the buffer with `get_unchecked`: ```rust let entry = self.map.find_mut(hash, |header| { if header.hash != hash { return false; } // SAFETY: buffer is only appended to, and we correctly inserted values and offsets let existing_value = unsafe { self.buffer.get_unchecked(header.range()) }; value == existing_value }); ``` When that hash matches a **short** entry, `offset_or_inline` holds inlined bytes rather than an offset, so the range is nonsense and the read goes outside the buffer. The lookup for a short value does not have this problem, because it already compares `header.len` first, which is precisely what keeps it sound. This is reachable whenever a value of at most 8 bytes and a value longer than 8 bytes collide on the full 64-bit hash inside one map. Two things make that worth fixing rather than dismissing: - DataFusion hashes with `foldhash::fast::FixedState::default()`, a fixed compile-time seed, so a colliding pair is reproducible across runs and machines and can be found offline once. - The expected number of such pairs grows with the product of the short and long distinct counts, so the exposure rises with cardinality rather than staying constant. The map backs grouped aggregation on a single string or binary column (`GroupValuesBytes`) and `count(DISTINCT …)` over those types (`ArrowBytesSet`), so this is on the path of ordinary `GROUP BY` and `COUNT(DISTINCT)` queries. What happens on a collision depends on the inlined bytes read as an offset. If they happen to land inside the buffer, the comparison sees unrelated bytes, reports inequality, and the query is correct by luck. Otherwise the read is out of bounds, which is undefined behaviour: with UB checks enabled it aborts, and in a release build it reads unmapped or unrelated memory. The extended CI job builds the workspace with `force_hash_collisions`, where every value hashes to `0` so every short/long pair collides. Today `dataframe::test_grouping_with_alias` aborts there: ``` unsafe precondition(s) violated: slice::get_unchecked requires that the range is within the slice thread caused non-unwinding panic. aborting. ``` The check is asymmetric for a reason worth stating: the short-value lookup already compares `header.len`, so a long entry can never be mistaken for a short one. Only the long-value lookup was missing it. That single comparison guards two different failure modes, an out-of-bounds read in one direction and a false equality that merges two distinct values in the other, which is why the regression test exercises both orders. ## What changes are included in this PR? - Compare `header.len` on the long-value path, exactly as the short-value path already does. That rules out short entries, so `offset_or_inline` is only ever treated as an offset when it is one, and it doubles as a cheap early-out. It cannot reject a real match, since two equal long values have equal lengths. - Rewrite the `SAFETY` comment to state the invariant that actually holds. - Give `datafusion-physical-expr-common` the `force_hash_collisions` passthrough feature that `datafusion-common`, `datafusion-core`, `datafusion-physical-plan` and the two aggregate crates already have, so the path can be reached deterministically from a unit test. ## What is the testing strategy for this PR? `short_and_long_values_with_colliding_hashes` in `binary_map.rs`, gated on `force_hash_collisions`, inserts a short and a long value in both orders and asserts each distinct value is offered to the payload function exactly once. The values are chosen so that removing either length check actually fails: `"ab"` inlines to 0x6162, far past the end of the buffer, and a single NUL byte inlines to 0, which is the offset of a long value stored first. Mutation-checked both ways: removing the long path's check aborts on the precondition violation, removing the short path's check trips the assertion. With the fix, the whole `core_integration` suite passes under `--features force_hash_collisions` (1144 tests) where it previously aborted. Without the feature, `datafusion-physical-expr-common`, `datafusion-functions-aggregate` and the `dataframe` integration tests all pass, and clippy is clean in both feature configurations. ## Are there any user-facing changes?
1 parent 9a1ad7c commit a407990

2 files changed

Lines changed: 68 additions & 4 deletions

File tree

‎datafusion/physical-expr-common/Cargo.toml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,9 @@ default = []
4545
# Enables the `PhysicalExpr::to_proto` hook used by `datafusion-proto`.
4646
# Off by default so crates that never serialize plans pay nothing.
4747
proto = ["dep:datafusion-proto-models"]
48+
# Test-only: makes every value hash to the same bucket, so that code paths
49+
# which only run on a hash collision are reachable in tests.
50+
force_hash_collisions = ["datafusion-common/force_hash_collisions"]
4851

4952
[dependencies]
5053
arrow = { workspace = true }

‎datafusion/physical-expr-common/src/binary_map.rs‎

Lines changed: 65 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -511,12 +511,15 @@ where
511511
else {
512512
// Check if the value is already present in the set
513513
let entry = self.map.find_mut(hash, |header| {
514-
// compare value if hashes match
515-
if header.hash != hash {
514+
// Without the length test, a short entry colliding on
515+
// `hash` builds `range()` from its inlined bytes and
516+
// reads past the buffer: memory safety, not just a
517+
// wasted comparison.
518+
if header.hash != hash || header.len != value_len {
516519
return false;
517520
}
518-
// Need to compare the bytes in the buffer
519-
// SAFETY: buffer is only appended to, and we correctly inserted values and offsets
521+
// SAFETY: a long entry, so `offset_or_inline` indexes the
522+
// append-only `buffer`, written before the insert.
520523
let existing_value =
521524
unsafe { self.buffer.get_unchecked(header.range()) };
522525
value == existing_value
@@ -1103,6 +1106,64 @@ mod tests {
11031106
);
11041107
}
11051108

1109+
/// A short value and a long value whose hashes collide must not be confused,
1110+
/// whichever of them is inserted first.
1111+
///
1112+
/// A short value keeps its bytes inlined in the entry's `offset_or_inline`
1113+
/// field, while a long one keeps an offset into the value buffer there. An
1114+
/// entry is only examined when its hash matches, so the length has to be
1115+
/// compared before that field is interpreted at all. Each direction fails
1116+
/// differently without that check, and each case below is chosen to make it
1117+
/// actually fail rather than get away with it:
1118+
///
1119+
/// * long value looked up against a short entry: `header.range()` is built
1120+
/// out of the inlined bytes, so `"ab"` (0x6162) sends the read tens of
1121+
/// kilobytes past the end of the buffer;
1122+
/// * short value looked up against a long entry: the inlined bytes are
1123+
/// compared against that entry's offset, so a single NUL byte (inlined as
1124+
/// 0) matches the long value stored at offset 0 and the two distinct
1125+
/// values collapse into one.
1126+
///
1127+
/// `force_hash_collisions` makes every value collide, which is the only way
1128+
/// to reach either path deterministically.
1129+
#[cfg(feature = "force_hash_collisions")]
1130+
#[test]
1131+
fn short_and_long_values_with_colliding_hashes() {
1132+
let long = "a value too long to be inlined";
1133+
assert!(long.len() > SHORT_VALUE_LEN);
1134+
1135+
for [first, second] in [[long, "\0"], ["ab", long]] {
1136+
assert!(first.len() <= SHORT_VALUE_LEN || second.len() <= SHORT_VALUE_LEN);
1137+
// The second value is the one looked up against an entry of the
1138+
// other kind; the repeats must match their own entry.
1139+
let values: ArrayRef = Arc::new(StringArray::from(vec![
1140+
Some(first),
1141+
Some(second),
1142+
Some(first),
1143+
Some(second),
1144+
]));
1145+
1146+
let mut map = ArrowBytesMap::<i32, usize>::new(OutputType::Utf8);
1147+
let mut inserted = vec![];
1148+
let mut payloads = 0;
1149+
map.insert_if_new(
1150+
&values,
1151+
|value| {
1152+
inserted.push(value.expect("no nulls in this test").to_vec());
1153+
payloads += 1;
1154+
payloads
1155+
},
1156+
|_| {},
1157+
);
1158+
1159+
assert_eq!(
1160+
inserted,
1161+
vec![first.as_bytes().to_vec(), second.as_bytes().to_vec()],
1162+
"inserting {first:?} then {second:?}"
1163+
);
1164+
}
1165+
}
1166+
11061167
// asserts that the set contains the expected strings, in the same order
11071168
fn assert_set<O: OffsetSizeTrait>(set: ArrowBytesSet<O>, expected: &[Option<&str>]) {
11081169
let strings = set.into_state();

0 commit comments

Comments
 (0)