]> wimlib.net Git - wimlib/blobdiff - src/modify.c
Increment real_refcnt for metadata lte's
[wimlib] / src / modify.c
index a8ac048ff8d68f280f3da79624782a415137fa4d..d5553dc8a2ba1838608f0df89fa2d897724a84fe 100644 (file)
@@ -32,6 +32,7 @@
 #include "dentry.h"
 #include "xml.h"
 #include "lookup_table.h"
+#include "timestamp.h"
 #include <sys/stat.h>
 #include <dirent.h>
 #include <string.h>
  * the WIM image. */
 #define WIMLIB_ADD_IMAGE_FLAG_ROOT 0x80000000
 
-void destroy_image_metadata(struct image_metadata *imd, struct lookup_table *lt)
+void destroy_image_metadata(struct image_metadata *imd,
+                           struct lookup_table *table)
 {
-       free_dentry_tree(imd->root_dentry, lt);
+       free_dentry_tree(imd->root_dentry, table);
        free_security_data(imd->security_data);
 
        /* Get rid of the lookup table entry for this image's metadata resource
         * */
-       if (lt)
-               lookup_table_remove(lt, imd->metadata_lte);
+       if (table) {
+               lookup_table_unlink(table, imd->metadata_lte);
+               free_lookup_table_entry(imd->metadata_lte);
+       }
 }
 
