Add the ClientEncryption.createEncryptedCollection helper method (#…

…1079)

Add the `ClientEncryption.createEncryptedCollection` helper method

JAVA-4679

config/spotbugs/exclude.xml

@@ -240,4 +240,18 @@
         <Method name="acquirePermitOrGetAvailableOpenedConnection"/>
         <Bug pattern="NS_NON_SHORT_CIRCUIT"/>
     </Match>
+
+    <!-- Can actually be null, but is not annotated as `@Nullable`. Annotating it as such causes warnings
+         in other places where `null` is not handled, see https://jira.mongodb.org/browse/JAVA-4861.
+         When the aforementioned ticket is done, it will be clear what to do with the warnings suppressed here. -->
+    <Match>
+        <Class name="com.mongodb.client.internal.ClientEncryptionImpl"/>
+        <Method name="createEncryptedCollection"/>
+        <Bug pattern="RCN_REDUNDANT_NULLCHECK_OF_NONNULL_VALUE"/>
+    </Match>
+    <Match>
+        <Class name="com.mongodb.reactivestreams.client.internal.vault.ClientEncryptionImpl"/>
+        <Method name="~.*createEncryptedCollection.*"/>
+        <Bug pattern="RCN_REDUNDANT_NULLCHECK_OF_NONNULL_VALUE"/>
+    </Match>
 </FindBugsFilter>

driver-core/src/main/com/mongodb/MongoUpdatedEncryptedFieldsException.java

@@ -0,0 +1,67 @@
+/*
+ * 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;
+
+import com.mongodb.annotations.Beta;
+import org.bson.BsonDocument;
+
+import static com.mongodb.assertions.Assertions.assertNotNull;
+
+/**
+ * An exception thrown by methods that may automatically create data encryption keys
+ * where needed based on the {@code encryptedFields} configuration.
+ *
+ * @since 4.9
+ */
+@Beta(Beta.Reason.SERVER)
+public final class MongoUpdatedEncryptedFieldsException extends MongoClientException {
+    private static final long serialVersionUID = 1;
+
+    private final BsonDocument encryptedFields;
+
+    /**
+     * Not part of the public API.
+     *
+     * @param encryptedFields The (partially) updated {@code encryptedFields} document,
+     * which allows users to infer which data keys are known to be created before the exception happened
+     * (see {@link #getEncryptedFields()} for more details).
+     * Reporting this back to a user may be helpful because creation of a data key includes persisting it in the key vault.
+     * @param msg The message.
+     * @param cause The cause.
+     */
+    public MongoUpdatedEncryptedFieldsException(final BsonDocument encryptedFields, final String msg, final Throwable cause) {
+        super(msg, assertNotNull(cause));
+        this.encryptedFields = assertNotNull(encryptedFields);
+    }
+
+    /**
+     * The {@code encryptedFields} document that allows inferring which data keys are <strong>known to be created</strong>
+     * before {@code this} exception happened by comparing this document with the original {@code encryptedFields} configuration.
+     * Creation of a data key includes persisting it in the key vault.
+     * <p>
+     * Note that the returned {@code encryptedFields} document is not guaranteed to contain information about all the data keys that
+     * may be created, only about those that the driver is certain about. For example, if persisting a data key times out,
+     * the driver does not know whether it can be considered created or not, and does not include the information about the key in
+     * the {@code encryptedFields} document. You can analyze whether the {@linkplain #getCause() cause} is a definite or indefinite
+     * error, and rely on the returned {@code encryptedFields} to be containing information on all created keys
+     * only if the error is definite.</p>
+     *
+     * @return The updated {@code encryptedFields} document.
+     */
+    public BsonDocument getEncryptedFields() {
+        return encryptedFields;
+    }
+}

driver-core/src/main/com/mongodb/client/model/CreateCollectionOptions.java

