]> wimlib.net Git - wimlib/blobdiff - src/ntfs-capture.c
ntfs capture: Store security descriptors in rbtree
[wimlib] / src / ntfs-capture.c
index 69123f0672c03540976a4f5b399d5fe4d89cda5d..a257cdb69d600041e155b432afb3260d8a048db0 100644 (file)
 #include <stdlib.h>
 #include <unistd.h>
 #include <errno.h>
+#include "rbtree.h"
 
 /* Structure that allows searching the security descriptors by SHA1 message
  * digest. */
 struct sd_set {
        struct wim_security_data *sd;
-       struct sd_node *root;
+       struct rb_root rb_root;
 };
 
 /* Binary tree node of security descriptors, indexed by the @hash field. */
 struct sd_node {
        int security_id;
        u8 hash[SHA1_HASH_SIZE];
-       struct sd_node *left;
-       struct sd_node *right;
+       struct rb_node rb_node;
 };
 
-static void free_sd_tree(struct sd_node *root)
+static void free_sd_tree(struct rb_node *node)
 {
-       if (root) {
-               free_sd_tree(root->left);
-               free_sd_tree(root->right);
-               FREE(root);
+       if (node) {
+               free_sd_tree(node->rb_left);
+               free_sd_tree(node->rb_right);
+               FREE(container_of(node, struct sd_node, rb_node));
        }
 }
 /* Frees a security descriptor index set. */
 static void destroy_sd_set(struct sd_set *sd_set)
 {
-       free_sd_tree(sd_set->root);
+       free_sd_tree(sd_set->rb_root.rb_node);
 }
 
 /* Inserts a a new node into the security descriptor index tree. */
-static void insert_sd_node(struct sd_node *new, struct sd_node *root)
+static void insert_sd_node(struct sd_set *set, struct sd_node *new)
 {
-       int cmp = hashes_cmp(new->hash, root->hash);
-       if (cmp < 0) {
-               if (root->left)
-                       insert_sd_node(new, root->left);
+       struct rb_root *root = &set->rb_root;
+       struct rb_node **p = &(root->rb_node);
+       struct rb_node *rb_parent = NULL;
+
+       while (*p) {
+               struct sd_node *this = container_of(*p, struct sd_node, rb_node);
+               int cmp = hashes_cmp(new->hash, this->hash);
+
+               rb_parent = *p;
+               if (cmp < 0)
+                       p = &((*p)->rb_left);
+               else if (cmp > 0)
+                       p = &((*p)->rb_right);
                else
-                       root->left = new;
-       } else if (cmp > 0) {
-               if (root->right)
-                       insert_sd_node(new, root->right);
-               else
-                       root->right = new;
-       } else {
-               wimlib_assert(0);
+                       wimlib_assert(0); /* Duplicate SHA1 message digest */
        }
+       rb_link_node(&new->rb_node, rb_parent, p);
+       rb_insert_color(&new->rb_node, root);
 }
 
-/* Returns the security ID of the security data having a SHA1 message digest of
- * @hash in the security descriptor index tree rooted at @root.
- *
- * If not found, return -1. */
-static int lookup_sd(const u8 hash[SHA1_HASH_SIZE], struct sd_node *root)
+/* Returns the index of the security descriptor having a SHA1 message digest of
+ * @hash.  If not found, return -1. */
+static int lookup_sd(struct sd_set *set, const u8 hash[SHA1_HASH_SIZE])
 {
-       int cmp;
-       if (!root)
-               return -1;
-       cmp = hashes_cmp(hash, root->hash);
-       if (cmp < 0)
-               return lookup_sd(hash, root->left);
-       else if (cmp > 0)
-               return lookup_sd(hash, root->right);
-       else
-               return root->security_id;
+       struct rb_node *node = set->rb_root.rb_node;
+
+       while (node) {
+               struct sd_node *sd_node = container_of(node, struct sd_node, rb_node);
+               int cmp = hashes_cmp(hash, sd_node->hash);
+               if (cmp < 0)
+                       node = node->rb_left;
+               else if (cmp > 0)
+                       node = node->rb_right;
+               else
+                       return sd_node->security_id;
+       }
+       return -1;
 }
 
 /*
@@ -134,10 +139,11 @@ static int sd_set_add_sd(struct sd_set *sd_set, const char descriptor[],
 
        sha1_buffer((const u8*)descriptor, size, hash);
 
-       security_id = lookup_sd(hash, sd_set->root);
-       if (security_id >= 0)
+       security_id = lookup_sd(sd_set, hash);
+       if (security_id >= 0) /* Identical descriptor already exists */
                return security_id;
 
+       /* Need to add a new security descriptor */
        new = MALLOC(sizeof(*new));
        if (!new)
                goto out;
@@ -149,11 +155,8 @@ static int sd_set_add_sd(struct sd_set *sd_set, const char descriptor[],
 
        memcpy(descr_copy, descriptor, size);
        new->security_id = sd->num_entries;
-       new->left = NULL;
-       new->right = NULL;
        copy_hash(new->hash, hash);
 
-
        descriptors = REALLOC(sd->descriptors,
                              (sd->num_entries + 1) * sizeof(sd->descriptors[0]));
        if (!descriptors)
@@ -169,11 +172,7 @@ static int sd_set_add_sd(struct sd_set *sd_set, const char descriptor[],
        sd->num_entries++;
        DEBUG("There are now %d security descriptors", sd->num_entries);
        sd->total_length += size + sizeof(sd->sizes[0]);
-
-       if (sd_set->root)
-               insert_sd_node(new, sd_set->root);
-       else
-               sd_set->root = new;
+       insert_sd_node(sd_set, new);
        return new->security_id;
 out_free_descr:
        FREE(descr_copy);
@@ -340,7 +339,7 @@ static int capture_ntfs_streams(struct wim_dentry *dentry, ntfs_inode *ni,
                                        lte->resource_entry.size = data_size;
                                }
                                ntfs_loc = NULL;
-                               DEBUG("Add resource for `%s' (size = %zu)",
+                               DEBUG("Add resource for `%s' (size = %"PRIu64")",
                                      dentry->file_name_utf8,
                                      lte->resource_entry.original_size);
                                copy_hash(lte->hash, attr_hash);
@@ -360,7 +359,7 @@ static int capture_ntfs_streams(struct wim_dentry *dentry, ntfs_inode *ni,
                } else {
                        /* Named data stream.  Put the reference to it in the
                         * alternate data stream entries */
-                       struct ads_entry *new_ads_entry;
+                       struct wim_ads_entry *new_ads_entry;
                        size_t stream_name_utf8_len;
 
                        ret = utf16_to_utf8((const char*)attr_record_name(actx->attr),
@@ -575,13 +574,16 @@ static int build_dentry_tree_ntfs_recursive(struct wim_dentry **root_p,
                        if (ret != 0)
                                return ret;
                } else {
+                       if (
                #ifdef ENODATA
-                       if (errno != ENODATA) {
+                           errno != ENODATA &&
+               #endif
+                           errno != EMLINK
+                           ) {
                                ERROR_WITH_ERRNO("Error getting DOS name "
                                                 "of `%s'", path);
                                return WIMLIB_ERR_NTFS_3G;
                        }
-               #endif
                }
        }
 
@@ -669,16 +671,22 @@ int build_dentry_tree_ntfs(struct wim_dentry **root_p,
 {
        ntfs_volume *vol;
        ntfs_inode *root_ni;
-       int ret = 0;
+       int ret;
        struct sd_set sd_set = {
                .sd = sd,
-               .root = NULL,
+               .rb_root = {NULL},
        };
        ntfs_volume **ntfs_vol_p = extra_arg;
 
        DEBUG("Mounting NTFS volume `%s' read-only", device);
 
+#ifdef HAVE_NTFS_MNT_RDONLY
+       /* NTFS-3g 2013 */
+       vol = ntfs_mount(device, NTFS_MNT_RDONLY);
+#else
+       /* NTFS-3g 2011, 2012 */
        vol = ntfs_mount(device, MS_RDONLY);
+#endif
        if (!vol) {
                ERROR_WITH_ERRNO("Failed to mount NTFS volume `%s' read-only",
                                 device);
@@ -705,6 +713,7 @@ int build_dentry_tree_ntfs(struct wim_dentry **root_p,
        char *path = MALLOC(32768);
        if (!path) {
                ERROR("Could not allocate memory for NTFS pathname");
+               ret = WIMLIB_ERR_NOMEM;
                goto out_cleanup;
        }