Avoid leak in DrainFilter when a drop panics

2 files changed, 50 insertions, 2 deletions

diff --git a/src/liballoc/collections/linked_list.rs b/src/liballoc/collections/linked_list.rs
index b88ca8a0fb0..463ec67d668 100644
--- a/src/liballoc/collections/linked_list.rs
+++ b/src/liballoc/collections/linked_list.rs
@@ -1565,7 +1565,21 @@ where
     F: FnMut(&mut T) -> bool,
 {
     fn drop(&mut self) {
-        self.for_each(drop);
+        struct DropGuard<'r, 'a, T, F>(&'r mut DrainFilter<'a, T, F>)
+            where F: FnMut(&mut T) -> bool;
+
+        impl<'r, 'a, T, F> Drop for DropGuard<'r, 'a, T, F>
+        where F: FnMut(&mut T) -> bool {
+            fn drop(&mut self) {
+                self.0.for_each(drop);
+            }
+        }
+
+        while let Some(item) = self.next() {
+            let guard = DropGuard(self);
+            drop(item);
+            mem::forget(guard);
+        }
     }
 }
 
diff --git a/src/liballoc/tests/linked_list.rs b/src/liballoc/tests/linked_list.rs
index b7736515b26..aa3f582a165 100644
--- a/src/liballoc/tests/linked_list.rs
+++ b/src/liballoc/tests/linked_list.rs
@@ -1,5 +1,5 @@
 use std::collections::LinkedList;
-use std::panic::catch_unwind;
+use std::panic::{catch_unwind, AssertUnwindSafe};
 
 #[test]
 fn test_basic() {
@@ -532,6 +532,40 @@ fn drain_filter_complex() {
 }
 
 #[test]
+fn drain_filter_drop_panic_leak() {
+    static mut DROPS: i32 = 0;
+
+    struct D(bool);
+
+    impl Drop for D {
+        fn drop(&mut self) {
+            unsafe {
+                DROPS += 1;
+            }
+
+            if self.0 {
+                panic!("panic in `drop`");
+            }
+        }
+    }
+
+    let mut q = LinkedList::new();
+    q.push_back(D(false));
+    q.push_back(D(false));
+    q.push_back(D(false));
+    q.push_back(D(false));
+    q.push_back(D(false));
+    q.push_front(D(false));
+    q.push_front(D(true));
+    q.push_front(D(false));
+
+    catch_unwind(AssertUnwindSafe(|| drop(q.drain_filter(|_| true)))).ok();
+
+    assert_eq!(unsafe { DROPS }, 8);
+    assert!(q.is_empty());
+}
+
+#[test]
 fn test_drop() {
     static mut DROPS: i32 = 0;
     struct Elem;