]> wimlib.net Git - wimlib/blobdiff - programs/imagex.c
Add --threads option to imagex-optimize
[wimlib] / programs / imagex.c
index dd1a1243900cff33ebf44859139da74fa083c8fa..372e115b64509d8b6a85ba72061038160c353aef 100644 (file)
@@ -23,8 +23,8 @@
  */
 
 #include "config.h"
-#include "wimlib_tchar.h"
 #include "wimlib.h"
+#include "wimlib_tchar.h"
 
 #include <ctype.h>
 #include <errno.h>
@@ -145,6 +145,7 @@ IMAGEX_PROGNAME" mountrw WIMFILE [IMAGE_NUM | IMAGE_NAME] DIRECTORY\n"
 [OPTIMIZE] =
 T(
 IMAGEX_PROGNAME" optimize WIMFILE [--check] [--recompress]\n"
+"                      [--threads=NUM_THREADS]\n"
 ),
 [SPLIT] =
 T(
@@ -262,6 +263,7 @@ static const struct option mount_options[] = {
 static const struct option optimize_options[] = {
        {T("check"),      no_argument, NULL, IMAGEX_CHECK_OPTION},
        {T("recompress"), no_argument, NULL, IMAGEX_RECOMPRESS_OPTION},
+       {T("threads"),    required_argument, NULL, IMAGEX_THREADS_OPTION},
        {NULL, 0, NULL, 0},
 };
 
@@ -369,22 +371,28 @@ file_get_size(const tchar *filename)
                return (off_t)-1;
 }
 
-static const tchar *default_capture_config =
-T(
-"[ExclusionList]\n"
-"\\$ntfs.log\n"
-"\\hiberfil.sys\n"
-"\\pagefile.sys\n"
-"\\System Volume Information\n"
-"\\RECYCLER\n"
-"\\Windows\\CSC\n"
-"\n"
-"[CompressionExclusionList]\n"
-"*.mp3\n"
-"*.zip\n"
-"*.cab\n"
-"\\WINDOWS\\inf\\*.pnf\n"
-);
+tchar pat_ntfs_log[]                  = T("/$ntfs.log");
+tchar pat_hiberfil_sys[]              = T("/hiberfil.sys");
+tchar pat_pagefile_sys[]              = T("/pagefile.sys");
+tchar pat_system_volume_information[] = T("/System Volume Information");
+tchar pat_recycler[]                  = T("/RECYCLER");
+tchar pat_windows_csc[]               = T("/Windows/CSC");
+
+tchar *default_pats[] = {
+       pat_ntfs_log,
+       pat_hiberfil_sys,
+       pat_pagefile_sys,
+       pat_system_volume_information,
+       pat_recycler,
+       pat_windows_csc,
+};
+
+static struct wimlib_capture_config default_capture_config = {
+       .exclusion_pats = {
+               .num_pats = sizeof(default_pats) / sizeof(*default_pats),
+               .pats = default_pats,
+       },
+};
 
 enum {
        PARSE_FILENAME_SUCCESS = 0,
@@ -505,6 +513,34 @@ is_comment_line(const tchar *line, size_t len)
        }
 }
 
