[PATCH v2 3/6] setupapi: Implement SetupQueryDrivesInDiskSpaceList.

[Vijay Kiran Kamuju]

From: Michael Müller

From: Michael Müller

Signed-off-by: Vijay Kiran Kamuju 
---
 dlls/setupapi/diskspace.c       | 85 +++++++++++++++++++++++++++++++++
 dlls/setupapi/stubs.c           | 18 -------
 dlls/setupapi/tests/diskspace.c | 70 +++++++++++++++++++++++++++
 3 files changed, 155 insertions(+), 18 deletions(-)

diff --git a/dlls/setupapi/diskspace.c b/dlls/setupapi/diskspace.c
index dcee29443d71..17873ddf9577 100644
--- a/dlls/setupapi/diskspace.c
+++ b/dlls/setupapi/diskspace.c
@@ -419,3 +419,88 @@ BOOL WINAPI SetupAddToDiskSpaceListA(HDSKSPC diskspace, PCSTR targetfile,
     if (targetfileW) HeapFree(GetProcessHeap(), 0, targetfileW);
     return ret;
 }
+
+/***********************************************************************
+ *      SetupQueryDrivesInDiskSpaceListW (SETUPAPI.@)
+ */
+BOOL WINAPI SetupQueryDrivesInDiskSpaceListW(HDSKSPC diskspace, PWSTR buffer, DWORD size, PDWORD required_size)
+{
+    struct space_list *list = diskspace;
+    struct file_entry *file;
+    DWORD cur_size = 1;
+    BOOL used[26];
+
+    TRACE("(%p, %p, %d, %p)\n", diskspace, buffer, size, required_size);
+
+    if (!diskspace)
+    {
+        SetLastError(ERROR_INVALID_PARAMETER);
+        return FALSE;
+    }
+
+    memset(&used, 0, sizeof(used));
+    LIST_FOR_EACH_ENTRY(file, &list->files, struct file_entry, entry)
+    {
+        int device;
+
+        /* UNC paths are not yet supported by this function */
+        if (towlower(file->path[0]) < 'a' || towlower(file->path[0]) > 'z' || file->path[1] != ':')
+            continue;
+
+        device = towlower(file->path[0]) - 'a';
+        if (used[device]) continue;
+
+        cur_size += 3;
+
+        if (buffer)
+        {
+            if (cur_size > size)
+            {
+                if (required_size) *required_size = cur_size;
+                SetLastError(ERROR_INSUFFICIENT_BUFFER);
+                return FALSE;
+            }
+            *buffer++ = towlower(file->path[0]);
+            *buffer++ = ':';
+            *buffer++ = 0;
+        }
+
+        used[device] = TRUE;
+    }
+
+    if (buffer && size) *buffer = 0;
+    if (required_size)  *required_size = cur_size;
+    return TRUE;
+}
+
+/***********************************************************************
+ *      SetupQueryDrivesInDiskSpaceListA (SETUPAPI.@)
+ */
+BOOL WINAPI SetupQueryDrivesInDiskSpaceListA(HDSKSPC diskspace, PSTR buffer, DWORD size, PDWORD required_size)
+{
+    WCHAR *bufferW = NULL;
+    BOOL ret;
+    int i;
+
+    if (buffer && size)
+    {
+        bufferW = HeapAlloc(GetProcessHeap(), 0, size * sizeof(WCHAR));
+        if (!bufferW)
+        {
+            SetLastError(ERROR_NOT_ENOUGH_MEMORY);
+            return FALSE;
+        }
+    }
+
+    ret = SetupQueryDrivesInDiskSpaceListW(diskspace, bufferW ? bufferW : (WCHAR *)buffer,
+                                           size, required_size);
+
+    if (bufferW)
+    {
+        for (i = 0; i < size; i++)
+            buffer[i] = bufferW[i];
+        HeapFree(GetProcessHeap(), 0, bufferW);
+    }
+
+    return ret;
+}
diff --git a/dlls/setupapi/stubs.c b/dlls/setupapi/stubs.c
index 47dca6c15052..80c3a7cfd9f7 100644
--- a/dlls/setupapi/stubs.c
+++ b/dlls/setupapi/stubs.c
@@ -230,24 +230,6 @@ INT WINAPI SetupPromptReboot( HSPFILEQ file_queue, HWND owner, BOOL scan_only )
     return 0;
 }

-/***********************************************************************
- *      SetupQueryDrivesInDiskSpaceListA (SETUPAPI.@)
- */
-BOOL WINAPI SetupQueryDrivesInDiskSpaceListA(HDSKSPC disk_space, PSTR return_buffer, DWORD return_buffer_size, PDWORD required_size)
-{
-    FIXME("%p, %p, %d, %p: stub\n", disk_space, return_buffer, return_buffer_size, required_size);
-    return FALSE;
-}
-
-/***********************************************************************
- *      SetupQueryDrivesInDiskSpaceListW (SETUPAPI.@)
- */
-BOOL WINAPI SetupQueryDrivesInDiskSpaceListW(HDSKSPC disk_space, PWSTR return_buffer, DWORD return_buffer_size, PDWORD required_size)
-{
-    FIXME("%p, %p, %d, %p: stub\n", disk_space, return_buffer, return_buffer_size, required_size);
-    return FALSE;
-}
-
 /***********************************************************************
  *      SetupAddToSourceListA (SETUPAPI.@)
  */
