1 files changed, 123 insertions, 21 deletions

diff --git a/src/util/hash_table.c b/src/util/hash_table.c
index 1811ee7432f..4a6cedc2b4b 100644
--- a/src/util/hash_table.c
+++ b/src/util/hash_table.c
@@ -427,8 +427,7 @@ _mesa_hash_table_rehash(struct hash_table *ht, unsigned new_size_index)
 }
 
 static struct hash_entry *
-hash_table_insert(struct hash_table *ht, uint32_t hash,
-                  const void *key, void *data)
+hash_table_get_entry(struct hash_table *ht, uint32_t hash, const void *key)
 {
    struct hash_entry *available_entry = NULL;
 
@@ -469,11 +468,8 @@ hash_table_insert(struct hash_table *ht, uint32_t hash,
        */
       if (!entry_is_deleted(ht, entry) &&
           entry->hash == hash &&
-          ht->key_equals_function(key, entry->key)) {
-         entry->key = key;
-         entry->data = data;
+          ht->key_equals_function(key, entry->key))
          return entry;
-      }
 
       hash_address += double_hash;
       if (hash_address >= size)
@@ -484,8 +480,6 @@ hash_table_insert(struct hash_table *ht, uint32_t hash,
       if (entry_is_deleted(ht, available_entry))
          ht->deleted_entries--;
       available_entry->hash = hash;
-      available_entry->key = key;
-      available_entry->data = data;
       ht->entries++;
       return available_entry;
    }
@@ -496,6 +490,20 @@ hash_table_insert(struct hash_table *ht, uint32_t hash,
    return NULL;
 }
 
+static struct hash_entry *
+hash_table_insert(struct hash_table *ht, uint32_t hash,
+                  const void *key, void *data)
+{
+   struct hash_entry *entry = hash_table_get_entry(ht, hash, key);
+
+   if (entry) {
+      entry->key = key;
+      entry->data = data;
+   }
+
+   return entry;
+}
+
 /**
  * Inserts the key with the given hash into the table.
  *
@@ -659,13 +667,18 @@ _mesa_hash_u32(const void *key)
 uint32_t
 _mesa_hash_string(const void *_key)
 {
+   return _mesa_hash_string_with_length(_key, strlen((const char *)_key));
+}
+
+uint32_t
+_mesa_hash_string_with_length(const void *_key, unsigned length)
+{
    uint32_t hash = 0;
    const char *key = _key;
-   size_t len = strlen(key);
 #if defined(_WIN64) || defined(__x86_64__)
-   hash = (uint32_t)XXH64(key, len, hash);
+   hash = (uint32_t)XXH64(key, length, hash);
 #else
-   hash = XXH32(key, len, hash);
+   hash = XXH32(key, length, hash);
 #endif
    return hash;
 }
@@ -764,22 +777,54 @@ key_u64_equals(const void *a, const void *b)
 
 #define FREED_KEY_VALUE 0
 
+static void _mesa_hash_table_u64_delete_keys(void *data)
+{
+   struct hash_table_u64 *ht = ralloc_parent(data);
+
+   _mesa_hash_table_u64_clear(ht);
+}
+
 struct hash_table_u64 *
 _mesa_hash_table_u64_create(void *mem_ctx)
 {
    STATIC_ASSERT(FREED_KEY_VALUE != DELETED_KEY_VALUE);
    struct hash_table_u64 *ht;
 
-   ht = CALLOC_STRUCT(hash_table_u64);
+   ht = rzalloc(mem_ctx, struct hash_table_u64);
    if (!ht)
       return NULL;
 
    if (sizeof(void *) == 8) {
-      ht->table = _mesa_hash_table_create(mem_ctx, _mesa_hash_pointer,
+      ht->table = _mesa_hash_table_create(ht, _mesa_hash_pointer,
                                           _mesa_key_pointer_equal);
    } else {
-      ht->table = _mesa_hash_table_create(mem_ctx, key_u64_hash,
+      ht->table = _mesa_hash_table_create(ht, key_u64_hash,
                                           key_u64_equals);
+
+      /* Allocate a ralloc sub-context which takes the u64 hash table
+       * as a parent and attach a destructor to it so we can free the
+       * hash_key_u64 objects that were allocated by
+       * _mesa_hash_table_u64_insert().
+       *
+       * The order of creation of this sub-context is crucial: it needs
+       * to happen after the _mesa_hash_table_create() call to guarantee
+       * that the destructor is called before ht->table and its children
+       * are freed, otherwise the _mesa_hash_table_u64_clear() call in the
+       * destructor leads to a use-after-free situation.
+       */
+      if (ht->table) {
+         void *dummy_ctx = ralloc_context(ht);
+
+         /* If we can't allocate a sub-context, free the hash table
+          * immediately and return NULL to avoid future leaks.
+          */
+         if (!dummy_ctx) {
+            ralloc_free(ht);
+            return NULL;
+         }
+
+         ralloc_set_destructor(dummy_ctx, _mesa_hash_table_u64_delete_keys);
+      }
    }
 
    if (ht->table)
