setupapi: Rewrite DiskSpaceList logic using lists.

dlls/setupapi/diskspace.c

 /*
  * SetupAPI DiskSpace functions
  *
+ * Copyright 2016 Michael Müller
  * Copyright 2004 CodeWeavers (Aric Stewart)
  *
  * This library is free software; you can redistribute it and/or
@@ -27,69 +28,50 @@
 #include "winnls.h"
 #include "winreg.h"
 #include "setupapi.h"
+#include "wine/list.h"
 #include "wine/debug.h"
+#include "setupapi_private.h"
 
 WINE_DEFAULT_DEBUG_CHANNEL(setupapi);
 
-typedef struct {
-    WCHAR   lpzName[20];
-    LONGLONG dwFreeSpace;
-    LONGLONG dwWantedSpace;
-} DRIVE_ENTRY, *LPDRIVE_ENTRY;
+struct file_entry
+{
+    struct list entry;
+    WCHAR *path;
+    UINT operation;
+    LONGLONG size;
+};
+
+struct space_list
+{
+    struct list files;
+    UINT flags;
+};
 
-typedef struct {
-    DWORD   dwDriveCount;
-    DRIVE_ENTRY Drives[26];
-} DISKSPACELIST, *LPDISKSPACELIST;
 
 
 /***********************************************************************
  *		SetupCreateDiskSpaceListW  (SETUPAPI.@)
  */
-HDSKSPC WINAPI SetupCreateDiskSpaceListW(PVOID Reserved1, DWORD Reserved2, UINT Flags)
+HDSKSPC WINAPI SetupCreateDiskSpaceListW(PVOID reserved1, DWORD reserved2, UINT flags)
 {
-    WCHAR drives[255];
-    DWORD rc;
-    WCHAR *ptr;
-    LPDISKSPACELIST list=NULL;
+    struct space_list *list;
 
-    TRACE("(%p, %lu, 0x%08x)\n", Reserved1, Reserved2, Flags);
+    TRACE("(%p, %lu, 0x%08x)\n", reserved1, reserved2, flags);
 
-    if (Reserved1 || Reserved2 || Flags & ~SPDSL_IGNORE_DISK)
+    if (reserved1 || reserved2 || flags & ~SPDSL_IGNORE_DISK)
     {
         SetLastError(ERROR_INVALID_PARAMETER);
         return NULL;
     }
 
-    rc = GetLogicalDriveStringsW(255,drives);
-
-    if (rc == 0)
-        return NULL;
-
-    list = HeapAlloc(GetProcessHeap(),0,sizeof(DISKSPACELIST));
-
-    list->dwDriveCount = 0;
-    
-    ptr = drives;
-    
-    while (*ptr)
+    list = HeapAlloc(GetProcessHeap(), 0, sizeof(*list));
+    if (list)
     {
-        DWORD type = GetDriveTypeW(ptr);
-        if (type == DRIVE_FIXED)
-        {
-            DWORD clusters;
-            DWORD sectors;
-            DWORD bytes;
-            DWORD total;
-            lstrcpyW(list->Drives[list->dwDriveCount].lpzName,ptr);
-            GetDiskFreeSpaceW(ptr,&sectors,&bytes,&clusters,&total);
-            list->Drives[list->dwDriveCount].dwFreeSpace = clusters * sectors *
-                                                           bytes;
-            list->Drives[list->dwDriveCount].dwWantedSpace = 0;
-            list->dwDriveCount++;
-        }
-       ptr += lstrlenW(ptr) + 1;
+        list->flags = flags;
+        list_init(&list->files);
     }
+
     return list;
 }
 
