From b837a58ccfe076e8ec4eca5102127da61cd0c5b3 Mon Sep 17 00:00:00 2001
From: JustForFun88 <100504524+JustForFun88@users.noreply.github.com>
Date: Wed, 7 Sep 2022 11:29:08 +0500
Subject: [PATCH] Fix leak of memory if dropping function panic

---
 src/map.rs     | 29 +++++++++++++++++++++++++++++
 src/raw/mod.rs | 26 ++++++++++++++++++++++----
 2 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/src/map.rs b/src/map.rs
index 6a1e18c23b..2957c694ba 100644
--- a/src/map.rs
+++ b/src/map.rs
@@ -8379,4 +8379,33 @@ mod test_map {
 
         map2.clone_from(&map1);
     }
+
+    // Test that map do not leak memory if dropping function panic
+    #[test]
+    #[should_panic = "panic in drop"]
+    fn test_panic_in_drop() {
+        #[derive(Clone)]
+        struct CheckedDrop {
+            panic_in_drop: bool,
+        }
+        impl Drop for CheckedDrop {
+            fn drop(&mut self) {
+                if self.panic_in_drop {
+                    panic!("panic in drop");
+                }
+            }
+        }
+        const DISARMED: CheckedDrop = CheckedDrop {
+            panic_in_drop: false,
+        };
+        const ARMED: CheckedDrop = CheckedDrop {
+            panic_in_drop: true,
+        };
+
+        let mut map1 = HashMap::new();
+        map1.insert(1, DISARMED);
+        map1.insert(2, DISARMED);
+        map1.insert(3, ARMED);
+        map1.insert(4, DISARMED);
+    }
 }
diff --git a/src/raw/mod.rs b/src/raw/mod.rs
index 211b818a5f..0fc9bf0629 100644
--- a/src/raw/mod.rs
+++ b/src/raw/mod.rs
@@ -1803,8 +1803,17 @@ unsafe impl<#[may_dangle] T, A: Allocator + Clone> Drop for RawTable<T, A> {
     fn drop(&mut self) {
         if !self.table.is_empty_singleton() {
             unsafe {
-                self.drop_elements();
-                self.free_buckets();
+                // Guard that provide deallocation of memory even if any panics occurs during dropping
+                let mut self_ = guard(self, |self_| {
+                    self_.free_buckets();
+                });
+
+                // This may panic but in any case the scope guard will deallocate memory of
+                // the table, leaking any elements that were not dropped yet if panic occurs.
+                //
+                // This leak is unavoidable: we can't try dropping more elements
+                // since this could lead to another panic and abort the process.
+                self_.drop_elements();
             }
         }
     }
@@ -1815,8 +1824,17 @@ impl<T, A: Allocator + Clone> Drop for RawTable<T, A> {
     fn drop(&mut self) {
         if !self.table.is_empty_singleton() {
             unsafe {
-                self.drop_elements();
-                self.free_buckets();
+                // Guard that provide deallocation of memory even if any panics occurs during dropping
+                let mut self_ = guard(self, |self_| {
+                    self_.free_buckets();
+                });
+
+                // This may panic but in any case the scope guard will deallocate memory of
+                // the table, leaking any elements that were not dropped yet if panic occurs.
+                //
+                // This leak is unavoidable: we can't try dropping more elements
+                // since this could lead to another panic and abort the process.
+                self_.drop_elements();
             }
         }
     }