@@ -797,7 +842,7 @@ _mesa_hash_table_u64_delete_key(struct hash_entry *entry)
    struct hash_key_u64 *_key = (struct hash_key_u64 *)entry->key;
 
    if (_key)
-      free(_key);
+      FREE(_key);
 }
 
 void
@@ -807,6 +852,8 @@ _mesa_hash_table_u64_clear(struct hash_table_u64 *ht)
       return;
 
    _mesa_hash_table_clear(ht->table, _mesa_hash_table_u64_delete_key);
+   ht->freed_key_data = NULL;
+   ht->deleted_key_data = NULL;
 }
 
 void
@@ -814,10 +861,7 @@ _mesa_hash_table_u64_destroy(struct hash_table_u64 *ht)
 {
    if (!ht)
       return;
-
-   _mesa_hash_table_u64_clear(ht);
-   _mesa_hash_table_destroy(ht->table, NULL);
-   free(ht);
+   ralloc_free(ht);
 }
 
 void
@@ -843,7 +887,19 @@ _mesa_hash_table_u64_insert(struct hash_table_u64 *ht, uint64_t key,
          return;
       _key->value = key;
 
-      _mesa_hash_table_insert(ht->table, _key, data);
+      struct hash_entry *entry =
+         hash_table_get_entry(ht->table, key_u64_hash(_key), _key);
+
+      if (!entry) {
+         FREE(_key);
+         return;
+      }
+
+      entry->data = data;
+      if (!entry_is_present(ht->table, entry))
+         entry->key = _key;
+      else
+         FREE(_key);
    }
 }
 
@@ -901,6 +957,52 @@ _mesa_hash_table_u64_remove(struct hash_table_u64 *ht, uint64_t key)
       struct hash_key *_key = (struct hash_key *)entry->key;
 
       _mesa_hash_table_remove(ht->table, entry);
-      free(_key);
+      FREE(_key);
    }
 }
+
+
+/*
+ * Iterates in order ("freed key", "deleted key", regular entries...)
+ */
+struct hash_entry_u64
+_mesa_hash_table_u64_next_entry(struct hash_table_u64 *ht,
+                                struct hash_entry_u64 *ent)
+{
+   /* First entry: freed key */
+   if (!ent && ht->freed_key_data) {
+      return (struct hash_entry_u64){
+         .key = FREED_KEY_VALUE,
+         .data = ht->freed_key_data,
+      };
+   }
+
+   /* Second entry: deleted key */
+   if ((!ent || ent->key == FREED_KEY_VALUE) && ht->deleted_key_data) {
+      return (struct hash_entry_u64){
+         .key = DELETED_KEY_VALUE,
+         .data = ht->deleted_key_data,
+      };
+   }
+
+   /* All other entries: regular */
+   struct hash_entry *next =
+      _mesa_hash_table_next_entry(ht->table, ent ? ent->_entry : NULL);
+
+   if (!next)
+      return (struct hash_entry_u64){.data = NULL};
+
+   uint64_t key;
+   if (sizeof(void *) == 8) {
+      key = (uintptr_t)next->key;
+   } else {
+      const struct hash_key_u64 *_key = next->key;
+      key = _key->value;
+   }
+
+   return (struct hash_entry_u64){
+      .key = key,
+      .data = next->data,
+      ._entry = next,
+   };
+}