@@ -105,32 +87,58 @@ HDSKSPC WINAPI SetupCreateDiskSpaceListA(PVOID Reserved1, DWORD Reserved2, UINT
 /***********************************************************************
  *		SetupDuplicateDiskSpaceListW  (SETUPAPI.@)
  */
-HDSKSPC WINAPI SetupDuplicateDiskSpaceListW(HDSKSPC DiskSpace, PVOID Reserved1, DWORD Reserved2, UINT Flags)
+HDSKSPC WINAPI SetupDuplicateDiskSpaceListW(HDSKSPC diskspace, PVOID reserved1, DWORD reserved2, UINT flags)
 {
-    DISKSPACELIST *list_copy, *list_original = DiskSpace;
+    struct space_list *list_copy, *list = diskspace;
+    struct file_entry *file, *file_copy;
+
+    TRACE("(%p, %p, %lu, %u)\n", diskspace, reserved1, reserved2, flags);
 
-    if (Reserved1 || Reserved2 || Flags)
+    if (reserved1 || reserved2 || flags)
     {
         SetLastError(ERROR_INVALID_PARAMETER);
         return NULL;
     }
 
-    if (!DiskSpace)
+    if (!diskspace)
     {
         SetLastError(ERROR_INVALID_HANDLE);
         return NULL;
     }
 
-    list_copy = HeapAlloc(GetProcessHeap(), 0, sizeof(DISKSPACELIST));
+    list_copy = HeapAlloc(GetProcessHeap(), 0, sizeof(*list_copy));
     if (!list_copy)
     {
         SetLastError(ERROR_NOT_ENOUGH_MEMORY);
         return NULL;
     }
 
-    *list_copy = *list_original;
+    list_copy->flags = list->flags;
+    list_init(&list_copy->files);
+
+    LIST_FOR_EACH_ENTRY(file, &list->files, struct file_entry, entry)
+    {
+        file_copy = HeapAlloc(GetProcessHeap(), 0, sizeof(*file_copy));
+        if (!file_copy) goto error;
+
+        file_copy->path = strdupW(file->path);
+        if (!file_copy->path)
+        {
+            HeapFree(GetProcessHeap(), 0, file_copy);
+            goto error;
+        }
+
+        file_copy->operation = file->operation;
+        file_copy->size = file->size;
+        list_add_head(&list_copy->files, &file->entry);
+    }
 
     return list_copy;
+
+error:
+    SetLastError(ERROR_NOT_ENOUGH_MEMORY);
+    SetupDestroyDiskSpaceList(list_copy);
+    return NULL;
 }
 
 /***********************************************************************
@@ -155,55 +163,51 @@ BOOL WINAPI SetupAddInstallSectionToDiskSpaceListA(HDSKSPC DiskSpace,
 /***********************************************************************
 *		SetupQuerySpaceRequiredOnDriveW  (SETUPAPI.@)
 */
-BOOL WINAPI SetupQuerySpaceRequiredOnDriveW(HDSKSPC DiskSpace,
-                        LPCWSTR DriveSpec, LONGLONG *SpaceRequired,
-                        PVOID Reserved1, UINT Reserved2)
+BOOL WINAPI SetupQuerySpaceRequiredOnDriveW(HDSKSPC diskspace,
+                        LPCWSTR drivespec, LONGLONG *required,
+                        PVOID reserved1, UINT reserved2)
 {
-    WCHAR *driveW;
-    unsigned int i;
-    LPDISKSPACELIST list = DiskSpace;
-    BOOL rc = FALSE;
-    static const WCHAR bkslsh[]= {'\\',0};
+    struct space_list *list = diskspace;
+    struct file_entry *file;
+    LONGLONG sum = 0;
 
-    if (!DiskSpace)
+    TRACE("(%p, %s, %p, %p, %u)\n", diskspace, debugstr_w(drivespec), required, reserved1, reserved2);
+
+    if (!diskspace)
     {
         SetLastError(ERROR_INVALID_HANDLE);
         return FALSE;
     }
 
-    if (!DriveSpec)
+    if (!drivespec || !drivespec[0])
     {
-        SetLastError(ERROR_INVALID_PARAMETER);
+        SetLastError(drivespec ? ERROR_INVALID_DRIVE : ERROR_INVALID_DRIVE);
         return FALSE;
     }
 
-    driveW = HeapAlloc(GetProcessHeap(), 0, (lstrlenW(DriveSpec) + 2) * sizeof(WCHAR));
-    if (!driveW)
+    if (!required)
     {
-        SetLastError(ERROR_NOT_ENOUGH_MEMORY);
+        SetLastError(ERROR_INVALID_PARAMETER);
         return FALSE;
     }
 
-    lstrcpyW(driveW,DriveSpec);
-    lstrcatW(driveW,bkslsh);
-
-    TRACE("Looking for drive %s\n",debugstr_w(driveW));
- 
-    for (i = 0; i < list->dwDriveCount; i++)
+    if (towlower(drivespec[0]) < 'a' || towlower(drivespec[0]) > 'z' ||
+        drivespec[1] != ':' || drivespec[2] != 0)
     {
-        TRACE("checking drive %s\n",debugstr_w(list->Drives[i].lpzName));
-        if (wcscmp(driveW,list->Drives[i].lpzName)==0)
-        {
-            rc = TRUE;
-            *SpaceRequired = list->Drives[i].dwWantedSpace;
-            break;
-        }
+        FIXME("UNC paths not yet supported (%s)\n", debugstr_w(drivespec));
+        SetLastError((GetVersion() & 0x80000000) ? ERROR_INVALID_DRIVE : ERROR_INVALID_PARAMETER);
+        return FALSE;
     }
 
-    HeapFree(GetProcessHeap(), 0, driveW);
+    LIST_FOR_EACH_ENTRY(file, &list->files, struct file_entry, entry)
+    {
+        if (towlower(file->path[0]) == towlower(drivespec[0]) &&
+            file->path[1] == ':' && file->path[2] == '\\')
+            sum += file->size;
+    }
 
-    if (!rc) SetLastError(ERROR_INVALID_DRIVE);
-    return rc;
+    *required = sum;
+    return TRUE;
 }
 
 /***********************************************************************
@@ -253,10 +257,25 @@ BOOL WINAPI SetupQuerySpaceRequiredOnDriveA(HDSKSPC DiskSpace,
 /***********************************************************************
 *		SetupDestroyDiskSpaceList  (SETUPAPI.@)
 */
-BOOL WINAPI SetupDestroyDiskSpaceList(HDSKSPC DiskSpace)
+BOOL WINAPI SetupDestroyDiskSpaceList(HDSKSPC diskspace)
 {
-    LPDISKSPACELIST list = DiskSpace;
-    HeapFree(GetProcessHeap(),0,list);
+    struct space_list *list = diskspace;
+    struct file_entry *file, *file2;
+
+    if (!diskspace)
+    {
+        SetLastError(ERROR_INVALID_PARAMETER);
+        return FALSE;
+    }
+
+    LIST_FOR_EACH_ENTRY_SAFE(file, file2, &list->files, struct file_entry, entry)
+    {
+        HeapFree(GetProcessHeap(), 0, file->path);
+        list_remove(&file->entry);
+        HeapFree(GetProcessHeap(), 0, file);
+    }
+
+    HeapFree(GetProcessHeap(), 0, list);
     return TRUE; 
 }