+ if (file_has_streams(file)) {
+ ret = validate_streams_and_compute_total_length(file, &n);
+ if (ret)
+ return ret;
+ inode_size += n;
+ }
+
+ /* To save memory, we allocate the ntfs_dentry's and ntfs_stream's in
+ * the same memory block as their ntfs_inode. */
+ ni = CALLOC(1, inode_size);
+ if (!ni)
+ return WIMLIB_ERR_NOMEM;
+
+ ni->ino = file->FileReferenceNumber;
+ ni->attributes = info->BasicInformation.FileAttributes;
+ ni->creation_time = info->BasicInformation.CreationTime;
+ ni->last_write_time = info->BasicInformation.LastWriteTime;
+ ni->last_access_time = info->BasicInformation.LastAccessTime;
+ ni->security_id = info->SecurityId;
+
+ p = FIRST_DENTRY(ni);
+
+ if (!NTFS_IS_ROOT_FILE(file->FileReferenceNumber))
+ p = load_name_information(file, ni, p);
+
+ if (file_has_streams(file))
+ p = load_stream_information(file, ni, p);
+
+ wimlib_assert((u8 *)p - (u8 *)ni == inode_size);
+
+ ntfs_inode_map_add_inode(inode_map, ni);
+ return 0;
+}
+
+/*
+ * Quickly find all files on an NTFS volume by using FSCTL_QUERY_FILE_LAYOUT to
+ * scan the MFT. The NTFS volume is specified by the NT namespace path @path.
+ * For each file, allocate an 'ntfs_inode' structure for each file and add it to
+ * 'inode_map' keyed by inode number. Include NTFS special files such as
+ * $Bitmap (they will be removed later).
+ */
+static int
+load_files_from_mft(const wchar_t *path, struct ntfs_inode_map *inode_map)
+{
+ HANDLE h = NULL;
+ QUERY_FILE_LAYOUT_INPUT in = (QUERY_FILE_LAYOUT_INPUT) {
+ .NumberOfPairs = 0,
+ .Flags = QUERY_FILE_LAYOUT_RESTART |
+ QUERY_FILE_LAYOUT_INCLUDE_NAMES |
+ QUERY_FILE_LAYOUT_INCLUDE_STREAMS |
+ QUERY_FILE_LAYOUT_INCLUDE_EXTENTS |
+ QUERY_FILE_LAYOUT_INCLUDE_EXTRA_INFO |
+ QUERY_FILE_LAYOUT_INCLUDE_STREAMS_WITH_NO_CLUSTERS_ALLOCATED,
+ .FilterType = QUERY_FILE_LAYOUT_FILTER_TYPE_NONE,
+ };
+ const size_t outsize = 32768;
+ QUERY_FILE_LAYOUT_OUTPUT *out = NULL;
+ int ret;
+ NTSTATUS status;
+
+ status = winnt_open(path, wcslen(path),
+ FILE_READ_DATA | FILE_READ_ATTRIBUTES, &h);
+ if (!NT_SUCCESS(status)) {
+ ret = -1; /* Silently try standard recursive scan instead */
+ goto out;
+ }
+
+ out = MALLOC(outsize);
+ if (!out) {
+ ret = WIMLIB_ERR_NOMEM;
+ goto out;
+ }
+
+ while (NT_SUCCESS(status = winnt_fsctl(h, FSCTL_QUERY_FILE_LAYOUT,
+ &in, sizeof(in),
+ out, outsize, NULL)))
+ {
+ const FILE_LAYOUT_ENTRY *file =
+ (const void *)out + out->FirstFileOffset;
+ for (;;) {
+ ret = load_one_file(file, inode_map);
+ if (ret)
+ goto out;
+ if (file->NextFileOffset == 0)
+ break;
+ file = (const void *)file + file->NextFileOffset;
+ }
+ in.Flags &= ~QUERY_FILE_LAYOUT_RESTART;
+ }
+
+ /* Normally, FSCTL_QUERY_FILE_LAYOUT fails with STATUS_END_OF_FILE after
+ * all files have been enumerated. */
+ if (status != STATUS_END_OF_FILE) {
+ if (status == STATUS_INVALID_DEVICE_REQUEST /* old OS */ ||
+ status == STATUS_INVALID_PARAMETER /* not root directory */ ) {
+ /* Silently try standard recursive scan instead */
+ ret = -1;
+ } else {
+ winnt_error(status,
+ L"Error enumerating files on volume \"%ls\"",
+ path);
+ /* Try standard recursive scan instead */
+ ret = WIMLIB_ERR_UNSUPPORTED;
+ }
+ goto out;
+ }
+ ret = 0;
+out:
+ FREE(out);
+ NtClose(h);
+ return ret;
+}
+
+/* Build the list of child dentries for each inode in @map. This is done by
+ * iterating through each name of each inode and adding it to its parent's
+ * children list. Note that every name should have a parent, i.e. should belong
+ * to some directory. The root directory does not have any names. */
+static int
+build_children_lists(struct ntfs_inode_map *map, struct ntfs_inode **root_ret)
+{
+ struct ntfs_inode *ni;
+
+ avl_tree_for_each_in_order(ni, map->root, struct ntfs_inode, index_node)
+ {
+ struct ntfs_dentry *nd;
+ u32 n;
+
+ if (NTFS_IS_ROOT_FILE(ni->ino)) {
+ *root_ret = ni;
+ continue;
+ }
+
+ n = ni->num_aliases;
+ nd = FIRST_DENTRY(ni);
+ for (;;) {
+ struct ntfs_inode *parent;
+
+ parent = ntfs_inode_map_lookup(map, nd->parent_ino);
+ if (unlikely(!parent)) {
+ ERROR("Parent inode 0x%016"PRIx64" of"
+ "directory entry \"%ls\" (inode "
+ "0x%016"PRIx64") was missing from the "
+ "MFT listing!",
+ nd->parent_ino, nd->name, ni->ino);
+ return WIMLIB_ERR_UNSUPPORTED;
+ }
+ nd->next_child = parent->first_child;
+ parent->first_child = nd;
+ if (!--n)
+ break;
+ nd = NEXT_DENTRY(nd);
+ }
+ }
+ return 0;
+}
+
+struct security_map_node {
+ struct avl_tree_node index_node;
+ u32 disk_security_id;
+ u32 wim_security_id;
+};
+
+/* 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);
+ 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;
+ }
+
+ root->num_aliases = 1;
+
+ ret = generate_wim_structures_recursive(root_ret, path, path_nchars,
+ L"", false, root, ctx,
+ &inode_map, &security_map);
+out:
+ ntfs_inode_map_destroy(&inode_map);
+ security_map_destroy(&security_map);
+ return ret;
+}
+
+#endif /* ENABLE_FAST_MFT_SCAN */
+
+/*----------------------------------------------------------------------------*
+ * Entry point for directory tree scans on Windows *
+ *----------------------------------------------------------------------------*/
+
+#define WINDOWS_NT_MAX_PATH 32768
+
+int
+win32_build_dentry_tree(struct wim_dentry **root_ret,
+ const wchar_t *root_disk_path,
+ struct capture_params *params)
+{
+ wchar_t *path = NULL;
+ struct winnt_scan_ctx ctx = { .params = params };
+ UNICODE_STRING ntpath;
+ size_t ntpath_nchars;
+ HANDLE h = NULL;
+ NTSTATUS status;
+ int ret;