Tweak reading of concat runs

[wimlib] / src / win32_capture.c

diff --git a/src/win32_capture.c b/src/win32_capture.c
index 308c0677ccb90b02c0a3b6f7eb31a51e7bc9558a..eb8bd7f9891f68ac038c964acc6a0746918d5446 100644 (file)
--- a/src/win32_capture.c
+++ b/src/win32_capture.c
@@ -58,11 +58,13 @@ int
 read_win32_file_prefix(const struct wim_lookup_table_entry *lte,
                       u64 size,
                       consume_data_callback_t cb,
+                      u32 in_chunk_size,
                       void *ctx_or_buf,
                       int _ignored_flags)
 {
        int ret = 0;
        void *out_buf;
+       bool out_buf_malloced;
        u64 bytes_remaining;
 
        HANDLE hFile = win32_open_existing_file(lte->file_on_disk,
@@ -73,16 +75,27 @@ read_win32_file_prefix(const struct wim_lookup_table_entry *lte,
                return WIMLIB_ERR_OPEN;
        }
 
-       if (cb)
-               out_buf = alloca(WIM_CHUNK_SIZE);
-       else
+       out_buf_malloced = false;
+       if (cb) {
+               if (in_chunk_size <= STACK_MAX) {
+                       out_buf = alloca(in_chunk_size);
+               } else {
+                       out_buf = MALLOC(in_chunk_size);
+                       if (out_buf == NULL) {
+                               ret = WIMLIB_ERR_NOMEM;
+                               goto out_close_handle;
+                       }
+                       out_buf_malloced = true;
+               }
+       } else {
                out_buf = ctx_or_buf;
+       }
 
        bytes_remaining = size;
        while (bytes_remaining) {
                DWORD bytesToRead, bytesRead;
 
-               bytesToRead = min(WIM_CHUNK_SIZE, bytes_remaining);
+               bytesToRead = min(in_chunk_size, bytes_remaining);
                if (!ReadFile(hFile, out_buf, bytesToRead, &bytesRead, NULL) ||
                    bytesRead != bytesToRead)
                {
@@ -101,6 +114,9 @@ read_win32_file_prefix(const struct wim_lookup_table_entry *lte,
                        out_buf += bytesRead;
                }
        }
+       if (out_buf_malloced)
+               FREE(out_buf);
+out_close_handle:
        CloseHandle(hFile);
        return ret;
 }
@@ -112,6 +128,7 @@ struct win32_encrypted_read_ctx {
        void *buf;
        size_t buf_filled;
        u64 bytes_remaining;
+       u32 in_chunk_size;
 };
 
 static DWORD WINAPI
@@ -119,6 +136,7 @@ win32_encrypted_export_cb(unsigned char *_data, void *_ctx, unsigned long len)
 {
        const void *data = _data;
        struct win32_encrypted_read_ctx *ctx = _ctx;
+       u32 in_chunk_size = ctx->in_chunk_size;
        int ret;
 
        DEBUG("len = %lu", len);
@@ -130,7 +148,7 @@ win32_encrypted_export_cb(unsigned char *_data, void *_ctx, unsigned long len)
                                             len);
                while (bytes_to_buffer) {
                        size_t bytes_to_copy_to_buf =
-                               min(bytes_to_buffer, WIM_CHUNK_SIZE - ctx->buf_filled);
+                               min(bytes_to_buffer, in_chunk_size - ctx->buf_filled);
 
                        memcpy(ctx->buf + ctx->buf_filled, data,
                               bytes_to_copy_to_buf);
@@ -138,7 +156,7 @@ win32_encrypted_export_cb(unsigned char *_data, void *_ctx, unsigned long len)
                        data += bytes_to_copy_to_buf;
                        bytes_to_buffer -= bytes_to_copy_to_buf;
 
-                       if (ctx->buf_filled == WIM_CHUNK_SIZE ||
+                       if (ctx->buf_filled == in_chunk_size ||
                            ctx->buf_filled == ctx->bytes_remaining)
                        {
                                ret = (*ctx->read_prefix_cb)(ctx->buf,
@@ -168,6 +186,7 @@ int
 read_win32_encrypted_file_prefix(const struct wim_lookup_table_entry *lte,
                                 u64 size,
                                 consume_data_callback_t cb,
+                                u32 in_chunk_size,
                                 void *ctx_or_buf,
                                 int _ignored_flags)
 {
@@ -183,7 +202,7 @@ read_win32_encrypted_file_prefix(const struct wim_lookup_table_entry *lte,
        export_ctx.read_prefix_ctx_or_buf = ctx_or_buf;
        export_ctx.wimlib_err_code = 0;
        if (cb) {
-               export_ctx.buf = MALLOC(WIM_CHUNK_SIZE);
+               export_ctx.buf = MALLOC(in_chunk_size);
                if (!export_ctx.buf)
                        return WIMLIB_ERR_NOMEM;
        } else {
@@ -952,10 +971,10 @@ win32_capture_stream(const wchar_t *path,
                ret = win32_get_encrypted_file_size(path, &encrypted_size);
                if (ret)
                        goto out_free_spath;
-               lte->resource_entry.original_size = encrypted_size;
+               lte->size = encrypted_size;
        } else {
                lte->resource_location = RESOURCE_IN_FILE_ON_DISK;
-               lte->resource_entry.original_size = (u64)dat->StreamSize.QuadPart;
+               lte->size = (u64)dat->StreamSize.QuadPart;
        }
 
        u32 stream_id;