+static ssize_t
+text_file_count_lines(tchar **contents_p, size_t *nchars_p)
+{
+       ssize_t nlines = 0;
+       tchar *contents = *contents_p;
+       size_t nchars = *nchars_p;
+       size_t i;
+
+       for (i = 0; i < nchars; i++)
+               if (contents[i] == T('\n'))
+                       nlines++;
+
+       /* Handle last line not terminated by a newline */
+       if (nchars != 0 && contents[nchars - 1] != T('\n')) {
+               contents = realloc(contents, (nchars + 1) * sizeof(tchar));
+               if (!contents) {
+                       imagex_error(T("Out of memory!"));
+                       return -1;
+               }
+               contents[nchars] = T('\n');
+               *contents_p = contents;
+               nchars++;
+               nlines++;
+       }
+       *nchars_p = nchars;
+       return nlines;
+}
+
 /* Parses a file in the source list format.  (See the man page for
  * 'wimlib-imagex capture' for details on this format and the meaning.)
  *
@@ -526,37 +562,24 @@ static struct wimlib_capture_source *
 parse_source_list(tchar **source_list_contents_p, size_t source_list_nchars,
                  size_t *nsources_ret)
 {
-       size_t nlines;
+       ssize_t nlines;
        tchar *p;
        struct wimlib_capture_source *sources;
        size_t i, j;
-       tchar *source_list_contents = *source_list_contents_p;
 
-       nlines = 0;
-       for (i = 0; i < source_list_nchars; i++)
-               if (source_list_contents[i] == T('\n'))
-                       nlines++;
-
-       /* Handle last line not terminated by a newline */
-       if (source_list_nchars != 0 &&
-           source_list_contents[source_list_nchars - 1] != T('\n'))
-       {
-               source_list_contents = realloc(source_list_contents,
-                                              (source_list_nchars + 1) * sizeof(tchar));
-               if (!source_list_contents)
-                       goto oom;
-               source_list_contents[source_list_nchars] = T('\n');
-               *source_list_contents_p = source_list_contents;
-               source_list_nchars++;
-               nlines++;
-       }
+       nlines = text_file_count_lines(source_list_contents_p,
+                                      &source_list_nchars);
+       if (nlines < 0)
+               return NULL;
 
        /* Always allocate at least 1 slot, just in case the implementation of
         * calloc() returns NULL if 0 bytes are requested. */
        sources = calloc(nlines ?: 1, sizeof(*sources));