-/* 
+/*
  * Recursively builds a dentry tree from a directory tree on disk, outside the
  * WIM file.
  *
@@ -63,10 +67,10 @@ void destroy_image_metadata(struct image_metadata *imd, struct lookup_table *lt)
  *             modified if successful.  NULL if the file or directory was
  *             excluded from capture.
  *
- * @root_disk_path:  The path to the root of the directory tree on disk. 
+ * @root_disk_path:  The path to the root of the directory tree on disk.
  *
  * @lookup_table: The lookup table for the WIM file.  For each file added to the
- *             dentry tree being built, an entry is added to the lookup table, 
+ *             dentry tree being built, an entry is added to the lookup table,
  *             unless an identical stream is already in the lookup table.
  *             These lookup table entries that are added point to the path of
  *             the file on disk.
@@ -92,7 +96,8 @@ static int build_dentry_tree(struct dentry **root_ret,
                             struct lookup_table *lookup_table,
                             struct wim_security_data *sd,
                             const struct capture_config *config,
-                            int add_flags,
+                            int add_image_flags,
+                            wimlib_progress_func_t progress_func,
                             void *extra_arg)
 {
        struct stat root_stbuf;
@@ -100,35 +105,46 @@ static int build_dentry_tree(struct dentry **root_ret,
        int (*stat_fn)(const char *restrict, struct stat *restrict);
        struct dentry *root;
        const char *filename;
+       struct inode *inode;
 
        if (exclude_path(root_disk_path, config, true)) {
-               if (add_flags & WIMLIB_ADD_IMAGE_FLAG_ROOT) {
+               if (add_image_flags & WIMLIB_ADD_IMAGE_FLAG_ROOT) {
                        ERROR("Cannot exclude the root directory from capture");
                        return WIMLIB_ERR_INVALID_CAPTURE_CONFIG;
                }
-               if (add_flags & WIMLIB_ADD_IMAGE_FLAG_VERBOSE)
-                       printf("Excluding file `%s' from capture\n",
-                              root_disk_path);
+               if ((add_image_flags & WIMLIB_ADD_IMAGE_FLAG_VERBOSE)
+                   && progress_func)
+               {
+                       union wimlib_progress_info info;
+                       info.scan.cur_path = root_disk_path;
+                       info.scan.excluded = true;
+                       progress_func(WIMLIB_PROGRESS_MSG_SCAN_DENTRY, &info);
+               }
                *root_ret = NULL;
                return 0;
        }
 
+       if ((add_image_flags & WIMLIB_ADD_IMAGE_FLAG_VERBOSE)
+           && progress_func)
+       {
+               union wimlib_progress_info info;
+               info.scan.cur_path = root_disk_path;
+               info.scan.excluded = false;
+               progress_func(WIMLIB_PROGRESS_MSG_SCAN_DENTRY, &info);
+       }
 
-       if (add_flags & WIMLIB_ADD_IMAGE_FLAG_DEREFERENCE)
+       if (add_image_flags & WIMLIB_ADD_IMAGE_FLAG_DEREFERENCE)
                stat_fn = stat;
        else
                stat_fn = lstat;
 
-       if (add_flags & WIMLIB_ADD_IMAGE_FLAG_VERBOSE)
-               printf("Scanning `%s'\n", root_disk_path);
-
        ret = (*stat_fn)(root_disk_path, &root_stbuf);
        if (ret != 0) {
                ERROR_WITH_ERRNO("Failed to stat `%s'", root_disk_path);
                return WIMLIB_ERR_STAT;
        }
 
-       if ((add_flags & WIMLIB_ADD_IMAGE_FLAG_ROOT) && 
+       if ((add_image_flags & WIMLIB_ADD_IMAGE_FLAG_ROOT) &&
              !S_ISDIR(root_stbuf.st_mode)) {
                ERROR("`%s' is not a directory", root_disk_path);
                return WIMLIB_ERR_NOTDIR;
@@ -140,7 +156,7 @@ static int build_dentry_tree(struct dentry **root_ret,
                return WIMLIB_ERR_SPECIAL_FILE;
        }
 
-       if (add_flags & WIMLIB_ADD_IMAGE_FLAG_ROOT)
+       if (add_image_flags & WIMLIB_ADD_IMAGE_FLAG_ROOT)
                filename = "";
        else
                filename = path_basename(root_disk_path);
@@ -149,77 +165,33 @@ static int build_dentry_tree(struct dentry **root_ret,
        if (!root)
                return WIMLIB_ERR_NOMEM;
 
-       stbuf_to_inode(&root_stbuf, root->d_inode);
-       add_flags &= ~WIMLIB_ADD_IMAGE_FLAG_ROOT;
-       root->d_inode->resolved = true;
-
-       if (dentry_is_directory(root)) { /* Archiving a directory */
-               DIR *dir;
-               struct dirent *p;
-               struct dentry *child;
+       inode = root->d_inode;
+
+#ifdef HAVE_STAT_NANOSECOND_PRECISION
+       inode->creation_time = timespec_to_wim_timestamp(&root_stbuf.st_mtim);
+       inode->last_write_time = timespec_to_wim_timestamp(&root_stbuf.st_mtim);
+       inode->last_access_time = timespec_to_wim_timestamp(&root_stbuf.st_atim);
+#else
+       inode->creation_time = unix_timestamp_to_wim(root_stbuf.st_mtime);
+       inode->last_write_time = unix_timestamp_to_wim(root_stbuf.st_mtime);
+       inode->last_access_time = unix_timestamp_to_wim(root_stbuf.st_atime);
+#endif
+       if (sizeof(ino_t) >= 8)
+               inode->ino = (u64)root_stbuf.st_ino;
+       else
+               inode->ino = (u64)root_stbuf.st_ino |
+                                  ((u64)root_stbuf.st_dev << ((sizeof(ino_t) * 8) & 63));
 
-               dir = opendir(root_disk_path);
-               if (!dir) {
-                       ERROR_WITH_ERRNO("Failed to open the directory `%s'",
-                                        root_disk_path);
-                       ret = WIMLIB_ERR_OPEN;
-                       goto out;
-               }
+       add_image_flags &= ~WIMLIB_ADD_IMAGE_FLAG_ROOT;
+       inode->resolved = true;
 
