liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-11 07:06:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fixed integer wraparound leading to buffer overflow with parallel HNSW index build

Browse Source 
Reported-by: chungkn from OneMount Group <kimngocchung.k2a@gmail.com>
This commit is contained in:
Andrew Kane
2026-02-25 10:35:14 -08:00
parent f5841f07fd
commit b7e680d41a
2 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGELOG.md
View File

@@ -1,5 +1,6 @@
## 0.8.2 (unreleased) ## 0.8.2 (unreleased)
- Fixed buffer overflow with parallel HNSW index build
- Improved `install` target on Windows - Improved `install` target on Windows
- Fixed `Index Searches` in `EXPLAIN` output for Postgres 18 - Fixed `Index Searches` in `EXPLAIN` output for Postgres 18

15
src/hnswbuild.c
View File

@@ -77,6 +77,8 @@
#define PARALLEL_KEY_HNSW_AREA UINT64CONST(0xA000000000000002) #define PARALLEL_KEY_HNSW_AREA UINT64CONST(0xA000000000000002)
#define PARALLEL_KEY_QUERY_TEXT UINT64CONST(0xA000000000000003) #define PARALLEL_KEY_QUERY_TEXT UINT64CONST(0xA000000000000003)
#define HNSW_MAX_GRAPH_MEMORY (SIZE_MAX / 2)
/* /*
* Create the metapage * Create the metapage
*/ */
@@ -489,6 +491,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
LWLock *flushLock = &graph->flushLock; LWLock *flushLock = &graph->flushLock;
char *base = buildstate->hnswarea; char *base = buildstate->hnswarea;
Datum value; Datum value;
Size memoryMargin;
/* Form index value */ /* Form index value */
if (!HnswFormIndexValue(&value, values, isnull, buildstate->typeInfo, support)) if (!HnswFormIndexValue(&value, values, isnull, buildstate->typeInfo, support))
@@ -497,6 +500,9 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
/* Get datum size */ /* Get datum size */
valueSize = VARSIZE_ANY(DatumGetPointer(value)); valueSize = VARSIZE_ANY(DatumGetPointer(value));
/* In a parallel build, add a margin so allocations never fail */
memoryMargin = base == NULL ? 0 : 1024 * 1024;
/* Ensure graph not flushed when inserting */ /* Ensure graph not flushed when inserting */
LWLockAcquire(flushLock, LW_SHARED); LWLockAcquire(flushLock, LW_SHARED);
@@ -518,7 +524,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
* Check that we have enough memory available for the new element now that * Check that we have enough memory available for the new element now that
* we have the allocator lock, and flush pages if needed. * we have the allocator lock, and flush pages if needed.
*/ */
if (graph->memoryUsed >= graph->memoryTotal) if (graph->memoryUsed + memoryMargin >= graph->memoryTotal)
{ {
LWLockRelease(&graph->allocatorLock); LWLockRelease(&graph->allocatorLock);
@@ -611,7 +617,7 @@ InitGraph(HnswGraph * graph, char *base, Size memoryTotal)
HnswPtrStore(base, graph->head, (HnswElement) NULL); HnswPtrStore(base, graph->head, (HnswElement) NULL);
HnswPtrStore(base, graph->entryPoint, (HnswElement) NULL); HnswPtrStore(base, graph->entryPoint, (HnswElement) NULL);
graph->memoryUsed = 0; graph->memoryUsed = 0;
graph->memoryTotal = memoryTotal; graph->memoryTotal = Min(memoryTotal, HNSW_MAX_GRAPH_MEMORY);
graph->flushed = false; graph->flushed = false;
graph->indtuples = 0; graph->indtuples = 0;
SpinLockInit(&graph->lock); SpinLockInit(&graph->lock);
@@ -940,6 +946,8 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
if (esthnswarea > estother) if (esthnswarea > estother)
esthnswarea -= estother; esthnswarea -= estother;
esthnswarea = Min(esthnswarea, HNSW_MAX_GRAPH_MEMORY);
shm_toc_estimate_chunk(&pcxt->estimator, esthnswarea); shm_toc_estimate_chunk(&pcxt->estimator, esthnswarea);
shm_toc_estimate_keys(&pcxt->estimator, 2); shm_toc_estimate_keys(&pcxt->estimator, 2);
@@ -982,8 +990,7 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
snapshot); snapshot);
hnswarea = (char *) shm_toc_allocate(pcxt->toc, esthnswarea); hnswarea = (char *) shm_toc_allocate(pcxt->toc, esthnswarea);
/* Report less than allocated so never fails */ InitGraph(&hnswshared->graphData, hnswarea, esthnswarea);
InitGraph(&hnswshared->graphData, hnswarea, esthnswarea - 1024 * 1024);
/* /*
* Avoid base address for relptr for Postgres < 14.5 * Avoid base address for relptr for Postgres < 14.5