From 99174b786a85040bd75c533e9ed97613b91627b4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Michael=20M=C3=BCller?= <michael@fds-team.de>
Date: Fri, 4 Mar 2016 04:53:00 +0100
Subject: [PATCH] setupapi: ImplementSetupAddSectionToDiskSpaceList.

---
 dlls/setupapi/diskspace.c        | 130 +++++++++++++++++++++++++-
 dlls/setupapi/queue.c            |   2 +-
 dlls/setupapi/setupapi.spec      |   4 +-
 dlls/setupapi/setupapi_private.h |   2 +
 dlls/setupapi/tests/diskspace.c  | 193 +++++++++++++++++++++++++++++++++++++++
 5 files changed, 327 insertions(+), 4 deletions(-)

diff --git a/dlls/setupapi/diskspace.c b/dlls/setupapi/diskspace.c
index af36dd4..73e7fbd 100644
--- a/dlls/setupapi/diskspace.c
+++ b/dlls/setupapi/diskspace.c
@@ -65,6 +65,23 @@ static LONGLONG get_file_size(WCHAR *path)
     return size.QuadPart;
 }
 
+static BOOL get_size_from_inf(HINF layoutinf, WCHAR *filename, LONGLONG *size)
+{
+    static const WCHAR SourceDisksFiles[]  = {'S','o','u','r','c','e','D','i','s','k','s','F','i','l','e','s',0};
+    INFCONTEXT context;
+    WCHAR buffer[20];
+
+    if (!SetupFindFirstLineW(layoutinf, SourceDisksFiles, filename, &context))
+        return FALSE;
+
+    if (!SetupGetStringFieldW(&context, 3, buffer, sizeof(buffer), NULL))
+        return FALSE;
+
+    /* FIXME: is there a atollW ? */
+    *size = atolW(buffer);
+    return TRUE;
+}
+
 /***********************************************************************
  *		SetupCreateDiskSpaceListW  (SETUPAPI.@)
  */
