From 15f5f03b008f82388d4ea24d2989e63883ee8878 Mon Sep 17 00:00:00 2001 From: yhmo Date: Fri, 23 Oct 2026 16:40:44 +0800 Subject: [PATCH] fix: escape iterator pk/id values and remove sparse vector stdout dump The query and search iterators interpolated VARCHAR primary-key and ID values into filter-expression string literals without escaping. A crafted value containing quotes or backslashes could break out of the literal and inject expression clauses that the server evaluates under the requesting user identity, bypassing client-side filters and enabling cross-scope data reads. Escape backslash, single/double quote, CR, LF and TAB per the Milvus expression grammar before interpolation, keeping benign values unchanged. Also stop printing raw sparse-vector JSON to stdout when parsing fails. Signed-off-by: yhmo --- src/impl/types/QueryIteratorImpl.cpp | 2 +- src/impl/types/SearchIteratorImpl.cpp | 2 +- src/impl/utils/DmlUtils.cpp | 1 - src/impl/utils/DqlUtils.h | 38 ++++++++++++ test/it/v1/TestQueryIterator.cpp | 86 ++++++++++++++++++++++++++ test/it/v1/TestSearchIterator.cpp | 88 +++++++++++++++++++++++++++ test/ut/utils/TestDqlUtils.cpp | 26 ++++++++ 7 files changed, 240 insertions(+), 3 deletions(-) diff --git a/src/impl/types/QueryIteratorImpl.cpp b/src/impl/types/QueryIteratorImpl.cpp index 55b5ea24..4b37a508 100644 --- a/src/impl/types/QueryIteratorImpl.cpp +++ b/src/impl/types/QueryIteratorImpl.cpp @@ -187,7 +187,7 @@ QueryIteratorImpl::setupNextFilter() { std::string pk_name = args_.PkSchema().Name(); std::string iter_filter; if (args_.PkSchema().FieldDataType() == DataType::VARCHAR) { - iter_filter = pk_name + " > " + "\"" + next_id_ + "\""; + iter_filter = pk_name + " > " + "\"" + EscapeExprStringValue(next_id_) + "\""; } else { iter_filter = pk_name + " > " + next_id_; } diff --git a/src/impl/types/SearchIteratorImpl.cpp b/src/impl/types/SearchIteratorImpl.cpp index dad91ddc..a8750e34 100644 --- a/src/impl/types/SearchIteratorImpl.cpp +++ b/src/impl/types/SearchIteratorImpl.cpp @@ -473,7 +473,7 @@ SearchIteratorImpl::filteredDuplicatedResultFilter() { if (args_.PkSchema().FieldDataType() == DataType::INT64) { filtered_ids_str << id; } else { - filtered_ids_str << "'" << id << "'"; + filtered_ids_str << "'" << EscapeExprStringValue(id) << "'"; } } diff --git a/src/impl/utils/DmlUtils.cpp b/src/impl/utils/DmlUtils.cpp index a9f90fa0..e7d08b05 100644 --- a/src/impl/utils/DmlUtils.cpp +++ b/src/impl/utils/DmlUtils.cpp @@ -269,7 +269,6 @@ EncodeSparseFloatVector(const SparseFloatVecFieldData::ElementT& sparse) { Status ParseSparseFloatVector(const nlohmann::json& obj, const std::string& field_name, std::map& pairs) { if (!obj.is_object()) { - std::cout << obj << std::endl; std::string err_msg = "Value type should be a dict for field: " + field_name; return {StatusCode::INVALID_ARGUMENT, err_msg}; } diff --git a/src/impl/utils/DqlUtils.h b/src/impl/utils/DqlUtils.h index f94dfa0c..4917e75f 100644 --- a/src/impl/utils/DqlUtils.h +++ b/src/impl/utils/DqlUtils.h @@ -35,6 +35,44 @@ namespace milvus { class SearchResponse; +/** + * @brief Escape a server-returned value for safe interpolation into a Milvus filter-expression + * string literal (single- or double-quoted). Milvus expression literals terminate on an unescaped + * quote/backslash and cannot contain raw line breaks, so values such as VARCHAR primary keys must + * be escaped before being embedded, otherwise a crafted value could inject expression clauses. + */ +inline std::string +EscapeExprStringValue(const std::string& value) { + std::string escaped; + escaped.reserve(value.size()); + for (const char c : value) { + switch (c) { + case '\\': + escaped += "\\\\"; + break; + case '"': + escaped += "\\\""; + break; + case '\'': + escaped += "\\'"; + break; + case '\r': + escaped += "\\r"; + break; + case '\n': + escaped += "\\n"; + break; + case '\t': + escaped += "\\t"; + break; + default: + escaped += c; + break; + } + } + return escaped; +} + Status CreateMilvusFieldData(const proto::schema::FieldData& proto_data, size_t offset, size_t count, FieldDataPtr& field_data); diff --git a/test/it/v1/TestQueryIterator.cpp b/test/it/v1/TestQueryIterator.cpp index e7790962..6cbea6b0 100644 --- a/test/it/v1/TestQueryIterator.cpp +++ b/test/it/v1/TestQueryIterator.cpp @@ -229,3 +229,89 @@ TEST_F(MilvusMockedTest, QueryIteratorUsesExplicitDatabaseForSchemaAndQuery) { EXPECT_TRUE(milvus::SchemaCache::GetInstance().Get(endpoint, database_name, collection_name, desc)); EXPECT_FALSE(milvus::SchemaCache::GetInstance().Get(endpoint, "default", collection_name, desc)); } + +TEST_F(MilvusMockedTest, QueryIteratorEscapesVarcharPkInPagingExpression) { + milvus::ConnectParam connect_param{"127.0.0.1", server_.ListenPort()}; + client_->Connect(connect_param); + + const std::string collection_name = "varchar_pk_col"; + milvus::CollectionSchema collection_schema(collection_name); + collection_schema.AddField(milvus::FieldSchema("pk", milvus::DataType::VARCHAR, "pk", true, false)); + collection_schema.AddField(milvus::FieldSchema("age", milvus::DataType::INT64)); + + // attack payload planted as a VARCHAR primary key value + const std::string injection_pk = "x\" or 1==1 or \"a\"==\"a"; + + size_t page_query_count = 0; + + EXPECT_CALL(service_, DescribeCollection( + _, Property(&DescribeCollectionRequest::collection_name, collection_name), _)) + .WillOnce([&](::grpc::ServerContext*, const DescribeCollectionRequest*, DescribeCollectionResponse* response) { + response->set_collectionid(100); + milvus::ConvertCollectionSchema(collection_schema, *response->mutable_schema()); + return ::grpc::Status{}; + }); + + EXPECT_CALL(service_, Query(_, _, _)) + .WillRepeatedly([&](::grpc::ServerContext*, const QueryRequest* request, QueryResults* response) { + // the session-setup query uses limit=1 + for (const auto& pair : request->query_params()) { + if (pair.key() == milvus::LIMIT && pair.value() == "1") { + response->set_session_ts(999999); + return ::grpc::Status{}; + } + } + + // the first paging page uses the plain user filter; every later page must + // keep the crafted pk inside the string literal: the quote after 'x' is + // escaped (`x\"`), so the raw `x" or` breakout never occurs + ++page_query_count; + if (page_query_count == 1) { + EXPECT_EQ(request->expr(), "age >= 0"); + } else { + EXPECT_EQ(request->expr(), " ( age >= 0 ) and pk > \"x\\\" or 1==1 or \\\"a\\\"==\\\"a\""); + EXPECT_NE(request->expr().find("x\\\" or"), std::string::npos); + EXPECT_EQ(request->expr().find("x\" or"), std::string::npos); + } + + // return one page with the crafted pk as the trailing value + auto pk_data = std::make_shared( + "pk", std::vector{"alpha", injection_pk}); + milvus::FieldDataSchema pk_bridge(pk_data, nullptr); + milvus::proto::schema::FieldData proto_pk; + auto status = milvus::CreateProtoFieldData(pk_bridge, proto_pk); + EXPECT_TRUE(status.IsOk()); + response->mutable_fields_data()->Add(std::move(proto_pk)); + + auto age_data = std::make_shared("age", std::vector{1, 2}); + milvus::FieldDataSchema age_bridge(age_data, nullptr); + milvus::proto::schema::FieldData proto_age; + status = milvus::CreateProtoFieldData(age_bridge, proto_age); + EXPECT_TRUE(status.IsOk()); + response->mutable_fields_data()->Add(std::move(proto_age)); + + return ::grpc::Status{}; + }); + + milvus::QueryIteratorArguments arguments; + arguments.SetBatchSize(2); + arguments.SetLimit(4); + arguments.SetCollectionName(collection_name); + arguments.SetFilter("age >= 0"); + arguments.SetConsistencyLevel(milvus::ConsistencyLevel::BOUNDED); + arguments.AddOutputField("pk"); + arguments.AddOutputField("age"); + + milvus::QueryIteratorPtr iterator; + auto status = client_->QueryIterator(arguments, iterator); + EXPECT_TRUE(status.IsOk()); + + // first page returns the crafted pk, second page re-queries with the escaped expression + milvus::QueryResults batch_results; + status = iterator->Next(batch_results); + EXPECT_TRUE(status.IsOk()); + EXPECT_EQ(batch_results.GetRowCount(), 2); + + status = iterator->Next(batch_results); + EXPECT_TRUE(status.IsOk()); +} diff --git a/test/it/v1/TestSearchIterator.cpp b/test/it/v1/TestSearchIterator.cpp index 1c65e40b..e9759f55 100644 --- a/test/it/v1/TestSearchIterator.cpp +++ b/test/it/v1/TestSearchIterator.cpp @@ -302,3 +302,91 @@ TEST_F(MilvusMockedTest, SearchIteratorV2PinsProbeTimestampForSessionConsistency milvus::CollectionTsCache::GetInstance().Invalidate(endpoint, "default", collection_name); } + +TEST_F(MilvusMockedTest, SearchIteratorEscapesVarcharIdsInNotInFilter) { + milvus::ConnectParam connect_param{"127.0.0.1", server_.ListenPort()}; + client_->Connect(connect_param); + + const std::string collection_name = "varchar_pk_search_col"; + milvus::CollectionSchema collection_schema(collection_name); + collection_schema.AddField(milvus::FieldSchema("pk", milvus::DataType::VARCHAR, "pk", true, false)); + collection_schema.AddField(milvus::FieldSchema("f16_vector", milvus::DataType::FLOAT16_VECTOR).WithDimension(16)); + + // attack payload planted as a VARCHAR primary key value + const std::string injection_id = "a' or '1'=='1"; + + EXPECT_CALL(service_, + DescribeCollection(_, Property(&DescribeCollectionRequest::collection_name, collection_name), _)) + .WillOnce([&](::grpc::ServerContext*, const DescribeCollectionRequest*, DescribeCollectionResponse* response) { + response->set_collectionid(100); + milvus::ConvertCollectionSchema(collection_schema, *response->mutable_schema()); + return ::grpc::Status{}; + }); + + auto build_page = [&](SearchResults* response, const std::string& pk_value, float score) { + response->mutable_status()->set_code(milvus::proto::common::ErrorCode::Success); + auto* results = response->mutable_results(); + results->set_top_k(1); + results->set_num_queries(1); + results->set_primary_field_name("pk"); + auto pk_data = std::make_shared("pk", std::vector{pk_value}); + milvus::FieldDataSchema bridge(pk_data, nullptr); + milvus::proto::schema::FieldData proto_field; + auto status = milvus::CreateProtoFieldData(bridge, proto_field); + EXPECT_TRUE(status.IsOk()); + results->mutable_fields_data()->Add(std::move(proto_field)); + results->mutable_ids()->mutable_str_id()->add_data(pk_value); + results->mutable_topks()->Add(1); + results->mutable_scores()->Add(score); + }; + + size_t search_call_count = 0; + EXPECT_CALL(service_, Search(_, _, _)) + .WillRepeatedly([&](::grpc::ServerContext*, const SearchRequest* request, SearchResults* response) { + ++search_call_count; + if (search_call_count == 1) { + // V2 compatibility probe: empty token makes the SDK fall back to the V1 iterator + response->mutable_results()->mutable_search_iterator_v2_results()->set_token(""); + return ::grpc::Status{}; + } + if (search_call_count == 2) { + // the init page uses the plain user filter and returns the crafted id at the tail + EXPECT_EQ(request->dsl(), "age >= 0"); + build_page(response, injection_id, 0.5f); + return ::grpc::Status{}; + } + // the fill search must keep the crafted id inside the single-quoted literal: + // `a' or '1'=='1` becomes `a\' or \'1\'==\'1`, so no raw breakout occurs + EXPECT_EQ(request->dsl(), "age >= 0 and pk not in ['a\\' or \\'1\\'==\\'1']"); + EXPECT_NE(request->dsl().find("a\\' or"), std::string::npos); + EXPECT_EQ(request->dsl().find("a' or"), std::string::npos); + build_page(response, "beta", 0.6f); + return ::grpc::Status{}; + }); + + milvus::SearchIteratorArguments arguments; + arguments.SetBatchSize(1); + arguments.SetLimit(2); + arguments.SetCollectionName(collection_name); + arguments.SetFilter("age >= 0"); + arguments.SetConsistencyLevel(milvus::ConsistencyLevel::STRONG); + arguments.SetMetricType(milvus::MetricType::COSINE); + arguments.AddOutputField("pk"); + std::vector vector(milvus::T_DIMENSION, 1.0f); + auto status = arguments.AddFloat16Vector("f16_vector", vector); + EXPECT_TRUE(status.IsOk()); + + milvus::SearchIteratorPtr iterator; + status = client_->SearchIterator(arguments, iterator); + EXPECT_TRUE(status.IsOk()); + + milvus::SingleResult first_page; + status = iterator->Next(first_page); + EXPECT_TRUE(status.IsOk()); + EXPECT_EQ(first_page.GetRowCount(), 1); + + milvus::SingleResult second_page; + status = iterator->Next(second_page); + EXPECT_TRUE(status.IsOk()); + EXPECT_EQ(second_page.GetRowCount(), 1); +} diff --git a/test/ut/utils/TestDqlUtils.cpp b/test/ut/utils/TestDqlUtils.cpp index dcf3abd6..719bdba7 100644 --- a/test/ut/utils/TestDqlUtils.cpp +++ b/test/ut/utils/TestDqlUtils.cpp @@ -112,6 +112,32 @@ TEST_F(DqlUtilsTest, FillSearchResponseExtraInfoResetsStaleMetadataTest) { EXPECT_FLOAT_EQ(response.CacheHitRatio(), -1.0f); } +TEST_F(DqlUtilsTest, EscapeExprStringValueTest) { + // benign values pass through unchanged + EXPECT_EQ(milvus::EscapeExprStringValue(""), ""); + EXPECT_EQ(milvus::EscapeExprStringValue("hello"), "hello"); + EXPECT_EQ(milvus::EscapeExprStringValue("hello world"), "hello world"); + + // quotes and backslashes are escaped so a crafted value cannot break out of + // a Milvus expression string literal (single- or double-quoted) + EXPECT_EQ(milvus::EscapeExprStringValue("a\"b"), "a\\\"b"); + EXPECT_EQ(milvus::EscapeExprStringValue("a'b"), "a\\'b"); + EXPECT_EQ(milvus::EscapeExprStringValue("a\\b"), "a\\\\b"); + EXPECT_EQ(milvus::EscapeExprStringValue("a\\\"b"), "a\\\\\\\"b"); + + // control characters are escaped per the expression grammar + EXPECT_EQ(milvus::EscapeExprStringValue("a\tb"), "a\\tb"); + EXPECT_EQ(milvus::EscapeExprStringValue("a\nb"), "a\\nb"); + EXPECT_EQ(milvus::EscapeExprStringValue("a\rb"), "a\\rb"); + + // an injection payload is fully contained: no raw quote remains to close the + // literal early and no standalone `or 1==1` clause is formed + EXPECT_EQ(milvus::EscapeExprStringValue("x\" or 1==1 or \"a\"==\"a"), + "x\\\" or 1==1 or \\\"a\\\"==\\\"a"); + EXPECT_EQ(milvus::EscapeExprStringValue("a' or '1'=='1"), + "a\\' or \\'1\\'==\\'1"); +} + TEST_F(DqlUtilsTest, DeduceGuaranteeTimestampTest) { auto& cache = milvus::CollectionTsCache::GetInstance(); cache.Clear();