-       if (!sources)
-               goto oom;
-       p = source_list_contents;
+       if (!sources) {
+               imagex_error(T("out of memory"));
+               return NULL;
+       }
+       p = *source_list_contents_p;
        j = 0;
        for (i = 0; i < nlines; i++) {
                /* XXX: Could use rawmemchr() here instead, but it may not be
@@ -575,9 +598,144 @@ parse_source_list(tchar **source_list_contents_p, size_t source_list_nchars,
        }
        *nsources_ret = j;
        return sources;
-oom:
-       imagex_error(T("out of memory"));
-       return NULL;
+}
+
+
+enum capture_config_section {
+       CAPTURE_CONFIG_NO_SECTION,
+       CAPTURE_CONFIG_EXCLUSION_SECTION,
+       CAPTURE_CONFIG_EXCLUSION_EXCEPTION_SECTION,
+       CAPTURE_CONFIG_IGNORE_SECTION,
+};
+
+enum {
+       CAPTURE_CONFIG_INVALID_SECTION,
+       CAPTURE_CONFIG_CHANGED_SECTION,
+       CAPTURE_CONFIG_SAME_SECTION,
+};
+
+static int
+check_config_section(tchar *line, size_t len,
+                    enum capture_config_section *cur_section)
+{
+       while (istspace(*line))
+               line++;
+
+       if (*line != T('['))
+               return CAPTURE_CONFIG_SAME_SECTION;
+
+       line++;
+       tchar *endbrace = tstrrchr(line, T(']'));
+       if (!endbrace)
+               return CAPTURE_CONFIG_SAME_SECTION;
+
+       if (!tmemcmp(line, T("ExclusionList"), endbrace - line)) {
+               *cur_section = CAPTURE_CONFIG_EXCLUSION_SECTION;
+       } else if (!tmemcmp(line, T("ExclusionException"), endbrace - line)) {
+               *cur_section = CAPTURE_CONFIG_EXCLUSION_EXCEPTION_SECTION;
+       } else if (!tmemcmp(line, T("CompressionExclusionList"), endbrace - line)) {
+               *cur_section = CAPTURE_CONFIG_IGNORE_SECTION;
+               tfputs(T("WARNING: Ignoring [CompressionExclusionList] section "
+                        "of capture config file\n"),
+                      stderr);
+       } else if (!tmemcmp(line, T("AlignmentList"), endbrace - line)) {
+               *cur_section = CAPTURE_CONFIG_IGNORE_SECTION;
+               tfputs(T("WARNING: Ignoring [AlignmentList] section "
+                        "of capture config file\n"),
+                      stderr);
+       } else {
+               imagex_error(T("Invalid capture config file section \"%"TS"\""),
+                            line - 1);
+               return CAPTURE_CONFIG_INVALID_SECTION;
+       }
+       return CAPTURE_CONFIG_CHANGED_SECTION;
+}
+
+
+static bool
+pattern_list_add_pattern(struct wimlib_pattern_list *pat_list,
+                        tchar *pat)
+{
+       if (pat_list->num_pats == pat_list->num_allocated_pats) {
+               tchar **pats;
+               size_t num_allocated_pats = pat_list->num_pats + 8;
+
+               pats = realloc(pat_list->pats,
+                              num_allocated_pats * sizeof(pat_list->pats[0]));
+               if (!pats) {
+                       imagex_error(T("Out of memory!"));
+                       return false;
+               }
+               pat_list->pats = pats;
+               pat_list->num_allocated_pats = num_allocated_pats;
+       }
+       pat_list->pats[pat_list->num_pats++] = pat;
+       return true;
+}
+
+static bool
+parse_capture_config_line(tchar *line, size_t len,
+                         enum capture_config_section *cur_section,
+                         struct wimlib_capture_config *config)
+{
+       tchar *filename;
+       int ret;
+
+       ret = check_config_section(line, len, cur_section);
+       if (ret == CAPTURE_CONFIG_INVALID_SECTION)
+               return false;
+       if (ret == CAPTURE_CONFIG_CHANGED_SECTION)
+               return true;
+
+       switch (*cur_section) {
+       case CAPTURE_CONFIG_NO_SECTION:
+               imagex_error(T("Line \"%"TS"\" is not in a section "
+                              "(such as [ExclusionList]"), line);
+               return false;
+       case CAPTURE_CONFIG_EXCLUSION_SECTION:
+               if (parse_filename(&line, &len, &filename) != PARSE_FILENAME_SUCCESS)
+                       return false;
+               return pattern_list_add_pattern(&config->exclusion_pats,
+                                               filename);
+       case CAPTURE_CONFIG_EXCLUSION_EXCEPTION_SECTION:
+               if (parse_filename(&line, &len, &filename) != PARSE_FILENAME_SUCCESS)
+                       return false;
+               return pattern_list_add_pattern(&config->exclusion_exception_pats,
+                                               filename);
+       case CAPTURE_CONFIG_IGNORE_SECTION:
+               return true;
+       }
+       return false;
+}
+
+static int
+parse_capture_config(tchar **contents_p, size_t nchars,
+                    struct wimlib_capture_config *config)
+{
+       ssize_t nlines;
+       tchar *p;
+       size_t i;
+       enum capture_config_section cur_section;
+
+       memset(config, 0, sizeof(*config));
+
+       nlines = text_file_count_lines(contents_p, &nchars);
+       if (nlines < 0)
+               return -1;
+
+       cur_section = CAPTURE_CONFIG_NO_SECTION;
+       p = *contents_p;
+       for (i = 0; i < nlines; i++) {
+               tchar *endp = tmemchr(p, T('\n'), nchars);
+               size_t len = endp - p + 1;
+               *endp = T('\0');
+               if (!is_comment_line(p, len))
+                       if (!parse_capture_config_line(p, len, &cur_section, config))
+                               return -1;
+               p = endp + 1;
+
+       }
+       return 0;
 }
 
 /* Reads the contents of a file into memory. */
@@ -1112,14 +1270,14 @@ imagex_capture_or_append(int argc, tchar **argv)
 {
        int c;
        int open_flags = 0;
-       int add_image_flags = 0;
+       int add_image_flags = WIMLIB_ADD_IMAGE_FLAG_EXCLUDE_VERBOSE;
        int write_flags = 0;
        int compression_type = WIMLIB_COMPRESSION_TYPE_XPRESS;
        const tchar *wimfile;
        const tchar *name;
        const tchar *desc;
        const tchar *flags_element = NULL;
-       WIMStruct *w = NULL;
+       WIMStruct *w;
        int ret;
        int cur_image;
        int cmd = tstrcmp(argv[0], T("append")) ? CAPTURE : APPEND;
@@ -1130,13 +1288,13 @@ imagex_capture_or_append(int argc, tchar **argv)
        tchar *source_copy;
 
        const tchar *config_file = NULL;
-       tchar *config_str = NULL;
-       size_t config_len;
+       tchar *config_str;
+       struct wimlib_capture_config *config = NULL;
 
        bool source_list = false;
        size_t source_list_nchars;
-       tchar *source_list_contents = NULL;
-       bool capture_sources_malloced = false;
+       tchar *source_list_contents;
+       bool capture_sources_malloced;
        struct wimlib_capture_source *capture_sources;
        size_t num_sources;
 
@@ -1231,7 +1389,7 @@ imagex_capture_or_append(int argc, tchar **argv)
                                                    &num_sources);
                if (!capture_sources) {
                        ret = -1;
-                       goto out;
+                       goto out_free_source_list_contents;
                }
                capture_sources_malloced = true;
        } else {
@@ -1242,14 +1400,25 @@ imagex_capture_or_append(int argc, tchar **argv)
                capture_sources[0].wim_target_path = NULL;
                capture_sources[0].reserved = 0;
                num_sources = 1;
+               capture_sources_malloced = false;
+               source_list_contents = NULL;
        }
 
        if (config_file) {
+               size_t config_len;
+
                config_str = file_get_text_contents(config_file, &config_len);
                if (!config_str) {
                        ret = -1;
-                       goto out;
+                       goto out_free_capture_sources;
                }
+
+               config = alloca(sizeof(*config));
+               ret = parse_capture_config(&config_str, config_len, config);
+               if (ret)
+                       goto out_free_config;
+       } else {
+               config = &default_capture_config;
        }
 
        if (cmd == APPEND)
