]> wimlib.net Git - wimlib/blobdiff - src/security.c
Remove buffer_io.h
[wimlib] / src / security.c
index 089067dc433f5152cd6cd4fd15261744d56eb48d..628b93584d1dfff77727ecc9f845c85c2b6a6152 100644 (file)
 /*
  * security.c
  *
- * Read the security data from the WIM.  Doing anything with the security data
- * is not yet implemented other than printing some information about it.
+ * Read and write the per-WIM-image table of security descriptors.
  */
 
 /*
- * Copyright (C) 2012 Eric Biggers
+ * Copyright (C) 2012, 2013 Eric Biggers
  *
  * This file is part of wimlib, a library for working with WIM files.
  *
  * wimlib is free software; you can redistribute it and/or modify it under the
- * terms of the GNU Lesser General Public License as published by the Free
- * Software Foundation; either version 2.1 of the License, or (at your option)
+ * terms of the GNU General Public License as published by the Free
+ * Software Foundation; either version 3 of the License, or (at your option)
  * any later version.
  *
  * wimlib is distributed in the hope that it will be useful, but WITHOUT ANY
  * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
- * A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
+ * A PARTICULAR PURPOSE. See the GNU General Public License for more
  * details.
  *
- * You should have received a copy of the GNU Lesser General Public License
+ * You should have received a copy of the GNU General Public License
  * along with wimlib; if not, see http://www.gnu.org/licenses/.
  */
 
-#include "wimlib_internal.h"
-#include "io.h"
-#include "security.h"
+#ifdef HAVE_CONFIG_H
+#  include "config.h"
+#endif
+
+#include "wimlib/assert.h"
+#include "wimlib/endianness.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 _ACE_HEADER {
+       /* enum ace_type, specifies what type of ACE this is.  */
+       u8 type;
+
+       /* bitwise OR of the inherit ACE flags #defined above */
+       u8 flags;
+
+       /* Size of the access control entry. */
+       le16 size;
+} _packed_attribute ACE_HEADER;
+
+/* Grants rights to a user or group */
+typedef struct _ACCESS_ALLOWED_ACE {
+       ACE_HEADER hdr;
+       le32 mask;
+       le32 sid_start;
+} _packed_attribute ACCESS_ALLOWED_ACE;
+
+/* Denies rights to a user or group */
+typedef struct _ACCESS_DENIED_ACE {
+       ACE_HEADER hdr;
+       le32 mask;
+       le32 sid_start;
+} _packed_attribute ACCESS_DENIED_ACE;
+
+typedef struct _SYSTEM_AUDIT_ACE {
+       ACE_HEADER hdr;
+       le32 mask;
+       le32 sid_start;
+} _packed_attribute SYSTEM_AUDIT_ACE;
+
+
+/* Header of an access control list. */
+typedef struct _ACL {
+       /* ACL_REVISION or ACL_REVISION_DS */
+       u8 revision;
+
+       /* padding */
+       u8 sbz1;
+
+       /* Total size of the ACL, including all access control entries */
+       le16 acl_size;
+
+       /* Number of access control entry structures that follow the ACL
+        * structure. */
+       le16 ace_count;
+
+       /* padding */
+       le16 sbz2;
+} _packed_attribute ACL;
+
+/* A structure used to identify users or groups. */
+typedef struct _SID {
+
+       /* example: 0x1 */
+       u8  revision;
+       u8  sub_authority_count;
+
+       /* Identifies the authority that issued the SID.  Can be, but does not
+        * have to be, one of enum sid_authority_value */
+       u8  identifier_authority[6];
+
+       le32 sub_authority[];
+} _packed_attribute SID;
+
+typedef struct _SECURITY_DESCRIPTOR_RELATIVE  {
+       /* Example: 0x1 */
+       u8 revision;
+       /* Example: 0x0 */
+       u8 sbz1;
+
+       /* Example: 0x4149 */
+       le16 security_descriptor_control;
+
+       /* Offset of a SID structure in the security descriptor. */
+       /* Example: 0x14 */
+       le32 owner_offset;
+
+       /* Offset of a SID structure in the security descriptor. */
+       /* Example: 0x24 */
+       le32 group_offset;
+
+       /* Offset of an ACL structure in the security descriptor. */
+       /* System ACL. */
+       /* Example: 0x00 */
+       le32 sacl_offset;
+
+       /* Offset of an ACL structure in the security descriptor. */
+       /* Discretionary ACL. */
+       /* Example: 0x34 */
+       le32 dacl_offset;
+} _packed_attribute SECURITY_DESCRIPTOR_RELATIVE;
+
+struct wim_security_data_disk {
+       le32 total_length;
+       le32 num_entries;
+       le64 sizes[];
+} _packed_attribute;
+
+/*
+ * This is a hack to work around a problem in libntfs-3g.  libntfs-3g validates
+ * security descriptors with a function named ntfs_valid_descr().
+ * ntfs_valid_descr() considers a security descriptor that ends in a SACL
+ * (Sysetm Access Control List) with no ACE's (Access Control Entries) to be
+ * invalid.  However, a security descriptor like this exists in the Windows 7
+ * install.wim.  Here, security descriptors matching this pattern are modified
+ * to have no SACL.  This should make no difference since the SACL had no
+ * entries anyway; however this ensures that that the security descriptors pass
+ * the validation in libntfs-3g.
+ */
+static void
+empty_sacl_fixup(SECURITY_DESCRIPTOR_RELATIVE *descr, size_t *size_p)
+{
+       /* No-op if no NTFS-3g support, or if NTFS-3g is version 2013 or later
+        * */
+#if defined(WITH_NTFS_3G) && !defined(HAVE_NTFS_MNT_RDONLY)
+       if (*size_p >= sizeof(SECURITY_DESCRIPTOR_RELATIVE)) {
+               u32 sacl_offset = le32_to_cpu(descr->sacl_offset);
+               if (sacl_offset == *size_p - sizeof(ACL)) {
+                       descr->sacl_offset = cpu_to_le32(0);
+                       *size_p -= sizeof(ACL);
+               }
+       }
+#endif
+}
 