diff --git a/dlls/setupapi/tests/diskspace.c b/dlls/setupapi/tests/diskspace.c
index bb8c7b78463b..60087d11f689 100644
--- a/dlls/setupapi/tests/diskspace.c
+++ b/dlls/setupapi/tests/diskspace.c
@@ -628,6 +628,75 @@ static void test_SetupAddToDiskSpaceListA(void)
        "Expected SetupDestroyDiskSpaceList to succeed\n");
 }

+static void test_SetupQueryDrivesInDiskSpaceListA(void)
+{
+    char buffer[MAX_PATH];
+    HDSKSPC handle;
+    DWORD size;
+    BOOL ret;
+
+    handle = SetupCreateDiskSpaceListA(NULL, 0, SPDSL_IGNORE_DISK);
+    ok(handle != NULL,"Expected SetupCreateDiskSpaceListA to return a valid handle\n");
+
+    ret = SetupQueryDrivesInDiskSpaceListA(handle, NULL, 0, NULL);
+    ok(ret, "Expected SetupQueryDrivesInDiskSpaceListA to succeed\n");
+
+    size = 0;
+    ret = SetupQueryDrivesInDiskSpaceListA(handle, NULL, 0, &size);
+    ok(ret, "Expected SetupQueryDrivesInDiskSpaceListA to succeed\n");
+    ok(size == 1, "Expected size 1, got %u\n", size);
+
+    ret = SetupAddToDiskSpaceListA(handle, "F:\random-file.dat", 0, FILEOP_COPY, 0, 0);
+    ok(ret, "Expected SetupAddToDiskSpaceListA to succeed\n");
+
+    ret = SetupAddToDiskSpaceListA(handle, "G:\random-file.dat", 0, FILEOP_COPY, 0, 0);
+    ok(ret, "Expected SetupAddToDiskSpaceListA to succeed\n");
+
+    ret = SetupAddToDiskSpaceListA(handle, "G:\random-file2.dat", 0, FILEOP_COPY, 0, 0);
+    ok(ret, "Expected SetupAddToDiskSpaceListA to succeed\n");
+
+    ret = SetupAddToDiskSpaceListA(handle, "X:\random-file.dat", 0, FILEOP_COPY, 0, 0);
+    ok(ret, "Expected SetupAddToDiskSpaceListA to succeed\n");
+
+    size = 0;
+    ret = SetupQueryDrivesInDiskSpaceListA(handle, NULL, 0, &size);
+    ok(ret, "Expected SetupQueryDrivesInDiskSpaceListA to succeed\n");
+    ok(size == 10, "Expected size 10, got %u\n", size);
+
+    size = 0;
+    ret = SetupQueryDrivesInDiskSpaceListA(handle, buffer, 0, &size);
+    ok(!ret, "Expected SetupQueryDrivesInDiskSpaceListA to fail\n");
+    ok(size == 4, "Expected size 4, got %u\n", size);
+    ok(GetLastError() == ERROR_INSUFFICIENT_BUFFER,
+       "Expected GetLastError() to return ERROR_INSUFFICIENT_BUFFER, got %u\n", GetLastError());
+
+    size = 0;
+    ret = SetupQueryDrivesInDiskSpaceListA(handle, buffer, 4, &size);
+    ok(!ret, "Expected SetupQueryDrivesInDiskSpaceListA to fail\n");
+    ok(size == 7, "Expected size 7, got %u\n", size);
+    ok(GetLastError() == ERROR_INSUFFICIENT_BUFFER,
+       "Expected GetLastError() to return ERROR_INSUFFICIENT_BUFFER, got %u\n", GetLastError());
+
+    size = 0;
+    ret = SetupQueryDrivesInDiskSpaceListA(handle, buffer, 7, &size);
+    ok(!ret, "Expected SetupQueryDrivesInDiskSpaceListA to fail\n");
+    ok(size == 10, "Expected size 10, got %u\n", size);
+    ok(GetLastError() == ERROR_INSUFFICIENT_BUFFER,
+       "Expected GetLastError() to return ERROR_INSUFFICIENT_BUFFER, got %u\n", GetLastError());
+
+    size = 0;
+    memset(buffer, 0xff, sizeof(buffer));
+    ret = SetupQueryDrivesInDiskSpaceListA(handle, buffer, sizeof(buffer), &size);
+    ok(ret, "Expected SetupQueryDrivesInDiskSpaceListA to succeed\n");
+    ok(size == 10, "Expected size 10, got %u\n", size);
+    ok(!memcmp("f:\0g:\0x:\0\0", buffer, 10), "Device list does not match\n");
+
+    memset(buffer, 0xff, sizeof(buffer));
+    ret = SetupQueryDrivesInDiskSpaceListA(handle, buffer, sizeof(buffer), NULL);
+    ok(ret, "Expected SetupQueryDrivesInDiskSpaceListA to succeed\n");
+    ok(!memcmp("f:\0g:\0x:\0\0", buffer, 10), "Device list does not match\n");
+}
+
 START_TEST(diskspace)
 {
     is_win9x = !SetupCreateDiskSpaceListW((void *)0xdeadbeef, 0xdeadbeef, 0) &&
@@ -639,4 +708,5 @@ START_TEST(diskspace)
     test_SetupQuerySpaceRequiredOnDriveA();
     test_SetupQuerySpaceRequiredOnDriveW();
     test_SetupAddToDiskSpaceListA();
+    test_SetupQueryDrivesInDiskSpaceListA();
 }