+/* Map from disk security IDs to WIM security IDs */
+struct security_map {
+ struct avl_tree_node *root;
+};
+
+#define SECURITY_MAP_NODE(node) \
+ avl_tree_entry((node), struct security_map_node, index_node)
+
+static int
+_avl_cmp_security_map_nodes(const struct avl_tree_node *node1,
+ const struct avl_tree_node *node2)
+{
+ return cmp_u32(SECURITY_MAP_NODE(node1)->disk_security_id,
+ SECURITY_MAP_NODE(node2)->disk_security_id);
+}
+
+static s32
+security_map_lookup(struct security_map *map, u32 disk_security_id)
+{
+ struct security_map_node tmp;
+ const struct avl_tree_node *res;
+
+ if (disk_security_id == 0) /* No on-disk security ID; uncacheable */
+ return -1;
+
+ tmp.disk_security_id = disk_security_id;
+ res = avl_tree_lookup_node(map->root, &tmp.index_node,
+ _avl_cmp_security_map_nodes);
+ if (!res)
+ return -1;
+ return SECURITY_MAP_NODE(res)->wim_security_id;
+}
+
+static int
+security_map_insert(struct security_map *map, u32 disk_security_id,
+ u32 wim_security_id)
+{
+ struct security_map_node *node;
+
+ if (disk_security_id == 0) /* No on-disk security ID; uncacheable */
+ return 0;
+
+ node = MALLOC(sizeof(*node));
+ if (!node)
+ return WIMLIB_ERR_NOMEM;
+
+ node->disk_security_id = disk_security_id;
+ node->wim_security_id = wim_security_id;
+ avl_tree_insert(&map->root, &node->index_node,
+ _avl_cmp_security_map_nodes);
+ return 0;
+}
+
+static void
+security_map_destroy(struct security_map *map)
+{
+ struct security_map_node *node;
+
+ avl_tree_for_each_in_postorder(node, map->root,
+ struct security_map_node, index_node)
+ FREE(node);
+}
+
+/*
+ * Turn our temporary NTFS structures into the final WIM structures:
+ *
+ * ntfs_inode => wim_inode
+ * ntfs_dentry => wim_dentry
+ * ntfs_stream => wim_inode_stream
+ *
+ * This also handles things such as exclusions and issuing progress messages.
+ * It's similar to winnt_build_dentry_tree_recursive(), but this is much faster
+ * because almost all information we need is already loaded in memory in the
+ * ntfs_* structures. However, in some cases we still fall back to
+ * winnt_build_dentry_tree_recursive() and/or opening the file.
+ */
+static int
+generate_wim_structures_recursive(struct wim_dentry **root_ret,
+ wchar_t *path, size_t path_nchars,
+ const wchar_t *filename, bool is_primary_name,
+ struct ntfs_inode *ni,
+ struct winnt_scan_ctx *ctx,
+ struct ntfs_inode_map *inode_map,
+ struct security_map *security_map)
+{
+ int ret = 0;
+ struct wim_dentry *root = NULL;
+ struct wim_inode *inode = NULL;
+ const struct ntfs_stream *ns;
+
+ /* Completely ignore NTFS special files. */
+ if (NTFS_IS_SPECIAL_FILE(ni->ino))
+ goto out;
+
+ /* Fall back to a recursive scan for unhandled cases. Reparse points,
+ * in particular, can't be properly handled here because a commonly used
+ * filter driver (WOF) hides reparse points from regular filesystem APIs
+ * but not from FSCTL_QUERY_FILE_LAYOUT. */
+ if (ni->attributes & (FILE_ATTRIBUTE_REPARSE_POINT |
+ FILE_ATTRIBUTE_ENCRYPTED))
+ {
+ ret = winnt_build_dentry_tree_recursive(&root,
+ NULL,
+ path,
+ path_nchars,
+ path,
+ path_nchars,
+ filename,
+ ctx);
+ goto out;
+ }
+
+ /* Test for exclusion based on path. */
+ ret = try_exclude(path, ctx->params);
+ if (unlikely(ret < 0)) /* Excluded? */
+ goto out_progress;
+ if (unlikely(ret > 0)) /* Error? */
+ goto out;
+
+ /* Create the WIM dentry and possibly a new WIM inode */
+ ret = inode_table_new_dentry(ctx->params->inode_table, filename,
+ ni->ino, ctx->params->capture_root_dev,
+ false, &root);
+ if (ret)
+ goto out;
+
+ inode = root->d_inode;
+
+ /* Set the short name if needed. */
+ if (is_primary_name && *ni->short_name) {
+ size_t nbytes = wcslen(ni->short_name) * sizeof(wchar_t);
+ root->d_short_name = memdup(ni->short_name,
+ nbytes + sizeof(wchar_t));
+ if (!root->d_short_name) {
+ ret = WIMLIB_ERR_NOMEM;
+ goto out;
+ }
+ root->d_short_name_nbytes = nbytes;
+ }
+
+ if (inode->i_nlink > 1) { /* Already seen this inode? */
+ ret = 0;
+ goto out_progress;
+ }
+
+ /* The file attributes and timestamps were cached from the MFT. */
+ inode->i_attributes = ni->attributes;
+ inode->i_creation_time = ni->creation_time;
+ inode->i_last_write_time = ni->last_write_time;
+ inode->i_last_access_time = ni->last_access_time;
+
+ /* Set the security descriptor if needed. */
+ if (!(ctx->params->add_flags & WIMLIB_ADD_FLAG_NO_ACLS)) {
+ /* Look up the WIM security ID that corresponds to the on-disk
+ * security ID. */
+ s32 wim_security_id =
+ security_map_lookup(security_map, ni->security_id);
+ if (likely(wim_security_id >= 0)) {
+ /* The mapping for this security ID is already cached.*/
+ inode->i_security_id = wim_security_id;
+ } else {
+ HANDLE h;
+ NTSTATUS status;
+
+ /* Create a mapping for this security ID and insert it
+ * into the security map. */
+
+ status = winnt_open(path, path_nchars,
+ READ_CONTROL |
+ ACCESS_SYSTEM_SECURITY, &h);
+ if (!NT_SUCCESS(status)) {
+ winnt_error(status, L"Can't open \"%ls\" to "
+ "read security descriptor",
+ printable_path(path));
+ ret = WIMLIB_ERR_OPEN;
+ goto out;
+ }
+ ret = winnt_load_security_descriptor(h, inode, path, ctx);
+ (*func_NtClose)(h);
+ if (ret)
+ goto out;
+
+ ret = security_map_insert(security_map, ni->security_id,
+ inode->i_security_id);
+ if (ret)
+ goto out;
+ }
+ }
+
+ /* Add data streams based on the cached information from the MFT. */
+ ns = FIRST_STREAM(ni);
+ for (u32 i = 0; i < ni->num_streams; i++) {
+ struct windows_file *windows_file;
+
+ /* Reference the stream by path if it's a named data stream, or
+ * if the volume doesn't support "open by file ID", or if the
+ * application hasn't explicitly opted in to "open by file ID".
+ * Otherwise, only save the inode number (file ID). */
+ if (*ns->name ||
+ !(ctx->vol_flags & FILE_SUPPORTS_OPEN_BY_FILE_ID) ||
+ !(ctx->params->add_flags & WIMLIB_ADD_FLAG_FILE_PATHS_UNNEEDED))
+ {
+ windows_file = alloc_windows_file(path,
+ path_nchars,
+ ns->name,
+ wcslen(ns->name),
+ ctx->snapshot,
+ false);
+ } else {
+ windows_file = alloc_windows_file_for_file_id(ni->ino,
+ path,
+ ctx->params->capture_root_nchars + 1,
+ ctx->snapshot);
+ }
+
+ ret = add_stream(inode, windows_file, ns->size,
+ STREAM_TYPE_DATA, ns->name,
+ ctx->params->unhashed_blobs);
+ if (ret)
+ goto out;
+ ns = NEXT_STREAM(ns);
+ }
+
+ set_sort_key(inode, ni->starting_lcn);
+
+ /* If processing a directory, then recurse to its children. In this
+ * version there is no need to go to disk, as we already have the list
+ * of children cached from the MFT. */
+ if (inode_is_directory(inode)) {
+ const struct ntfs_dentry *nd = ni->first_child;
+
+ while (nd != NULL) {
+ const size_t name_len = wcslen(nd->name);
+ wchar_t *p = path + path_nchars;
+ struct wim_dentry *child;
+ const struct ntfs_dentry *next = nd->next_child;
+
+ if (*(p - 1) != L'\\')
+ *p++ = L'\\';
+ p = wmempcpy(p, nd->name, name_len);
+ *p = '\0';
+
+ ret = generate_wim_structures_recursive(
+ &child,
+ path,
+ p - path,
+ nd->name,
+ nd->is_primary,
+ (void *)nd - nd->offset_from_inode,
+ ctx,
+ inode_map,
+ security_map);
+
+ path[path_nchars] = L'\0';
+
+ if (ret)
+ goto out;
+
+ attach_scanned_tree(root, child, ctx->params->blob_table);
+ nd = next;
+ }
+ }
+
+out_progress:
+ ctx->params->progress.scan.cur_path = path;
+ if (likely(root))
+ ret = do_capture_progress(ctx->params, WIMLIB_SCAN_DENTRY_OK, inode);
+ else
+ ret = do_capture_progress(ctx->params, WIMLIB_SCAN_DENTRY_EXCLUDED, NULL);
+out:
+ if (--ni->num_aliases == 0) {
+ /* Memory usage optimization: when we don't need the ntfs_inode
+ * (and its names and streams) anymore, free it. */
+ ntfs_inode_map_remove(inode_map, ni);
+ }
+ if (unlikely(ret)) {
+ free_dentry_tree(root, ctx->params->blob_table);
+ root = NULL;
+ }
+ *root_ret = root;
+ return ret;
+}
+
+static int
+winnt_build_dentry_tree_fast(struct wim_dentry **root_ret, wchar_t *path,
+ size_t path_nchars, struct winnt_scan_ctx *ctx)
+{
+ struct ntfs_inode_map inode_map = { .root = NULL };
+ struct security_map security_map = { .root = NULL };
+ struct ntfs_inode *root = NULL;
+ bool adjust_path;
+ int ret;
+
+ adjust_path = (path[path_nchars - 1] == L'\\');
+ if (adjust_path)
+ path[path_nchars - 1] = L'\0';
+
+ ret = load_files_from_mft(path, &inode_map);
+
+ if (adjust_path)
+ path[path_nchars - 1] = L'\\';
+
+ if (ret)
+ goto out;
+
+ ret = build_children_lists(&inode_map, &root);
+ if (ret)
+ goto out;
+
+ if (!root) {
+ ERROR("The MFT listing for volume \"%ls\" did not include a "
+ "root directory!", path);
+ ret = WIMLIB_ERR_UNSUPPORTED;
+ goto out;