Skip to content

Commit da46146

Browse files
erinzhu001MongoDB Bot
erinzhu001
authored and
MongoDB Bot
committed
SERVER-101843 Add test with payload for $encStrEndsWith (#35269)
GitOrigin-RevId: b15948ad1ea30f304b574518314c14a610b18c20
1 parent 06d1b0c commit da46146

File tree

1 file changed

+113
-41
lines changed

1 file changed

+113
-41
lines changed

‎src/mongo/db/query/fle/text_search_predicate_test.cpp

+113-41
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,6 @@ BSONObj generateFFP(StringData path, Value value, EncryptionPlaceholderContext c
196196
BSONObj doc = BSON("value" << value);
197197
auto element = doc.firstElement();
198198
auto cdrValue = ConstDataRange(element.value(), element.value() + element.valuesize());
199-
auto fpp = FLEClientCrypto::serializeFindPayloadV2(indexKeyAndId, userKeyAndId, element, 0);
200199

201200
// Tokens are made according to fle_tokens.h.
202201
auto collectionToken = CollectionsLevel1Token::deriveFrom(indexKey.data);
@@ -231,6 +230,24 @@ BSONObj generateFFP(StringData path, Value value, EncryptionPlaceholderContext c
231230
escTextPrefixDerivedFromDataToken,
232231
serverTextPrefixDerivedFromDataToken});
233232
} break;
233+
case EncryptionPlaceholderContext::kTextSuffixComparison: {
234+
auto edcTextSuffixToken = EDCTextSuffixToken::deriveFrom(edcToken);
235+
auto edcTextSuffixDerivedFromDataToken =
236+
EDCTextSuffixDerivedFromDataToken::deriveFrom(edcTextSuffixToken, cdrValue);
237+
238+
auto escTextSuffixToken = ESCTextSuffixToken::deriveFrom(escToken);
239+
auto escTextSuffixDerivedFromDataToken =
240+
ESCTextSuffixDerivedFromDataToken::deriveFrom(escTextSuffixToken, cdrValue);
241+
242+
auto serverTextSuffixToken = ServerTextSuffixToken::deriveFrom(serverToken);
243+
auto serverTextSuffixDerivedFromDataToken =
244+
ServerTextSuffixDerivedFromDataToken::deriveFrom(serverTextSuffixToken, cdrValue);
245+
246+
payload.getTokenSets().setSuffixTokens(
247+
TextSuffixFindTokenSet{edcTextSuffixDerivedFromDataToken,
248+
escTextSuffixDerivedFromDataToken,
249+
serverTextSuffixDerivedFromDataToken});
250+
} break;
234251
default:
235252
MONGO_UNIMPLEMENTED_TASSERT(10172502);
236253
}
@@ -246,8 +263,7 @@ template <typename T>
246263
std::unique_ptr<Expression> makeEncStrStartsWith(ExpressionContext* const expCtx,
247264
StringData path,
248265
T value) {
249-
// A BSONObj corresponds to tests using an actual FindTextPayload, while a Value is used
250-
// otherwise.
266+
// A BSONObj corresponds to using an actual FindTextPayload, while a Value is used otherwise.
251267
static_assert(std::is_same_v<T, BSONObj> || std::is_same_v<T, Value>);
252268

253269
auto fieldpath = ExpressionFieldPath::createPathFromString(
@@ -264,14 +280,25 @@ std::unique_ptr<Expression> makeEncStrStartsWith(ExpressionContext* const expCtx
264280
}
265281
}
266282

267-
std::unique_ptr<Expression> makeEncStrEndsWithAggExpr(ExpressionContext* const expCtx,
268-
StringData path,
269-
Value value) {
283+
template <typename T>
284+
std::unique_ptr<Expression> makeEncStrEndsWith(ExpressionContext* const expCtx,
285+
StringData path,
286+
T value) {
287+
// A BSONObj corresponds to using an actual FindTextPayload, while a Value is used otherwise.
288+
static_assert(std::is_same_v<T, BSONObj> || std::is_same_v<T, Value>);
289+
270290
auto fieldpath = ExpressionFieldPath::createPathFromString(
271291
expCtx, path.toString(), expCtx->variablesParseState);
272-
auto textExpr = make_intrusive<ExpressionConstant>(expCtx, value);
273-
return std::make_unique<ExpressionEncStrEndsWith>(
274-
expCtx, std::move(fieldpath), std::move(textExpr));
292+
293+
if constexpr (std::is_same_v<T, BSONObj>) {
294+
return std::make_unique<ExpressionEncStrEndsWith>(
295+
expCtx,
296+
std::move(fieldpath),
297+
ExpressionConstant::parse(expCtx, value.firstElement(), expCtx->variablesParseState));
298+
} else {
299+
return std::make_unique<ExpressionEncStrEndsWith>(
300+
expCtx, std::move(fieldpath), make_intrusive<ExpressionConstant>(expCtx, value));
301+
}
275302
}
276303

277304
std::unique_ptr<Expression> makeEncStrContainsAggExpr(ExpressionContext* const expCtx,
@@ -322,6 +349,34 @@ void runEncStrStartsWithTest(T& predicate,
322349
ASSERT_BSONOBJ_EQ(expectedResult, actualBson);
323350
}
324351

352+
template <typename T>
353+
void runEncStrEndsWithTest(T& predicate,
354+
ExpressionContext* const expCtx,
355+
StringData fieldName,
356+
StringData valueSD,
357+
std::vector<PrfBlock> tags,
358+
BSONObj expectedResult) {
359+
static_assert(std::is_same_v<T, MockTextSearchPredicate> ||
360+
std::is_same_v<T, MockTextSearchPredicateWithFFP>);
361+
362+
std::unique_ptr<Expression> input = nullptr;
363+
Value value = Value(valueSD);
364+
365+
if constexpr (std::is_same_v<T, MockTextSearchPredicateWithFFP>) {
366+
BSONObj ffp =
367+
generateFFP(fieldName, value, EncryptionPlaceholderContext::kTextSuffixComparison);
368+
predicate.setEncryptedTags(Value(ffp.firstElement()), tags);
369+
input = makeEncStrEndsWith(expCtx, fieldName, ffp);
370+
} else {
371+
predicate.setEncryptedTags(valueSD, std::move(tags));
372+
input = makeEncStrEndsWith(expCtx, fieldName, value);
373+
}
374+
375+
auto actual = predicate.rewrite(input.get());
376+
auto actualBson = actual->serialize().getDocument().toBson();
377+
ASSERT_BSONOBJ_EQ(expectedResult, actualBson);
378+
}
379+
325380
TEST_F(TextSearchPredicateRewriteTest, Enc_Starts_With_NoFFP_Expr) {
326381
std::unique_ptr<Expression> input =
327382
makeEncStrStartsWith(_mock.getExpressionContext(), "ssn"_sd, Value("5"_sd));
@@ -381,42 +436,59 @@ TEST_F(TextSearchPredicateRewriteTestWithFFP, Enc_Starts_With_FFP_Expr) {
381436

382437
TEST_F(TextSearchPredicateRewriteTest, Enc_Str_Ends_With_NoFFP_Expr) {
383438
std::unique_ptr<Expression> input =
384-
makeEncStrEndsWithAggExpr(_mock.getExpressionContext(), "ssn"_sd, Value("5"_sd));
439+
makeEncStrEndsWith(_mock.getExpressionContext(), "ssn"_sd, Value("5"_sd));
385440
ASSERT_EQ(_predicate.rewrite(input.get()), nullptr);
386441
}
387442

388443
TEST_F(TextSearchPredicateRewriteTest, Enc_Str_Ends_With_Expr) {
389-
std::unique_ptr<Expression> input =
390-
makeEncStrEndsWithAggExpr(_mock.getExpressionContext(), "ssn"_sd, Value("hello"_sd));
391-
std::vector<PrfBlock> tags = {{1}, {2}};
392-
393-
_predicate.setEncryptedTags("hello", tags);
394-
395-
auto actual = _predicate.rewrite(input.get());
396-
auto actualBson = actual->serialize().getDocument().toBson();
444+
auto expectedResult = fromjson(R"({
445+
"$or": [
446+
{
447+
"$in": [
448+
{
449+
"$const": {"$binary":{"base64":"AQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=","subType":"0"}}
450+
},
451+
"$__safeContent__"
452+
]
453+
},
454+
{
455+
"$in": [
456+
{
457+
"$const": {"$binary":{"base64":"AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=","subType":"0"}}
458+
},
459+
"$__safeContent__"
460+
]
461+
}
462+
]
463+
})");
464+
runEncStrEndsWithTest(
465+
_predicate, _mock.getExpressionContext(), "ssn"_sd, "hello"_sd, {{1}, {2}}, expectedResult);
466+
}
397467

398-
ASSERT_BSONOBJ_EQ_AUTO( // NOLINT
399-
R"({
400-
"$or": [
401-
{
402-
"$in": [
403-
{
404-
"$const": {"$binary":{"base64":"AQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=","subType":"0"}}
405-
},
406-
"$__safeContent__"
407-
]
408-
},
409-
{
410-
"$in": [
411-
{
412-
"$const": {"$binary":{"base64":"AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=","subType":"0"}}
413-
},
414-
"$__safeContent__"
415-
]
416-
}
417-
]
418-
})",
419-
actualBson);
468+
// Tests same as above with an actual FindTextPayload.
469+
TEST_F(TextSearchPredicateRewriteTestWithFFP, Enc_Str_Ends_With_Expr) {
470+
auto expectedResult = fromjson(R"({
471+
"$or": [
472+
{
473+
"$in": [
474+
{
475+
"$const": {"$binary":{"base64":"AQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=","subType":"0"}}
476+
},
477+
"$__safeContent__"
478+
]
479+
},
480+
{
481+
"$in": [
482+
{
483+
"$const": {"$binary":{"base64":"AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=","subType":"0"}}
484+
},
485+
"$__safeContent__"
486+
]
487+
}
488+
]
489+
})");
490+
runEncStrEndsWithTest(
491+
_predicate, _mock.getExpressionContext(), "ssn"_sd, "hello"_sd, {{1}, {2}}, expectedResult);
420492
}
421493

422494
TEST_F(TextSearchPredicateRewriteTest, Enc_Str_Contains_NoFFP_Expr) {
@@ -531,7 +603,7 @@ TEST_F(TextSearchPredicateCollScanRewriteTest, Enc_Str_Starts_With_Expr) {
531603

532604
TEST_F(TextSearchPredicateCollScanRewriteTest, Enc_Str_Ends_With_Expr) {
533605
std::unique_ptr<Expression> input =
534-
makeEncStrEndsWithAggExpr(_mock.getExpressionContext(), "ssn"_sd, Value("hello"_sd));
606+
makeEncStrEndsWith(_mock.getExpressionContext(), "ssn"_sd, Value("hello"_sd));
535607

536608
// Serialize the expression before any rewrites occur to validate later.
537609
auto expressionPreRewrite = input->serialize().getDocument().toBson();

0 commit comments

Comments
 (0)