diff --git a/weed/s3api/s3vectors/handlers.go b/weed/s3api/s3vectors/handlers.go index 7a0ef0e279..b8ede3c827 100644 --- a/weed/s3api/s3vectors/handlers.go +++ b/weed/s3api/s3vectors/handlers.go @@ -1301,12 +1301,25 @@ func toOutputVectors(vectors []*vector_pb.Vector, returnData, returnMetadata boo return out } +// vectorKeysBounds enforces the AWS 1..max key list window; GetVectors is +// capped lower than a PutVectors or DeleteVectors batch. +func vectorKeysBounds(keys []string, max int) error { + if len(keys) == 0 || len(keys) > max { + return fmt.Errorf("keys must be between 1 and %d", max) + } + return nil +} + func (s *Server) handleGetVectors(w http.ResponseWriter, r *http.Request) { var req GetVectorsRequest if err := readJSON(r, &req); err != nil { writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) return } + if err := vectorKeysBounds(req.Keys, vectorsMaxGetKeys); err != nil { + writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) + return + } arns, err := s.resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn) if err != nil { writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) @@ -1409,6 +1422,7 @@ func (s *Server) handleListVectors(w http.ResponseWriter, r *http.Request) { // pages like a plain one does; vectorsMaxTopK is the AWS QueryVectors limit. const ( vectorsDefaultPage = 500 + vectorsMaxGetKeys = 100 vectorsMaxPutBatch = 500 vectorsMaxPage = 1000 vectorsMaxListPage = 500 @@ -1469,6 +1483,10 @@ func (s *Server) handleDeleteVectors(w http.ResponseWriter, r *http.Request) { writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) return } + if err := vectorKeysBounds(req.Keys, vectorsMaxPutBatch); err != nil { + writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) + return + } arns, err := s.resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn) if err != nil { writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) diff --git a/weed/s3api/s3vectors/handlers_test.go b/weed/s3api/s3vectors/handlers_test.go index 045e51ae51..85bd4fd5c7 100644 --- a/weed/s3api/s3vectors/handlers_test.go +++ b/weed/s3api/s3vectors/handlers_test.go @@ -27,6 +27,8 @@ type fakeVectors struct { rows []*vector_pb.Vector created []*vector_pb.CreateIndexRequest puts []*vector_pb.PutVectorsRequest + gets []*vector_pb.GetVectorsRequest + deletes []*vector_pb.DeleteVectorsRequest } func newFakeVectors(n int) *fakeVectors { @@ -73,9 +75,15 @@ func (f *fakeVectors) CreateIndex(ctx context.Context, req *vector_pb.CreateInde } func (f *fakeVectors) GetVectors(ctx context.Context, req *vector_pb.GetVectorsRequest, opts ...grpc.CallOption) (*vector_pb.GetVectorsResponse, error) { + f.gets = append(f.gets, req) return &vector_pb.GetVectorsResponse{}, nil } +func (f *fakeVectors) DeleteVectors(ctx context.Context, req *vector_pb.DeleteVectorsRequest, opts ...grpc.CallOption) (*vector_pb.DeleteVectorsResponse, error) { + f.deletes = append(f.deletes, req) + return &vector_pb.DeleteVectorsResponse{}, nil +} + func (f *fakeVectors) PutVectors(ctx context.Context, req *vector_pb.PutVectorsRequest, opts ...grpc.CallOption) (*vector_pb.PutVectorsResponse, error) { f.puts = append(f.puts, req) return &vector_pb.PutVectorsResponse{}, nil @@ -907,3 +915,55 @@ func TestDeleteVectorBucketPolicyRemovesLegacyDoc(t *testing.T) { t.Fatalf("empty shared policy should be removed: %s", doc) } } + +func keys(n int) []string { + out := make([]string, n) + for i := range out { + out[i] = "k" + } + return out +} + +func TestVectorKeysBounds(t *testing.T) { + for _, tc := range []struct { + name string + keys []string + max int + ok bool + }{ + {"empty", nil, vectorsMaxGetKeys, false}, + {"one", []string{"a"}, vectorsMaxGetKeys, true}, + {"get full", keys(vectorsMaxGetKeys), vectorsMaxGetKeys, true}, + {"get over", keys(vectorsMaxGetKeys + 1), vectorsMaxGetKeys, false}, + {"delete full", keys(vectorsMaxPutBatch), vectorsMaxPutBatch, true}, + {"delete over", keys(vectorsMaxPutBatch + 1), vectorsMaxPutBatch, false}, + } { + if err := vectorKeysBounds(tc.keys, tc.max); (err == nil) != tc.ok { + t.Fatalf("%s: accepted=%v, want %v", tc.name, err == nil, tc.ok) + } + } +} + +// An out-of-window key list is rejected before the engine sees it. +func TestVectorKeysBoundsRejectsBeforeEngine(t *testing.T) { + f := newFakeVectors(0) + s := &Server{vectors: f} + for _, tc := range []struct { + name string + handler http.HandlerFunc + payload interface{} + }{ + {"get empty", s.handleGetVectors, &GetVectorsRequest{}}, + {"get over", s.handleGetVectors, &GetVectorsRequest{Keys: keys(vectorsMaxGetKeys + 1)}}, + {"delete empty", s.handleDeleteVectors, &DeleteVectorsRequest{}}, + {"delete over", s.handleDeleteVectors, &DeleteVectorsRequest{Keys: keys(vectorsMaxPutBatch + 1)}}, + } { + w := postJSONAs(t, "", tc.handler, tc.payload) + if w.Code != http.StatusBadRequest || !strings.Contains(w.Body.String(), "ValidationException") { + t.Fatalf("%s: %d %s, want 400 ValidationException", tc.name, w.Code, w.Body.String()) + } + } + if len(f.gets)+len(f.deletes) != 0 { + t.Fatalf("engine saw %d get/%d delete calls, want none", len(f.gets), len(f.deletes)) + } +}