@@ -1257,8 +1426,8 @@ imagex_capture_or_append(int argc, tchar **argv)
                                      imagex_progress_func);
        else
                ret = wimlib_create_new_wim(compression_type, &w);
-       if (ret != 0)
-               goto out;
+       if (ret)
+               goto out_free_config;
 
        if (!source_list) {
                struct stat stbuf;
@@ -1274,7 +1443,7 @@ imagex_capture_or_append(int argc, tchar **argv)
                                imagex_error_with_errno(T("Failed to stat "
                                                          "\"%"TS"\""), source);
                                ret = -1;
-                               goto out;
+                               goto out_wimlib_free;
                        }
                }
        }
@@ -1282,12 +1451,11 @@ imagex_capture_or_append(int argc, tchar **argv)
        win32_acquire_capture_privileges();
 #endif
 
-       ret = wimlib_add_image_multisource(w, capture_sources,
-                                          num_sources, name,
-                                          (config_str ? config_str :
-                                               default_capture_config),
-                                          (config_str ? config_len :
-                                               tstrlen(default_capture_config)),
+       ret = wimlib_add_image_multisource(w,
+                                          capture_sources,
+                                          num_sources,
+                                          name,
+                                          config,
                                           add_image_flags,
                                           imagex_progress_func);
        if (ret != 0)
@@ -1319,12 +1487,19 @@ out_release_privs:
 #ifdef __WIN32__
        win32_release_capture_privileges();
 #endif
-out:
+out_wimlib_free:
        wimlib_free(w);
-       free(config_str);
-       free(source_list_contents);
+out_free_config:
+       if (config != NULL && config != &default_capture_config) {
+               free(config->exclusion_pats.pats);
+               free(config->exclusion_exception_pats.pats);
+               free(config_str);
+       }
+out_free_capture_sources:
        if (capture_sources_malloced)
                free(capture_sources);
+out_free_source_list_contents:
+       free(source_list_contents);
        return ret;
 }
 
@@ -2081,6 +2256,7 @@ imagex_optimize(int argc, tchar **argv)
        const tchar *wimfile;
        off_t old_size;
        off_t new_size;
+       unsigned num_threads = 0;
 
        for_opt(c, optimize_options) {
                switch (c) {
@@ -2091,6 +2267,11 @@ imagex_optimize(int argc, tchar **argv)
                case IMAGEX_RECOMPRESS_OPTION:
                        write_flags |= WIMLIB_WRITE_FLAG_RECOMPRESS;
                        break;
+               case IMAGEX_THREADS_OPTION:
+                       num_threads = parse_num_threads(optarg);
+                       if (num_threads == UINT_MAX)
+                               return -1;
+                       break;
                default:
                        usage(OPTIMIZE);
                        return -1;
@@ -2118,7 +2299,8 @@ imagex_optimize(int argc, tchar **argv)
        else
                tprintf(T("%"PRIu64" KiB\n"), old_size >> 10);
 
-       ret = wimlib_overwrite(w, write_flags, 0, imagex_progress_func);
+       ret = wimlib_overwrite(w, write_flags, num_threads,
+                              imagex_progress_func);
 
        if (ret == 0) {
                new_size = file_get_size(argv[0]);