@@ -45,6 +45,32 @@ public class CreateCollectionOptions {
     private ClusteredIndexOptions clusteredIndexOptions;
     private Bson encryptedFields;
 
+    public CreateCollectionOptions() {
+    }
+
+    /**
+     * A shallow copy constructor.
+     *
+     * @param options The options to copy.
+     *
+     * @since 4.9
+     */
+    public CreateCollectionOptions(final CreateCollectionOptions options) {
+        notNull("options", options);
+        maxDocuments = options.maxDocuments;
+        capped = options.capped;
+        sizeInBytes = options.sizeInBytes;
+        storageEngineOptions = options.storageEngineOptions;
+        indexOptionDefaults = options.indexOptionDefaults;
+        validationOptions = options.validationOptions;
+        collation = options.collation;
+        expireAfterSeconds = options.expireAfterSeconds;
+        timeSeriesOptions = options.timeSeriesOptions;
+        changeStreamPreAndPostImagesOptions = options.changeStreamPreAndPostImagesOptions;
+        clusteredIndexOptions = options.clusteredIndexOptions;
+        encryptedFields = options.encryptedFields;
+    }
+
     /**
      * Gets the maximum number of documents allowed in a capped collection.
      *

driver-core/src/main/com/mongodb/client/model/CreateEncryptedCollectionParams.java

@@ -0,0 +1,83 @@
+/*
+ * 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.client.model;
+
+import com.mongodb.annotations.Beta;
+import com.mongodb.client.model.vault.DataKeyOptions;
+import com.mongodb.lang.Nullable;
+import org.bson.BsonDocument;
+
+import static com.mongodb.assertions.Assertions.notNull;
+
+/**
+ * Auxiliary parameters for creating an encrypted collection.
+ *
+ * @since 4.9
+ */
+@Beta(Beta.Reason.SERVER)
+public final class CreateEncryptedCollectionParams {
+    private final String kmsProvider;
+    @Nullable
+    private BsonDocument masterKey;
+
+    /**
+     * @param kmsProvider The name of the KMS provider.
+     */
+    public CreateEncryptedCollectionParams(final String kmsProvider) {
+        this.kmsProvider = notNull("kmsProvider", kmsProvider);
+        masterKey = null;
+    }
+
+    /**
+     * The name of the KMS provider.
+     *
+     * @return The name of the KMS provider.
+     */
+    public String getKmsProvider() {
+        return kmsProvider;
+    }
+
+    /**
+     * Sets the {@linkplain DataKeyOptions#getMasterKey() master key} for creating a data key.
+     *
+     * @param masterKey The master key for creating a data key.
+     * @return {@code this}.
+     */
+    public CreateEncryptedCollectionParams masterKey(@Nullable final BsonDocument masterKey) {
+        this.masterKey = masterKey;
+        return this;
+    }
+
+    /**
+     * The {@linkplain DataKeyOptions#getMasterKey() master key} for creating a data key.
+     * The default is {@code null}.
+     *
+     * @return The master key for creating a data key.
+     */
+    @Nullable
+    public BsonDocument getMasterKey() {
+        return masterKey;
+    }
+
+    @Override
+    public String toString() {
+        return "CreateEncryptedCollectionParams{"
+                + ", kmsProvider=" + kmsProvider
+                + ", masterKey=" + masterKey
+                + '}';
+    }
+}

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/vault/ClientEncryptionImpl.java

@@ -17,9 +17,14 @@
 package com.mongodb.reactivestreams.client.internal.vault;
 
 import com.mongodb.ClientEncryptionSettings;
+import com.mongodb.MongoClientSettings;
+import com.mongodb.MongoConfigurationException;
 import com.mongodb.MongoNamespace;
+import com.mongodb.MongoUpdatedEncryptedFieldsException;
 import com.mongodb.ReadConcern;
 import com.mongodb.WriteConcern;
+import com.mongodb.client.model.CreateCollectionOptions;
+import com.mongodb.client.model.CreateEncryptedCollectionParams;
 import com.mongodb.client.model.Filters;
 import com.mongodb.client.model.UpdateOneModel;
 import com.mongodb.client.model.Updates;
@@ -32,22 +37,30 @@
 import com.mongodb.reactivestreams.client.MongoClient;
 import com.mongodb.reactivestreams.client.MongoClients;
 import com.mongodb.reactivestreams.client.MongoCollection;
+import com.mongodb.reactivestreams.client.MongoDatabase;
 import com.mongodb.reactivestreams.client.internal.crypt.Crypt;
 import com.mongodb.reactivestreams.client.internal.crypt.Crypts;
 import com.mongodb.reactivestreams.client.vault.ClientEncryption;
 import org.bson.BsonArray;
 import org.bson.BsonBinary;
 import org.bson.BsonDocument;
+import org.bson.BsonNull;
 import org.bson.BsonString;
 import org.bson.BsonValue;
+import org.bson.codecs.configuration.CodecRegistry;
 import org.bson.conversions.Bson;
 import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 
+import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.internal.capi.MongoCryptHelper.validateRewrapManyDataKeyOptions;
+import static java.lang.String.format;
 import static java.util.Arrays.asList;
 import static java.util.Collections.singletonList;
 
@@ -183,6 +196,79 @@ public Publisher<RewrapManyDataKeyResult> rewrapManyDataKey(final Bson filter, f
                         }));
     }
 
