[Windows] disk_partitions: use wide-char APIs to fix UnicodeDecodeError (#1959)

ANSI GetLogicalDriveStrings/FindFirstVolumeMountPoint/GetVolumeInformation
return bytes in the system ANSI code page; feeding them to Py_BuildValue
"(ssss)" decodes as strict UTF-8 and crashes on non-UTF8 systems (e.g.
cp1251). Switch to the W family + PyUnicode_FromWideChar, and make
psutil_QueryDosDevice wide for the same reason.

Signed-off-by: Alex Chen <l46983284@gmail.com>

Author: l46983284-cpu

psutil/arch/windows/disk.c

@@ -6,7 +6,6 @@
 
 #include <Python.h>
 #include <windows.h>
-#include <tchar.h>
 
 #include "../../arch/all/init.h"
 
@@ -221,21 +220,25 @@ psutil_disk_io_counters(PyObject *self, PyObject *args) {
 
 PyObject *
 psutil_disk_partitions(PyObject *self, PyObject *args) {
-    DWORD num_bytes;
-    char drive_strings[255];
-    char *drive_letter = drive_strings;
-    char mp_buf[MAX_PATH];
-    char mp_path[MAX_PATH];
+    DWORD num_chars;
+    wchar_t drive_strings[255];
+    wchar_t *drive_letter = drive_strings;
+    wchar_t mp_buf[MAX_PATH];
+    wchar_t mp_path[MAX_PATH];
     int all = 0;
     int type;
     int ret;
     unsigned int old_mode = 0;
     char opts[50];
     HANDLE mp_h;
     BOOL mp_flag = TRUE;
-    char fs_type[MAX_PATH + 1] = {0};
+    wchar_t fs_type[MAX_PATH + 1] = {0};
     DWORD pflags = 0;
     DWORD lpMaximumComponentLength = 0;  // max file name
+    PyObject *py_drive = NULL;
+    PyObject *py_mount = NULL;
+    PyObject *py_fs = NULL;
+    PyObject *py_opts = NULL;
     PyObject *py_retlist = PyList_New(0);
 
     if (py_retlist == NULL) {
@@ -250,11 +253,11 @@ psutil_disk_partitions(PyObject *self, PyObject *args) {
         goto error;
 
     Py_BEGIN_ALLOW_THREADS
-    num_bytes = GetLogicalDriveStrings(254, drive_letter);
+    num_chars = GetLogicalDriveStringsW(254, drive_letter);
     Py_END_ALLOW_THREADS
 
-    if (num_bytes == 0) {
-        psutil_oserror();
+    if (num_chars == 0) {
+        psutil_oserror_wsyscall("GetLogicalDriveStringsW");
         goto error;
     }
 
@@ -263,7 +266,7 @@ psutil_disk_partitions(PyObject *self, PyObject *args) {
         fs_type[0] = 0;
 
         Py_BEGIN_ALLOW_THREADS
-        type = GetDriveType(drive_letter);
+        type = GetDriveTypeW(drive_letter);
         Py_END_ALLOW_THREADS
 
         // by default we only show hard drives and cd-roms
@@ -276,7 +279,7 @@ psutil_disk_partitions(PyObject *self, PyObject *args) {
             // floppy disk: skip it by default as it introduces a
             // considerable slowdown.
             if ((type == DRIVE_REMOVABLE)
-                && (strcmp(drive_letter, "A:\\") == 0))
+                && (wcscmp(drive_letter, L"A:\\") == 0))
             {
                 goto next;
             }
@@ -285,8 +288,8 @@ psutil_disk_partitions(PyObject *self, PyObject *args) {
         // May spin up a removable drive or go over the wire for a
         // network one, so do it without the GIL.
         Py_BEGIN_ALLOW_THREADS
-        ret = GetVolumeInformation(
-            (LPCTSTR)drive_letter,
+        ret = GetVolumeInformationW(
+            drive_letter,
             NULL,  // we don't want the volume name
             0,
             NULL,
@@ -320,40 +323,64 @@ psutil_disk_partitions(PyObject *self, PyObject *args) {
             str_append(opts, sizeof(opts), ",");
         str_append(opts, sizeof(opts), psutil_get_drive_type(type));
 
+        // Convert the strings which will go into the result tuples.
+        py_drive = PyUnicode_FromWideChar(drive_letter, wcslen(drive_letter));
+        if (py_drive == NULL)
+            goto error;
+        py_fs = PyUnicode_FromWideChar(fs_type, wcslen(fs_type));
+        if (py_fs == NULL)
+            goto error;
+        // opts holds pure ASCII, so plain UTF-8 decoding is safe.
+        py_opts = PyUnicode_FromString(opts);
+        if (py_opts == NULL)
+            goto error;
+
         // Check for mount points on this volume and add/get info
         // (checks first to know if we can even have mount points)
         if ((ret != 0) && (pflags & FILE_SUPPORTS_REPARSE_POINTS)) {
             Py_BEGIN_ALLOW_THREADS
-            mp_h = FindFirstVolumeMountPoint(drive_letter, mp_buf, MAX_PATH);
+            mp_h = FindFirstVolumeMountPointW(drive_letter, mp_buf, MAX_PATH);
             Py_END_ALLOW_THREADS
 
             if (mp_h != INVALID_HANDLE_VALUE) {
                 mp_flag = TRUE;
                 while (mp_flag) {
-                    // Append full mount path with drive letter
-                    str_copy(
-                        mp_path, sizeof(mp_path), drive_letter
-                    );  // initialize
-                    str_append(
-                        mp_path, sizeof(mp_path), mp_buf
-                    );  // append mount point
+                    // Append full mount path with drive letter.
+                    mp_path[0] = L'\0';
+                    wcsncat(mp_path, drive_letter, _ARRAYSIZE(mp_path) - 1);
+                    wcsncat(
+                        mp_path,
+                        mp_buf,
+                        _ARRAYSIZE(mp_path) - wcslen(mp_path) - 1
+                    );
+
+                    py_mount = PyUnicode_FromWideChar(
+                        mp_path, wcslen(mp_path)
+                    );
+                    if (py_mount == NULL) {
+                        FindVolumeMountPointClose(mp_h);
+                        goto error;
+                    }
 
                     if (!pylist_append_fmt(
                             py_retlist,
-                            "(ssss)",
-                            drive_letter,
-                            mp_path,
-                            fs_type,  // typically "NTFS"
-                            opts
+                            "(OOOO)",
+                            py_drive,
+                            py_mount,
+                            py_fs,  // typically "NTFS"
+                            py_opts
                         ))
                     {
                         FindVolumeMountPointClose(mp_h);
                         goto error;
                     }
+                    Py_CLEAR(py_mount);
 
                     // Continue looking for more mount points
                     Py_BEGIN_ALLOW_THREADS
-                    mp_flag = FindNextVolumeMountPoint(mp_h, mp_buf, MAX_PATH);
+                    mp_flag = FindNextVolumeMountPointW(
+                        mp_h, mp_buf, MAX_PATH
+                    );
                     Py_END_ALLOW_THREADS
                 }
                 FindVolumeMountPointClose(mp_h);
@@ -362,25 +389,33 @@ psutil_disk_partitions(PyObject *self, PyObject *args) {
 
         if (!pylist_append_fmt(
                 py_retlist,
-                "(ssss)",
-                drive_letter,
-                drive_letter,
-                fs_type,  // either FAT, FAT32, NTFS, HPFS, CDFS, UDF or NWFS
-                opts
+                "(OOOO)",
+                py_drive,
+                py_drive,
+                py_fs,  // either FAT, FAT32, NTFS, HPFS, CDFS, UDF or NWFS
+                py_opts
             ))
         {
             goto error;
         }
         goto next;
 
     next:
-        drive_letter = strchr(drive_letter, 0) + 1;
+        Py_CLEAR(py_opts);
+        Py_CLEAR(py_fs);
+        Py_CLEAR(py_mount);
+        Py_CLEAR(py_drive);
+        drive_letter = wcschr(drive_letter, 0) + 1;
     }
 
     SetErrorMode(old_mode);
     return py_retlist;
 
 error:
+    Py_XDECREF(py_opts);
+    Py_XDECREF(py_fs);
+    Py_XDECREF(py_mount);
+    Py_XDECREF(py_drive);
     SetErrorMode(old_mode);
     Py_DECREF(py_retlist);
     return NULL;
@@ -392,23 +427,32 @@ psutil_disk_partitions(PyObject *self, PyObject *args) {
 // If no match is found return an empty string.
 PyObject *
 psutil_QueryDosDevice(PyObject *self, PyObject *args) {
-    LPCTSTR lpDevicePath;
-    TCHAR d = TEXT('A');
-    TCHAR szBuff[5];
+    PyObject *py_device_path;
+    PyObject *py_ret = NULL;
+    wchar_t *device_path;
+    wchar_t d = L'A';
 
-    if (!PyArg_ParseTuple(args, "s", &lpDevicePath))
+    if (!PyArg_ParseTuple(args, "U", &py_device_path))
         return NULL;
 
-    while (d <= TEXT('Z')) {
-        TCHAR szDeviceName[3] = {d, TEXT(':'), TEXT('\0')};
-        TCHAR szTarget[512] = {0};
-        if (QueryDosDevice(szDeviceName, szTarget, 511) != 0) {
-            if (_tcscmp(lpDevicePath, szTarget) == 0) {
-                _stprintf_s(szBuff, _countof(szBuff), TEXT("%c:"), d);
-                return PyUnicode_FromString(szBuff);
+    device_path = PyUnicode_AsWideCharString(py_device_path, NULL);
+    if (device_path == NULL)
+        return NULL;
+
+    while (d <= L'Z') {
+        wchar_t szDeviceName[3] = {d, L':', L'\0'};
+        wchar_t szTarget[512] = {0};
+        if (QueryDosDeviceW(szDeviceName, szTarget, 511) != 0) {
+            if (wcscmp(device_path, szTarget) == 0) {
+                py_ret = PyUnicode_FromWideChar(
+                    szDeviceName, wcslen(szDeviceName)
+                );
+                PyMem_Free(device_path);
+                return py_ret;
             }
         }
         d++;
     }
+    PyMem_Free(device_path);
     return PyUnicode_FromString("");
 }