From ba79ac713deb154896fb5fa2247f9887952f5071 Mon Sep 17 00:00:00 2001
From: Matt Brubeck <mbrubeck@limpet.net>
Date: Thu, 7 Jan 2021 21:28:46 -0800
Subject: [PATCH] Fix potential buffer overflow in `insert_many`

Fixes #252.
---
 src/lib.rs   | 48 +++++++++++++++++++++++++++---------------------
 src/tests.rs | 13 +++++++++++++
 2 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 0241aef..29a43e1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1009,7 +1009,7 @@ impl<A: Array> SmallVec<A> {
     /// Insert multiple elements at position `index`, shifting all following elements toward the
     /// back.
     pub fn insert_many<I: IntoIterator<Item = A::Item>>(&mut self, index: usize, iterable: I) {
-        let iter = iterable.into_iter();
+        let mut iter = iterable.into_iter();
         if index == self.len() {
             return self.extend(iter);
         }
@@ -1019,11 +1019,13 @@ impl<A: Array> SmallVec<A> {
         assert!(index + lower_size_bound >= index); // Protect against overflow
         self.reserve(lower_size_bound);
 
+        let mut num_added = 0;
+        let old_len = self.len();
+        assert!(index <= old_len);
+
         unsafe {
-            let old_len = self.len();
-            assert!(index <= old_len);
             let start = self.as_mut_ptr();
-            let mut ptr = start.add(index);
+            let ptr = start.add(index);
 
             // Move the trailing elements.
             ptr::copy(ptr, ptr.add(lower_size_bound), old_len - index);
@@ -1036,26 +1038,16 @@ impl<A: Array> SmallVec<A> {
                 len: old_len + lower_size_bound,
             };
 
-            let mut num_added = 0;
-            for element in iter {
-                let mut cur = ptr.add(num_added);
-                if num_added >= lower_size_bound {
-                    // Iterator provided more elements than the hint.  Move trailing items again.
-                    self.reserve(1);
-                    let start = self.as_mut_ptr();
-                    ptr = start.add(index);
-                    cur = ptr.add(num_added);
-                    ptr::copy(cur, cur.add(1), old_len - index);
-
-                    guard.start = start;
-                    guard.len += 1;
-                    guard.skip.end += 1;
-                }
+            while num_added < lower_size_bound {
+                let element = match iter.next() {
+                    Some(x) => x,
+                    None => break,
+                };
+                let cur = ptr.add(num_added);
                 ptr::write(cur, element);
                 guard.skip.start += 1;
                 num_added += 1;
             }
-            mem::forget(guard);
 
             if num_added < lower_size_bound {
                 // Iterator provided fewer elements than the hint
@@ -1066,12 +1058,26 @@ impl<A: Array> SmallVec<A> {
                 );
             }
 
+            // There are no more duplicate or uninitialized slots, so the guard is not needed.
             self.set_len(old_len + num_added);
+            mem::forget(guard);
+        }
+
+        unsafe {
+            for element in iter {
+                // Iterator provided more elements than the hint.  Move trailing items again.
+                self.reserve(1);
+                let cur = self.as_mut_ptr().add(index).add(num_added);
+                ptr::copy(cur, cur.add(1), old_len - index);
+                ptr::write(cur, element);
+                self.set_len(self.len() + 1);
+                num_added += 1;
+            }
         }
 
         struct DropOnPanic<T> {
             start: *mut T,
-            skip: Range<usize>,
+            skip: Range<usize>, // Space we copied-out-of, but haven't written-to yet.
             len: usize,
         }
 
diff --git a/src/tests.rs b/src/tests.rs
index 0452ae8..19f6da8 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -905,3 +905,16 @@ fn empty_macro() {
 fn zero_size_items() {
     SmallVec::<[(); 0]>::new().push(());
 }
+
+#[test]
+fn test_insert_many_overflow() {
+    let mut v: SmallVec<[u8; 1]> = SmallVec::new();
+    v.push(123);
+
+    // Prepare an iterator with small lower bound
+    let iter = (0u8..5).filter(|n| n % 2 == 0);
+    assert_eq!(iter.size_hint().0, 0);
+
+    v.insert_many(0, iter);
+    assert_eq!(&*v, &[0, 2, 4, 123]);
+}