-               /* Buffer for names of files in directory. */
-               size_t len = strlen(root_disk_path);
-               char name[len + 1 + FILENAME_MAX + 1];
-               memcpy(name, root_disk_path, len);
-               name[len] = '/';
-
-               /* Create a dentry for each entry in the directory on disk, and recurse
-                * to any subdirectories. */
-               while (1) {
-                       errno = 0;
-                       p = readdir(dir);
-                       if (p == NULL) {
-                               if (errno) {
-                                       ret = WIMLIB_ERR_READ;
-                                       ERROR_WITH_ERRNO("Error reading the "
-                                                        "directory `%s'",
-                                                        root_disk_path);
-                               }
-                               break;
-                       }
-                       if (p->d_name[0] == '.' && (p->d_name[1] == '\0'
-                             || (p->d_name[1] == '.' && p->d_name[2] == '\0')))
-                                       continue;
-                       strcpy(name + len + 1, p->d_name);
-                       ret = build_dentry_tree(&child, name, lookup_table,
-                                               NULL, config,
-                                               add_flags, NULL);
-                       if (ret != 0)
-                               break;
-                       if (child)
-                               link_dentry(child, root);
-               }
-               closedir(dir);
-       } else if (dentry_is_symlink(root)) { /* Archiving a symbolic link */
-               char deref_name_buf[4096];
-               ssize_t deref_name_len;
-               
-               deref_name_len = readlink(root_disk_path, deref_name_buf,
-                                         sizeof(deref_name_buf) - 1);
-               if (deref_name_len >= 0) {
-                       deref_name_buf[deref_name_len] = '\0';
-                       DEBUG("Read symlink `%s'", deref_name_buf);
-                       ret = inode_set_symlink(root->d_inode, deref_name_buf,
-                                               lookup_table, NULL);
-               } else {
-                       ERROR_WITH_ERRNO("Failed to read target of "
-                                        "symbolic link `%s'", root_disk_path);
-                       ret = WIMLIB_ERR_READLINK;
-               }
-       } else { /* Archiving a regular file */
+       if (S_ISREG(root_stbuf.st_mode)) { /* Archiving a regular file */
 
                struct lookup_table_entry *lte;
                u8 hash[SHA1_HASH_SIZE];
 
+               inode->attributes = FILE_ATTRIBUTE_NORMAL;
+
                /* Empty files do not have to have a lookup table entry. */
                if (root_stbuf.st_size == 0)
                        goto out;
@@ -259,6 +231,97 @@ static int build_dentry_tree(struct dentry **root_ret,
                        lookup_table_insert(lookup_table, lte);
                }
                root->d_inode->lte = lte;
+       } else if (S_ISDIR(root_stbuf.st_mode)) { /* Archiving a directory */
+
+               inode->attributes = FILE_ATTRIBUTE_DIRECTORY;
+
+               DIR *dir;
+               struct dirent entry, *result;
+               struct dentry *child;
+
+               dir = opendir(root_disk_path);
+               if (!dir) {
+                       ERROR_WITH_ERRNO("Failed to open the directory `%s'",
+                                        root_disk_path);
+                       ret = WIMLIB_ERR_OPEN;
+                       goto out;
+               }
+
+               /* Buffer for names of files in directory. */
+               size_t len = strlen(root_disk_path);
+               char name[len + 1 + FILENAME_MAX + 1];
+               memcpy(name, root_disk_path, len);
+               name[len] = '/';
+
+               /* Create a dentry for each entry in the directory on disk, and recurse
+                * to any subdirectories. */
+               while (1) {
+                       errno = 0;
+                       ret = readdir_r(dir, &entry, &result);
+                       if (ret != 0) {
+                               ret = WIMLIB_ERR_READ;
+                               ERROR_WITH_ERRNO("Error reading the "
+                                                "directory `%s'",
+                                                root_disk_path);
+                               break;
+                       }
+                       if (result == NULL)
+                               break;
+                       if (result->d_name[0] == '.' && (result->d_name[1] == '\0'
+                             || (result->d_name[1] == '.' && result->d_name[2] == '\0')))
+                                       continue;
+                       strcpy(name + len + 1, result->d_name);
+                       ret = build_dentry_tree(&child, name, lookup_table,
+                                               NULL, config, add_image_flags,
+                                               progress_func, NULL);
+                       if (ret != 0)
+                               break;
+                       if (child)
+                               dentry_add_child(root, child);
+               }
+               closedir(dir);
+       } else { /* Archiving a symbolic link */
+               inode->attributes = FILE_ATTRIBUTE_REPARSE_POINT;
+               inode->reparse_tag = WIM_IO_REPARSE_TAG_SYMLINK;
+
+               /* The idea here is to call readlink() to get the UNIX target of
+                * the symbolic link, then turn the target into a reparse point
+                * data buffer that contains a relative or absolute symbolic
+                * link (NOT a junction point or *full* path symbolic link with
+                * drive letter).
+                */
+
+               char deref_name_buf[4096];
+               ssize_t deref_name_len;
+
+               deref_name_len = readlink(root_disk_path, deref_name_buf,
+                                         sizeof(deref_name_buf) - 1);
+               if (deref_name_len >= 0) {
+                       deref_name_buf[deref_name_len] = '\0';
+                       DEBUG("Read symlink `%s'", deref_name_buf);
+                       ret = inode_set_symlink(root->d_inode, deref_name_buf,
+                                               lookup_table, NULL);
+                       if (ret == 0) {
+                               /*
+                                * Unfortunately, Windows seems to have the
+                                * concept of "file" symbolic links as being
+                                * different from "directory" symbolic links...
+                                * so FILE_ATTRIBUTE_DIRECTORY needs to be set
+                                * on the symbolic link if the *target* of the
+                                * symbolic link is a directory.
+                                */
+                               struct stat stbuf;
+                               if (stat(root_disk_path, &stbuf) == 0 &&
+                                   S_ISDIR(stbuf.st_mode))
+                               {
+                                       inode->attributes |= FILE_ATTRIBUTE_DIRECTORY;
+                               }
+                       }
+               } else {
+                       ERROR_WITH_ERRNO("Failed to read target of "
+                                        "symbolic link `%s'", root_disk_path);
+                       ret = WIMLIB_ERR_READLINK;
+               }
        }
 out:
        if (ret == 0)
@@ -300,14 +363,14 @@ static int allocate_lte_if_needed(struct dentry *dentry, void *arg)
                                dest_lte = clone_lookup_table_entry(src_lte);
                                if (!dest_lte)
                                        return WIMLIB_ERR_NOMEM;
-                               list_add_tail(&dest_lte->list, lte_list_head);
+                               list_add_tail(&dest_lte->staging_list, lte_list_head);
                        }
                }
        }
        return 0;
 }
 
