From e1dd41643db125f29a5c595c2062bea427f35b7b Mon Sep 17 00:00:00 2001
From: Daniel Widdis <widdis@gmail.com>
Date: Mon, 30 May 2022 11:53:26 -0700
Subject: [PATCH] Handle race condition in PdhUtil#PdhEnumObjectItems

---
 CHANGES.md                                    |  1 +
 .../com/sun/jna/platform/win32/PdhUtil.java   | 39 ++++++++++++++-----
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 470456a22a..ef9cf0154a 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -13,6 +13,7 @@ Features
 Bug Fixes
 ---------
 * [#1438](https://github.com/java-native-access/jna/pull/1438): Handle arrays in structures with differing size - [@matthiasblaesing](https://github.com/matthiasblaesing).
+* [#1442](https://github.com/java-native-access/jna/issues/1442): Handle race condition in `c.s.j.p.win32.PdhUtil#PdhEnumObjectItems` - [@dbwiddis](https://github.com/dbwiddis).
 
 Important Changes
 -----------------
diff --git a/contrib/platform/src/com/sun/jna/platform/win32/PdhUtil.java b/contrib/platform/src/com/sun/jna/platform/win32/PdhUtil.java
index 5f53f4f28e..ebfa8e7f60 100644
--- a/contrib/platform/src/com/sun/jna/platform/win32/PdhUtil.java
+++ b/contrib/platform/src/com/sun/jna/platform/win32/PdhUtil.java
@@ -162,7 +162,9 @@ public static PdhEnumObjectItems PdhEnumObjectItems(String szDataSource, String
         List<String> counters = new ArrayList<String>();
         List<String> instances = new ArrayList<String>();
 
-        // Call once to get string lengths
+        // Call once to get counter and instance string lengths
+        // If zero on input and the object exists, the function returns PDH_MORE_DATA
+        // and sets these parameters to the required buffer size.
         DWORDByReference pcchCounterListLength = new DWORDByReference(new DWORD(0));
         DWORDByReference pcchInstanceListLength = new DWORDByReference(new DWORD(0));
         int result = Pdh.INSTANCE.PdhEnumObjectItems(szDataSource, szMachineName, szObjectName, null, pcchCounterListLength, null,
@@ -174,16 +176,35 @@ public static PdhEnumObjectItems PdhEnumObjectItems(String szDataSource, String
         Memory mszCounterList = null;
         Memory mszInstanceList = null;
 
-        if (pcchCounterListLength.getValue().intValue() > 0) {
-            mszCounterList = new Memory(pcchCounterListLength.getValue().intValue() * CHAR_TO_BYTES);
-        }
+        // A race condition may exist for some instance lists such as processes or
+        // threads which may result in PDH_MORE_DATA.
+        do {
+            if (pcchCounterListLength.getValue().intValue() > 0) {
+                mszCounterList = new Memory(pcchCounterListLength.getValue().intValue() * CHAR_TO_BYTES);
+            }
 
-        if (pcchInstanceListLength.getValue().intValue() > 0) {
-            mszInstanceList = new Memory(pcchInstanceListLength.getValue().intValue() * CHAR_TO_BYTES);
-        }
+            if (pcchInstanceListLength.getValue().intValue() > 0) {
+                mszInstanceList = new Memory(pcchInstanceListLength.getValue().intValue() * CHAR_TO_BYTES);
+            }
 
-        result = Pdh.INSTANCE.PdhEnumObjectItems(szDataSource, szMachineName, szObjectName, mszCounterList,
-                pcchCounterListLength, mszInstanceList, pcchInstanceListLength, dwDetailLevel, 0);
+            result = Pdh.INSTANCE.PdhEnumObjectItems(szDataSource, szMachineName, szObjectName, mszCounterList,
+                    pcchCounterListLength, mszInstanceList, pcchInstanceListLength, dwDetailLevel, 0);
+            if (result == Pdh.PDH_MORE_DATA) {
+                // If the specified size on input is greater than zero but less than the
+                // required size, you should not rely on the returned size to reallocate the
+                // buffer.
+                if (mszCounterList != null) {
+                    long tooSmallSize = mszCounterList.size() / CHAR_TO_BYTES;
+                    pcchCounterListLength.setValue(new DWORD(tooSmallSize + 1024));
+                    mszCounterList.close();
+                }
+                if (mszInstanceList != null) {
+                    long tooSmallSize = mszInstanceList.size() / CHAR_TO_BYTES;
+                    pcchInstanceListLength.setValue(new DWORD(tooSmallSize + 1024));
+                    mszInstanceList.close();
+                }
+            }
+        } while (result == Pdh.PDH_MORE_DATA);
 
         if(result != WinError.ERROR_SUCCESS) {
             throw new PdhException(result);