Skip to content

Commit b15ae6d

Browse files
arjunr2arjunr2
committed
Fix: Bug in flat storage variant computation
1 parent 1b94378 commit b15ae6d

File tree

2 files changed

+64
-63
lines changed

2 files changed

+64
-63
lines changed

‎crates/environ/src/component/types.rs‎

Lines changed: 56 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -373,6 +373,10 @@ impl ComponentTypes {
373373
/// not encodable in flat types, the values are all lowered to memory, implied by
374374
/// the empty storage
375375
pub fn flat_types_storage(&self, ty: &InterfaceType, limit: usize) -> FlatTypesStorage {
376+
assert!(
377+
limit <= MAX_FLAT_TYPES,
378+
"limit exceeding maximum flat types not allowed"
379+
);
376380
self.flat_types_storage_inner(ty, limit)
377381
.unwrap_or_else(|| FlatTypesStorage::new())
378382
}
@@ -385,7 +389,7 @@ impl ComponentTypes {
385389
// Helper routines
386390
let push = |storage: &mut FlatTypesStorage, t32: FlatType, t64: FlatType| -> bool {
387391
storage.push(t32, t64);
388-
(storage.len as usize) < limit
392+
(storage.len as usize) <= limit
389393
};
390394

391395
let push_discrim = |storage: &mut FlatTypesStorage| -> bool {
@@ -394,64 +398,62 @@ impl ComponentTypes {
394398

395399
let push_storage =
396400
|storage: &mut FlatTypesStorage, other: Option<FlatTypesStorage>| -> bool {
397-
if let Some(other) = other {
398-
let len = usize::from(storage.len);
399-
let other_len = usize::from(other.len);
400-
if len + other_len <= limit {
401-
storage.memory32[len..len + other_len]
402-
.copy_from_slice(&other.memory32[..other_len]);
403-
storage.memory64[len..len + other_len]
404-
.copy_from_slice(&other.memory64[..other_len]);
405-
storage.len += other.len;
406-
true
407-
} else {
408-
false
409-
}
410-
} else {
411-
false
412-
}
401+
other
402+
.and_then(|other| {
403+
let len = usize::from(storage.len);
404+
let other_len = usize::from(other.len);
405+
(len + other_len <= limit).then(|| {
406+
storage.memory32[len..len + other_len]
407+
.copy_from_slice(&other.memory32[..other_len]);
408+
storage.memory64[len..len + other_len]
409+
.copy_from_slice(&other.memory64[..other_len]);
410+
storage.len += other.len;
411+
})
412+
})
413+
.is_some()
413414
};
414415

416+
// Case is broken down as:
417+
// * None => No field
418+
// * Some(None) => Invalid storage (overflow)
419+
// * Some(storage) => Valid storage
415420
let push_storage_variant_case =
416421
|storage: &mut FlatTypesStorage, case: Option<Option<FlatTypesStorage>>| -> bool {
417-
if let Some(case) = case {
418-
if let Some(case) = case {
419-
// Discriminant will make size[case] = limit overshoot
420-
if case.len as usize >= limit {
421-
false
422-
} else {
423-
// Skip 1 for discriminant
424-
let dst = storage
425-
.memory32
426-
.iter_mut()
427-
.zip(&mut storage.memory64)
428-
.skip(1);
429-
for (i, ((t32, t64), (dst32, dst64))) in case
430-
.memory32
431-
.iter()
432-
.take(case.len as usize)
433-
.zip(case.memory64.iter())
434-
.zip(dst)
435-
.enumerate()
436-
{
437-
if i + 1 < usize::from(storage.len) {
438-
// Populated Index
439-
dst32.join(*t32);
440-
dst64.join(*t64);
441-
} else {
442-
// New Index
443-
storage.len += 1;
444-
*dst32 = *t32;
445-
*dst64 = *t64;
422+
match case {
423+
None => true,
424+
Some(case) => {
425+
case.and_then(|case| {
426+
// Discriminant will make size[case] = limit overshoot
427+
((1 + case.len as usize) <= limit).then(|| {
428+
// Skip 1 for discriminant
429+
let dst = storage
430+
.memory32
431+
.iter_mut()
432+
.zip(&mut storage.memory64)
433+
.skip(1);
434+
for (i, ((t32, t64), (dst32, dst64))) in case
435+
.memory32
436+
.iter()
437+
.take(case.len as usize)
438+
.zip(case.memory64.iter())
439+
.zip(dst)
440+
.enumerate()
441+
{
442+
if i + 1 < usize::from(storage.len) {
443+
// Populated Index
444+
dst32.join(*t32);
445+
dst64.join(*t64);
446+
} else {
447+
// New Index
448+
storage.len += 1;
449+
*dst32 = *t32;
450+
*dst64 = *t64;
451+
}
446452
}
447-
}
448-
true
449-
}
450-
} else {
451-
true
453+
})
454+
})
455+
.is_some()
452456
}
453-
} else {
454-
false
455457
}
456458
};
457459

@@ -509,7 +511,7 @@ impl ComponentTypes {
509511
}
510512
InterfaceType::Option(i) => {
511513
push_discrim(storage)
512-
&& push_storage_variant_case(storage, Some(None))
514+
&& push_storage_variant_case(storage, None)
513515
&& push_storage_variant_case(
514516
storage,
515517
Some(self.flat_types_storage_inner(&self[*i].ty, limit)),

‎crates/wasmtime/src/runtime/rr/hooks/component_hooks.rs‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -41,16 +41,9 @@ where
4141
{
4242
let _ = (args, id, func_idx, type_idx);
4343
#[cfg(feature = "rr-component")]
44-
let component = id.get(store.0).component();
45-
#[cfg(feature = "rr-component")]
46-
let types = component.types();
47-
#[cfg(all(feature = "rr-component", feature = "rr-validate"))]
48-
let flat_results = types.flat_types_storage(
49-
&InterfaceType::Tuple(types[type_idx].results),
50-
MAX_FLAT_RESULTS,
51-
);
52-
#[cfg(feature = "rr-component")]
5344
{
45+
let component = id.get(store.0).component();
46+
let types = component.types();
5447
let checksum = *component.checksum();
5548
let flat_params = types.flat_types_storage(
5649
&InterfaceType::Tuple(types[type_idx].params),
@@ -63,6 +56,12 @@ where
6356
let result = wasm_call(store);
6457
#[cfg(all(feature = "rr-component", feature = "rr-validate"))]
6558
{
59+
let component = id.get(store.0).component();
60+
let types = component.types();
61+
let flat_results = types.flat_types_storage(
62+
&InterfaceType::Tuple(types[type_idx].results),
63+
MAX_FLAT_RESULTS,
64+
);
6665
let result = result.map(|_| RRFuncArgVals::from_raw_slice(args, flat_results.iter32()));
6766
store.0.record_event_validation(|| {
6867
WasmFuncReturnEvent(ResultEvent::from_anyhow_result(&result))

0 commit comments

Comments
 (0)