Fixing OOM crash in base::ProcessMetrics::GetWorkingSetKBytes

base/process/process_metrics_win.cc

Index: base/process/process_metrics_win.cc
diff --git a/base/process/process_metrics_win.cc b/base/process/process_metrics_win.cc
index fe4b84988f1a2736da6fef1b9a3c5ca0552f5ba0..e1a5bfecad99431f46038ba797fe54ec405e0454 100644
--- a/base/process/process_metrics_win.cc
+++ b/base/process/process_metrics_win.cc
@@ -11,6 +11,7 @@
 #include <winternl.h>
 
 #include <algorithm>
+#include <memory>
 
 #include "base/logging.h"
 #include "base/memory/ptr_util.h"
@@ -148,41 +149,43 @@ bool ProcessMetrics::GetWorkingSetKBytes(WorkingSetKBytes* ws_usage) const {
   DCHECK(ws_usage);
   memset(ws_usage, 0, sizeof(*ws_usage));
 
-  DWORD number_of_entries = 4096;  // Just a guess.
-  PSAPI_WORKING_SET_INFORMATION* buffer = NULL;
+  // Call QueryWorkingSet once to get number of items.
+  PSAPI_WORKING_SET_INFORMATION initial_buffer;
+  QueryWorkingSet(process_, &initial_buffer, sizeof(initial_buffer));
+  if (GetLastError() != ERROR_BAD_LENGTH)
+    return false;
+
+  DWORD number_of_entries = static_cast<DWORD>(initial_buffer.NumberOfEntries);
+
+  std::unique_ptr<PSAPI_WORKING_SET_INFORMATION> buffer;
   int retries = 5;
   for (;;) {
+    // Maybe some entries are being added right now. Increase the buffer to
+    // take that into account.
+    number_of_entries = static_cast<DWORD>(number_of_entries * 1.1);
+
     DWORD buffer_size = sizeof(PSAPI_WORKING_SET_INFORMATION) +
                         (number_of_entries * sizeof(PSAPI_WORKING_SET_BLOCK));
 
-    // if we can't expand the buffer, don't leak the previous
-    // contents or pass a NULL pointer to QueryWorkingSet
-    PSAPI_WORKING_SET_INFORMATION* new_buffer =
-        reinterpret_cast<PSAPI_WORKING_SET_INFORMATION*>(
-            realloc(buffer, buffer_size));
-    if (!new_buffer) {
-      free(buffer);
-      return false;
-    }
-    buffer = new_buffer;
+    // Allocate the buffer.
+    DCHECK(!buffer.get());
+    buffer.reset(reinterpret_cast<PSAPI_WORKING_SET_INFORMATION*>(
+        new uint8_t[buffer_size]));

[brucedawson, 2016/09/07 00:41:12]

This seems slightly dodgy, using unique_ptr to hold a pointer that has had
reinterpet_cast run on it. It means that the new and delete calls will have
mismatched sizes.

I'd suggest a vector<uint8_t> except that then we get the zero initialization
behavior that would waste memory. Sigh...

Would be best to have std::unique_ptr<uint8_t> and cast that result when you use
it, except that even then you are misusing unique_ptr which always does a scalar
delete - not legal for an array.

 
     // Call the function once to get number of items
-    if (QueryWorkingSet(process_, buffer, buffer_size))
+    if (QueryWorkingSet(process_, buffer.get(), buffer_size))
       break;  // Success
 
-    if (GetLastError() != ERROR_BAD_LENGTH) {
-      free(buffer);
+    if (GetLastError() != ERROR_BAD_LENGTH)
       return false;
-    }
 
+    // Need to grow the number of entries.
     number_of_entries = static_cast<DWORD>(buffer->NumberOfEntries);
-
-    // Maybe some entries are being added right now. Increase the buffer to
-    // take that into account.
-    number_of_entries = static_cast<DWORD>(number_of_entries * 1.25);
+    // Free the current buffer before allocating one with larger size.
+    buffer.reset();
 
     if (--retries == 0) {
-      free(buffer);  // If we're looping, eventually fail.
+      // If we're looping, eventually fail.
       return false;
     }
   }
@@ -205,7 +208,6 @@ bool ProcessMetrics::GetWorkingSetKBytes(WorkingSetKBytes* ws_usage) const {
   ws_usage->priv = ws_private * PAGESIZE_KB;
   ws_usage->shareable = ws_shareable * PAGESIZE_KB;
   ws_usage->shared = ws_shared * PAGESIZE_KB;
-  free(buffer);
   return true;
 }