Race Condition / Use-After-Free in HNSW cursor handling

// Proof-of-concept: demonstrates a race between cursor iteration and freeing in the pre-fix API // NOTE: This PoC targets the pre-fix behavior where hnsw_cursor_next(index, cursor) and // hnsw_cursor_free(index, cursor) operate without proper synchronization. // The code is a conceptual demonstration and may require adaptation to a real test harness. #include <pthread.h> #include <unistd.h> #include "hnsw.h" static HNSW *g_index = NULL; static hnswCursor *g_cursor = NULL; void* iter_thread(void* arg) { // Initialize a cursor and iterate through the index. g_cursor = hnsw_cursor_init(g_index); if (!g_cursor) return NULL; hnswNode *node; while ((node = hnsw_cursor_next(g_index, g_cursor)) != NULL) { // Simulate work on node (void)node; usleep(100); // small delay to widen the race window } // Free the cursor (without synchronization in pre-fix code) hnsw_cursor_free(g_index, g_cursor); g_cursor = NULL; return NULL; } void* free_thread(void* arg) { // Sleep a bit to ensure the iterator has started usleep(50); if (g_cursor) { // Potential race: freeing the cursor while iter_thread is using it hnsw_cursor_free(g_index, g_cursor); g_cursor = NULL; } return NULL; } int main() { // Initialize index (details omitted). In a real test, populate with some nodes. g_index = hnsw_create_index(/* params */); pthread_t t1, t2; pthread_create(&t1, NULL, iter_thread, NULL); pthread_create(&t2, NULL, free_thread, NULL); pthread_join(t1, NULL); pthread_join(t2, NULL); // Cleanup index hnsw_free_index(g_index); return 0; } // Expected outcome prior to the fix: a race between hnsw_cursor_next() and hnsw_cursor_free() // may cause use-after-free or memory corruption when the cursor is freed while being iterated.

diff --git a/hnsw.c b/hnsw.c index f09e0497a8b..1c29b75c136 100644 --- a/hnsw.c +++ b/hnsw.c @@ -2237,36 +2237,58 @@ int hnsw_deserialize_index(HNSW *index) { * * The function returns NULL on out of memory. */ hnswCursor *hnsw_cursor_init(HNSW *index) { + if (pthread_rwlock_wrlock(&index->global_lock) != 0) return NULL; hnswCursor *cursor = hmalloc(sizeof(*cursor)); - if (cursor == NULL) return NULL; + if (cursor == NULL) { + pthread_rwlock_unlock(&index->global_lock); + return NULL; + } + cursor->index = index; cursor->next = index->cursors; cursor->current = index->head; index->cursors = cursor; + pthread_rwlock_unlock(&index->global_lock); return cursor; } /* Free the cursor. Can be called both at the end of the iteration, when * hnsw_cursor_next() returned NULL, or before. */ -void hnsw_cursor_free(HNSW *index, hnswCursor *cursor) { - hnswCursor *x = index->cursors; +void hnsw_cursor_free(hnswCursor *cursor) { + pthread_rwlock_wrlock(&cursor->index->global_lock); // Best effort. + hnswCursor *x = cursor->index->cursors; hnswCursor *prev = NULL; while(x) { if (x == cursor) { if (prev) prev->next = cursor->next; else - index->cursors = cursor->next; + cursor->index->cursors = cursor->next; hfree(cursor); - return; + break; } x = x->next; } + pthread_rwlock_unlock(&cursor->index->global_lock); +} + +/* Acquire a lock to use the cursor. Returns 1 if the lock was acquired + * with success, otherwise zero is returned. The returned element is + * protected after calling hnsw_cursor_next() for all the time required to + * access it, then hnsw_cursor_release_lock() should be called in order + * to unlock the HNSW index. */ +int hnsw_cursor_acquire_lock(hnswCursor *cursor) { + return pthread_rwlock_rdlock(&cursor->index->global_lock) == 0; +} + +/* Release the cursor lock, see hnsw_cursor_acquire_lock() top comment + * for more information. */ +void hnsw_cursor_release_lock(hnswCursor *cursor) { + pthread_rwlock_unlock(&cursor->index->global_lock); } /* Return the next element of the HNSW. See hnsw_cursor_init() for * the guarantees of the function. */ -hnswNode *hnsw_cursor_next(HNSW *index, hnswCursor *cursor) { - (void)index; // Unused but future proof to have. +hnswNode *hnsw_cursor_next(hnswCursor *cursor) { hnswNode *ret = cursor->current; if (ret) cursor->current = ret->next; return ret; diff --git a/hnsw.h b/hnsw.h index 69c3153b365..48af67d788e 100644 --- a/hnsw.h +++ b/hnsw.h @@ -65,12 +65,15 @@ typedef struct hnswNode { hnswNodeLayer layers[]; } hnswNode; +struct HNSW; + /* It is possible to navigate an HNSW with a cursor that guarantees * visiting all the elements that remain in the HNSW from the start to the * end of the process (but not the new ones, so that the process will * eventually finish). Check hnsw_cursor_init(), hnsw_cursor_next() and * hnsw_cursor_free(). */ typedef struct hnswCursor { + struct HNSW *index; // Reference to the index of this cursor. hnswNode *current; // Element to report when hnsw_cursor_next() is called. struct hnswCursor *next; // Next cursor active. } hnswCursor; @@ -156,8 +159,10 @@ uint32_t hnsw_quants_bytes(HNSW *index); /* Cursors. */ hnswCursor *hnsw_cursor_init(HNSW *index); -void hnsw_cursor_free(HNSW *index, hnswCursor *cursor); -hnswNode *hnsw_cursor_next(HNSW *index, hnswCursor *cursor); +void hnsw_cursor_free(hnswCursor *cursor); +hnswNode *hnsw_cursor_next(hnswCursor *cursor); +int hnsw_cursor_acquire_lock(hnswCursor *cursor); +void hnsw_cursor_release_lock(hnswCursor *cursor); /* Allocator selection. */ void hnsw_set_allocator(void (*free_ptr)(void*), void *(*malloc_ptr)(size_t),