diff --git a/query/vector/vector_test.go b/query/vector/vector_test.go index ba1b62bd680..4d578214a55 100644 --- a/query/vector/vector_test.go +++ b/query/vector/vector_test.go @@ -485,6 +485,63 @@ func TestSimilarToOptionsIntegration(t *testing.T) { }) } +// TestSimilarToNonPositiveNeighbors is the end-to-end guard for a crash reachable +// from an unvalidated similar_to number-of-neighbors argument. A count <= 0 used to +// panic the query goroutine and take down the Alpha; the worker/task.go boundary +// check must now turn it into an ordinary query error while the server keeps serving. +func TestSimilarToNonPositiveNeighbors(t *testing.T) { + const pred = "vnonpos" + dropPredicate(pred) + t.Cleanup(func() { dropPredicate(pred) }) + + setSchema(fmt.Sprintf(vectorSchemaWithIndex, pred, "4", "euclidean")) + + rdf := `<0x1> "[0,0]" . + <0x2> "[1,0]" . + <0x3> "[2,0]" .` + require.NoError(t, addTriplesToCluster(rdf)) + + // Both k=0 and k=-1 parse fine through DQL, so they reach the worker guard. + for _, k := range []int{0, -1} { + query := fmt.Sprintf(`{ + results(func: similar_to(%s, %d, "[0,0]")) { + uid + } + }`, pred, k) + _, err := processQuery(context.Background(), t, query) + require.Errorf(t, err, "similar_to with k=%d should return an error, not crash", k) + require.ErrorContains(t, err, "number of neighbors") + } + + // The load-bearing assertion: a valid query after the bad ones must succeed. + // On the pre-fix code the first non-positive query panics and crashes the + // Alpha, so this only passes if the server stayed up. + query := fmt.Sprintf(`{ + results(func: similar_to(%s, 3, "[0,0]")) { + uid + } + }`, pred) + resp := processQueryNoErr(t, query) + + var result struct { + Data struct { + Results []struct { + UID string `json:"uid"` + } `json:"results"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(resp), &result)) + require.Len(t, result.Data.Results, 3) + + expected := map[string]struct{}{"0x1": {}, "0x2": {}, "0x3": {}} + for _, r := range result.Data.Results { + _, ok := expected[r.UID] + require.Truef(t, ok, "unexpected uid %s", r.UID) + delete(expected, r.UID) + } + require.Empty(t, expected) +} + func TestVectorInQueryArgument(t *testing.T) { dropPredicate("vtest") setSchema(fmt.Sprintf(vectorSchemaWithIndex, "vtest", "4", "euclidean")) diff --git a/tok/hnsw/persistent_hnsw.go b/tok/hnsw/persistent_hnsw.go index 864e7e98637..01ca3265d53 100644 --- a/tok/hnsw/persistent_hnsw.go +++ b/tok/hnsw/persistent_hnsw.go @@ -512,6 +512,13 @@ func (ph *persistentHNSW[T]) SearchWithPath( start := time.Now().UnixMilli() r = index.NewSearchPathResult() + // Mirror the clamp already applied on the SearchWithOptions path: a negative + // maxResults must not reach the bottom-layer search, where it would drive a + // negative slice bound in addPathNode. + if maxResults < 0 { + maxResults = 0 + } + // 0-profile_vector_entry var startVec []T entry, err := ph.PickStartNode(ctx, c, &startVec) diff --git a/tok/hnsw/persistent_hnsw_test.go b/tok/hnsw/persistent_hnsw_test.go index 9a221f50fcf..f7667deb261 100644 --- a/tok/hnsw/persistent_hnsw_test.go +++ b/tok/hnsw/persistent_hnsw_test.go @@ -14,6 +14,7 @@ import ( c "github.com/dgraph-io/dgraph/v25/tok/constraints" "github.com/dgraph-io/dgraph/v25/tok/index" opt "github.com/dgraph-io/dgraph/v25/tok/options" + "github.com/stretchr/testify/require" "golang.org/x/exp/slices" ) @@ -920,3 +921,30 @@ func TestSearchReturnsCorrectOrderForAllMetrics(t *testing.T) { }) } } + +// TestAddPathNodeNonPositiveMaxResults is a regression test for a crash reachable +// from an unvalidated similar_to number-of-neighbors argument. A non-positive +// maxResults made effectiveMaxLen negative and panicked addPathNode: a negative +// slice bound (slr.neighbors[:-1]) for a negative value, and an empty-slice index +// (slr.neighbors[0]) once the neighbors were truncated to zero. Both must now be +// handled without panicking. +func TestAddPathNodeNonPositiveMaxResults(t *testing.T) { + simType := GetSimType[float64](Euclidean, 64) + for _, maxResults := range []int{-1, 0} { + slr := newLayerResult[float64](0) + slr.setFirstPathNode(persistentHeapElement[float64]{value: 0.1, index: 1}) + + require.NotPanicsf(t, func() { + slr.addPathNode(persistentHeapElement[float64]{value: 0.2, index: 2}, simType, maxResults) + }, "addPathNode panicked with maxResults=%d", maxResults) + + // A non-positive limit clamps effectiveMaxLen to 0, so the neighbor set + // truncates to empty and the bottom-of-path append is skipped. Pin that + // resulting state, not just the absence of a panic: neighbors is empty and + // the path is left exactly as setFirstPathNode set it. + require.Emptyf(t, slr.neighbors, + "neighbors should truncate to empty for maxResults=%d", maxResults) + require.Equalf(t, []uint64{1}, slr.path, + "path should be unchanged for maxResults=%d", maxResults) + } +} diff --git a/tok/hnsw/search_layer.go b/tok/hnsw/search_layer.go index c234b55f73c..ee829a8f50e 100644 --- a/tok/hnsw/search_layer.go +++ b/tok/hnsw/search_layer.go @@ -57,11 +57,16 @@ func (slr *searchLayerResult[T]) addPathNode( slr.filtered++ } effectiveMaxLen := maxResults + slr.filtered + // A non-positive maxResults would make effectiveMaxLen negative and panic the + // slice expression below; clamp it so the truncation degrades to an empty set. + if effectiveMaxLen < 0 { + effectiveMaxLen = 0 + } if len(slr.neighbors) > effectiveMaxLen { slr.neighbors = slr.neighbors[:effectiveMaxLen] } - if slr.neighbors[0].index == n.index { + if len(slr.neighbors) > 0 && slr.neighbors[0].index == n.index { slr.path = append(slr.path, slr.neighbors[0].index) } } diff --git a/worker/task.go b/worker/task.go index a59bff64fcc..8c42b7cd791 100644 --- a/worker/task.go +++ b/worker/task.go @@ -361,6 +361,10 @@ func (qs *queryState) handleValuePostings(ctx context.Context, args funcArgs) er if err != nil { return fmt.Errorf("invalid value for number of neighbors: %s", q.SrcFunc.Args[0]) } + if numNeighbors < 1 { + return fmt.Errorf( + "invalid value for number of neighbors, must be a positive integer: %d", numNeighbors) + } cspec, err := pickFactoryCreateSpec(ctx, args.q.Attr) if err != nil { return err