@@ -165,7 +182,118 @@ HDSKSPC WINAPI SetupDuplicateDiskSpaceListA(HDSKSPC DiskSpace, PVOID Reserved1,
 }
 
 /***********************************************************************
- *		SetupAddInstallSectionToDiskSpaceListA  (SETUPAPI.@)
+ *      SetupAddSectionToDiskSpaceListW  (SETUPAPI.@)
+ */
+BOOL WINAPI SetupAddSectionToDiskSpaceListW(HDSKSPC diskspace, HINF hinf, HINF hlist,
+                                            PCWSTR section, UINT operation, PVOID reserved1,
+                                            UINT reserved2)
+{
+    static const WCHAR sepW[] = {'\\',0};
+    WCHAR dest[MAX_PATH], src[MAX_PATH], *dest_dir, *full_path;
+    INFCONTEXT context;
+    BOOL ret = FALSE;
+
+    TRACE("(%p, %p, %p, %s, %u, %p, %u)\n", diskspace, hinf, hlist, debugstr_w(section),
+                                            operation, reserved1, reserved2);
+
+    if (!diskspace)
+    {
+        SetLastError(ERROR_INVALID_HANDLE);
+        return FALSE;
+    }
+
+    if (!section)
+    {
+        SetLastError(ERROR_INVALID_PARAMETER);
+        return FALSE;
+    }
+
+    if (!hlist) hlist = hinf;
+
+    if (!SetupFindFirstLineW(hlist, section, NULL, &context))
+    {
+        SetLastError(ERROR_SECTION_NOT_FOUND);
+        return FALSE;
+    }
+
+    dest_dir = get_destination_dir(hinf, section);
+    if (!dest_dir)
+    {
+        SetLastError(ERROR_NOT_ENOUGH_MEMORY);
+        return FALSE;
+    }
+
+    do
+    {
+        LONGLONG filesize;
+        int path_size;
+        BOOL tmp_ret;
+
+        if (!SetupGetStringFieldW(&context, 1, dest, sizeof(dest) / sizeof(WCHAR), NULL))
+            goto end;
+        if (!SetupGetStringFieldW(&context, 2, src, sizeof(src) / sizeof(WCHAR), NULL))
+            *src = 0;
+        if (!get_size_from_inf(hinf, src[0] ? src : dest, &filesize))
+            goto end;
+
+        path_size = strlenW(dest_dir) + strlenW(dest) + 2;
+        full_path = HeapAlloc(GetProcessHeap(), 0, path_size * sizeof(WCHAR));
+        if (!full_path)
+        {
+            SetLastError(ERROR_NOT_ENOUGH_MEMORY);
+            goto end;
+        }
+
+        strcpyW(full_path, dest_dir);
+        strcatW(full_path, sepW);
+        strcatW(full_path, dest);
+
+        tmp_ret = SetupAddToDiskSpaceListW(diskspace, full_path, filesize, operation, 0, 0);
+        HeapFree(GetProcessHeap(), 0, full_path);
+        if (!tmp_ret) goto end;
+    }
+    while (SetupFindNextLine(&context, &context));
+
+    ret = TRUE;
+
+end:
+    HeapFree(GetProcessHeap(), 0, dest_dir);
+    return ret;
+}
+
+/***********************************************************************
+ *      SetupAddInstallSectionToDiskSpaceListA  (SETUPAPI.@)
+ */
+BOOL WINAPI SetupAddSectionToDiskSpaceListA(HDSKSPC diskspace, HINF hinf, HINF hlist,
+                                            PCSTR section, UINT operation, PVOID reserved1,
+                                            UINT reserved2)
+{
+    LPWSTR sectionW = NULL;
+    DWORD len;
+    BOOL ret;
+
+    if (section)
+    {
+        len = MultiByteToWideChar(CP_ACP, 0, section, -1, NULL, 0);
+
+        sectionW = HeapAlloc(GetProcessHeap(), 0, len * sizeof(WCHAR));
+        if (!sectionW)
+        {
+            SetLastError(ERROR_NOT_ENOUGH_MEMORY);
+            return FALSE;
+        }
+
+        MultiByteToWideChar(CP_ACP, 0, section, -1, sectionW, len);
+    }
+
+    ret = SetupAddSectionToDiskSpaceListW(diskspace, hinf, hlist, sectionW, operation,
+                                          reserved1, reserved2);
+    if (sectionW) HeapFree(GetProcessHeap(), 0, sectionW);
+    return ret;
+}
+
+/***********************************************************************
+ *      SetupAddInstallSectionToDiskSpaceListW  (SETUPAPI.@)
  */
 BOOL WINAPI SetupAddInstallSectionToDiskSpaceListA(HDSKSPC DiskSpace, 
                         HINF InfHandle, HINF LayoutInfHandle, 
diff --git a/dlls/setupapi/queue.c b/dlls/setupapi/queue.c
index c065c6e..b50e8cf 100644
--- a/dlls/setupapi/queue.c
+++ b/dlls/setupapi/queue.c
@@ -336,7 +336,7 @@ static void get_source_info( HINF hinf, const WCHAR *src_file, SP_FILE_COPY_PARA
  *
  * Retrieve the destination dir for a given section.
  */
-static WCHAR *get_destination_dir( HINF hinf, const WCHAR *section )
+WCHAR *get_destination_dir( HINF hinf, const WCHAR *section )
 {
     static const WCHAR Dest[] = {'D','e','s','t','i','n','a','t','i','o','n','D','i','r','s',0};
     static const WCHAR Def[]  = {'D','e','f','a','u','l','t','D','e','s','t','D','i','r',0};
diff --git a/dlls/setupapi/setupapi.spec b/dlls/setupapi/setupapi.spec
index 1d18870..dd41e16 100644
--- a/dlls/setupapi/setupapi.spec
+++ b/dlls/setupapi/setupapi.spec
@@ -244,8 +244,8 @@
 @ stub SetArrayToMultiSzValue
 @ stdcall SetupAddInstallSectionToDiskSpaceListA(long long long str ptr long)
 @ stub SetupAddInstallSectionToDiskSpaceListW
-@ stub SetupAddSectionToDiskSpaceListA
-@ stub SetupAddSectionToDiskSpaceListW
+@ stdcall SetupAddSectionToDiskSpaceListA(long long long str long ptr long)
+@ stdcall SetupAddSectionToDiskSpaceListW(long long long wstr long ptr long)
 @ stdcall SetupAddToDiskSpaceListA(long str int64 long ptr long)
 @ stdcall SetupAddToDiskSpaceListW(long wstr int64 long ptr long)
 @ stdcall SetupAddToSourceListA(long str)
diff --git a/dlls/setupapi/setupapi_private.h b/dlls/setupapi/setupapi_private.h
index f4685ab..4df0623 100644
--- a/dlls/setupapi/setupapi_private.h
+++ b/dlls/setupapi/setupapi_private.h
@@ -92,6 +92,8 @@ extern const WCHAR *DIRID_get_string( int dirid ) DECLSPEC_HIDDEN;
 extern const WCHAR *PARSER_get_inf_filename( HINF hinf ) DECLSPEC_HIDDEN;
 extern WCHAR *PARSER_get_dest_dir( INFCONTEXT *context ) DECLSPEC_HIDDEN;
 
+extern WCHAR *get_destination_dir( HINF hinf, const WCHAR *section );
+
 /* support for Ascii queue callback functions */
 
 struct callback_WtoA_context
diff --git a/dlls/setupapi/tests/diskspace.c b/dlls/setupapi/tests/diskspace.c
index 793d3a1..f40506c 100644
--- a/dlls/setupapi/tests/diskspace.c
+++ b/dlls/setupapi/tests/diskspace.c
@@ -32,6 +32,8 @@
 
 static BOOL is_win9x;
 
+#define STD_HEADER "[Version]\r\nSignature=\"$CHICAGO$\"\r\n"
+
 static inline const char* debugstr_longlong(ULONGLONG ll)
 {
     static char string[17];
@@ -42,6 +44,18 @@ static inline const char* debugstr_longlong(ULONGLONG ll)
     return string;
 }
 
+/* create a new file with specified contents and open it */
+static HINF inf_open_file_content(const char * tmpfilename, const char *data, UINT *err_line)
+{
+    DWORD res;
+    HANDLE handle = CreateFileA(tmpfilename, GENERIC_READ|GENERIC_WRITE,
+                                FILE_SHARE_READ|FILE_SHARE_WRITE, NULL, CREATE_ALWAYS, 0, 0);
+    if (handle == INVALID_HANDLE_VALUE) return 0;
+    if (!WriteFile( handle, data, strlen(data), &res, NULL )) trace( "write error\n" );
+    CloseHandle( handle );
+    return SetupOpenInfFileA( tmpfilename, 0, INF_STYLE_WIN4, err_line );
+}
+
 static void test_SetupCreateDiskSpaceListA(void)
 {
     HDSKSPC ret;
@@ -753,6 +767,184 @@ static void test_SetupQueryDrivesInDiskSpaceListA(void)
     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");
+
+    ok(SetupDestroyDiskSpaceList(handle),
+       "Expected SetupDestroyDiskSpaceList to succeed\n");
+}
+
+struct device_usage
+{
+    const char *dev;
+    LONGLONG usage;
+};
+
+struct section
+{
+    const char *name;
+    UINT fileop;
+    BOOL result;
+    DWORD error_code;
+};
+
+static const struct
+{
+    const char *data;
+    struct section sections[2];
+    const char *devices;
+    int device_length;
+    struct device_usage usage[2];
+}
+section_test[] =
+{
+    /* 0 */
+    {STD_HEADER "[a]\ntest,,,\n[SourceDisksFiles]\ntest=1,,4096\r\n",
+     {{"a", FILEOP_COPY, TRUE, 0}, {NULL, 0, TRUE, 0}},
+     "c:\00", sizeof("c:\00"), {{"c:", 4096}, {NULL, 0}}},
+    /* 1 */
+    {STD_HEADER "[a]\ntest,,,\n[SourceDisksFiles]\ntest=1,,4096\r\n",
+     {{"a", FILEOP_DELETE, TRUE, 0}, {NULL, 0, TRUE, 0}},
+     "c:\00", sizeof("c:\00"), {{"c:", 0}, {NULL, 0}}},
+    /* 2 */
+    {STD_HEADER "[a]\ntest,,,\n\r\n",
+     {{"a", FILEOP_COPY, FALSE, ERROR_LINE_NOT_FOUND}, {NULL, 0, TRUE, 0}},
+     "", sizeof(""), {{NULL, 0}, {NULL, 0}}},
+    /* 3 */
+    {STD_HEADER "[a]\ntest,,,\n[SourceDisksFiles]\ntest=1,,4096\n[DestinationDirs]\nDefaultDestDir=-1,F:\\test\r\n",
+     {{"a", FILEOP_COPY, TRUE, 0}, {NULL, 0, TRUE, 0}},
+     "f:\00", sizeof("f:\00"), {{"f:", 4096}, {NULL, 0}}},
+    /* 4 */
+    {STD_HEADER "[a]\ntest,test2,,\n[SourceDisksFiles]\ntest2=1,,4096\r\n",
+     {{"a", FILEOP_COPY, TRUE, 0}, {NULL, 0, TRUE, 0}},
+     "c:\00", sizeof("c:\00"), {{"c:", 4096}, {NULL, 0}}},
+    /* 5 */
+    {STD_HEADER "[a]\ntest,,,\n[SourceDisksFiles]\ntest=1,,4096\r\n",
+     {{"b", FILEOP_COPY, FALSE, ERROR_SECTION_NOT_FOUND}, {NULL, 0, TRUE, 0}},
+     "", sizeof(""), {{NULL, 0}, {NULL, 0}}},
+    /* 6 */
+    {STD_HEADER "[a]\ntest,,,\n[b]\ntest,,,\n[SourceDisksFiles]\ntest=1,,4096\r\n",
+     {{"a", FILEOP_COPY, TRUE, 0}, {"b", FILEOP_COPY, TRUE, 0}},
+     "c:\00", sizeof("c:\00"), {{"c:", 4096}, {NULL, 0}}},
+    /* 7 */
+    {STD_HEADER "[a]\ntest,,,\n[b]\ntest,,,\n[SourceDisksFiles]\ntest=1,,4096\n[DestinationDirs]\nb=-1,F:\\test\r\n",
+     {{"a", FILEOP_COPY, TRUE, 0}, {"b", FILEOP_COPY, TRUE, 0}},
+     "c:\00f:\00", sizeof("c:\00f:\00"), {{"c:", 4096}, {"f:", 4096}}},
+    /* 8 */
+    {STD_HEADER "[a]\ntest,test1,,\n[b]\ntest,test2,,\n[SourceDisksFiles]\ntest1=1,,4096\ntest2=1,,8192\r\n",
+     {{"a", FILEOP_COPY, TRUE, 0}, {"b", FILEOP_COPY, TRUE, 0}},
+     "c:\00", sizeof("c:\00"), {{"c:", 8192}, {NULL, 0}}},
+    /* 9 */
+    {STD_HEADER "[a]\ntest1,test,,\n[b]\ntest2,test,,\n[SourceDisksFiles]\ntest=1,,4096\r\n",
+     {{"a", FILEOP_COPY, TRUE, 0}, {"b", FILEOP_COPY, TRUE, 0}},
+     "c:\00", sizeof("c:\00"), {{"c:", 8192}, {NULL, 0}}},
+};
+
+static void test_SetupAddSectionToDiskSpaceListA(void)
+{
+    char tmp[MAX_PATH];
+    char tmpfilename[MAX_PATH];
+    char buffer[MAX_PATH];
+    HDSKSPC diskspace;
+    UINT err_line;
+    LONGLONG space;
+    BOOL ret;
+    int i, j;
+    HINF inf;
+
+    if (!GetTempPathA(MAX_PATH, tmp))
+    {
+        win_skip("GetTempPath failed with error %d\n", GetLastError());
+        return;
+    }
+
+    if (!GetTempFileNameA(tmp, "inftest", 0, tmpfilename))
+    {
+        win_skip("GetTempFileNameA failed with error %d\n", GetLastError());
+        return;
+    }
+
+    inf = inf_open_file_content(tmpfilename, STD_HEADER "[a]\ntest,,,\n[SourceDisksFiles]\ntest=1,,4096\r\n", &err_line);
+    ok(!!inf, "Failed to open inf file (%d, line %d)\n", GetLastError(), err_line);
+
+    diskspace = SetupCreateDiskSpaceListA(NULL, 0, SPDSL_IGNORE_DISK);
+    ok(diskspace != NULL,"Expected SetupCreateDiskSpaceListA to return a valid handle\n");
+
+    ret = SetupAddSectionToDiskSpaceListA(diskspace, NULL, NULL, "a", FILEOP_COPY, 0, 0);
+    ok(!ret, "Expected SetupAddSectionToDiskSpaceListA to fail\n");
+    ok(GetLastError() == ERROR_SECTION_NOT_FOUND, "Expected ERROR_SECTION_NOT_FOUND as error, got %u\n",
+       GetLastError());
+
+    ret = SetupAddSectionToDiskSpaceListA(NULL, inf, NULL, "a", FILEOP_COPY, 0, 0);
+    ok(!ret, "Expected SetupAddSectionToDiskSpaceListA to fail\n");
+    ok(GetLastError() == ERROR_INVALID_HANDLE, "Expected ERROR_INVALID_HANDLE as error, got %u\n",
+       GetLastError());
+
+    ret = SetupAddSectionToDiskSpaceListA(NULL, inf, NULL, "b", FILEOP_COPY, 0, 0);
+    ok(!ret, "Expected SetupAddSectionToDiskSpaceListA to fail\n");
+    ok(GetLastError() == ERROR_INVALID_HANDLE, "Expected ERROR_INVALID_HANDLE as error, got %u\n",
+       GetLastError());
+
+    ret = SetupAddSectionToDiskSpaceListA(diskspace, inf, NULL, "a", 0, 0, 0);
+    ok(ret, "Expected SetupAddSectionToDiskSpaceListA to succeed (%u)\n", GetLastError());
+
+    ok(SetupDestroyDiskSpaceList(diskspace),
+       "Expected SetupDestroyDiskSpaceList to succeed\n");
+
+    for (i = 0; i < sizeof(section_test) / sizeof(section_test[0]); i++)
+    {
+        err_line = 0;
+
+        inf = inf_open_file_content(tmpfilename, section_test[i].data, &err_line);
+        ok(!!inf, "test %d: Failed to open inf file (%d, line %d)\n", i, GetLastError(), err_line);
+        if (!inf) continue;
+
+        diskspace = SetupCreateDiskSpaceListA(NULL, 0, SPDSL_IGNORE_DISK);
+        ok(diskspace != NULL, "Expected SetupCreateDiskSpaceListA to return a valid handle\n");
+
+        for (j = 0; j < 2; j++)
+        {
+            const struct section *section = &section_test[i].sections[j];
+            if (!section->name)
+                continue;
+
+            SetLastError(0xdeadbeef);
+            ret = SetupAddSectionToDiskSpaceListA(diskspace, inf, NULL, section->name, section->fileop, 0, 0);
+            if (section->result)
+                ok(ret, "test %d: Expected adding section %d to succeed (%u)\n", i, j, GetLastError());
+            else
+            {
+                ok(!ret, "test %d: Expected adding section %d to fail\n", i, j);
+                ok(GetLastError() == section->error_code, "test %d: Expected %u as error, got %u\n",
+                   i, section->error_code, GetLastError());
+            }
+        }
+
+        memset(buffer, 0x0, sizeof(buffer));
+        ret = SetupQueryDrivesInDiskSpaceListA(diskspace, buffer, sizeof(buffer), NULL);
+        ok(ret, "test %d: Expected SetupQueryDrivesInDiskSpaceListA to succeed (%u)\n", i, GetLastError());
+        ok(!memcmp(section_test[i].devices, buffer, section_test[i].device_length),
+           "test %d: Device list (%s) does not match\n", i, buffer);
+
+        for (j = 0; j < 2; j++)
+        {
+            const struct device_usage *usage = &section_test[i].usage[j];
+            if (!usage->dev)
+                continue;
+
+            space = 0;
+            ret = SetupQuerySpaceRequiredOnDriveA(diskspace, usage->dev, &space, NULL, 0);
+            ok(ret, "test %d: Expected SetupQuerySpaceRequiredOnDriveA to succeed for device %s (%u)\n",
+               i, usage->dev, GetLastError());
+            ok(space == usage->usage, "test %d: Expected size %u for device %s, got %u\n",
+               i, (DWORD)usage->usage, usage->dev, (DWORD)space);
+        }
+
+        ok(SetupDestroyDiskSpaceList(diskspace),
+           "Expected SetupDestroyDiskSpaceList to succeed\n");
+
+        SetupCloseInfFile(inf);
+    }
+
+    DeleteFileA(tmpfilename);
 }
 
 START_TEST(diskspace)
@@ -767,4 +959,5 @@ START_TEST(diskspace)
     test_SetupQuerySpaceRequiredOnDriveW();
     test_SetupAddToDiskSpaceListA();
     test_SetupQueryDrivesInDiskSpaceListA();
+    test_SetupAddSectionToDiskSpaceListA();
 }
-- 
1.9.1

