do _Streams/_Storages enumeration on the internal data structures

4 files changed, 105 insertions, 124 deletions

diff --git a/libmsi/database.c b/libmsi/database.c
index 38ea337..1ba8bf6 100644
--- a/libmsi/database.c
+++ b/libmsi/database.c
@@ -440,6 +440,56 @@ static void cache_infile_structure( LibmsiDatabase *db )
     IEnumSTATSTG_Release(stgenum);
 }
 
+unsigned msi_enum_db_streams(LibmsiDatabase *db,
+                             unsigned (*fn)(const WCHAR *, IStream *, void *),
+                             void *opaque)
+{
+    unsigned r;
+    LibmsiStream *stream, *stream2;
+
+    LIST_FOR_EACH_ENTRY_SAFE( stream, stream2, &db->streams, LibmsiStream, entry )
+    {
+        IStream *stm;
+
+        if (stream->stg != db->infile) continue;
+
+        stm = stream->stm;
+        IStream_AddRef(stm);
+        r = fn( stream->name, stm, opaque);
+        IStream_Release(stm);
+
+        if (r) {
+            return r;
+        }
+    }
+
+    return LIBMSI_RESULT_SUCCESS;
+}
+
+unsigned msi_enum_db_storages(LibmsiDatabase *db,
+                              unsigned (*fn)(const WCHAR *, IStorage *, void *),
+                              void *opaque)
+{
+    unsigned r;
+    LibmsiStorage *storage, *storage2;
+
+    LIST_FOR_EACH_ENTRY_SAFE( storage, storage2, &db->storages, LibmsiStorage, entry )
+    {
+        IStorage *stg;
+
+        stg = storage->stg;
+        IStorage_AddRef(stg);
+        r = fn( storage->name, stg, opaque);
+        IStorage_Release(stg);
+
+        if (r) {
+            return r;
+        }
+    }
+
+    return LIBMSI_RESULT_SUCCESS;
+}
+
 unsigned msi_clone_open_stream( LibmsiDatabase *db, IStorage *stg, const WCHAR *name, IStream **stm )
 {
     IStream *stream;
diff --git a/libmsi/msipriv.h b/libmsi/msipriv.h
index 8737403..653a694 100644
--- a/libmsi/msipriv.h
+++ b/libmsi/msipriv.h
@@ -359,9 +359,11 @@ unsigned msi_create_stream( LibmsiDatabase *db, const WCHAR *stname, IStream *st
 extern unsigned msi_get_raw_stream( LibmsiDatabase *, const WCHAR *, IStream **);
 extern unsigned msi_clone_open_stream( LibmsiDatabase *, IStorage *, const WCHAR *, IStream ** );
 void msi_destroy_stream( LibmsiDatabase *, const WCHAR * );
+extern unsigned msi_enum_db_streams(LibmsiDatabase *, unsigned (*fn)(const WCHAR *, IStream *, void *), void *);
 unsigned msi_create_storage( LibmsiDatabase *db, const WCHAR *stname, IStream *stm );
 unsigned msi_open_storage( LibmsiDatabase *db, const WCHAR *stname );
 void msi_destroy_storage( LibmsiDatabase *db, const WCHAR *stname );
+extern unsigned msi_enum_db_storages(LibmsiDatabase *, unsigned (*fn)(const WCHAR *, IStorage *, void *), void *);
 extern unsigned _libmsi_database_open_query(LibmsiDatabase *, const WCHAR *, LibmsiQuery **);
 extern unsigned _libmsi_query_open( LibmsiDatabase *, LibmsiQuery **, const WCHAR *, ... );
 typedef unsigned (*record_func)( LibmsiRecord *, void *);
diff --git a/libmsi/storages.c b/libmsi/storages.c
index 8d88a17..6da08ed 100644
--- a/libmsi/storages.c
+++ b/libmsi/storages.c
@@ -319,74 +319,39 @@ static const LibmsiViewOps storages_ops =
     NULL,
 };
 
-static int add_storages_to_table(LibmsiStorageView *sv)
+static unsigned add_storage_to_table(const WCHAR *name, IStorage *stg, void *opaque)
 {
-    STORAGE *storage = NULL;
-    IEnumSTATSTG *stgenum = NULL;
-    STATSTG stat;
-    HRESULT hr;
-    unsigned r;
-    unsigned count = 0, size;
-
-    hr = IStorage_EnumElements(sv->db->infile, 0, NULL, 0, &stgenum);
-    if (FAILED(hr))
-        return -1;
+    LibmsiStorageView *sv = (LibmsiStorageView *)opaque;
+    STORAGE *storage;
 
-    sv->max_storages = 1;
-    sv->storages = msi_alloc(sizeof(STORAGE *));
-    if (!sv->storages)
-        return -1;
+    storage = create_storage(sv, name);
+    if (!storage)
+        return LIBMSI_RESULT_NOT_ENOUGH_MEMORY;
 
-    while (true)
+    if (!storages_set_table_size(sv, ++sv->num_rows))
     {
-        size = 0;
-        hr = IEnumSTATSTG_Next(stgenum, 1, &stat, &size);
-        if (FAILED(hr) || !size)
-            break;
-
-        if (stat.type != STGTY_STORAGE)
-        {
-            CoTaskMemFree(stat.pwcsName);
-            continue;
-        }
-
-        TRACE("enumerated storage %s\n", debugstr_w(stat.pwcsName));
-
-        r = msi_open_storage(sv->db, stat.pwcsName);
-        if (r != LIBMSI_RESULT_SUCCESS)
-        {
-            count = -1;
-            CoTaskMemFree(stat.pwcsName);
-            break;
-        }
-
-        storage = create_storage(sv, stat.pwcsName);
-        if (!storage)
-        {
-            count = -1;
-            CoTaskMemFree(stat.pwcsName);
-            break;
-        }
-
-        CoTaskMemFree(stat.pwcsName);
+        msi_free(storage);
+        return LIBMSI_RESULT_NOT_ENOUGH_MEMORY;
+    }
 
-        if (!storages_set_table_size(sv, ++count))
-        {
-            count = -1;
-            break;
-        }
+    sv->storages[sv->num_rows - 1] = storage;
+    return LIBMSI_RESULT_SUCCESS;
+}
 
-        sv->storages[count - 1] = storage;
-    }
+static unsigned add_storages_to_table(LibmsiStorageView *sv)
+{
+    sv->max_storages = 1;
+    sv->storages = msi_alloc_zero(sizeof(STORAGE *));
+    if (!sv->storages)
+        return LIBMSI_RESULT_NOT_ENOUGH_MEMORY;
 
-    IEnumSTATSTG_Release(stgenum);
-    return count;
+    return msi_enum_db_storages(sv->db, add_storage_to_table, sv);
 }
 
 unsigned storages_view_create(LibmsiDatabase *db, LibmsiView **view)
 {
     LibmsiStorageView *sv;
-    int rows;
+    unsigned r;
 
     TRACE("(%p, %p)\n", db, view);
 
@@ -397,13 +362,12 @@ unsigned storages_view_create(LibmsiDatabase *db, LibmsiView **view)
     sv->view.ops = &storages_ops;
     sv->db = db;
 
-    rows = add_storages_to_table(sv);
-    if (rows < 0)
+    r = add_storages_to_table(sv);
+    if (r)
     {
         msi_free( sv );
-        return LIBMSI_RESULT_FUNCTION_FAILED;
+        return r;
     }
-    sv->num_rows = rows;
 
     *view = (LibmsiView *)sv;
 
diff --git a/libmsi/streams.c b/libmsi/streams.c
index 3a8ad62..2b379c9 100644
--- a/libmsi/streams.c
+++ b/libmsi/streams.c
@@ -84,6 +84,9 @@ static STREAM *create_stream(LibmsiStreamsView *sv, const WCHAR *name, bool enco
 
     stream->str_index = _libmsi_add_string(sv->db->strings, name, -1, 1, StringNonPersistent);
     stream->stream = stm;
+    if (stream->stream)
+        IStream_AddRef(stm);
+
     return stream;
 }
 
@@ -167,7 +170,9 @@ static unsigned streams_view_set_row(LibmsiView *view, unsigned row, LibmsiRecor
 
 done:
     if (r != LIBMSI_RESULT_SUCCESS)
+    {
         msi_free(stream);
+    }
 
     msi_free(name);
 
@@ -353,78 +358,39 @@ static const LibmsiViewOps streams_ops =
     NULL,
 };
 
-static int add_streams_to_table(LibmsiStreamsView *sv)
+static unsigned add_stream_to_table(const WCHAR *name, IStream *stm, void *opaque)
 {
-    IEnumSTATSTG *stgenum = NULL;
-    STATSTG stat;
-    STREAM *stream = NULL;
-    HRESULT hr;
-    unsigned r, count = 0, size;
-
-    hr = IStorage_EnumElements(sv->db->infile, 0, NULL, 0, &stgenum);
-    if (FAILED(hr))
-        return -1;
+    LibmsiStreamsView *sv = (LibmsiStreamsView *)opaque;
+    STREAM *stream;
 
-    sv->max_streams = 1;
-    sv->streams = msi_alloc_zero(sizeof(STREAM *));
-    if (!sv->streams)
-        return -1;
+    stream = create_stream(sv, name, true, stm);
+    if (!stream)
+        return LIBMSI_RESULT_NOT_ENOUGH_MEMORY;
 
-    while (true)
+    if (!streams_set_table_size(sv, ++sv->num_rows))
     {
-        size = 0;
-        hr = IEnumSTATSTG_Next(stgenum, 1, &stat, &size);
-        if (FAILED(hr) || !size)
-            break;
-
-        if (stat.type != STGTY_STREAM)
-        {
-            CoTaskMemFree(stat.pwcsName);
-            continue;
-        }
-
-        /* table streams are not in the _Streams table */
-        if (*stat.pwcsName == 0x4840)
-        {
-            CoTaskMemFree(stat.pwcsName);
-            continue;
-        }
-
-        stream = create_stream(sv, stat.pwcsName, true, NULL);
-        if (!stream)
-        {
-            count = -1;
-            CoTaskMemFree(stat.pwcsName);
-            break;
-        }
-
-        r = msi_get_raw_stream(sv->db, stat.pwcsName, &stream->stream);
-        CoTaskMemFree(stat.pwcsName);
-
-        if (r != LIBMSI_RESULT_SUCCESS)
-        {
-            WARN("unable to get stream %u\n", r);
-            count = -1;
-            break;
-        }
+        msi_free(stream);
+        return LIBMSI_RESULT_NOT_ENOUGH_MEMORY;
+    }
 
-        if (!streams_set_table_size(sv, ++count))
-        {
-            count = -1;
-            break;
-        }
+    sv->streams[sv->num_rows - 1] = stream;
+    return LIBMSI_RESULT_SUCCESS;
+}
 
-        sv->streams[count - 1] = stream;
-    }
+static unsigned add_streams_to_table(LibmsiStreamsView *sv)
+{
+    sv->max_streams = 1;
+    sv->streams = msi_alloc_zero(sizeof(STREAM *));
+    if (!sv->streams)
+        return LIBMSI_RESULT_NOT_ENOUGH_MEMORY;
 
-    IEnumSTATSTG_Release(stgenum);
-    return count;
+    return msi_enum_db_streams(sv->db, add_stream_to_table, sv);
 }
 
 unsigned streams_view_create(LibmsiDatabase *db, LibmsiView **view)
 {
     LibmsiStreamsView *sv;
-    int rows;
+    unsigned r;
 
     TRACE("(%p, %p)\n", db, view);
 
@@ -434,13 +400,12 @@ unsigned streams_view_create(LibmsiDatabase *db, LibmsiView **view)
 
     sv->view.ops = &streams_ops;
     sv->db = db;
-    rows = add_streams_to_table(sv);
-    if (rows < 0)
+    r = add_streams_to_table(sv);
+    if (r)
     {
         msi_free( sv );
-        return LIBMSI_RESULT_FUNCTION_FAILED;
+        return r;
     }
-    sv->num_rows = rows;
 
     *view = (LibmsiView *)sv;