X-Git-Url: https://wimlib.net/git/?p=wimlib;a=blobdiff_plain;f=src%2Fresource.c;h=f54ea259bb5e91c79fd52e8b2cd37f9874b12897;hp=69a531fbce354309867ca3bc2cc42689517ab983;hb=4f8059f2d0a74a9922128b162d9c9343b305999c;hpb=ebd6c0ec0ff47ac18af4ef918fd78fb8d9f19540 diff --git a/src/resource.c b/src/resource.c index 69a531fb..f54ea259 100644 --- a/src/resource.c +++ b/src/resource.c @@ -26,53 +26,40 @@ #include "dentry.h" #include "lookup_table.h" #include "buffer_io.h" -#include "lzx.h" -#include "xpress.h" #include "sha1.h" +#ifdef __WIN32__ +# include "win32.h" +#endif + #include #include #include #include +#include -#ifdef WITH_NTFS_3G -#include -#include -#include -#include +#ifdef HAVE_ALLOCA_H +# include #endif /* - * Reads all or part of a compressed resource into an in-memory buffer. - * - * @fp: The FILE* for the WIM file. - * @resource_compressed_size: The compressed size of the resource. - * @resource_uncompressed_size: The uncompressed size of the resource. - * @resource_offset: The offset of the start of the resource from - * the start of the stream @fp. - * @resource_ctype: The compression type of the resource. - * @len: The number of bytes of uncompressed data to read from - * the resource. - * @offset: The offset of the bytes to read within the uncompressed - * resource. - * @contents_len: An array into which the uncompressed data is written. - * It must be at least @len bytes long. + * Reads all or part of a compressed WIM resource. * * Returns zero on success, nonzero on failure. */ -static int read_compressed_resource(FILE *fp, u64 resource_compressed_size, - u64 resource_uncompressed_size, - u64 resource_offset, int resource_ctype, - u64 len, u64 offset, u8 contents_ret[]) +static int +read_compressed_resource(int in_fd, + u64 resource_compressed_size, + u64 resource_uncompressed_size, + u64 resource_offset, + int resource_ctype, + u64 len, + u64 offset, + consume_data_callback_t cb, + void *ctx_or_buf) { + int ret; - DEBUG2("comp size = %"PRIu64", uncomp size = %"PRIu64", " - "res offset = %"PRIu64"", - resource_compressed_size, - resource_uncompressed_size, - resource_offset); - DEBUG2("resource_ctype = %s, len = %"PRIu64", offset = %"PRIu64"", - wimlib_get_compression_type_string(resource_ctype), len, offset); /* Trivial case */ if (len == 0) return 0; @@ -80,9 +67,9 @@ static int read_compressed_resource(FILE *fp, u64 resource_compressed_size, int (*decompress)(const void *, unsigned, void *, unsigned); /* Set the appropriate decompress function. */ if (resource_ctype == WIMLIB_COMPRESSION_TYPE_LZX) - decompress = lzx_decompress; + decompress = wimlib_lzx_decompress; else - decompress = xpress_decompress; + decompress = wimlib_xpress_decompress; /* The structure of a compressed resource consists of a table of chunk * offsets followed by the chunks themselves. Each chunk consists of @@ -96,7 +83,8 @@ static int read_compressed_resource(FILE *fp, u64 resource_compressed_size, * follows the chunk table and therefore must have an offset of 0. */ - /* Calculate how many chunks the resource conists of in its entirety. */ + /* Calculate how many chunks the resource consists of in its entirety. + * */ u64 num_chunks = (resource_uncompressed_size + WIM_CHUNK_SIZE - 1) / WIM_CHUNK_SIZE; /* As mentioned, the first chunk has no entry in the chunk table. */ @@ -123,9 +111,22 @@ static int read_compressed_resource(FILE *fp, u64 resource_compressed_size, if (end_chunk != num_chunks - 1) num_needed_chunks++; - /* Declare the chunk table. It will only contain offsets for the chunks - * that are actually needed for this read. */ - u64 chunk_offsets[num_needed_chunks]; + /* Allocate the chunk table. It will only contain offsets for the + * chunks that are actually needed for this read. */ + u64 *chunk_offsets; + bool chunk_offsets_malloced; + if (num_needed_chunks < 1000) { + chunk_offsets = alloca(num_needed_chunks * sizeof(u64)); + chunk_offsets_malloced = false; + } else { + chunk_offsets = malloc(num_needed_chunks * sizeof(u64)); + if (!chunk_offsets) { + ERROR("Failed to allocate chunk table " + "with %"PRIu64" entries", num_needed_chunks); + return WIMLIB_ERR_NOMEM; + } + chunk_offsets_malloced = true; + } /* Set the implicit offset of the first chunk if it is included in the * needed chunks. @@ -161,20 +162,17 @@ static int read_compressed_resource(FILE *fp, u64 resource_compressed_size, /* Skip over unneeded chunk table entries. */ u64 file_offset_of_needed_chunk_entries = resource_offset + start_table_idx * chunk_entry_size; - if (fseeko(fp, file_offset_of_needed_chunk_entries, SEEK_SET) != 0) { - ERROR_WITH_ERRNO("Failed to seek to byte %"PRIu64" to read " - "chunk table of compressed resource", - file_offset_of_needed_chunk_entries); - return WIMLIB_ERR_READ; - } /* Number of bytes we need to read from the chunk table. */ size_t size = num_needed_chunk_entries * chunk_entry_size; - u8 chunk_tab_buf[size]; + /* Read the raw data into the end of the chunk_offsets array to + * avoid allocating another array. */ + void *chunk_tab_buf = (void*)&chunk_offsets[num_needed_chunks] - size; - if (fread(chunk_tab_buf, 1, size, fp) != size) - goto err; + if (full_pread(in_fd, chunk_tab_buf, size, + file_offset_of_needed_chunk_entries) != size) + goto read_error; /* Now fill in chunk_offsets from the entries we have read in * chunk_tab_buf. */ @@ -196,18 +194,16 @@ static int read_compressed_resource(FILE *fp, u64 resource_compressed_size, /* Done with the chunk table now. We must now seek to the first chunk * that is needed for the read. */ - u64 file_offset_of_first_needed_chunk = resource_offset + - chunk_table_size + chunk_offsets[0]; - if (fseeko(fp, file_offset_of_first_needed_chunk, SEEK_SET) != 0) { - ERROR_WITH_ERRNO("Failed to seek to byte %"PRIu64" to read " - "first chunk of compressed resource", - file_offset_of_first_needed_chunk); - return WIMLIB_ERR_READ; - } + u64 cur_read_offset = resource_offset + chunk_table_size + chunk_offsets[0]; /* Pointer to current position in the output buffer for uncompressed - * data. */ - u8 *out_p = (u8*)contents_ret; + * data. Alternatively, if using a callback function, we repeatedly + * fill a temporary buffer to feed data into the callback function. */ + u8 *out_p; + if (cb) + out_p = alloca(WIM_CHUNK_SIZE); + else + out_p = ctx_or_buf; /* Buffer for compressed data. While most compressed chunks will have a * size much less than WIM_CHUNK_SIZE, WIM_CHUNK_SIZE - 1 is the maximum @@ -215,15 +211,11 @@ static int read_compressed_resource(FILE *fp, u64 resource_compressed_size, * happen to compress to more than the uncompressed size (i.e. a * sequence of random bytes) are always stored uncompressed. But this seems * to be the case in M$'s WIM files, even though it is undocumented. */ - u8 compressed_buf[WIM_CHUNK_SIZE - 1]; - + void *compressed_buf = alloca(WIM_CHUNK_SIZE - 1); /* Decompress all the chunks. */ for (u64 i = start_chunk; i <= end_chunk; i++) { - DEBUG2("Chunk %"PRIu64" (start %"PRIu64", end %"PRIu64").", - i, start_chunk, end_chunk); - /* Calculate the sizes of the compressed chunk and of the * uncompressed chunk. */ unsigned compressed_chunk_size; @@ -258,11 +250,6 @@ static int read_compressed_resource(FILE *fp, u64 resource_compressed_size, uncompressed_chunk_size = WIM_CHUNK_SIZE; } - DEBUG2("compressed_chunk_size = %u, " - "uncompressed_chunk_size = %u", - compressed_chunk_size, uncompressed_chunk_size); - - /* Figure out how much of this chunk we actually need to read */ u64 start_offset; if (i == start_chunk) @@ -275,106 +262,93 @@ static int read_compressed_resource(FILE *fp, u64 resource_compressed_size, else end_offset = WIM_CHUNK_SIZE - 1; - u64 partial_chunk_size = end_offset + 1 - start_offset; - bool is_partial_chunk = (partial_chunk_size != - uncompressed_chunk_size); - - DEBUG2("start_offset = %u, end_offset = %u", start_offset, - end_offset); - DEBUG2("partial_chunk_size = %u", partial_chunk_size); + unsigned partial_chunk_size = end_offset + 1 - start_offset; + bool is_partial_chunk = (partial_chunk_size != uncompressed_chunk_size); /* This is undocumented, but chunks can be uncompressed. This * appears to always be the case when the compressed chunk size * is equal to the uncompressed chunk size. */ if (compressed_chunk_size == uncompressed_chunk_size) { - /* Probably an uncompressed chunk */ - - if (start_offset != 0) { - if (fseeko(fp, start_offset, SEEK_CUR) != 0) { - ERROR_WITH_ERRNO("Uncompressed partial " - "chunk fseek() error"); - return WIMLIB_ERR_READ; - } + /* Uncompressed chunk */ + if (full_pread(in_fd, + cb ? out_p + start_offset : out_p, + partial_chunk_size, + cur_read_offset + start_offset) != partial_chunk_size) + { + goto read_error; } - if (fread(out_p, 1, partial_chunk_size, fp) != - partial_chunk_size) - goto err; } else { /* Compressed chunk */ - int ret; /* Read the compressed data into compressed_buf. */ - if (fread(compressed_buf, 1, compressed_chunk_size, - fp) != compressed_chunk_size) - goto err; + if (full_pread(in_fd, + compressed_buf, + compressed_chunk_size, + cur_read_offset) != compressed_chunk_size) + { + goto read_error; + } - /* For partial chunks we must buffer the uncompressed - * data because we don't need all of it. */ - if (is_partial_chunk) { + /* For partial chunks and when writing directly to a + * buffer, we must buffer the uncompressed data because + * we don't need all of it. */ + if (is_partial_chunk && !cb) { u8 uncompressed_buf[uncompressed_chunk_size]; ret = decompress(compressed_buf, - compressed_chunk_size, - uncompressed_buf, - uncompressed_chunk_size); - if (ret != 0) - return WIMLIB_ERR_DECOMPRESSION; + compressed_chunk_size, + uncompressed_buf, + uncompressed_chunk_size); + if (ret) { + ret = WIMLIB_ERR_DECOMPRESSION; + goto out; + } memcpy(out_p, uncompressed_buf + start_offset, - partial_chunk_size); + partial_chunk_size); } else { ret = decompress(compressed_buf, - compressed_chunk_size, - out_p, - uncompressed_chunk_size); - if (ret != 0) - return WIMLIB_ERR_DECOMPRESSION; + compressed_chunk_size, + out_p, + uncompressed_chunk_size); + if (ret) { + ret = WIMLIB_ERR_DECOMPRESSION; + goto out; + } } } - - /* Advance the pointer into the uncompressed output data by the - * number of uncompressed bytes that were written. */ - out_p += partial_chunk_size; - } - - return 0; - -err: - if (feof(fp)) - ERROR("Unexpected EOF in compressed file resource"); - else - ERROR_WITH_ERRNO("Error reading compressed file resource"); - return WIMLIB_ERR_READ; -} - -/* - * Reads uncompressed data from an open file stream. - */ -int read_uncompressed_resource(FILE *fp, u64 offset, u64 len, - u8 contents_ret[]) -{ - if (fseeko(fp, offset, SEEK_SET) != 0) { - ERROR("Failed to seek to byte %"PRIu64" of input file " - "to read uncompressed resource (len = %"PRIu64")", - offset, len); - return WIMLIB_ERR_READ; - } - if (fread(contents_ret, 1, len, fp) != len) { - if (feof(fp)) { - ERROR("Unexpected EOF in uncompressed file resource"); + if (cb) { + /* Feed the data to the callback function */ + ret = cb(out_p + start_offset, + partial_chunk_size, ctx_or_buf); + if (ret) + goto out; } else { - ERROR("Failed to read %"PRIu64" bytes from " - "uncompressed resource at offset %"PRIu64, - len, offset); + /* No callback function provided; we are writing + * directly to a buffer. Advance the pointer into this + * buffer by the number of uncompressed bytes that were + * written. */ + out_p += partial_chunk_size; } - return WIMLIB_ERR_READ; + cur_read_offset += compressed_chunk_size; } - return 0; + + ret = 0; +out: + if (chunk_offsets_malloced) + FREE(chunk_offsets); + return ret; + +read_error: + ERROR_WITH_ERRNO("Error reading compressed file resource"); + ret = WIMLIB_ERR_READ; + goto out; } /* Reads the contents of a struct resource_entry, as represented in the on-disk * format, from the memory pointed to by @p, and fills in the fields of @entry. * A pointer to the byte after the memory read at @p is returned. */ -const u8 *get_resource_entry(const u8 *p, struct resource_entry *entry) +const void * +get_resource_entry(const void *p, struct resource_entry *entry) { u64 size; u8 flags; @@ -404,7 +378,8 @@ const u8 *get_resource_entry(const u8 *p, struct resource_entry *entry) /* Copies the struct resource_entry @entry to the memory pointed to by @p in the * on-disk format. A pointer to the byte after the memory written at @p is * returned. */ -u8 *put_resource_entry(u8 *p, const struct resource_entry *entry) +void * +put_resource_entry(void *p, const struct resource_entry *entry) { p = put_u56(p, entry->size); p = put_u8(p, entry->flags); @@ -413,201 +388,235 @@ u8 *put_resource_entry(u8 *p, const struct resource_entry *entry) return p; } -#ifdef WITH_FUSE -static FILE *wim_get_fp(WIMStruct *w) +static int +read_partial_wim_resource(const struct wim_lookup_table_entry *lte, + u64 size, + consume_data_callback_t cb, + void *ctx_or_buf, + int flags, + u64 offset) { - pthread_mutex_lock(&w->fp_tab_mutex); - FILE *fp; - - wimlib_assert(w->filename != NULL); + WIMStruct *wim; + int in_fd; + int ret; - for (size_t i = 0; i < w->num_allocated_fps; i++) { - if (w->fp_tab[i]) { - fp = w->fp_tab[i]; - w->fp_tab[i] = NULL; - goto out; + wimlib_assert(lte->resource_location == RESOURCE_IN_WIM); + + wim = lte->wim; + in_fd = wim->in_fd; + + if (lte->resource_entry.flags & WIM_RESHDR_FLAG_COMPRESSED && + !(flags & WIMLIB_RESOURCE_FLAG_RAW)) + { + ret = read_compressed_resource(in_fd, + lte->resource_entry.size, + lte->resource_entry.original_size, + lte->resource_entry.offset, + wimlib_get_compression_type(wim), + size, + offset, + cb, + ctx_or_buf); + } else { + offset += lte->resource_entry.offset; + if (cb) { + /* Send data to callback function */ + u8 buf[min(WIM_CHUNK_SIZE, size)]; + while (size) { + size_t bytes_to_read = min(WIM_CHUNK_SIZE, size); + size_t bytes_read = full_pread(in_fd, buf, + bytes_to_read, offset); + if (bytes_read != bytes_to_read) + goto read_error; + ret = cb(buf, bytes_read, ctx_or_buf); + if (ret) + goto out; + size -= bytes_read; + offset += bytes_read; + } + } else { + /* Send data directly to a buffer */ + if (full_pread(in_fd, ctx_or_buf, size, offset) != size) + goto read_error; } + ret = 0; } - DEBUG("Opening extra file descriptor to `%s'", w->filename); - fp = fopen(w->filename, "rb"); - if (!fp) - ERROR_WITH_ERRNO("Failed to open `%s'", w->filename); + goto out; +read_error: + ERROR_WITH_ERRNO("Error reading data from WIM"); + ret = WIMLIB_ERR_READ; out: - pthread_mutex_unlock(&w->fp_tab_mutex); - return fp; + if (ret) { + if (errno == 0) + errno = EIO; + } + return ret; } -static int wim_release_fp(WIMStruct *w, FILE *fp) + +int +read_partial_wim_resource_into_buf(const struct wim_lookup_table_entry *lte, + size_t size, u64 offset, void *buf) { - int ret = 0; - FILE **fp_tab; + return read_partial_wim_resource(lte, size, NULL, buf, 0, offset); +} - pthread_mutex_lock(&w->fp_tab_mutex); +static int +read_wim_resource_prefix(const struct wim_lookup_table_entry *lte, + u64 size, + consume_data_callback_t cb, + void *ctx_or_buf, + int flags) +{ + return read_partial_wim_resource(lte, size, cb, ctx_or_buf, flags, 0); +} - for (size_t i = 0; i < w->num_allocated_fps; i++) { - if (w->fp_tab[i] == NULL) { - w->fp_tab[i] = fp; - goto out; - } - } - fp_tab = REALLOC(w->fp_tab, sizeof(FILE*) * (w->num_allocated_fps + 4)); - if (!fp_tab) { - ret = WIMLIB_ERR_NOMEM; - goto out; +#ifndef __WIN32__ +static int +read_file_on_disk_prefix(const struct wim_lookup_table_entry *lte, + u64 size, + consume_data_callback_t cb, + void *ctx_or_buf, + int _ignored_flags) +{ + const tchar *filename = lte->file_on_disk; + int ret; + int fd; + size_t bytes_read; + + fd = open(filename, O_RDONLY); + if (fd < 0) { + ERROR_WITH_ERRNO("Can't open \"%"TS"\"", filename); + return WIMLIB_ERR_OPEN; } - w->fp_tab = fp_tab; - memset(&w->fp_tab[w->num_allocated_fps], 0, 4 * sizeof(FILE*)); - w->fp_tab[w->num_allocated_fps] = fp; - w->num_allocated_fps += 4; -out: - pthread_mutex_unlock(&w->fp_tab_mutex); + if (cb) { + /* Send data to callback function */ + u8 buf[min(WIM_CHUNK_SIZE, size)]; + size_t bytes_to_read; + while (size) { + bytes_to_read = min(WIM_CHUNK_SIZE, size); + bytes_read = full_read(fd, buf, bytes_to_read); + if (bytes_read != bytes_to_read) + goto read_error; + ret = cb(buf, bytes_read, ctx_or_buf); + if (ret) + goto out_close; + size -= bytes_read; + } + } else { + /* Send data directly to a buffer */ + bytes_read = full_read(fd, ctx_or_buf, size); + if (bytes_read != size) + goto read_error; + } + ret = 0; + goto out_close; +read_error: + ERROR_WITH_ERRNO("Error reading \"%"TS"\"", filename); + ret = WIMLIB_ERR_READ; +out_close: + close(fd); return ret; } -#endif +#endif /* !__WIN32__ */ -/* - * Reads some data from the resource corresponding to a WIM lookup table entry. - * - * @lte: The WIM lookup table entry for the resource. - * @buf: Buffer into which to write the data. - * @size: Number of bytes to read. - * @offset: Offset at which to start reading the resource. - * - * Returns zero on success, nonzero on failure. - */ -int read_wim_resource(const struct wim_lookup_table_entry *lte, u8 buf[], - size_t size, u64 offset, int flags) +static int +read_buffer_prefix(const struct wim_lookup_table_entry *lte, + u64 size, consume_data_callback_t cb, + void *ctx_or_buf, int _ignored_flags) { - int ctype; - int ret = 0; - FILE *fp; - - /* We shouldn't be allowing read over-runs in any part of the library. - * */ - if (flags & WIMLIB_RESOURCE_FLAG_RAW) - wimlib_assert(offset + size <= lte->resource_entry.size); - else - wimlib_assert(offset + size <= lte->resource_entry.original_size); - - switch (lte->resource_location) { - case RESOURCE_IN_WIM: - /* The resource is in a WIM file, and its WIMStruct is given by - * the lte->wim member. The resource may be either compressed - * or uncompressed. */ - wimlib_assert(lte->wim != NULL); - - #ifdef WITH_FUSE - if (flags & WIMLIB_RESOURCE_FLAG_MULTITHREADED) { - fp = wim_get_fp(lte->wim); - if (!fp) - return WIMLIB_ERR_OPEN; - } else - #endif - { - wimlib_assert(!(flags & WIMLIB_RESOURCE_FLAG_MULTITHREADED)); - wimlib_assert(lte->wim->fp != NULL); - fp = lte->wim->fp; - } - - ctype = wim_resource_compression_type(lte); - - wimlib_assert(ctype != WIMLIB_COMPRESSION_TYPE_NONE || - (lte->resource_entry.original_size == - lte->resource_entry.size)); + const void *inbuf = lte->attached_buffer; + int ret; - if ((flags & WIMLIB_RESOURCE_FLAG_RAW) - || ctype == WIMLIB_COMPRESSION_TYPE_NONE) - ret = read_uncompressed_resource(fp, - lte->resource_entry.offset + offset, - size, buf); - else - ret = read_compressed_resource(fp, - lte->resource_entry.size, - lte->resource_entry.original_size, - lte->resource_entry.offset, - ctype, size, offset, buf); - #ifdef WITH_FUSE - if (flags & WIMLIB_RESOURCE_FLAG_MULTITHREADED) { - int ret2 = wim_release_fp(lte->wim, fp); - if (ret == 0) - ret = ret2; + if (cb) { + while (size) { + size_t chunk_size = min(WIM_CHUNK_SIZE, size); + ret = cb(inbuf, chunk_size, ctx_or_buf); + if (ret) + return ret; + size -= chunk_size; + inbuf += chunk_size; } - #endif - break; - case RESOURCE_IN_STAGING_FILE: - case RESOURCE_IN_FILE_ON_DISK: - /* The resource is in some file on the external filesystem and - * needs to be read uncompressed */ - wimlib_assert(lte->file_on_disk != NULL); - BUILD_BUG_ON(<e->file_on_disk != <e->staging_file_name); - /* Use existing file pointer if available; otherwise open one - * temporarily */ - if (lte->file_on_disk_fp) { - fp = lte->file_on_disk_fp; - } else { - fp = fopen(lte->file_on_disk, "rb"); - if (!fp) { - ERROR_WITH_ERRNO("Failed to open the file " - "`%s'", lte->file_on_disk); - ret = WIMLIB_ERR_OPEN; - break; - } - } - ret = read_uncompressed_resource(fp, offset, size, buf); - if (fp != lte->file_on_disk_fp) - fclose(fp); - break; -#if defined(__CYGWIN__) || defined(__WIN32__) - case RESOURCE_WIN32: - wimlib_assert(lte->file_on_disk_fp != NULL); - DEBUG("Calling win32_read_file()"); - ret = win32_read_file(lte->file_on_disk, lte->file_on_disk_fp, - offset, size, buf); - break; -#endif - case RESOURCE_IN_ATTACHED_BUFFER: - /* The resource is directly attached uncompressed in an - * in-memory buffer. */ - wimlib_assert(lte->attached_buffer != NULL); - memcpy(buf, lte->attached_buffer + offset, size); - break; -#ifdef WITH_NTFS_3G - case RESOURCE_IN_NTFS_VOLUME: - wimlib_assert(lte->ntfs_loc != NULL); - wimlib_assert(lte->attr != NULL); - if (lte->ntfs_loc->is_reparse_point) - offset += 8; - if (ntfs_attr_pread(lte->attr, offset, size, buf) != size) { - ERROR_WITH_ERRNO("Error reading NTFS attribute " - "at `%s'", - lte->ntfs_loc->path_utf8); - ret = WIMLIB_ERR_NTFS_3G; - } - break; -#endif - default: - wimlib_assert(0); - ret = -1; - break; + } else { + memcpy(ctx_or_buf, inbuf, size); } - return ret; + return 0; } +typedef int (*read_resource_prefix_handler_t)(const struct wim_lookup_table_entry *lte, + u64 size, + consume_data_callback_t cb, + void *ctx_or_buf, + int flags); + /* - * Reads all the data from the resource corresponding to a WIM lookup table - * entry. + * Read the first @size bytes from a generic "resource", which may be located in + * the WIM (compressed or uncompressed), in an external file, or directly in an + * in-memory buffer. * - * @lte: The WIM lookup table entry for the resource. - * @buf: Buffer into which to write the data. It must be at least - * wim_resource_size(lte) bytes long. + * Feed the data either to a callback function (cb != NULL, passing it + * ctx_or_buf), or write it directly into a buffer (cb == NULL, ctx_or_buf + * specifies the buffer, which must have room for @size bytes). * - * Returns 0 on success; nonzero on failure. + * When using a callback function, it is called with chunks up to 32768 bytes in + * size until the resource is exhausted. + * + * If the resource is located in a WIM file, @flags can be: + * * WIMLIB_RESOURCE_FLAG_THREADSAFE_READ if it must be safe to access the resource + * concurrently by multiple threads. + * * WIMLIB_RESOURCE_FLAG_RAW if the raw compressed data is to be supplied + * instead of the uncompressed data. + * Otherwise, the @flags are ignored. */ -int read_full_wim_resource(const struct wim_lookup_table_entry *lte, u8 buf[], - int flags) +int +read_resource_prefix(const struct wim_lookup_table_entry *lte, + u64 size, consume_data_callback_t cb, void *ctx_or_buf, + int flags) +{ + static const read_resource_prefix_handler_t handlers[] = { + [RESOURCE_IN_WIM] = read_wim_resource_prefix, + #ifndef __WIN32__ + [RESOURCE_IN_FILE_ON_DISK] = read_file_on_disk_prefix, + #endif + [RESOURCE_IN_ATTACHED_BUFFER] = read_buffer_prefix, + #ifdef WITH_FUSE + [RESOURCE_IN_STAGING_FILE] = read_file_on_disk_prefix, + #endif + #ifdef WITH_NTFS_3G + [RESOURCE_IN_NTFS_VOLUME] = read_ntfs_file_prefix, + #endif + #ifdef __WIN32__ + [RESOURCE_WIN32] = read_win32_file_prefix, + [RESOURCE_WIN32_ENCRYPTED] = read_win32_encrypted_file_prefix, + #endif + }; + wimlib_assert(lte->resource_location < ARRAY_LEN(handlers) + && handlers[lte->resource_location] != NULL); + return handlers[lte->resource_location](lte, size, cb, ctx_or_buf, flags); +} + +int +read_full_resource_into_buf(const struct wim_lookup_table_entry *lte, + void *buf) +{ + return read_resource_prefix(lte, wim_resource_size(lte), NULL, buf, 0); +} + +struct extract_ctx { + SHA_CTX sha_ctx; + consume_data_callback_t extract_chunk; + void *extract_chunk_arg; +}; + +static int +extract_chunk_sha1_wrapper(const void *chunk, size_t chunk_size, + void *_ctx) { - return read_wim_resource(lte, buf, wim_resource_size(lte), 0, flags); + struct extract_ctx *ctx = _ctx; + + sha1_update(&ctx->sha_ctx, chunk, chunk_size); + return ctx->extract_chunk(chunk, chunk_size, ctx->extract_chunk_arg); } /* Extracts the first @size bytes of a WIM resource to somewhere. In the @@ -616,77 +625,48 @@ int read_full_wim_resource(const struct wim_lookup_table_entry *lte, u8 buf[], * * @extract_chunk is a function that is called to extract each chunk of the * resource. */ -int extract_wim_resource(const struct wim_lookup_table_entry *lte, - u64 size, - extract_chunk_func_t extract_chunk, - void *extract_chunk_arg) +int +extract_wim_resource(const struct wim_lookup_table_entry *lte, + u64 size, + consume_data_callback_t extract_chunk, + void *extract_chunk_arg) { - u64 bytes_remaining = size; - u8 buf[min(WIM_CHUNK_SIZE, bytes_remaining)]; - u64 offset = 0; - int ret = 0; - u8 hash[SHA1_HASH_SIZE]; - bool check_hash = (size == wim_resource_size(lte)); - SHA_CTX ctx; - - if (check_hash) - sha1_init(&ctx); - - while (bytes_remaining) { - u64 to_read = min(bytes_remaining, sizeof(buf)); - ret = read_wim_resource(lte, buf, to_read, offset, 0); - if (ret != 0) - return ret; - if (check_hash) - sha1_update(&ctx, buf, to_read); - ret = extract_chunk(buf, to_read, offset, extract_chunk_arg); - if (ret != 0) { - ERROR_WITH_ERRNO("Error extracting WIM resource"); - return ret; - } - bytes_remaining -= to_read; - offset += to_read; - } - if (check_hash) { - sha1_final(hash, &ctx); - if (!hashes_equal(hash, lte->hash)) { - #ifdef ENABLE_ERROR_MESSAGES - ERROR("Invalid checksum on the following WIM resource:"); - print_lookup_table_entry(lte, stderr); - #endif - return WIMLIB_ERR_INVALID_RESOURCE_HASH; - } - } - return 0; -} - -/* Write @n bytes from @buf to the file descriptor @fd, retrying on internupt - * and on short writes. - * - * Returns short count and set errno on failure. */ -static ssize_t full_write(int fd, const void *buf, size_t n) -{ - const char *p = buf; - ssize_t ret; - ssize_t total = 0; - - while (total != n) { - ret = write(fd, p, n); - if (ret < 0) { - if (errno == EINTR) - continue; - else - break; + int ret; + if (size == wim_resource_size(lte)) { + /* Do SHA1 */ + struct extract_ctx ctx; + ctx.extract_chunk = extract_chunk; + ctx.extract_chunk_arg = extract_chunk_arg; + sha1_init(&ctx.sha_ctx); + ret = read_resource_prefix(lte, size, + extract_chunk_sha1_wrapper, + &ctx, 0); + if (ret == 0) { + u8 hash[SHA1_HASH_SIZE]; + sha1_final(hash, &ctx.sha_ctx); + if (!hashes_equal(hash, lte->hash)) { + #ifdef ENABLE_ERROR_MESSAGES + ERROR("Invalid SHA1 message digest " + "on the following WIM resource:"); + print_lookup_table_entry(lte, stderr); + if (lte->resource_location == RESOURCE_IN_WIM) + ERROR("The WIM file appears to be corrupt!"); + #endif + ret = WIMLIB_ERR_INVALID_RESOURCE_HASH; + } } - total += ret; - p += ret; + } else { + /* Don't do SHA1 */ + ret = read_resource_prefix(lte, size, extract_chunk, + extract_chunk_arg, 0); } - return total; + return ret; } -int extract_wim_chunk_to_fd(const u8 *buf, size_t len, u64 offset, void *arg) +static int +extract_wim_chunk_to_fd(const void *buf, size_t len, void *_fd_p) { - int fd = *(int*)arg; + int fd = *(int*)_fd_p; ssize_t ret = full_write(fd, buf, len); if (ret < len) { ERROR_WITH_ERRNO("Error writing to file descriptor"); @@ -696,6 +676,36 @@ int extract_wim_chunk_to_fd(const u8 *buf, size_t len, u64 offset, void *arg) } } +int +extract_wim_resource_to_fd(const struct wim_lookup_table_entry *lte, + int fd, u64 size) +{ + return extract_wim_resource(lte, size, extract_wim_chunk_to_fd, &fd); +} + + +static int +sha1_chunk(const void *buf, size_t len, void *ctx) +{ + sha1_update(ctx, buf, len); + return 0; +} + +/* Calculate the SHA1 message digest of a stream. */ +int +sha1_resource(struct wim_lookup_table_entry *lte) +{ + int ret; + SHA_CTX sha_ctx; + + sha1_init(&sha_ctx); + ret = read_resource_prefix(lte, wim_resource_size(lte), + sha1_chunk, &sha_ctx, 0); + if (ret == 0) + sha1_final(lte->hash, &sha_ctx); + return ret; +} + /* * Copies the file resource specified by the lookup table entry @lte from the * input WIM to the output WIM that has its FILE * given by @@ -706,21 +716,18 @@ int extract_wim_chunk_to_fd(const u8 *buf, size_t len, u64 offset, void *arg) * * (This function is confusing and should be refactored somehow.) */ -int copy_resource(struct wim_lookup_table_entry *lte, void *wim) +int +copy_resource(struct wim_lookup_table_entry *lte, void *wim) { WIMStruct *w = wim; int ret; - if ((lte->resource_entry.flags & WIM_RESHDR_FLAG_METADATA) && - !w->write_metadata) - return 0; - - ret = write_wim_resource(lte, w->out_fp, + ret = write_wim_resource(lte, w->out_fd, wim_resource_compression_type(lte), <e->output_resource_entry, 0); - if (ret != 0) - return ret; - lte->out_refcnt = lte->refcnt; - lte->part_number = w->hdr.part_number; - return 0; + if (ret == 0) { + lte->out_refcnt = lte->refcnt; + lte->part_number = w->hdr.part_number; + } + return ret; }