-#ifdef ENABLE_SECURITY_DATA
+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.
+/*
+ * Reads the security data from the metadata resource of a WIM image.
  *
  * @metadata_resource: An array that contains the uncompressed metadata
- *                             resource for the WIM file.
- * @metadata_resource_len:     The length of @metadata_resource.
- * @sd_p:      A pointer to a pointer wim_security_data structure that will be filled
- *             in with a pointer to a new wim_security_data structure on success.
+ *                             resource for the WIM image.
+ * @metadata_resource_len:     The length of @metadata_resource.  It must be at
+ *                             least 8 bytes.
+ * @sd_ret:    A pointer to a pointer to a wim_security_data structure that
+ *             will be filled in with a pointer to a new wim_security_data
+ *             structure containing the security data on success.
  *
  * Note: There is no `offset' argument because the security data is located at
  * the beginning of the metadata resource.
+ *
+ * Possible errors include:
+ *     WIMLIB_ERR_NOMEM
+ *     WIMLIB_ERR_INVALID_SECURITY_DATA
  */
-int read_security_data(const u8 metadata_resource[], 
-                      u64 metadata_resource_len, struct wim_security_data **sd_p)
+int
+read_wim_security_data(const u8 metadata_resource[], size_t metadata_resource_len,
+                      struct wim_security_data **sd_ret)
 {
        struct wim_security_data *sd;
-       const u8 *p;
+       int ret;
+       u64 total_len;
        u64 sizes_size;
+       u64 size_no_descriptors;
+       const struct wim_security_data_disk *sd_disk;
+       const u8 *p;
 
-       if (metadata_resource_len < 8) {
-               ERROR("Not enough space in %"PRIu64"-byte file resource for "
-                               "security data!\n", metadata_resource_len);
-               return WIMLIB_ERR_INVALID_RESOURCE_SIZE;
-       }
-       sd = MALLOC(sizeof(struct wim_security_data));
+       wimlib_assert(metadata_resource_len >= 8);
+
+       sd = new_wim_security_data();
        if (!sd)
-               return WIMLIB_ERR_NOMEM;
-       p = get_u32(metadata_resource, &sd->total_length);
-       p = get_u32(p, &sd->num_entries);
+               goto out_of_memory;
+
+       sd_disk = (const struct wim_security_data_disk*)metadata_resource;
+       sd->total_length = le32_to_cpu(sd_disk->total_length);
+       sd->num_entries = le32_to_cpu(sd_disk->num_entries);
+
+       DEBUG("Reading security data: num_entries=%u, total_length=%u",
+             sd->num_entries, sd->total_length);
+
+       /* Length field of 0 is a special case that really means length
+        * of 8. */
+       if (sd->total_length == 0)
+               sd->total_length = 8;
+
+       /* The security_id field of each dentry is a signed 32-bit integer, so
+        * the possible indices into the security descriptors table are 0
+        * through 0x7fffffff.  Which means 0x80000000 security descriptors
+        * maximum.  Not like you should ever have anywhere close to that many
+        * security descriptors! */
+       if (sd->num_entries > 0x80000000)
+               goto out_invalid_sd;
 
        /* Verify the listed total length of the security data is big enough to
         * include the sizes array, verify that the file data is big enough to
-        * include it as well, then allocate the array of sizes. */
-       sizes_size = sd->num_entries * sizeof(u64);
+        * include it as well, then allocate the array of sizes.
+        *
+        * Note: The total length of the security data must fit in a 32-bit
+        * integer, even though each security descriptor size is a 64-bit
+        * integer.  This is stupid, and we need to be careful not to actually
+        * let the security descriptor sizes be over 0xffffffff.  */
+       if (sd->total_length > metadata_resource_len)
+               goto out_invalid_sd;
 
-       DEBUG("Reading security data with %u entries\n", sd->num_entries);
+       sizes_size = (u64)sd->num_entries * sizeof(u64);
+       size_no_descriptors = 8 + sizes_size;
+       if (size_no_descriptors > sd->total_length)
+               goto out_invalid_sd;
 
-       if (sd->num_entries == 0) {
-               FREE(sd);
-               return 0;
-       }
+       total_len = size_no_descriptors;
 
-       u64 size_no_descriptors = 8 + sizes_size;
-       if (size_no_descriptors > sd->total_length) {
-               ERROR("Security data total length of %"PRIu64" is too short because\n"
-                               "there must be at least %"PRIu64" bytes of security "
-                               "data!\n", sd->total_length, 
-                               8 + sizes_size);
-               FREE(sd);
-               return WIMLIB_ERR_INVALID_RESOURCE_SIZE;
-       }
-       if (size_no_descriptors > metadata_resource_len) {
-               ERROR("File resource of %"PRIu64" bytes is not big enough\n"
-                               "to hold security data of at least %"PRIu64" "
-                               "bytes!\n", metadata_resource_len, size_no_descriptors);
-               FREE(sd);
-               return WIMLIB_ERR_INVALID_RESOURCE_SIZE;
-       }
+       /* Return immediately if no security descriptors. */
+       if (sd->num_entries == 0)
+               goto out_align_total_length;
+
+       /* Allocate a new buffer for the sizes array */
        sd->sizes = MALLOC(sizes_size);
-       if (!sd->sizes) {
-               FREE(sd);
-               return WIMLIB_ERR_NOMEM;
+       if (!sd->sizes)
+               goto out_of_memory;
+
+       /* Copy the sizes array into the new buffer */
+       for (u32 i = 0; i < sd->num_entries; i++) {
+               sd->sizes[i] = le64_to_cpu(sd_disk->sizes[i]);
+               if (sd->sizes[i] > 0xffffffff)
+                       goto out_invalid_sd;
        }
 
-       /* Copy the sizes array in from the file data. */
-       p = get_bytes(p, sizes_size, sd->sizes);
-       array_to_le64(sd->sizes, sd->num_entries);
+       p = (const u8*)sd_disk + size_no_descriptors;
 
-       /* Allocate the array of pointers to descriptors, and read them in. */
-       sd->descriptors = CALLOC(sd->num_entries, sizeof(u8*));
-       if (!sd->descriptors) {
-               FREE(sd);
-               FREE(sd->sizes);
-               return WIMLIB_ERR_NOMEM;
-       }
-       u64 total_len = size_no_descriptors;
+       /* Allocate the array of pointers to the security descriptors, then read
+        * them into separate buffers. */
+       sd->descriptors = CALLOC(sd->num_entries, sizeof(sd->descriptors[0]));
+       if (!sd->descriptors)
+               goto out_of_memory;
 
-       for (uint i = 0; i < sd->num_entries; i++) {
+       for (u32 i = 0; i < sd->num_entries; i++) {
+               if (sd->sizes[i] == 0)
+                       continue;
                total_len += sd->sizes[i];
-               if (total_len > sd->total_length) {
-                       ERROR("Security data total length of %"PRIu64" is too "
-                                       "short because there are at least %"PRIu64" "
-                                       "bytes of security data!\n", 
-                                       sd->total_length, total_len);
-                       free_security_data(sd);
-                       return WIMLIB_ERR_INVALID_RESOURCE_SIZE;
-               }
-               if (total_len > metadata_resource_len) {
-                       ERROR("File resource of %"PRIu64" bytes is not big enough "
-                                       "to hold security data of at least %"PRIu64" "
-                                       "bytes!\n", metadata_resource_len, total_len);
-                       free_security_data(sd);
-                       return WIMLIB_ERR_INVALID_RESOURCE_SIZE;
-               }
-               sd->descriptors[i] = MALLOC(sd->sizes[i]);
-               if (!sd->descriptors[i]) {
-                       free_security_data(sd);
-                       return WIMLIB_ERR_NOMEM;
-               }
-               p = get_bytes(p, sd->sizes[i], sd->descriptors[i]);
+               if (total_len > (u64)sd->total_length)
+                       goto out_invalid_sd;
+               sd->descriptors[i] = memdup(p, sd->sizes[i]);
+               if (!sd->descriptors[i])
+                       goto out_of_memory;
+               p += sd->sizes[i];
+               empty_sacl_fixup((SECURITY_DESCRIPTOR_RELATIVE*)sd->descriptors[i],
+                                &sd->sizes[i]);
        }
-       sd->refcnt = 1;
-       *sd_p = sd;
-       return 0;
+out_align_total_length:
+       total_len = (total_len + 7) & ~7;
+       sd->total_length = (sd->total_length + 7) & ~7;
+       if (total_len != sd->total_length) {
+               WARNING("Expected WIM security data total length of "
+                       "%u bytes, but calculated %u bytes",
+                       sd->total_length, (unsigned)total_len);
+       }
+       *sd_ret = sd;
+       ret = 0;
+       goto out;
+out_invalid_sd:
+       ERROR("WIM security data is invalid!");
+       ret = WIMLIB_ERR_INVALID_SECURITY_DATA;
+       goto out_free_sd;
+out_of_memory:
+       ERROR("Out of memory while reading WIM security data!");
+       ret = WIMLIB_ERR_NOMEM;
+out_free_sd:
+       free_wim_security_data(sd);
+out:
+       return ret;
 }
 
-/* 
- * Writes security data to an in-memory buffer.
+/*
+ * Writes the security data for a WIM image to an in-memory buffer.
  */
-u8 *write_security_data(const struct wim_security_data *sd, u8 *p)
+u8 *
+write_wim_security_data(const struct wim_security_data * restrict sd,
+                       u8 * restrict p)
 {
-       if (sd) {
-               DEBUG("Writing security data (total_length = %u, "
-                               "num_entries = %u)\n", sd->total_length, 
-                               sd->num_entries);
-               u8 *orig_p = p;
-               p = put_u32(p, sd->total_length);
-               p = put_u32(p, sd->num_entries);
+       DEBUG("Writing security data (total_length = %"PRIu32", num_entries "
+             "= %"PRIu32")", sd->total_length, sd->num_entries);
 
-               for (uint i = 0; i < sd->num_entries; i++)
-                       p = put_u64(p, sd->sizes[i]);
+       u8 *orig_p = p;
+       struct wim_security_data_disk *sd_disk = (struct wim_security_data_disk*)p;
+       u32 num_entries = sd->num_entries;
 
-               for (uint i = 0; i < sd->num_entries; i++)
-                       p = put_bytes(p, sd->sizes[i], sd->descriptors[i]);
+       sd_disk->total_length = cpu_to_le32(sd->total_length);
+       sd_disk->num_entries = cpu_to_le32(num_entries);
 
-               wimlib_assert(p - orig_p <= sd->total_length);
+       for (u32 i = 0; i < num_entries; i++)
+               sd_disk->sizes[i] = cpu_to_le64(sd->sizes[i]);
 
-               DEBUG("Successfully wrote security data.\n");
-               return orig_p + sd->total_length;
-       } else {
-               DEBUG("Writing security data (total_length = 8, "
-                               "num_entries = 0)\n");
-               p = put_u32(p, 8);
-               return put_u32(p, 0);
+       p = (u8*)&sd_disk->sizes[num_entries];
 
-       }
-}
+       for (u32 i = 0; i < num_entries; i++)
+               p = mempcpy(p, sd->descriptors[i], sd->sizes[i]);
+
+       while ((uintptr_t)p & 7)
+               *p++ = 0;
 
-/* XXX We don't actually do anything with the ACL's yet besides being able to
- * print a few things.  It seems it would be a lot of work to have comprehensive
- * support for all the weird flags and stuff, and Windows PE seems to be okay
- * running from a WIM file that doesn't have any security data at all...  */
+       wimlib_assert(p - orig_p == sd->total_length);
 
-static void print_acl(const u8 *p)
+       DEBUG("Successfully wrote security data.");
+       return p;
+}
+
+static void
+print_acl(const ACL *acl, const tchar *type, size_t max_size)
 {
-       ACL *acl = (ACL*)p;
-       TO_LE16(acl->acl_size);
-       TO_LE16(acl->acl_count);
-       printf("    [ACL]\n");
-       printf("    Revision = %u\n", acl->revision);
-       printf("    ACL Size = %u\n", acl->acl_size);
-       printf("    ACE Count = %u\n", acl->ace_count);
-
-       p += sizeof(ACL);
-       for (uint i = 0; i < acl->ace_count; i++) {
-               ACEHeader *hdr = (ACEHeader*)p;
-               printf("        [ACE]\n");
-               printf("        ACE type  = %d\n", hdr->type);
-               printf("        ACE flags = 0x%x\n", hdr->flags);
-               printf("        ACE size  = %u\n", hdr->size);
-               AccessAllowedACE *aaa = (AccessAllowedACE*)hdr;
-               printf("        ACE mask = %x\n", to_le32(aaa->mask));
-               printf("        SID start = %u\n", to_le32(aaa->sid_start));
-               p += hdr->size;
+       const u8 *p;
+
+       if (max_size < sizeof(ACL))
+               return;
+
+       u8 revision = acl->revision;
+       u16 acl_size = le16_to_cpu(acl->acl_size);
+       u16 ace_count = le16_to_cpu(acl->ace_count);
+
+       tprintf(T("    [%"TS" ACL]\n"), type);
+       tprintf(T("    Revision = %u\n"), revision);
+       tprintf(T("    ACL Size = %u\n"), acl_size);
+       tprintf(T("    ACE Count = %u\n"), ace_count);
+
+       p = (const u8*)acl + sizeof(ACL);
+       for (u16 i = 0; i < ace_count; i++) {
+               if (max_size < p + sizeof(ACCESS_ALLOWED_ACE) - (const u8*)acl)
+                       break;
+               const ACCESS_ALLOWED_ACE *aaa = (const ACCESS_ALLOWED_ACE*)p;
+               tprintf(T("        [ACE]\n"));
+               tprintf(T("        ACE type  = %d\n"), aaa->hdr.type);
+               tprintf(T("        ACE flags = 0x%x\n"), aaa->hdr.flags);
+               tprintf(T("        ACE size  = %u\n"), le16_to_cpu(aaa->hdr.size));
+               tprintf(T("        ACE mask = %x\n"), le32_to_cpu(aaa->mask));
+               tprintf(T("        SID start = %u\n"), le32_to_cpu(aaa->sid_start));
+               p += le16_to_cpu(aaa->hdr.size);
        }
+       tputchar(T('\n'));
 }
 
-static void print_sid(const u8 *p)
+static void
+print_sid(const SID *sid, const tchar *type, size_t max_size)
 {
-       SID *sid = (SID*)p;
-       printf("    [SID]\n");
-       printf("    Revision = %u\n", sid->revision);
-       printf("    Subauthority count = %u\n", sid->sub_authority_count);
-       printf("    Identifier authority = ");
-       print_byte_field(sid->identifier_authority, sizeof(sid->identifier_authority));
-       putchar('\n');
-       for (uint i = 0; i < sid->sub_authority_count; i++)
-               printf("    Subauthority %u = %u\n", i, to_le32(sid->sub_authority[i]));
+       if (max_size < sizeof(SID))
+               return;
+
+       tprintf(T("    [%"TS" SID]\n"), type);
+       tprintf(T("    Revision = %u\n"), sid->revision);
+       tprintf(T("    Subauthority count = %u\n"), sid->sub_authority_count);
+       tprintf(T("    Identifier authority = "));
+       print_byte_field(sid->identifier_authority,
+                        sizeof(sid->identifier_authority), stdout);
+       tputchar(T('\n'));
+       if (max_size < sizeof(SID) + (size_t)sid->sub_authority_count * sizeof(u32))
+               return;
+       for (u8 i = 0; i < sid->sub_authority_count; i++) {
+               tprintf(T("    Subauthority %u = %u\n"),
+                       i, le32_to_cpu(sid->sub_authority[i]));
+       }
+       tputchar(T('\n'));
 }
 
-static void print_security_descriptor(const u8 *p, u64 size)
+static void
+print_security_descriptor(const SECURITY_DESCRIPTOR_RELATIVE *descr,
+                         size_t size)
 {
-       SecurityDescriptor *sd = (SecurityDescriptor*)p;
-       TO_LE16(sd->security_descriptor_control);
-       TO_LE32(sd->owner_offset);
-       TO_LE32(sd->group_offset);
-       TO_LE32(sd->sacl_offset);
-       TO_LE32(sd->dacl_offset);
-       printf("Revision = %u\n", sd->revision);
-       printf("Security Descriptor Control = %u\n", sd->security_descriptor_control);
-       printf("Owner offset = %u\n", sd->owner_offset);
-       printf("Group offset = %u\n", sd->group_offset);
-       printf("System ACL offset = %u\n", sd->sacl_offset);
-       printf("Discretionary ACL offset = %u\n", sd->dacl_offset);
-
-       if (sd->owner_offset != 0)
-               print_sid(p + sd->owner_offset);
-       if (sd->group_offset != 0)
-               print_sid(p + sd->group_offset);
-       if (sd->sacl_offset != 0)
-               print_acl(p + sd->sacl_offset);
-       if (sd->dacl_offset != 0)
-               print_acl(p + sd->dacl_offset);
+       u8 revision      = descr->revision;
+       u16 control      = le16_to_cpu(descr->security_descriptor_control);
+       u32 owner_offset = le32_to_cpu(descr->owner_offset);
+       u32 group_offset = le32_to_cpu(descr->group_offset);
+       u32 dacl_offset  = le32_to_cpu(descr->dacl_offset);
+       u32 sacl_offset  = le32_to_cpu(descr->sacl_offset);
+
+       tprintf(T("Revision = %u\n"), revision);
+       tprintf(T("Security Descriptor Control = %#x\n"), control);
+       tprintf(T("Owner offset = %u\n"), owner_offset);
+       tprintf(T("Group offset = %u\n"), group_offset);
+       tprintf(T("Discretionary ACL offset = %u\n"), dacl_offset);
+       tprintf(T("System ACL offset = %u\n"), sacl_offset);
+
+       if (owner_offset != 0 && owner_offset <= size)
+               print_sid((const SID*)((const u8*)descr + owner_offset),
+                         T("Owner"), size - owner_offset);
+
+       if (group_offset != 0 && group_offset <= size)
+               print_sid((const SID*)((const u8*)descr + group_offset),
+                         T("Group"), size - group_offset);
+
+       if (dacl_offset != 0 && dacl_offset <= size)
+               print_acl((const ACL*)((const u8*)descr + dacl_offset),
+                         T("Discretionary"), size - dacl_offset);
+
+       if (sacl_offset != 0 && sacl_offset <= size)
+               print_acl((const ACL*)((const u8*)descr + sacl_offset),
+                         T("System"), size - sacl_offset);
 }
 
-/* 
+/*
  * Prints the security data for a WIM file.
  */
-void print_security_data(const struct wim_security_data *sd)
+void
+print_wim_security_data(const struct wim_security_data *sd)
 {
-       puts("[SECURITY DATA]");
-       if (sd) {
-               printf("Length            = %u bytes\n", sd->total_length);
-               printf("Number of Entries = %u\n", sd->num_entries);
-
-               u64 num_entries = (u64)sd->num_entries;
-               for (u64 i = 0; i < num_entries; i++) {
-                       printf("[SecurityDescriptor %"PRIu64", "
-                                       "length = %"PRIu64"]\n", 
-                                       i, sd->sizes[i]);
-                       print_security_descriptor(sd->descriptors[i], 
-                                                 sd->sizes[i]);
-                       putchar('\n');
-               }
-       } else {
-               puts("Length            = 8 bytes\n"
-                    "Number of Entries = 0");
+       tputs(T("[SECURITY DATA]"));
+       tprintf(T("Length            = %"PRIu32" bytes\n"), sd->total_length);
+       tprintf(T("Number of Entries = %"PRIu32"\n"), sd->num_entries);
+
+       for (u32 i = 0; i < sd->num_entries; i++) {
+               tprintf(T("[SECURITY_DESCRIPTOR_RELATIVE %"PRIu32", length = %"PRIu64"]\n"),
+                       i, sd->sizes[i]);
+               print_security_descriptor((const SECURITY_DESCRIPTOR_RELATIVE*)sd->descriptors[i],
+                                         sd->sizes[i]);
+               tputchar(T('\n'));
        }
-       putchar('\n');
+       tputchar(T('\n'));
 }
 
-void free_security_data(struct wim_security_data *sd)
+void
+free_wim_security_data(struct wim_security_data *sd)
 {
-       if (!sd)
-               return;
-       wimlib_assert(sd->refcnt >= 1);
-       if (sd->refcnt == 1) {
+       if (sd) {
                u8 **descriptors = sd->descriptors;
-               u32 num_entries = sd->num_entries;
-
+               u32 num_entries  = sd->num_entries;
                if (descriptors)
                        while (num_entries--)
                                FREE(*descriptors++);
                FREE(sd->sizes);
                FREE(sd->descriptors);
                FREE(sd);
-       } else {
-               sd->refcnt--;
        }
 }
 
-#endif
+struct sd_node {
+       int security_id;
+       u8 hash[SHA1_HASH_SIZE];
+       struct rb_node rb_node;
+};
+
+static void
+free_sd_tree(struct rb_node *node)
+{
+       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. */
+void
+destroy_sd_set(struct wim_sd_set *sd_set, bool rollback)
+{
+       if (rollback) {
+               struct wim_security_data *sd = sd_set->sd;
+               u8 **descriptors = sd->descriptors + sd_set->orig_num_entries;
+               u32 num_entries  = sd->num_entries - sd_set->orig_num_entries;
+               while (num_entries--)
+                       FREE(*descriptors++);
+               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 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);
+       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
+                       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 wim_sd_set *set, const u8 hash[SHA1_HASH_SIZE])
+{
+       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;
+}
+
+/*
+ * Adds a security descriptor to the indexed security descriptor set as well as
+ * the corresponding `struct wim_security_data', and returns the new security
+ * ID; or, if there is an existing security descriptor that is the same, return
+ * the security ID for it.  If a new security descriptor cannot be allocated,
+ * return -1.
+ */
+int
+sd_set_add_sd(struct wim_sd_set *sd_set, const char *descriptor, size_t size)
+{
+       u8 hash[SHA1_HASH_SIZE];
+       int security_id;
+       struct sd_node *new;
+       u8 **descriptors;
+       u64 *sizes;
+       u8 *descr_copy;
+       struct wim_security_data *sd;
+       bool bret;
+
+       sha1_buffer(descriptor, size, hash);
+
+       security_id = lookup_sd(sd_set, hash);
+       if (security_id >= 0) /* Identical descriptor already exists */
+               goto out;
+
+       /* Need to add a new security descriptor */
+       security_id = -1;
+
+       new = MALLOC(sizeof(*new));
+       if (!new)
+               goto out;
+       descr_copy = MALLOC(size);
+       if (!descr_copy)
+               goto out_free_node;
+
+       sd = sd_set->sd;
+
+       memcpy(descr_copy, descriptor, 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)
+               goto out_free_descr;
+       sd->descriptors = descriptors;
+       sizes = REALLOC(sd->sizes,
+                       (sd->num_entries + 1) * sizeof(sd->sizes[0]));
+       if (!sizes)
+               goto out_free_descr;
+       sd->sizes = sizes;
+       sd->descriptors[sd->num_entries] = descr_copy;
+       sd->sizes[sd->num_entries] = size;
+       sd->num_entries++;
+       DEBUG("There are now %u security descriptors", sd->num_entries);
+       bret = insert_sd_node(sd_set, new);
+       wimlib_assert(bret);
+       security_id = new->security_id;
+       goto out;
+out_free_descr:
+       FREE(descr_copy);
+out_free_node:
+       FREE(new);
+out:
+       return security_id;
+}
+
+/* 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 (u32 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;
+}