restructure to use safe_memcpy to avoid memory leak

Dobatymo · Dobatymo · commit 8a3c0d4465aa · 2024-04-25T12:18:06.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()
@@ -255,6 +256,44 @@ do {                                                                    \
 } while (0)
 #endif /* UNIX */
 
+#if defined(MS_WIN32) && !defined(DONT_USE_SEH)
+static DWORD
+HandlePageException(EXCEPTION_POINTERS *ptrs, EXCEPTION_RECORD *record)
+{
+    *record = *ptrs->ExceptionRecord;
+    if (ptrs->ExceptionRecord->ExceptionCode == EXCEPTION_IN_PAGE_ERROR) {
+        return EXCEPTION_EXECUTE_HANDLER;
+    }
+    return EXCEPTION_CONTINUE_SEARCH;
+}
+#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_byte_method(mmap_object *self,
                       PyObject *Py_UNUSED(ignored))
@@ -264,7 +303,14 @@ mmap_read_byte_method(mmap_object *self,
         PyErr_SetString(PyExc_ValueError, "read byte out of range");
         return NULL;
     }
-    return PyLong_FromLong((unsigned char)self->data[self->pos++]);
+    unsigned char dest;
+    if (safe_memcpy(dest, self->data + self->pos, 1)) {
+        self->pos++;
+        return PyLong_FromLong(dest);
+    }
+    else {
+        return NULL;
+    }
 }
 
 static PyObject *
@@ -291,17 +337,6 @@ mmap_read_line_method(mmap_object *self,
     return result;
 }
 
-#if defined(MS_WIN32) && !defined(DONT_USE_SEH)
-static DWORD HandlePageException(EXCEPTION_POINTERS *ptrs, EXCEPTION_RECORD *record)
-{
-    *record = *ptrs->ExceptionRecord;
-    if (ptrs->ExceptionRecord->ExceptionCode == EXCEPTION_IN_PAGE_ERROR) {
-        return EXCEPTION_EXECUTE_HANDLER;
-    }
-    return EXCEPTION_CONTINUE_SEARCH;
-}
-#endif
-
 static PyObject *
 mmap_read_method(mmap_object *self,
                  PyObject *args)
@@ -319,22 +354,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;
 }
 
