From 9526813f5bc088261201934d7d6e574231eb252e Mon Sep 17 00:00:00 2001
From: Andrew Paseltiner <apaseltiner@gmail.com>
Date: Fri, 18 Sep 2015 14:15:02 -0400
Subject: [PATCH] Avoid zero-sized leaf allocations in `BTreeMap`

When both the key and value types were zero-sized, `BTreeMap` previously
called `heap::allocate` with `size == 0` for leaf nodes, which is
undefined behavior, and jemalloc would attempt to read invalid memory,
crashing the process.

This avoids undefined behavior by allocating enough space to store one
edge in leaf nodes that would otherwise have `size == 0`. Although this
uses extra memory, maps with zero-sized key types that have sensible
implementations of the ordering traits can only contain a single
key-value pair (and therefore only a single leaf node), and maps with
key and value types that are both zero-sized have few uses, if any.

Furthermore, this is a temporary fix that will likely be unnecessary
once the `BTreeMap` implementation is rewritten to use parent pointers.

Closes #28493.
---
 src/libcollections/btree/node.rs    |  7 +++-
 src/libcollectionstest/btree/map.rs | 53 +++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/src/libcollections/btree/node.rs b/src/libcollections/btree/node.rs
index d8f8ca6eae59..f5088bf4646a 100644
--- a/src/libcollections/btree/node.rs
+++ b/src/libcollections/btree/node.rs
@@ -164,7 +164,12 @@ fn calculate_allocation_generic<K, V>(capacity: usize, is_leaf: bool) -> (usize,
     let (keys_size, keys_align) = (capacity * mem::size_of::<K>(), mem::align_of::<K>());
     let (vals_size, vals_align) = (capacity * mem::size_of::<V>(), mem::align_of::<V>());
     let (edges_size, edges_align) = if is_leaf {
-        (0, 1)
+        // allocate one edge to ensure that we don't pass size 0 to `heap::allocate`
+        if mem::size_of::<K>() == 0 && mem::size_of::<V>() == 0 {
+            (1, mem::align_of::<Node<K, V>>())
+        } else {
+            (0, 1)
+        }
     } else {
         ((capacity + 1) * mem::size_of::<Node<K, V>>(), mem::align_of::<Node<K, V>>())
     };
diff --git a/src/libcollectionstest/btree/map.rs b/src/libcollectionstest/btree/map.rs
index 62b46433da95..846353cc4e7c 100644
--- a/src/libcollectionstest/btree/map.rs
+++ b/src/libcollectionstest/btree/map.rs
@@ -294,6 +294,59 @@ fn test_extend_ref() {
     assert_eq!(a[&3], "three");
 }
 
+#[test]
+fn test_zst() {
+    let mut m = BTreeMap::new();
+    assert_eq!(m.len(), 0);
+
+    assert_eq!(m.insert((), ()), None);
+    assert_eq!(m.len(), 1);
+
+    assert_eq!(m.insert((), ()), Some(()));
+    assert_eq!(m.len(), 1);
+    assert_eq!(m.iter().count(), 1);
+
+    m.clear();
+    assert_eq!(m.len(), 0);
+
+    for _ in 0..100 {
+        m.insert((), ());
+    }
+
+    assert_eq!(m.len(), 1);
+    assert_eq!(m.iter().count(), 1);
+}
+
+// This test's only purpose is to ensure that zero-sized keys with nonsensical orderings
+// do not cause segfaults when used with zero-sized values. All other map behavior is
+// undefined.
+#[test]
+fn test_bad_zst() {
+    use std::cmp::Ordering;
+
+    struct Bad;
+
+    impl PartialEq for Bad {
+        fn eq(&self, _: &Self) -> bool { false }
+    }
+
+    impl Eq for Bad {}
+
+    impl PartialOrd for Bad {
+        fn partial_cmp(&self, _: &Self) -> Option<Ordering> { Some(Ordering::Less) }
+    }
+
+    impl Ord for Bad {
+        fn cmp(&self, _: &Self) -> Ordering { Ordering::Less }
+    }
+
+    let mut m = BTreeMap::new();
+
+    for _ in 0..100 {
+        m.insert(Bad, Bad);
+    }
+}
+
 mod bench {
     use std::collections::BTreeMap;
     use std::__rand::{Rng, thread_rng};