+    @Override
+    public Publisher<BsonDocument> createEncryptedCollection(final MongoDatabase database, final String collectionName,
+            final CreateCollectionOptions createCollectionOptions, final CreateEncryptedCollectionParams createEncryptedCollectionParams) {
+        notNull("collectionName", collectionName);
+        notNull("createCollectionOptions", createCollectionOptions);
+        notNull("createEncryptedCollectionParams", createEncryptedCollectionParams);
+        MongoNamespace namespace = new MongoNamespace(database.getName(), collectionName);
+        Bson rawEncryptedFields = createCollectionOptions.getEncryptedFields();
+        if (rawEncryptedFields == null) {
+            throw new MongoConfigurationException(format("`encryptedFields` is not configured for the collection %s.", namespace));
+        }
+        CodecRegistry codecRegistry = options.getKeyVaultMongoClientSettings() == null
+                ? MongoClientSettings.getDefaultCodecRegistry()
+                : options.getKeyVaultMongoClientSettings().getCodecRegistry();
+        BsonDocument encryptedFields = rawEncryptedFields.toBsonDocument(BsonDocument.class, codecRegistry);
+        BsonValue fields = encryptedFields.get("fields");
+        if (fields != null && fields.isArray()) {
+            String kmsProvider = createEncryptedCollectionParams.getKmsProvider();
+            DataKeyOptions dataKeyOptions = new DataKeyOptions();
+            BsonDocument masterKey = createEncryptedCollectionParams.getMasterKey();
+            if (masterKey != null) {
+                dataKeyOptions.masterKey(masterKey);
+            }
+            String keyIdBsonKey = "keyId";
+            return Mono.defer(() -> {
+                // `Mono.defer` results in `maybeUpdatedEncryptedFields` and `dataKeyMightBeCreated` (mutable state)
+                // being created once per `Subscriber`, which allows the produced `Mono` to support multiple `Subscribers`.
+                BsonDocument maybeUpdatedEncryptedFields = encryptedFields.clone();
+                AtomicBoolean dataKeyMightBeCreated = new AtomicBoolean();
+                Iterable<Mono<BsonDocument>> publishersOfUpdatedFields = () -> maybeUpdatedEncryptedFields.get("fields").asArray()
+                        .stream()
+                        .filter(BsonValue::isDocument)
+                        .map(BsonValue::asDocument)
+                        .filter(field -> field.containsKey(keyIdBsonKey))
+                        .filter(field -> Objects.equals(field.get(keyIdBsonKey), BsonNull.VALUE))
+                        // here we rely on the `createDataKey` publisher being cold, i.e., doing nothing until it is subscribed to
+                        .map(field -> Mono.fromDirect(createDataKey(kmsProvider, dataKeyOptions))
+                                // This is the closest we can do with reactive streams to setting the `dataKeyMightBeCreated` flag
+                                // immediately before calling `createDataKey`.
+                                .doOnSubscribe(subscription -> dataKeyMightBeCreated.set(true))
+                                .doOnNext(dataKeyId -> field.put(keyIdBsonKey, dataKeyId))
+                                .map(dataKeyId -> field)
+                        )
+                        .iterator();
+                // `Flux.concat` ensures that data keys are created / fields are updated sequentially one by one
+                Flux<BsonDocument> publisherOfUpdatedFields = Flux.concat(publishersOfUpdatedFields);
+                return publisherOfUpdatedFields
+                        // All write actions in `doOnNext` above happen-before the completion (`onComplete`/`onError`) signals
+                        // for this publisher, because all signals are serial. `thenEmpty` further guarantees that the completion signal
+                        // for this publisher happens-before the `onSubscribe` signal for the publisher passed to it
+                        // (the next publisher, which creates a collection).
+                        // `defer` defers calling `createCollection` until the next publisher is subscribed to.
+                        // Therefore, all write actions in `doOnNext` above happen-before the invocation of `createCollection`,
+                        // which means `createCollection` is guaranteed to observe all those write actions, i.e.,
+                        // it is guaranteed to observe the updated document via the `maybeUpdatedEncryptedFields` reference.
+                        //
+                        // Similarly, the `Subscriber` of the returned `Publisher` is guaranteed to observe all those write actions
+                        // via the `maybeUpdatedEncryptedFields` reference, which is emitted as a result of `thenReturn`.
+                        .thenEmpty(Mono.defer(() -> Mono.fromDirect(database.createCollection(collectionName,
+                                new CreateCollectionOptions(createCollectionOptions).encryptedFields(maybeUpdatedEncryptedFields))))
+                        )
+                        .onErrorMap(e -> dataKeyMightBeCreated.get(), e ->
+                                new MongoUpdatedEncryptedFieldsException(maybeUpdatedEncryptedFields,
+                                        format("Failed to create %s.", namespace), e)
+                        )
+                        .thenReturn(maybeUpdatedEncryptedFields);
+            });
+        } else {
+            return Mono.fromDirect(database.createCollection(collectionName, createCollectionOptions))
+                    .thenReturn(encryptedFields);
+        }
+    }
+
     @Override
     public void close() {
         keyVaultClient.close();

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/vault/ClientEncryption.java

@@ -16,14 +16,19 @@
 
 package com.mongodb.reactivestreams.client.vault;
 
+import com.mongodb.AutoEncryptionSettings;
+import com.mongodb.MongoUpdatedEncryptedFieldsException;
 import com.mongodb.annotations.Beta;
+import com.mongodb.client.model.CreateCollectionOptions;
+import com.mongodb.client.model.CreateEncryptedCollectionParams;
 import com.mongodb.client.model.vault.DataKeyOptions;
 import com.mongodb.client.model.vault.EncryptOptions;
 import com.mongodb.client.model.vault.RewrapManyDataKeyOptions;
 import com.mongodb.client.model.vault.RewrapManyDataKeyResult;
 import com.mongodb.client.result.DeleteResult;
 import com.mongodb.lang.Nullable;
 import com.mongodb.reactivestreams.client.FindPublisher;
+import com.mongodb.reactivestreams.client.MongoDatabase;
 import org.bson.BsonBinary;
 import org.bson.BsonDocument;
 import org.bson.BsonValue;
@@ -187,6 +192,33 @@ public interface ClientEncryption extends Closeable {
      */
     Publisher<RewrapManyDataKeyResult> rewrapManyDataKey(Bson filter, RewrapManyDataKeyOptions options);
 
+    /**
+     * {@linkplain MongoDatabase#createCollection(String, CreateCollectionOptions) Create} a new collection with encrypted fields,
+     * automatically {@linkplain #createDataKey(String, DataKeyOptions) creating}
+     * new data encryption keys when needed based on the configured
+     * {@link CreateCollectionOptions#getEncryptedFields() encryptedFields}, which must be specified.
+     * This method does not modify the configured {@code encryptedFields} when creating new data keys,
+     * instead it creates a new configuration if needed.
+     *
+     * @param database The database to use for creating the collection.
+     * @param collectionName The name for the collection to create.
+     * @param createCollectionOptions Options for creating the collection.
+     * @param createEncryptedCollectionParams Auxiliary parameters for creating an encrypted collection.
+     * @return A publisher of the (potentially updated) {@code encryptedFields} configuration that was used to create the
+     * collection. A user may use this document to configure {@link AutoEncryptionSettings#getEncryptedFieldsMap()}.
+     * <p>
+     * {@linkplain org.reactivestreams.Subscriber#onError(Throwable) Signals} {@link MongoUpdatedEncryptedFieldsException}
+     * if an exception happens after creating at least one data key. This exception makes the updated {@code encryptedFields}
+     * {@linkplain MongoUpdatedEncryptedFieldsException#getEncryptedFields() available} to the caller.</p>
+     *
+     * @since 4.9
+     * @mongodb.server.release 6.0
+     * @mongodb.driver.manual reference/command/create Create Command
+     */
+    @Beta(Beta.Reason.SERVER)
+    Publisher<BsonDocument> createEncryptedCollection(MongoDatabase database, String collectionName,
+            CreateCollectionOptions createCollectionOptions, CreateEncryptedCollectionParams createEncryptedCollectionParams);
+
     @Override
     void close();
 }