-/* 
+/*
  * This function takes in a dentry that was previously located only in image(s)
  * in @src_wim, but now is being added to @dest_wim.  For each stream associated
  * with the dentry, if there is already a lookup table entry for that stream in
@@ -348,7 +411,8 @@ static int add_lte_to_dest_wim(struct dentry *dentry, void *arg)
 
                        next = lte_list_head->next;
                        list_del(next);
-                       dest_lte = container_of(next, struct lookup_table_entry, list);
+                       dest_lte = container_of(next, struct lookup_table_entry,
+                                               staging_list);
                        dest_lte->part_number = 1;
                        dest_lte->refcnt = 1;
                        wimlib_assert(hashes_equal(dest_lte->hash, src_lte->hash));
@@ -362,7 +426,7 @@ static int add_lte_to_dest_wim(struct dentry *dentry, void *arg)
 /*
  * Adds an image (given by its dentry tree) to the image metadata array of a WIM
  * file, adds an entry to the lookup table for the image metadata, updates the
- * image count in the header, and selects the new image. 
+ * image count in the header, and selects the new image.
  *
  * Does not update the XML data.
  *
@@ -389,15 +453,15 @@ static int add_new_dentry_tree(WIMStruct *w, struct dentry *root_dentry,
 
        if (!imd) {
                ERROR("Failed to allocate memory for new image metadata array");
-               return WIMLIB_ERR_NOMEM;
+               goto err;
        }
 
-       memcpy(imd, w->image_metadata, 
+       memcpy(imd, w->image_metadata,
               w->hdr.image_count * sizeof(struct image_metadata));
-       
+
        metadata_lte = new_lookup_table_entry();
        if (!metadata_lte)
-               goto out_free_imd;
+               goto err_free_imd;
 
        metadata_lte->resource_entry.flags = WIM_RESHDR_FLAG_METADATA;
        random_hash(metadata_lte->hash);
@@ -420,10 +484,9 @@ static int add_new_dentry_tree(WIMStruct *w, struct dentry *root_dentry,
        ret = select_wim_image(w, w->hdr.image_count);
        wimlib_assert(ret == 0);
        return ret;
-out_free_metadata_lte:
-       FREE(metadata_lte);
-out_free_imd:
+err_free_imd:
        FREE(imd);
+err:
        return WIMLIB_ERR_NOMEM;
 
 }
@@ -431,14 +494,15 @@ out_free_imd:
 /*
  * Copies an image, or all the images, from a WIM file, into another WIM file.
  */
