Refactor headers

[wimlib] / src / security.c

diff --git a/src/security.c b/src/security.c
index 37b9a063cfe5bec6d4793fe2586a54427f5654fc..0d0aa91ab3e33deb03f0f216637742880648a5f1 100644 (file)
--- a/src/security.c
+++ b/src/security.c
@@ -23,9 +23,16 @@
  * along with wimlib; if not, see http://www.gnu.org/licenses/.
  */
 
-#include "wimlib_internal.h"
-#include "buffer_io.h"
-#include "security.h"
+#ifdef HAVE_CONFIG_H
+#  include "config.h"
+#endif
+
+#include "wimlib/assert.h"
+#include "wimlib/buffer_io.h"
+#include "wimlib/error.h"
+#include "wimlib/security.h"
+#include "wimlib/sha1.h"
+#include "wimlib/util.h"
 
 /* At the start of each type of access control entry.  */
 typedef struct {
@@ -149,6 +156,12 @@ empty_sacl_fixup(u8 *descr, u64 *size_p)
 #endif
 }
 
+struct wim_security_data *
+new_wim_security_data(void)
+{
+       return CALLOC(1, sizeof(struct wim_security_data));
+}
+
 /*
  * Reads the security data from the metadata resource.
  *
@@ -186,7 +199,6 @@ read_security_data(const u8 metadata_resource[], u64 metadata_resource_len,
        }
        sd->sizes       = NULL;
        sd->descriptors = NULL;
-       sd->refcnt      = 1;
 
        p = metadata_resource;
        p = get_u32(p, &sd->total_length);
@@ -437,17 +449,14 @@ void
 free_security_data(struct wim_security_data *sd)
 {
        if (sd) {
-               wimlib_assert(sd->refcnt != 0);
-               if (--sd->refcnt == 0) {
-                       u8 **descriptors = sd->descriptors;
-                       u32 num_entries  = sd->num_entries;
-                       if (descriptors)
-                               while (num_entries--)
-                                       FREE(*descriptors++);
-                       FREE(sd->sizes);
-                       FREE(sd->descriptors);
-                       FREE(sd);
-               }
+               u8 **descriptors = sd->descriptors;
+               u32 num_entries  = sd->num_entries;
+               if (descriptors)
+                       while (num_entries--)
+                               FREE(*descriptors++);
+               FREE(sd->sizes);
+               FREE(sd->descriptors);
+               FREE(sd);
        }
 }
 
@@ -469,14 +478,20 @@ free_sd_tree(struct rb_node *node)
 
 /* Frees a security descriptor index set. */
 void
-destroy_sd_set(struct sd_set *sd_set)
+destroy_sd_set(struct wim_sd_set *sd_set, bool rollback)
 {
+       if (rollback) {
+               struct wim_security_data *sd = sd_set->sd;
+               for (s32 i = sd_set->orig_num_entries; i < sd->num_entries; i++)
+                       FREE(sd->descriptors[i]);
+               sd->num_entries = sd_set->orig_num_entries;
+       }
        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_set *set, struct sd_node *new)
+static bool
+insert_sd_node(struct wim_sd_set *set, struct sd_node *new)
 {
        struct rb_root *root = &set->rb_root;
        struct rb_node **p = &(root->rb_node);
@@ -492,16 +507,17 @@ insert_sd_node(struct sd_set *set, struct sd_node *new)
                else if (cmp > 0)
                        p = &((*p)->rb_right);
                else
-                       wimlib_assert(0); /* Duplicate SHA1 message digest */
+                       return false; /* Duplicate security descriptor */
        }
        rb_link_node(&new->rb_node, rb_parent, p);
        rb_insert_color(&new->rb_node, root);
+       return true;
 }
 
 /* Returns the index of the security descriptor having a SHA1 message digest of
  * @hash.  If not found, return -1. */
 int
-lookup_sd(struct sd_set *set, const u8 hash[SHA1_HASH_SIZE])
+lookup_sd(struct wim_sd_set *set, const u8 hash[SHA1_HASH_SIZE])
 {
        struct rb_node *node = set->rb_root.rb_node;
 
@@ -526,7 +542,7 @@ lookup_sd(struct sd_set *set, const u8 hash[SHA1_HASH_SIZE])
  * return -1.
  */
 int
-sd_set_add_sd(struct sd_set *sd_set, const char descriptor[], size_t size)
+sd_set_add_sd(struct wim_sd_set *sd_set, const char descriptor[], size_t size)
 {
        u8 hash[SHA1_HASH_SIZE];
        int security_id;
@@ -535,6 +551,7 @@ sd_set_add_sd(struct sd_set *sd_set, const char descriptor[], size_t size)
        u64 *sizes;
        u8 *descr_copy;
        struct wim_security_data *sd;
+       bool bret;
 
        sha1_buffer((const u8*)descriptor, size, hash);
 
@@ -556,6 +573,9 @@ sd_set_add_sd(struct sd_set *sd_set, const char descriptor[], size_t size)
        new->security_id = sd->num_entries;
        copy_hash(new->hash, hash);
 
+       /* There typically are only a few dozen security descriptors in a
+        * directory tree, so expanding the array of security descriptors by
+        * only 1 extra space each time should not be a problem. */
        descriptors = REALLOC(sd->descriptors,
                              (sd->num_entries + 1) * sizeof(sd->descriptors[0]));
        if (!descriptors)
@@ -571,7 +591,8 @@ sd_set_add_sd(struct sd_set *sd_set, const char descriptor[], size_t size)
        sd->num_entries++;
        DEBUG("There are now %d security descriptors", sd->num_entries);
        sd->total_length += size + sizeof(sd->sizes[0]);
-       insert_sd_node(sd_set, new);
+       bret = insert_sd_node(sd_set, new);
+       wimlib_assert(bret);
        return new->security_id;
 out_free_descr:
        FREE(descr_copy);
@@ -580,3 +601,38 @@ out_free_node:
 out:
        return -1;
 }
+
+/* Initialize a `struct sd_set' mapping from SHA1 message digests of security
+ * descriptors to indices into the security descriptors table of the WIM image
+ * (security IDs).  */
+int
+init_sd_set(struct wim_sd_set *sd_set, struct wim_security_data *sd)
+{
+       int ret;
+
+       sd_set->sd = sd;
+       sd_set->rb_root.rb_node = NULL;
+
+       /* Remember the original number of security descriptors so that newly
+        * added ones can be rolled back if needed. */
+       sd_set->orig_num_entries = sd->num_entries;
+       for (s32 i = 0; i < sd->num_entries; i++) {
+               struct sd_node *new;
+
+               new = MALLOC(sizeof(struct sd_node));
+               if (!new) {
+                       ret = WIMLIB_ERR_NOMEM;
+                       goto out_destroy_sd_set;
+               }
+               sha1_buffer(sd->descriptors[i], sd->sizes[i], new->hash);
+               new->security_id = i;
+               if (!insert_sd_node(sd_set, new))
+                       FREE(new); /* Ignore duplicate security descriptor */
+       }
+       ret = 0;
+       goto out;
+out_destroy_sd_set:
+       destroy_sd_set(sd_set, false);
+out:
+       return ret;
+}