JAVA-5767 Support $lookup in CSFLE and QE (#1638)

katcharov · web-flow · commit eeceb6323d9c · 2025-02-27T12:31:46.000-07:00
JAVA-5767
diff --git a/driver-core/src/test/resources/client-side-encryption-data/lookup/key-doc.json b/driver-core/src/test/resources/client-side-encryption-data/lookup/key-doc.json
@@ -0,0 +1,30 @@
+{
+    "_id": {
+        "$binary": {
+            "base64": "EjRWeBI0mHYSNBI0VniQEg==",
+            "subType": "04"
+        }
+    },
+    "keyMaterial": {
+        "$binary": {
+            "base64": "sHe0kz57YW7v8g9VP9sf/+K1ex4JqKc5rf/URX3n3p8XdZ6+15uXPaSayC6adWbNxkFskuMCOifDoTT+rkqMtFkDclOy884RuGGtUysq3X7zkAWYTKi8QAfKkajvVbZl2y23UqgVasdQu3OVBQCrH/xY00nNAs/52e958nVjBuzQkSb1T8pKJAyjZsHJ60+FtnfafDZSTAIBJYn7UWBCwQ==",
+            "subType": "00"
+        }
+    },
+    "creationDate": {
+        "$date": {
+            "$numberLong": "1648914851981"
+        }
+    },
+    "updateDate": {
+        "$date": {
+            "$numberLong": "1648914851981"
+        }
+    },
+    "status": {
+        "$numberInt": "0"
+    },
+    "masterKey": {
+        "provider": "local"
+    }
+}
diff --git a/driver-core/src/test/resources/client-side-encryption-data/lookup/schema-csfle.json b/driver-core/src/test/resources/client-side-encryption-data/lookup/schema-csfle.json
@@ -0,0 +1,19 @@
+{
+    "properties": {
+        "csfle": {
+            "encrypt": {
+                "keyId": [
+                    {
+                        "$binary": {
+                            "base64": "EjRWeBI0mHYSNBI0VniQEg==",
+                            "subType": "04"
+                        }
+                    }
+                ],
+                "bsonType": "string",
+                "algorithm": "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic"
+            }
+        }
+    },
+    "bsonType": "object"
+}
diff --git a/driver-core/src/test/resources/client-side-encryption-data/lookup/schema-csfle2.json b/driver-core/src/test/resources/client-side-encryption-data/lookup/schema-csfle2.json
@@ -0,0 +1,19 @@
+{
+    "properties": {
+        "csfle2": {
+            "encrypt": {
+                "keyId": [
+                    {
+                        "$binary": {
+                            "base64": "EjRWeBI0mHYSNBI0VniQEg==",
+                            "subType": "04"
+                        }
+                    }
+                ],
+                "bsonType": "string",
+                "algorithm": "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic"
+            }
+        }
+    },
+    "bsonType": "object"
+}
diff --git a/driver-core/src/test/resources/client-side-encryption-data/lookup/schema-qe.json b/driver-core/src/test/resources/client-side-encryption-data/lookup/schema-qe.json
@@ -0,0 +1,20 @@
+{
+    "escCollection": "enxcol_.qe.esc",
+    "ecocCollection": "enxcol_.qe.ecoc",
+    "fields": [
+        {
+            "keyId": {
+                "$binary": {
+                    "base64": "EjRWeBI0mHYSNBI0VniQEg==",
+                    "subType": "04"
+                }
+            },
+            "path": "qe",
+            "bsonType": "string",
+            "queries": {
+                "queryType": "equality",
+                "contention": 0
+            }
+        }
+    ]
+}
diff --git a/driver-core/src/test/resources/client-side-encryption-data/lookup/schema-qe2.json b/driver-core/src/test/resources/client-side-encryption-data/lookup/schema-qe2.json
@@ -0,0 +1,20 @@
+{
+    "escCollection": "enxcol_.qe2.esc",
+    "ecocCollection": "enxcol_.qe2.ecoc",
+    "fields": [
+        {
+            "keyId": {
+                "$binary": {
+                    "base64": "EjRWeBI0mHYSNBI0VniQEg==",
+                    "subType": "04"
+                }
+            },
+            "path": "qe2",
+            "bsonType": "string",
+            "queries": {
+                "queryType": "equality",
+                "contention": 0
+            }
+        }
+    ]
+}
diff --git a/driver-core/src/test/unit/com/mongodb/client/model/bulk/BaseClientDeleteOptionsTest.java b/driver-core/src/test/unit/com/mongodb/client/model/bulk/BaseClientDeleteOptionsTest.java
@@ -16,7 +16,7 @@
 
 package com.mongodb.client.model.bulk;
 
-import com.mongodb.MongoBaseInterfaceAssertions;
+import com.mongodb.testing.MongoBaseInterfaceAssertions;
 import org.junit.jupiter.api.Test;
 
 class BaseClientDeleteOptionsTest {
diff --git a/driver-core/src/test/unit/com/mongodb/client/model/bulk/BaseClientUpdateOptionsTest.java b/driver-core/src/test/unit/com/mongodb/client/model/bulk/BaseClientUpdateOptionsTest.java
@@ -16,7 +16,7 @@
 
 package com.mongodb.client.model.bulk;
 
-import com.mongodb.MongoBaseInterfaceAssertions;
+import com.mongodb.testing.MongoBaseInterfaceAssertions;
 import org.junit.jupiter.api.Test;
 
 class BaseClientUpdateOptionsTest {
diff --git a/driver-core/src/test/unit/com/mongodb/client/model/bulk/BaseClientUpsertableWriteModelOptionsTest.java b/driver-core/src/test/unit/com/mongodb/client/model/bulk/BaseClientUpsertableWriteModelOptionsTest.java
@@ -16,7 +16,7 @@
 
 package com.mongodb.client.model.bulk;
 
-import com.mongodb.MongoBaseInterfaceAssertions;
+import com.mongodb.testing.MongoBaseInterfaceAssertions;
 import org.junit.jupiter.api.Test;
 
 final class BaseClientUpsertableWriteModelOptionsTest {
diff --git a/driver-core/src/test/unit/com/mongodb/client/model/bulk/BaseClientWriteModelOptionsTest.java b/driver-core/src/test/unit/com/mongodb/client/model/bulk/BaseClientWriteModelOptionsTest.java
@@ -16,7 +16,7 @@
 
 package com.mongodb.client.model.bulk;
 
-import com.mongodb.MongoBaseInterfaceAssertions;
+import com.mongodb.testing.MongoBaseInterfaceAssertions;
 import org.junit.jupiter.api.Test;
 
 final class BaseClientWriteModelOptionsTest {
diff --git a/driver-core/src/test/unit/com/mongodb/testing/MongoAssertions.java b/driver-core/src/test/unit/com/mongodb/testing/MongoAssertions.java
@@ -0,0 +1,43 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.testing;
+
+import org.junit.jupiter.api.function.Executable;
+import org.opentest4j.AssertionFailedError;
+
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public final class MongoAssertions {
+
+    private MongoAssertions() {
+        //NOP
+    }
+
+    public static <T extends Throwable> void assertCause(
+            final Class<T> expectedCause, final String expectedMessageFragment, final Executable e) {
+        Throwable cause = assertThrows(Throwable.class, e);
+        while (cause.getCause() != null) {
+            cause = cause.getCause();
+        }
+        if (!cause.getMessage().contains(expectedMessageFragment)) {
+            throw new AssertionFailedError("Unexpected message: " + cause.getMessage(), cause);
+        }
+        if (!expectedCause.isInstance(cause)) {
+            throw new AssertionFailedError("Unexpected cause: " + cause.getClass(), assertThrows(Throwable.class, e));
+        }
+    }
+}
diff --git a/driver-core/src/test/unit/com/mongodb/testing/MongoBaseInterfaceAssertions.java b/driver-core/src/test/unit/com/mongodb/testing/MongoBaseInterfaceAssertions.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.mongodb;
+package com.mongodb.testing;
 
 import org.reflections.Reflections;
 
diff --git a/driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/CollectionInfoRetriever.java b/driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/CollectionInfoRetriever.java
@@ -20,7 +20,7 @@
 import com.mongodb.lang.Nullable;
 import com.mongodb.reactivestreams.client.MongoClient;
 import org.bson.BsonDocument;
-import reactor.core.publisher.Mono;
+import reactor.core.publisher.Flux;
 
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.reactivestreams.client.internal.TimeoutHelper.databaseWithTimeoutDeferred;
@@ -35,8 +35,8 @@ class CollectionInfoRetriever {
         this.client = notNull("client", client);
     }
 
-    public Mono<BsonDocument> filter(final String databaseName, final BsonDocument filter, @Nullable final Timeout operationTimeout) {
+    public Flux<BsonDocument> filter(final String databaseName, final BsonDocument filter, @Nullable final Timeout operationTimeout) {
         return databaseWithTimeoutDeferred(client.getDatabase(databaseName), TIMEOUT_ERROR_MESSAGE, operationTimeout)
-                .flatMap(database -> Mono.from(database.listCollections(BsonDocument.class).filter(filter).first()));
+                .flatMapMany(database -> Flux.from(database.listCollections(BsonDocument.class).filter(filter)));
     }
 }
diff --git a/driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/Crypt.java b/driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/Crypt.java
@@ -304,10 +304,8 @@ private void collInfo(final MongoCryptContext cryptContext,
         } else {
             collectionInfoRetriever.filter(databaseName, cryptContext.getMongoOperation(), operationTimeout)
                     .contextWrite(sink.contextView())
-                    .doOnSuccess(result -> {
-                        if (result != null) {
-                            cryptContext.addMongoOperationResult(result);
-                        }
+                    .doOnNext(result -> cryptContext.addMongoOperationResult(result))
+                    .doOnComplete(() -> {
                         cryptContext.completeMongoOperation();
                         executeStateMachineWithSink(cryptContext, databaseName, sink, operationTimeout);
                     })
diff --git a/driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/ClientSideEncryption25LookupProseTests.java b/driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/ClientSideEncryption25LookupProseTests.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.reactivestreams.client;
+
+import com.mongodb.ClientEncryptionSettings;
+import com.mongodb.MongoClientSettings;
+import com.mongodb.client.MongoClient;
+import com.mongodb.client.vault.ClientEncryption;
+import com.mongodb.reactivestreams.client.syncadapter.SyncClientEncryption;
+import com.mongodb.reactivestreams.client.syncadapter.SyncMongoClient;
+import com.mongodb.reactivestreams.client.vault.ClientEncryptions;
+
+public class ClientSideEncryption25LookupProseTests extends com.mongodb.client.ClientSideEncryption25LookupProseTests {
+
+    @Override
+    protected MongoClient createMongoClient(final MongoClientSettings settings) {
+        return new SyncMongoClient(MongoClients.create(settings));
+    }
+
+    @Override
+    protected ClientEncryption createClientEncryption(final ClientEncryptionSettings settings) {
+        return new SyncClientEncryption(ClientEncryptions.create(settings));
+    }
+
+}
diff --git a/driver-sync/src/main/com/mongodb/client/internal/CollectionInfoRetriever.java b/driver-sync/src/main/com/mongodb/client/internal/CollectionInfoRetriever.java
@@ -21,6 +21,9 @@
 import com.mongodb.lang.Nullable;
 import org.bson.BsonDocument;
 
+import java.util.ArrayList;
+import java.util.List;
+
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.client.internal.TimeoutHelper.databaseWithTimeout;
 
@@ -33,9 +36,10 @@ class CollectionInfoRetriever {
         this.client = notNull("client", client);
     }
 
-    @Nullable
-    public BsonDocument filter(final String databaseName, final BsonDocument filter, @Nullable final Timeout operationTimeout) {
-        return databaseWithTimeout(client.getDatabase(databaseName), TIMEOUT_ERROR_MESSAGE,
-                operationTimeout).listCollections(BsonDocument.class).filter(filter).first();
+    public List<BsonDocument> filter(final String databaseName, final BsonDocument filter, @Nullable final Timeout operationTimeout) {
+        return databaseWithTimeout(client.getDatabase(databaseName), TIMEOUT_ERROR_MESSAGE, operationTimeout)
+                .listCollections(BsonDocument.class)
+                .filter(filter)
+                .into(new ArrayList<>());
     }
 }
diff --git a/driver-sync/src/main/com/mongodb/client/internal/Crypt.java b/driver-sync/src/main/com/mongodb/client/internal/Crypt.java
@@ -41,6 +41,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
+import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
 
@@ -308,9 +309,10 @@ private void fetchCredentials(final MongoCryptContext cryptContext) {
 
     private void collInfo(final MongoCryptContext cryptContext, final String databaseName, @Nullable final Timeout operationTimeout) {
         try {
-            BsonDocument collectionInfo = assertNotNull(collectionInfoRetriever).filter(databaseName, cryptContext.getMongoOperation(), operationTimeout);
-            if (collectionInfo != null) {
-                cryptContext.addMongoOperationResult(collectionInfo);
+            List<BsonDocument> results = assertNotNull(collectionInfoRetriever)
+                    .filter(databaseName, cryptContext.getMongoOperation(), operationTimeout);
+            for (BsonDocument result : results) {
+                cryptContext.addMongoOperationResult(result);
             }
             cryptContext.completeMongoOperation();
         } catch (Throwable t) {
diff --git a/driver-sync/src/test/functional/com/mongodb/client/ClientSideEncryption25LookupProseTests.java b/driver-sync/src/test/functional/com/mongodb/client/ClientSideEncryption25LookupProseTests.java
diff --git a/driver-sync/src/test/functional/com/mongodb/internal/connection/OidcAuthenticationProseTests.java b/driver-sync/src/test/functional/com/mongodb/internal/connection/OidcAuthenticationProseTests.java
diff --git a/mongodb-crypt/build.gradle.kts b/mongodb-crypt/build.gradle.kts
diff --git a/mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java b/mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java
diff --git a/mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoCryptImpl.java b/mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoCryptImpl.java