-WIMLIBAPI int wimlib_export_image(WIMStruct *src_wim, 
-                                 int src_image, 
-                                 WIMStruct *dest_wim, 
-                                 const char *dest_name, 
-                                 const char *dest_description, 
-                                 int flags,
+WIMLIBAPI int wimlib_export_image(WIMStruct *src_wim,
+                                 int src_image,
+                                 WIMStruct *dest_wim,
+                                 const char *dest_name,
+                                 const char *dest_description,
+                                 int export_flags,
                                  WIMStruct **additional_swms,
-                                 unsigned num_additional_swms)
+                                 unsigned num_additional_swms,
+                                 wimlib_progress_func_t progress_func)
 {
        int i;
        int ret;
@@ -447,21 +511,18 @@ WIMLIBAPI int wimlib_export_image(WIMStruct *src_wim,
        struct wim_security_data *sd;
        struct lookup_table *joined_tab, *src_wim_tab_save;
 
-       if (!src_wim || !dest_wim)
-               return WIMLIB_ERR_INVALID_PARAM;
-
        if (dest_wim->hdr.total_parts != 1) {
                ERROR("Exporting an image to a split WIM is "
                      "unsupported");
                return WIMLIB_ERR_SPLIT_UNSUPPORTED;
        }
 
-       if (src_image == WIM_ALL_IMAGES) {
+       if (src_image == WIMLIB_ALL_IMAGES) {
                if (src_wim->hdr.image_count > 1) {
 
                        /* multi-image export. */
 
-                       if ((flags & WIMLIB_EXPORT_FLAG_BOOT) && 
+                       if ((export_flags & WIMLIB_EXPORT_FLAG_BOOT) &&
                              (src_wim->hdr.boot_idx == 0))
                        {
                                /* Specifying the boot flag on a multi-image
@@ -481,22 +542,25 @@ WIMLIBAPI int wimlib_export_image(WIMStruct *src_wim,
                                return WIMLIB_ERR_INVALID_PARAM;
                        }
                        for (i = 1; i <= src_wim->hdr.image_count; i++) {
-                               int export_flags = flags;
+                               int new_flags = export_flags;
 
                                if (i != src_wim->hdr.boot_idx)
-                                       export_flags &= ~WIMLIB_EXPORT_FLAG_BOOT;
+                                       new_flags &= ~WIMLIB_EXPORT_FLAG_BOOT;
 
-                               ret = wimlib_export_image(src_wim, i, dest_wim, 
+                               ret = wimlib_export_image(src_wim, i, dest_wim,
                                                          NULL, NULL,
-                                                         export_flags,
+                                                         new_flags,
                                                          additional_swms,
-                                                         num_additional_swms);
+                                                         num_additional_swms,
+                                                         progress_func);
                                if (ret != 0)
                                        return ret;
                        }
                        return 0;
+               } else if (src_wim->hdr.image_count == 1) {
+                       src_image = 1;
                } else {
-                       src_image = 1; 
+                       return 0;
                }
        }
 
@@ -576,9 +640,9 @@ WIMLIBAPI int wimlib_export_image(WIMStruct *src_wim,
        for_dentry_in_tree(root, add_lte_to_dest_wim, &wims);
        wimlib_assert(list_empty(&wims.lte_list_head));
 
-       if (flags & WIMLIB_EXPORT_FLAG_BOOT) {
+       if (export_flags & WIMLIB_EXPORT_FLAG_BOOT) {
                DEBUG("Setting boot_idx to %d", dest_wim->hdr.image_count);
-               dest_wim->hdr.boot_idx = dest_wim->hdr.image_count;
+               wimlib_set_boot_idx(dest_wim, dest_wim->hdr.image_count);
        }
        ret = 0;
        goto out;
@@ -588,7 +652,7 @@ out_xml_delete_image:
 out_free_ltes:
        {
                struct lookup_table_entry *lte, *tmp;
-               list_for_each_entry_safe(lte, tmp, &wims.lte_list_head, list)
+               list_for_each_entry_safe(lte, tmp, &wims.lte_list_head, staging_list)
                        free_lookup_table_entry(lte);
        }
 
@@ -600,8 +664,60 @@ out:
        return ret;
 }
 
-/* 
- * Deletes an image from the WIM. 
+static int image_run_full_verifications(WIMStruct *w)
+{
+       return for_dentry_in_tree(wim_root_dentry(w), verify_dentry, w);
+}
+
+static int lte_fix_refcnt(struct lookup_table_entry *lte, void *ctr)
+{
+       if (lte->refcnt != lte->real_refcnt) {
+               WARNING("The following lookup table entry has a reference "
+                       "count of %u, but", lte->refcnt);
+               WARNING("We found %u references to it",
+                       lte->real_refcnt);
+               print_lookup_table_entry(lte);
+               lte->refcnt = lte->real_refcnt;
+               ++*(unsigned long *)ctr;
+       }
+       return 0;
+}
+
+/* Ideally this would be unnecessary... however, the WIMs for Windows 8 are
+ * screwed up because some lookup table entries are referenced more times than
+ * their stated reference counts.  So theoretically, if we do the delete all the
+ * references to a stream and then remove it, it might still be referenced
+ * somewhere else... So, work around this problem by looking at ALL the images
+ * to re-calculate the reference count of EVERY lookup table entry. */
+int wim_run_full_verifications(WIMStruct *w)
+{
+       int ret;
+
+       for_lookup_table_entry(w->lookup_table, lte_zero_real_refcnt, NULL);
+       w->all_images_verified = true;
+       w->full_verification_in_progress = true;
+       ret = for_image(w, WIMLIB_ALL_IMAGES, image_run_full_verifications);
+       w->full_verification_in_progress = false;
+       if (ret == 0) {
+               unsigned long num_ltes_with_bogus_refcnt = 0;
+               for (int i = 0; i < w->hdr.image_count; i++)
+                       w->image_metadata[i].metadata_lte->real_refcnt++;
+               for_lookup_table_entry(w->lookup_table, lte_fix_refcnt,
+                                      &num_ltes_with_bogus_refcnt);
+               if (num_ltes_with_bogus_refcnt != 0) {
+                       WARNING("A total of %lu entries in the WIM's stream "
+                               "lookup table had to have\n"
+                               "          their reference counts fixed.",
+                               num_ltes_with_bogus_refcnt);
+               }
+       } else {
+               w->all_images_verified = false;
+       }
+       return ret;
+}
+
+/*
+ * Deletes an image from the WIM.
  */
 WIMLIBAPI int wimlib_delete_image(WIMStruct *w, int image)
 {
@@ -613,7 +729,7 @@ WIMLIBAPI int wimlib_delete_image(WIMStruct *w, int image)
                return WIMLIB_ERR_SPLIT_UNSUPPORTED;
        }
 
-       if (image == WIM_ALL_IMAGES) {
+       if (image == WIMLIB_ALL_IMAGES) {
                for (i = w->hdr.image_count; i >= 1; i--) {
                        ret = wimlib_delete_image(w, i);
                        if (ret != 0)
@@ -622,6 +738,12 @@ WIMLIBAPI int wimlib_delete_image(WIMStruct *w, int image)
                return 0;
        }
 
+       if (!w->all_images_verified) {
+               ret = wim_run_full_verifications(w);
+               if (ret != 0)
+                       return ret;
+       }
+
        DEBUG("Deleting image %d", image);
 
        /* Even if the dentry tree is not allocated, we must select it (and
@@ -651,10 +773,12 @@ WIMLIBAPI int wimlib_delete_image(WIMStruct *w, int image)
        else if (w->hdr.boot_idx > image)
                w->hdr.boot_idx--;
 
-       w->current_image = WIM_NO_IMAGE;
+       w->current_image = WIMLIB_NO_IMAGE;
 
        /* Remove the image from the XML information. */
        xml_delete_image(&w->wim_info, image);
+
+       w->deletion_occurred = true;
        return 0;
 }
 
@@ -742,7 +866,7 @@ static int init_capture_config(const char *_config_str, size_t config_len,
                FREE(config_str);
                return WIMLIB_ERR_NOMEM;
        }
-       
+
        memcpy(config_str, _config_str, config_len);
        next_p = config_str;
        config->config_str = config_str;
@@ -758,7 +882,7 @@ static int init_capture_config(const char *_config_str, size_t config_len,
                        ret = WIMLIB_ERR_INVALID_CAPTURE_CONFIG;
                        goto out_destroy;
                }
-               
+
                next_p = eol + 1;
                bytes_remaining -= (next_p - p);
                if (eol == p)
@@ -887,35 +1011,23 @@ bool exclude_path(const char *path, const struct capture_config *config,
                     && path[config->prefix_len] == '/')
                        path += config->prefix_len;
        }
-       return match_pattern(path, basename, &config->exclusion_list) && 
+       return match_pattern(path, basename, &config->exclusion_list) &&
                !match_pattern(path, basename, &config->exclusion_exception);
 
 }
 
-
-
-/*
- * Adds an image to the WIM, delegating the capture of the dentry tree and
- * security data to the function @capture_tree passed as a parameter.
- * Currently, @capture_tree may be build_dentry_tree() for capturing a "regular"
- * directory tree on disk, or build_dentry_tree_ntfs() for capturing a WIM image
- * directory from a NTFS volume using libntfs-3g.
- *
- * The @capture_tree function is also expected to create lookup table entries
- * for all the file streams it captures and insert them into @lookup_table,
- * being careful to look for identical entries that already exist and simply
- * increment the reference count for them rather than duplicating the entry.
- */
-int do_add_image(WIMStruct *w, const char *dir, const char *name,
-                const char *config_str, size_t config_len,
-                int flags,
-                int (*capture_tree)(struct dentry **, const char *,
-                                    struct lookup_table *, 
-                                    struct wim_security_data *,
-                                    const struct capture_config *,
-                                    int, void *),
-                void *extra_arg)
+WIMLIBAPI int wimlib_add_image(WIMStruct *w, const char *source,
+                              const char *name, const char *config_str,
+                              size_t config_len, int add_image_flags,
+                              wimlib_progress_func_t progress_func)
 {
+       int (*capture_tree)(struct dentry **, const char *,
+                           struct lookup_table *,
+                           struct wim_security_data *,
+                           const struct capture_config *,
+                           int, wimlib_progress_func_t, void *);
+       void *extra_arg;
+
        struct dentry *root_dentry = NULL;
        struct wim_security_data *sd;
        struct capture_config config;
@@ -923,13 +1035,31 @@ int do_add_image(WIMStruct *w, const char *dir, const char *name,
        struct hlist_head inode_list;
        int ret;
 
-       DEBUG("Adding dentry tree from directory or NTFS volume `%s'.", dir);
+       if (add_image_flags & WIMLIB_ADD_IMAGE_FLAG_NTFS) {
+#ifdef WITH_NTFS_3G
+               if (add_image_flags & (WIMLIB_ADD_IMAGE_FLAG_DEREFERENCE)) {
+                       ERROR("Cannot dereference files when capturing directly from NTFS");
+                       return WIMLIB_ERR_INVALID_PARAM;
+               }
+               capture_tree = build_dentry_tree_ntfs;
+               extra_arg = &w->ntfs_vol;
+#else
+               ERROR("wimlib was compiled without support for NTFS-3g, so\n"
+                     "        cannot capture a WIM image directly from a NTFS volume!");
+               return WIMLIB_ERR_UNSUPPORTED;
+#endif
+       } else {
+               capture_tree = build_dentry_tree;
+               extra_arg = NULL;
+       }
+
+       DEBUG("Adding dentry tree from directory or NTFS volume `%s'.", source);
 
        if (!name || !*name) {
                ERROR("Must specify a non-empty string for the image name");
                return WIMLIB_ERR_INVALID_PARAM;
        }
-       if (!dir) {
+       if (!source) {
                ERROR("Must specify the name of a directory or NTFS volume");
                return WIMLIB_ERR_INVALID_PARAM;
        }
@@ -951,7 +1081,7 @@ int do_add_image(WIMStruct *w, const char *dir, const char *name,
                config_str = default_config;
                config_len = strlen(default_config);
        }
-       ret = init_capture_config(config_str, config_len, dir, &config);
+       ret = init_capture_config(config_str, config_len, source, &config);
        if (ret != 0)
                return ret;
        print_capture_config(&config);
@@ -959,22 +1089,36 @@ int do_add_image(WIMStruct *w, const char *dir, const char *name,
        DEBUG("Allocating security data");
 
        sd = CALLOC(1, sizeof(struct wim_security_data));
-       if (!sd)
+       if (!sd) {
+               ret = WIMLIB_ERR_NOMEM;
                goto out_destroy_config;
+       }
        sd->total_length = 8;
        sd->refcnt = 1;
 
+       if (progress_func) {
+               union wimlib_progress_info progress;
+               progress.scan.source = source;
+               progress_func(WIMLIB_PROGRESS_MSG_SCAN_BEGIN, &progress);
+       }
+
        DEBUG("Building dentry tree.");
-       ret = (*capture_tree)(&root_dentry, dir, w->lookup_table, sd,
-                             &config, flags | WIMLIB_ADD_IMAGE_FLAG_ROOT,
-                             extra_arg);
+       ret = (*capture_tree)(&root_dentry, source, w->lookup_table, sd,
+                             &config, add_image_flags | WIMLIB_ADD_IMAGE_FLAG_ROOT,
+                             progress_func, extra_arg);
        destroy_capture_config(&config);
 
        if (ret != 0) {
-               ERROR("Failed to build dentry tree for `%s'", dir);
+               ERROR("Failed to build dentry tree for `%s'", source);
                goto out_free_security_data;
        }
 
+       if (progress_func) {
+               union wimlib_progress_info progress;
+               progress.scan.source = source;
+               progress_func(WIMLIB_PROGRESS_MSG_SCAN_END, &progress);
+       }
+
        DEBUG("Calculating full paths of dentries.");
        ret = for_dentry_in_tree(root_dentry, calculate_dentry_full_path, NULL);
        if (ret != 0)
@@ -1004,9 +1148,8 @@ int do_add_image(WIMStruct *w, const char *dir, const char *name,
        if (ret != 0)
                goto out_destroy_imd;
 
-       if (flags & WIMLIB_ADD_IMAGE_FLAG_BOOT)
+       if (add_image_flags & WIMLIB_ADD_IMAGE_FLAG_BOOT)
                wimlib_set_boot_idx(w, w->hdr.image_count);
-
        return 0;
 out_destroy_imd:
        destroy_image_metadata(&w->image_metadata[w->hdr.image_count - 1],
@@ -1021,14 +1164,3 @@ out_destroy_config:
        destroy_capture_config(&config);
        return ret;
 }
-
-/*
- * Adds an image to a WIM file from a directory tree on disk.
- */
-WIMLIBAPI int wimlib_add_image(WIMStruct *w, const char *dir, 
-                              const char *name, const char *config_str,
-                              size_t config_len, int flags)
-{
-       return do_add_image(w, dir, name, config_str, config_len, flags,
-                           build_dentry_tree, NULL);
-}