soundness issue in `IndexMap::insert` (`Pos::new`)

[japaric]

Spotted by @8rsrf7
Affected versions: 0.9.3 and probably earlier versions

heapless/src/index_map.rs

Lines 109 to 113 in 4e9d494

	nz: unsafe {
	NonZeroU32::new_unchecked(
	((u32::from(hash.0) << 16) + index as u32).wrapping_add(1),
	)
	},

^ This code probably assumes that the case index=u16::MAX and hash=u16::MAX is statistically impossible but it's not actually impossible.
Here's a proof of concept that triggers the issue

// src/main.rs
use std::hash::{BuildHasher, Hasher};
use std::sync::atomic::{self, AtomicU16};

use heapless::IndexMap;

fn main() {
    let mut map = IndexMap::<_, _, MyHasher, { 1 << 17 }>::default();
    for x in 0..=u16::MAX {
        map.insert(x, x).unwrap();
    }
    println!("{}", map.len());
}

#[derive(Default)]
struct MyHasher;

impl Hasher for MyHasher {
    fn finish(&self) -> u64 {
        static COUNT: AtomicU16 = AtomicU16::new(0);
        COUNT.fetch_add(1, atomic::Ordering::Relaxed).into()
    }

    fn write(&mut self, _bytes: &[u8]) {}
}

impl BuildHasher for MyHasher {
    type Hasher = MyHasher;

    fn build_hasher(&self) -> Self::Hasher {
        MyHasher
    }
}

$ cargo run
thread 'main' (33377) panicked at /tmp/heapless-0.9.3/src/index_map.rs:109:9:
(65535, 65535)

# Cargo.toml
[dependencies]
heapless.path = "../heapless-0.9.3"

with this patch applied to heapless-0.9.3 to get a panic instead of UB.

diff --git a/src/index_map.rs b/src/index_map.rs
index cd6e30b..80b74cb 100644
--- a/src/index_map.rs
+++ b/src/index_map.rs
@@ -105,12 +105,10 @@ pub struct Pos {

 impl Pos {
     fn new(index: usize, hash: HashValue) -> Self {
+        let val = ((u32::from(hash.0) << 16) + index as u32).wrapping_add(1);
+        assert!(val != 0, "{:?}", (index, hash.0));
         Self {
-            nz: unsafe {
-                NonZeroU32::new_unchecked(
-                    ((u32::from(hash.0) << 16) + index as u32).wrapping_add(1),
-                )
-            },
+            nz: unsafe { NonZeroU32::new_unchecked(val) },
         }
     }

One option to prevent this issue for certain would be to limit the capacity to <u16::MAX in the constructor with a (compile-time) assertion. Given that this library is commonly used on resource constrained devices, it's unlikely that people are using such large capacities.

[japaric]

with this patch applied to heapless-0.9.3 to get a panic instead of UB.

the patch is not need if using the dev profile. the debug assertion in NonZero::new_unchecked will detect the issue

$ cargo run
thread 'main' (38643) panicked at /home/jorge/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/heapless-0.9.3/src/index_map.rs:110:17:
unsafe precondition(s) violated: NonZero::new_unchecked requires the argument to be non-zero

This indicates a bug in the program. This Undefined Behavior check is optional, and cannot be relied on for safety.

[sgued]

This is indeed an issue.

I'm not a fan of the idea of preventing the use of N = u16::MAX, it seems like it's sufficiently low to be useful.

This makes me realise that just any N > u16::MAX will simply not work at all.

[sgued]

I see a possible fix, instead of having the assert!(N.is_power_of_two()) in a const block, we would instead encode that in a trait:

mod sealed {
    use core::num::{NonZeroU32, NonZeroU64};
    pub trait Sealed {
        type Pos;
    }
    pub struct PosU16 {
        nz: NonZeroU32,
    }

    pub struct PosU32 {
        nz: NonZeroU64,
    }
}

use sealed::*;

pub trait IndexMapSize: Sealed {}

#[non_exhaustive]
pub struct Const<const N: usize> {}

macro_rules! impl_const {
    ($($const:expr, $pos:ident);*;) => {$(
      impl Sealed for Const<$const> {
          type Pos = $pos;
      }

      impl IndexMapSize for Const<$const> {}
    )*};
}

impl_const! {
  {1 << 1}, PosU16;
  {1 << 2}, PosU16;
  {1 << 3}, PosU16;
  {1 << 4}, PosU16;
  {1 << 5}, PosU16;
  {1 << 6}, PosU16;
  {1 << 7}, PosU16;
  {1 << 8}, PosU16;
  {1 << 9}, PosU16;
  {1 << 10}, PosU16;
  {1 << 11}, PosU16;
  {1 << 12}, PosU16;
  {1 << 13}, PosU16;
  {1 << 14}, PosU16;
  {1 << 15}, PosU16;
  {1 << 16}, PosU32;
  {1 << 17}, PosU32;
  {1 << 18}, PosU32;
  {1 << 19}, PosU32;
  {1 << 20}, PosU32;
  {1 << 21}, PosU32;
  {1 << 22}, PosU32;
}

But I see two issues with that: this would be a breaking change and releasing this approach means that a View type could never be added to indexMap in the future.

[sgued]

After some messing around, I got a fix I'm satisfied with in #673. Thanks a lot for the report!