Reject non-positive result counts in the search endpoints (#1358)

Author: tomchop

core/web/apiv2/search.py

@@ -1,21 +1,26 @@
 from typing import Literal
 
 from fastapi import APIRouter, Request
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field
 
 from core.database_arango import ArangoYetiConnector
 
 # API endpoints
 router = APIRouter()
 
+# Both search endpoints bucket results per type, so this bounds each bucket
+# independently -- a request with no root_type can still return this many
+# results for every type it searches.
+MAX_RESULTS_PER_TYPE = 50
+
 
 class SearchRequest(BaseModel):
     """Global search request message."""
 
     model_config = ConfigDict(extra="forbid")
 
     query: str
-    count_per_type: int = 5
+    count_per_type: int = Field(default=5, ge=1, le=MAX_RESULTS_PER_TYPE)
 
 
 class SearchResultSection(BaseModel):
@@ -38,7 +43,12 @@ class SemanticSearchRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
     query: str
-    count: int = 10
+    # ChromaDB rejects a non-positive n_results, so an unconstrained count
+    # turns a client mistake into a server error. The ceiling bounds work
+    # rather than payload: count is overfetched into n_results (see
+    # _semantic_search_one_type) and each surviving candidate costs an ACL
+    # check and a database read, per type queried.
+    count: int = Field(default=10, ge=1, le=MAX_RESULTS_PER_TYPE)
     root_type: Literal["entity", "indicator", "dfiq"] | None = None
 
 

tests/apiv2/search.py

@@ -8,6 +8,7 @@
 from core.schemas import dfiq, entity, indicator, observable, rbac, roles, user
 from core.schemas.user import UserSensitive
 from core.web import webapp
+from core.web.apiv2.search import MAX_RESULTS_PER_TYPE
 
 client = TestClient(webapp.app)
 
@@ -165,6 +166,26 @@ def test_search_count_per_type_limits_results(self) -> None:
         self.assertEqual(sections["entity"]["total"], 3, data)
         self.assertEqual(len(sections["entity"]["results"]), 2, data)
 
+    def test_search_rejects_count_per_type_outside_the_allowed_range(self) -> None:
+        """A non-positive count_per_type used to be accepted and silently
+        return no results, which reads as "nothing matched" rather than
+        "your request was wrong"."""
+        for count_per_type in (0, -1, MAX_RESULTS_PER_TYPE + 1):
+            with self.subTest(count_per_type=count_per_type):
+                response = client.post(
+                    "/api/v2/search",
+                    json={"query": "test", "count_per_type": count_per_type},
+                )
+                self.assertEqual(response.status_code, 422, response.text)
+
+        for count_per_type in (1, MAX_RESULTS_PER_TYPE):
+            with self.subTest(count_per_type=count_per_type):
+                response = client.post(
+                    "/api/v2/search",
+                    json={"query": "test", "count_per_type": count_per_type},
+                )
+                self.assertEqual(response.status_code, 200, response.text)
+
 
 class searchRbacTest(unittest.TestCase):
     def setUp(self) -> None:

tests/core_tests/chromadb_test.py

@@ -8,6 +8,7 @@
 from core import database_arango
 from core.schemas import entity, rbac, roles, user
 from core.web import webapp
+from core.web.apiv2.search import MAX_RESULTS_PER_TYPE
 from plugins.analytics.public.chromadb_indexer import ChromaDBIndexer
 
 client = TestClient(webapp.app)
@@ -620,6 +621,30 @@ def test_documents_an_object_stops_producing_are_pruned(self, mock_get_client):
         }
         self.assertEqual(remaining, {"self", "approach:0"})
 
+    @mock.patch("core.chromadb_client.get_client")
+    def test_count_outside_the_allowed_range_is_rejected(self, mock_get_client):
+        """count is overfetched into ChromaDB's n_results, which raises on any
+        non-positive value -- unbounded, that surfaces as a 500 for what is a
+        client error. The upper bound caps the ACL checks and database reads a
+        single request can trigger."""
+        mock_get_client.return_value = self.chroma_client
+
+        for count in (0, -1, MAX_RESULTS_PER_TYPE + 1):
+            with self.subTest(count=count):
+                response = client.post(
+                    "/api/v2/search/semantic",
+                    json={"query": "russian actor", "count": count},
+                )
+                self.assertEqual(response.status_code, 422, response.json())
+
+        for count in (1, MAX_RESULTS_PER_TYPE):
+            with self.subTest(count=count):
+                response = client.post(
+                    "/api/v2/search/semantic",
+                    json={"query": "russian actor", "count": count},
+                )
+                self.assertEqual(response.status_code, 200, response.json())
+
 
 class SimilarityScoreTest(unittest.TestCase):
     def test_converts_squared_l2_distance_to_a_bounded_similarity(self):