From ab61f362accaa150aca82dab1040593e037d116f Mon Sep 17 00:00:00 2001
From: psrok1 <pawel.srokosz@cert.pl>
Date: Wed, 15 May 2024 17:33:05 +0200
Subject: [PATCH] Track references to mmap slices to release derived
 memoryviews when there is desire to release main MmapMemoryBuffer

---
 malduck/procmem/membuf.py | 43 +++++++++++++++++++++++++++++++++++----
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/malduck/procmem/membuf.py b/malduck/procmem/membuf.py
index 4f8e1b2..9f3245a 100644
--- a/malduck/procmem/membuf.py
+++ b/malduck/procmem/membuf.py
@@ -1,4 +1,5 @@
 import mmap
+import weakref
 from abc import ABC, abstractmethod
 from typing import Optional, Union
 
@@ -32,7 +33,6 @@ def __init__(
         self,
         buf: Union[bytes, bytearray, memoryview],
     ) -> None:
-        print("created ", self, id(self))
         if type(buf) is memoryview:
             self.buf = buf
         elif type(buf) in (bytearray, bytes):
@@ -49,13 +49,18 @@ def __setitem__(self, item: slice, value: bytes) -> None:
         if self.buf.readonly:
             # If buffer is read-only, make a copy (on write)
             patchable_buf = memoryview(bytearray(self.buf))
-            self.buf.release()
+            self.release()
             self.buf = patchable_buf
         self.buf[item] = value
 
     def __len__(self) -> int:
         return len(self.buf)
 
+    def _slice(
+        self, from_offset: Optional[int], to_offset: Optional[int]
+    ) -> memoryview:
+        return self.buf[from_offset:to_offset].toreadonly()
+
     def slice(
         self, from_offset: Optional[int] = None, to_offset: Optional[int] = None
     ) -> "MemoryBuffer":
@@ -66,10 +71,9 @@ def slice(
         changes. It means that changes on parent buffer may be seen in derived buffers,
         but not the other way.
         """
-        return PlainMemoryBuffer(self.buf[from_offset:to_offset].toreadonly())
+        return PlainMemoryBuffer(self._slice(from_offset, to_offset))
 
     def release(self) -> None:
-        print("released ", self, id(self))
         self.buf.release()
 
 
@@ -81,6 +85,7 @@ def __init__(
     ):
         self.opened_file = None
         self.mapped_buf = mapped_buf
+        self._slices: weakref.WeakSet = weakref.WeakSet()
         if mapped_buf is None and file_name is None:
             raise ValueError("Either file_name or map is required.")
         if file_name is not None:
@@ -101,6 +106,8 @@ def __init__(
                 self.opened_file = None
 
     def release(self) -> None:
+        for memory_slice in self._slices:
+            memory_slice.release()
         super().release()
         if self.mapped_buf is not None:
             self.mapped_buf.close()
@@ -108,3 +115,31 @@ def release(self) -> None:
         if self.opened_file is not None:
             self.opened_file.close()
             self.opened_file = None
+
+    def slice(
+        self, from_offset: Optional[int] = None, to_offset: Optional[int] = None
+    ) -> "MemoryBuffer":
+        return self.acquire_slice(self._slice(from_offset, to_offset))
+
+    def acquire_slice(self, buf: memoryview) -> "MemoryBuffer":
+        memory_slice = MmapSliceMemoryBuffer(buf, self)
+        self._slices.add(memory_slice)
+        return memory_slice
+
+    def release_slice(self, memory_slice: "MmapSliceMemoryBuffer") -> None:
+        self._slices.remove(memory_slice)
+
+
+class MmapSliceMemoryBuffer(PlainMemoryBuffer):
+    def __init__(self, buf: memoryview, parent: MmapMemoryBuffer):
+        super().__init__(buf)
+        self.parent = parent
+
+    def slice(
+        self, from_offset: Optional[int] = None, to_offset: Optional[int] = None
+    ) -> "MemoryBuffer":
+        return self.parent.acquire_slice(self._slice(from_offset, to_offset))
+
+    def release(self) -> None:
+        super().release()
+        self.parent.release_slice(self)