Windows: make uv_uptime() more robust

src/win/util.c

@@ -25,6 +25,7 @@
 #include <stdio.h>
 #include <string.h>
 #include <time.h>
+#include <wchar.h>
 
 #include "uv.h"
 #include "internal.h"
@@ -385,70 +386,103 @@ uv_err_t uv_resident_set_memory(size_t* rss) {
 
 
 uv_err_t uv_uptime(double* uptime) {
-  uv_err_t err;
-  PERF_DATA_BLOCK *dataBlock = NULL;
-  PERF_OBJECT_TYPE *objType;
-  PERF_COUNTER_DEFINITION *counterDef;
-  PERF_COUNTER_DEFINITION *counterDefUptime = NULL;
-  DWORD dataSize = 4096;
-  DWORD getSize;
-  LONG lError = ERROR_MORE_DATA;
-  uint64_t upsec;
-  unsigned int i;
-  BYTE *counterData;
+  BYTE stack_buffer[4096];
+  BYTE* malloced_buffer = NULL;
+  BYTE* buffer = (BYTE*) stack_buffer;
+  size_t buffer_size = sizeof(stack_buffer);
+  DWORD data_size;
 
-  *uptime = 0;
-  
-  while (lError == ERROR_MORE_DATA) {
-    if (dataBlock) {
-      free(dataBlock);
-    }
-    dataBlock = (PERF_DATA_BLOCK*)malloc(dataSize);
-    if (!dataBlock) {
-      uv_fatal_error(ERROR_OUTOFMEMORY, "malloc");
-    }
-    getSize = dataSize;
+  PERF_DATA_BLOCK* data_block;
+  PERF_OBJECT_TYPE* object_type;
+  PERF_COUNTER_DEFINITION* counter_definition;
 
-    lError = RegQueryValueExW(HKEY_PERFORMANCE_DATA, "2", NULL, NULL,
-                             (BYTE*)dataBlock, &getSize);
-    if (lError != ERROR_SUCCESS && getSize > 0) {
-      if (wcsncmp(dataBlock->Signature, "PERF", 4) == 0) {
-        break;
-      }
-    } else if (lError != ERROR_SUCCESS && lError != ERROR_MORE_DATA) {
-      err = uv__new_sys_error(GetLastError());
-      goto done;
+  DWORD i;
+
+  for (;;) {
+    LONG result;
+
+    data_size = (DWORD) buffer_size;
+    result = RegQueryValueExW(HKEY_PERFORMANCE_DATA,
+                              L"2",
+                              NULL,
+                              NULL,
+                              buffer,
+                              &data_size);
+    if (result == ERROR_SUCCESS) {
+      break;
+    } else if (result != ERROR_MORE_DATA) {
+      *uptime = 0;
+      return uv__new_sys_error(result);
     }
 
-    dataSize += 1024;
+    free(malloced_buffer);
+
+    buffer_size *= 2;
+    /* Don't let the buffer grow infinitely. */
+    if (buffer_size > 1 << 20) {
+      goto internalError;
+    }
+
+    buffer = malloced_buffer = (BYTE*) malloc(buffer_size);
+    if (malloced_buffer == NULL) {
+      *uptime = 0;
+      return uv__new_artificial_error(UV_ENOMEM);
+    }
   }
 
-  RegCloseKey(HKEY_PERFORMANCE_DATA);
+  if (data_size < sizeof(*data_block))
+    goto internalError;
 
-  objType = (PERF_OBJECT_TYPE*)((BYTE*)dataBlock + dataBlock->HeaderLength);
-  counterDef = (PERF_COUNTER_DEFINITION*)((BYTE*)objType + objType->HeaderLength);
+  data_block = (PERF_DATA_BLOCK*) buffer;
 
-  for (i = 0; i < objType->NumCounters; ++i) {
-    if (counterDef->CounterNameTitleIndex == 674) {
-      counterDefUptime = counterDef;
+  if (wmemcmp(data_block->Signature, L"PERF", 4) != 0)
+    goto internalError;
+
+  if (data_size < data_block->HeaderLength + sizeof(*object_type))
+    goto internalError;
+
+  object_type = (PERF_OBJECT_TYPE*) (buffer + data_block->HeaderLength);
+
+  if (object_type->NumInstances != PERF_NO_INSTANCES)
+    goto internalError;
+
+  counter_definition = (PERF_COUNTER_DEFINITION*) (buffer +
+      data_block->HeaderLength + object_type->HeaderLength);
+  for (i = 0; i < object_type->NumCounters; i++) {
+    if ((BYTE*) counter_definition + sizeof(*counter_definition) >
+        buffer + data_size) {
       break;
     }
-    counterDef = (PERF_COUNTER_DEFINITION*)((BYTE*)counterDef + counterDef->ByteLength);
+
+    if (counter_definition->CounterNameTitleIndex == 674 &&
+        counter_definition->CounterSize == sizeof(uint64_t)) {
+      if (counter_definition->CounterOffset + sizeof(uint64_t) > data_size ||
+          !(counter_definition->CounterType & PERF_OBJECT_TIMER)) {
+        goto internalError;
+      } else {
+        BYTE* address = (BYTE*) object_type + object_type->DefinitionLength +
+                        counter_definition->CounterOffset;
+        uint64_t value = *((uint64_t*) address);
+        *uptime = (double) (object_type->PerfTime.QuadPart - value) /
+                  (double) object_type->PerfFreq.QuadPart;
+        free(malloced_buffer);
+        return uv_ok_;
+      }
+    }
+
+    counter_definition = (PERF_COUNTER_DEFINITION*)
+        ((BYTE*) counter_definition + counter_definition->ByteLength);
   }
 
-  counterData = (BYTE*)objType + objType->DefinitionLength;
-  counterData += counterDefUptime->CounterOffset;
+  /* If we get here, the uptime value was not found. */
+  free(malloced_buffer);
+  *uptime = 0;
+  return uv__new_artificial_error(UV_ENOSYS);
 
-  upsec = *((uint64_t*)counterData);
-  *uptime = ((objType->PerfTime.QuadPart - upsec) / objType->PerfFreq.QuadPart);
-  err = uv_ok_;
-
-done:
-  if (dataBlock) {
-    free(dataBlock);
-  }
-
-  return err;
+ internalError:
+  free(malloced_buffer);
+  *uptime = 0;
+  return uv__new_artificial_error(UV_EIO);
 }