restructure to use safe_memcpy to avoid memory leak

Dobatymo · Dobatymo · commit 03a88169f3b6 · 2024-04-25T11:56:04.000+08:00
diff --git a/Modules/mmapmodule.c b/Modules/mmapmodule.c
@@ -27,6 +27,7 @@
 #include "pycore_bytesobject.h"   // _PyBytes_Find()
 #include "pycore_fileutils.h"     // _Py_stat_struct
 
+#include <stdbool.h>
 #include <stddef.h>               // offsetof()
 #ifndef MS_WINDOWS
 #  include <unistd.h>             // close()
@@ -302,6 +303,31 @@ static DWORD HandlePageException(EXCEPTION_POINTERS *ptrs, EXCEPTION_RECORD *rec
 }
 #endif
 
+bool safe_memcpy(void *restrict dest, const void *restrict src, size_t count) {
+#if defined(MS_WIN32) && !defined(DONT_USE_SEH)
+
+    // never fail for count 0
+    if (count == 0) {
+        return true;
+    }
+
+    EXCEPTION_RECORD record;
+    __try {
+        memcpy(dest, src, count);
+        return true;
+    }
+    __except (HandlePageException(GetExceptionInformation(), &record)) {
+        NTSTATUS status = record.ExceptionInformation[2];
+        ULONG code = LsaNtStatusToWinError(status);
+        PyErr_SetFromWindowsErr(code);
+        return false;
+    }
+#else
+    memcpy(dest, src, count);
+    return true;
+#endif
+}
+
 static PyObject *
 mmap_read_method(mmap_object *self,
                  PyObject *args)
@@ -319,22 +345,16 @@ mmap_read_method(mmap_object *self,
     if (num_bytes < 0 || num_bytes > remaining)
         num_bytes = remaining;
 
-#if defined(MS_WIN32) && !defined(DONT_USE_SEH)
-    EXCEPTION_RECORD record;
-    __try {
-        result = PyBytes_FromStringAndSize(&self->data[self->pos], num_bytes);
+    result = PyBytes_FromStringAndSize(NULL, num_bytes);
+    if (result == NULL) {
+        return NULL;
+    }
+    if (safe_memcpy(((PyBytesObject *) result)->ob_sval, &self->data[self->pos], num_bytes)) {
         self->pos += num_bytes;
     }
-    __except (HandlePageException(GetExceptionInformation(), &record)) {
-        NTSTATUS code = record.ExceptionInformation[2];
-        PyErr_SetFromWindowsErr(code);
-        result = NULL;
+    else {
+        Py_CLEAR(result);
     }
-#else
-    result = PyBytes_FromStringAndSize(&self->data[self->pos], num_bytes);
-    self->pos += num_bytes;
-#endif
-
     return result;
 }
 
