GODRIVER-2391 add encryptedFields to CreateCollection and Drop

mongo/client.go

@@ -68,12 +68,13 @@ type Client struct {
 	sessionPool     *session.Pool
 
 	// client-side encryption fields
-	keyVaultClientFLE *Client
-	keyVaultCollFLE   *Collection
-	mongocryptdFLE    *mcryptClient
-	cryptFLE          driver.Crypt
-	metadataClientFLE *Client
-	internalClientFLE *Client
+	keyVaultClientFLE  *Client
+	keyVaultCollFLE    *Collection
+	mongocryptdFLE     *mcryptClient
+	cryptFLE           driver.Crypt
+	metadataClientFLE  *Client
+	internalClientFLE  *Client
+	encryptedFieldsMap map[string]interface{}
 }
 
 // Connect creates a new Client and then initializes it using the Connect method. This is equivalent to calling
@@ -708,6 +709,7 @@ func (c *Client) configure(opts *options.ClientOptions) error {
 }
 
 func (c *Client) configureAutoEncryption(clientOpts *options.ClientOptions) error {
+	c.encryptedFieldsMap = clientOpts.AutoEncryptionOptions.EncryptedFieldsMap
 	if err := c.configureKeyVaultClientFLE(clientOpts); err != nil {
 		return err
 	}

mongo/collection.go

@@ -1734,6 +1734,68 @@ func (coll *Collection) Indexes() IndexView {
 // Drop drops the collection on the server. This method ignores "namespace not found" errors so it is safe to drop
 // a collection that does not exist on the server.
 func (coll *Collection) Drop(ctx context.Context) error {
+	// Follow Client-Side Encryption specification to check for encryptedFields.
+	// Drop does not have an encryptedFields option. See: GODRIVER-2413.
+	// Check for encryptedFields from the client EncryptedFieldsMap.
+	// Check for encryptedFields from the server if EncryptedFieldsMap is set.
+	ef := coll.db.getEncryptedFieldsFromMap(coll.name)
+	if ef == nil && coll.db.client.encryptedFieldsMap != nil {
+		var err error
+		if ef, err = coll.db.getEncryptedFieldsFromServer(ctx, coll.name); err != nil {
+			return err
+		}
+	}
+
+	if ef != nil {
+		return coll.dropEncryptedCollection(ctx, ef)
+	}
+
+	return coll.drop(ctx)
+}
+
+// dropEncryptedCollection drops a collection with EncryptedFields.
+func (coll *Collection) dropEncryptedCollection(ctx context.Context, ef interface{}) error {
+	efBSON, err := transformBsoncoreDocument(coll.registry, ef, true /* mapAllowed */, "encryptedFields")
+	if err != nil {
+		return fmt.Errorf("error transforming document: %v", err)
+	}
+
+	// Drop the data collection.
+	if err := coll.drop(ctx); err != nil {
+		return err
+	}
+	// Drop the three encryption-related, associated collections: `escCollection`, `eccCollection` and `ecocCollection`.
+	// Drop ESCCollection.
+	escCollection, err := getEncryptedStateCollectionName(efBSON, coll.name, "esc")
+	if err != nil {
+		return err
+	}
+	if err := coll.db.Collection(escCollection).drop(ctx); err != nil {
+		return err
+	}
+
+	// Drop ECCCollection.
+	eccCollection, err := getEncryptedStateCollectionName(efBSON, coll.name, "ecc")
+	if err != nil {
+		return err
+	}
+	if err := coll.db.Collection(eccCollection).drop(ctx); err != nil {
+		return err
+	}
+
+	// Drop ECOCCollection.
+	ecocCollection, err := getEncryptedStateCollectionName(efBSON, coll.name, "ecoc")
+	if err != nil {
+		return err
+	}
+	if err := coll.db.Collection(ecocCollection).drop(ctx); err != nil {
+		return err
+	}
+	return nil
+}
+
+// drop drops a collection without EncryptedFields.
+func (coll *Collection) drop(ctx context.Context) error {
 	if ctx == nil {
 		ctx = context.Background()
 	}

mongo/database.go

@@ -507,6 +507,158 @@ func (db *Database) Watch(ctx context.Context, pipeline interface{},
 //
 // For more information about the command, see https://www.mongodb.com/docs/manual/reference/command/create/.
 func (db *Database) CreateCollection(ctx context.Context, name string, opts ...*options.CreateCollectionOptions) error {
+	cco := options.MergeCreateCollectionOptions(opts...)
+	// Follow Client-Side Encryption specification to check for encryptedFields.
+	// Check for encryptedFields from create options.
+	ef := cco.EncryptedFields
+	// Check for encryptedFields from the client EncryptedFieldsMap.
+	if ef == nil {
+		ef = db.getEncryptedFieldsFromMap(name)
+	}
+	if ef != nil {
+		return db.createCollectionWithEncryptedFields(ctx, name, ef, opts...)
+	}
+
+	return db.createCollection(ctx, name, opts...)
+}
+
+// getEncryptedFieldsFromServer tries to get an "encryptedFields" document associated with collectionName by running the "listCollections" command.
+// Returns nil and no error if the listCollections command succeeds, but "encryptedFields" is not present.
+func (db *Database) getEncryptedFieldsFromServer(ctx context.Context, collectionName string) (interface{}, error) {
+	// Check if collection has an EncryptedFields configured server-side.
+	collSpecs, err := db.ListCollectionSpecifications(ctx, bson.D{{"name", collectionName}})
+	if err != nil {
+		return nil, err
+	}
+	if len(collSpecs) == 0 {
+		return nil, nil
+	}
+	if len(collSpecs) > 1 {
+		return nil, fmt.Errorf("expected 1 or 0 results from listCollections, got %v", len(collSpecs))
+	}
+	collSpec := collSpecs[0]
+	rawValue, err := collSpec.Options.LookupErr("encryptedFields")
+	if err == bsoncore.ErrElementNotFound {
+		return nil, nil
+	} else if err != nil {
+		return nil, err
+	}
+
+	encryptedFields, ok := rawValue.DocumentOK()
+	if !ok {
+		return nil, fmt.Errorf("expected encryptedFields of %v to be document, got %v", collectionName, rawValue.Type)
+	}
+
+	return encryptedFields, nil
+}
+
+// getEncryptedFieldsFromServer tries to get an "encryptedFields" document associated with collectionName by checking the client EncryptedFieldsMap.
+// Returns nil and no error if an EncryptedFieldsMap is not configured, or does not contain an entry for collectionName.
+func (db *Database) getEncryptedFieldsFromMap(collectionName string) interface{} {
+	// Check the EncryptedFieldsMap
+	efMap := db.client.encryptedFieldsMap
+	if efMap == nil {
+		return nil
+	}
+
+	namespace := db.name + "." + collectionName
+
+	ef, ok := efMap[namespace]
+	if ok {
+		return ef
+	}
+	return nil
+}
+
+// getEncryptedStateCollectionName returns the encrypted state collection name associated with dataCollectionName.
+func getEncryptedStateCollectionName(efBSON bsoncore.Document, dataCollectionName string, stateCollectionSuffix string) (string, error) {
+	if stateCollectionSuffix != "esc" && stateCollectionSuffix != "ecc" && stateCollectionSuffix != "ecoc" {
+		return "", fmt.Errorf("expected stateCollectionSuffix: esc, ecc, or ecoc. got %v", stateCollectionSuffix)
+	}
+	fieldName := stateCollectionSuffix + "Collection"
+	var val bsoncore.Value
+	var err error
+	if val, err = efBSON.LookupErr(fieldName); err != nil {
+		if err != bsoncore.ErrElementNotFound {
+			return "", err
+		}
+		// Return default name.
+		defaultName := "enxcol_." + dataCollectionName + "." + stateCollectionSuffix
+		return defaultName, nil
+	}
+
+	var stateCollectionName string
+	var ok bool
+	if stateCollectionName, ok = val.StringValueOK(); !ok {
+		return "", fmt.Errorf("expected string for '%v', got: %v", fieldName, val.Type)
+	}
+	return stateCollectionName, nil
+}
+
+// createCollectionWithEncryptedFields creates a collection with an EncryptedFields.
+func (db *Database) createCollectionWithEncryptedFields(ctx context.Context, name string, ef interface{}, opts ...*options.CreateCollectionOptions) error {
+	efBSON, err := transformBsoncoreDocument(db.registry, ef, true /* mapAllowed */, "encryptedFields")
+	if err != nil {
+		return fmt.Errorf("error transforming document: %v", err)
+	}
+
+	// Create the three encryption-related, associated collections: `escCollection`, `eccCollection` and `ecocCollection`.
+	// Create ESCCollection.
+	escCollection, err := getEncryptedStateCollectionName(efBSON, name, "esc")
+	if err != nil {
+		return err
+	}
+	if err := db.createCollection(ctx, escCollection); err != nil {
+		return err
+	}
+
+	// Create ECCCollection.
+	eccCollection, err := getEncryptedStateCollectionName(efBSON, name, "ecc")
+	if err != nil {
+		return err
+	}
+	if err := db.createCollection(ctx, eccCollection); err != nil {
+		return err
+	}
+
+	// Create ECOCCollection.
+	ecocCollection, err := getEncryptedStateCollectionName(efBSON, name, "ecoc")
+	if err != nil {
+		return err
+	}
+	if err := db.createCollection(ctx, ecocCollection); err != nil {
+		return err
+	}
+
+	// Create a data collection with the 'encryptedFields' option.
+	op, err := db.createCollectionOperation(name, opts...)
+	if err != nil {
+		return err
+	}
+
+	op.EncryptedFields(efBSON)
+	if err := db.executeCreateOperation(ctx, op); err != nil {
+		return err
+	}
+
+	// Create an index on the __safeContent__ field in the collection @collectionName.
+	if _, err := db.Collection(name).Indexes().CreateOne(ctx, IndexModel{Keys: bson.D{{"__safeContent__", 1}}}); err != nil {
+		return fmt.Errorf("error creating safeContent index: %v", err)
+	}
+
+	return nil
+}
+
+// createCollection creates a collection without EncryptedFields.
+func (db *Database) createCollection(ctx context.Context, name string, opts ...*options.CreateCollectionOptions) error {
+	op, err := db.createCollectionOperation(name, opts...)
+	if err != nil {
+		return err
+	}
+	return db.executeCreateOperation(ctx, op)
+}
+
+func (db *Database) createCollectionOperation(name string, opts ...*options.CreateCollectionOptions) (*operation.Create, error) {
 	cco := options.MergeCreateCollectionOptions(opts...)
 	op := operation.NewCreate(name).ServerAPI(db.client.serverAPI)
 
@@ -519,7 +671,7 @@ func (db *Database) CreateCollection(ctx context.Context, name string, opts ...*
 	if cco.ChangeStreamPreAndPostImages != nil {
 		csppi, err := transformBsoncoreDocument(db.registry, cco.ChangeStreamPreAndPostImages, true, "changeStreamPreAndPostImages")
 		if err != nil {
-			return err
+			return nil, err
 		}
 		op.ChangeStreamPreAndPostImages(csppi)
 	}
@@ -528,14 +680,14 @@ func (db *Database) CreateCollection(ctx context.Context, name string, opts ...*
 		if cco.DefaultIndexOptions.StorageEngine != nil {
 			storageEngine, err := transformBsoncoreDocument(db.registry, cco.DefaultIndexOptions.StorageEngine, true, "storageEngine")
 			if err != nil {
-				return err
+				return nil, err
 			}
 
 			doc = bsoncore.AppendDocumentElement(doc, "storageEngine", storageEngine)
 		}
 		doc, err := bsoncore.AppendDocumentEnd(doc, idx)
 		if err != nil {
-			return err
+			return nil, err
 		}
 
 		op.IndexOptionDefaults(doc)
@@ -549,7 +701,7 @@ func (db *Database) CreateCollection(ctx context.Context, name string, opts ...*
 	if cco.StorageEngine != nil {
 		storageEngine, err := transformBsoncoreDocument(db.registry, cco.StorageEngine, true, "storageEngine")
 		if err != nil {
-			return err
+			return nil, err
 		}
 		op.StorageEngine(storageEngine)
 	}
@@ -562,7 +714,7 @@ func (db *Database) CreateCollection(ctx context.Context, name string, opts ...*
 	if cco.Validator != nil {
 		validator, err := transformBsoncoreDocument(db.registry, cco.Validator, true, "validator")
 		if err != nil {
-			return err
+			return nil, err
 		}
 		op.Validator(validator)
 	}
@@ -582,13 +734,13 @@ func (db *Database) CreateCollection(ctx context.Context, name string, opts ...*
 
 		doc, err := bsoncore.AppendDocumentEnd(doc, idx)
 		if err != nil {
-			return err
+			return nil, err
 		}
 
 		op.TimeSeries(doc)
 	}
 
-	return db.executeCreateOperation(ctx, op)
+	return op, nil
 }
 
 // CreateView executes a create command to explicitly create a view on the server. See

mongo/integration/client_side_encryption_test.go

@@ -381,3 +381,82 @@ func TestClientSideEncryptionCustomCrypt(t *testing.T) {
 			"expected 2 calls to BypassAutoEncryption, got %v", cc.numBypassAutoEncryptionCalls)
 	})
 }
+
+func TestFLE2CreateCollection(t *testing.T) {
+	// FLE 2 is not supported on Standalone topology.
+	mtOpts := mtest.NewOptions().
+		MinServerVersion("6.0").
+		Enterprise(true).
+		CreateClient(false).
+		Topologies(mtest.ReplicaSet,
+			mtest.Sharded,
+			mtest.LoadBalanced,
+			mtest.ShardedReplicaSet)
+	mt := mtest.New(t, mtOpts)
+	defer mt.Close()
+
+	efJSON := `
+	{
+		"escCollection": "encryptedCollection.esc",
+		"eccCollection": "encryptedCollection.ecc",
+		"ecocCollection": "encryptedCollection.ecoc",
+		"fields": [
+		  {
+			"path": "firstName",
+			"bsonType": "string",
+			"keyId": {
+			  "$binary": {
+				"subType": "04",
+				"base64": "AAAAAAAAAAAAAAAAAAAAAA=="
+			  }
+			}
+		  }
+		]
+	  }
+	`
+	var efBSON bson.Raw
+	err := bson.UnmarshalExtJSON([]byte(efJSON), true /* canonical */, &efBSON)
+	assert.Nil(mt, err, "UnmarshalExtJSON error: %v", err)
+
+	// Test the behavior in the specification test fle2-CreateCollection.json: "CreateCollection from encryptedFields.".
+	// The Go driver does not support encryptedFields as an option to Drop. See: GODRIVER-2413.
+	mt.Run("CreateCollection from encryptedFields", func(mt *mtest.T) {
+		// Drop data and state collections to clean up from a prior test run.
+		{
+			err := mt.DB.Collection("coll").Drop(context.Background())
+			assert.Nil(mt, err, "error in Drop: %v", err)
+			err = mt.DB.Collection("encryptedCollection.esc").Drop(context.Background())
+			assert.Nil(mt, err, "error in Drop: %v", err)
+			err = mt.DB.Collection("encryptedCollection.ecc").Drop(context.Background())
+			assert.Nil(mt, err, "error in Drop: %v", err)
+			err = mt.DB.Collection("encryptedCollection.ecoc").Drop(context.Background())
+			assert.Nil(mt, err, "error in Drop: %v", err)
+		}
+
+		mt.DB.CreateCollection(context.Background(), "coll", options.CreateCollection().SetEncryptedFields(efBSON))
+
+		// Check expected collections and index exist.
+		{
+			got, err := mt.DB.ListCollectionNames(context.Background(), bson.D{{"name", "coll"}})
+			assert.Nil(mt, err, "error in ListCollectionNames")
+			assert.Equal(mt, got, []string{"coll"}, "expected ['coll'], got: %v", got)
+
+			got, err = mt.DB.ListCollectionNames(context.Background(), bson.D{{"name", "encryptedCollection.esc"}})
+			assert.Nil(mt, err, "error in ListCollectionNames")
+			assert.Equal(mt, got, []string{"encryptedCollection.esc"}, "expected ['encryptedCollection.esc'], got: %v", got)
+
+			got, err = mt.DB.ListCollectionNames(context.Background(), bson.D{{"name", "encryptedCollection.ecc"}})
+			assert.Nil(mt, err, "error in ListCollectionNames")
+			assert.Equal(mt, got, []string{"encryptedCollection.ecc"}, "expected ['encryptedCollection.ecc'], got: %v", got)
+
+			got, err = mt.DB.ListCollectionNames(context.Background(), bson.D{{"name", "encryptedCollection.ecoc"}})
+			assert.Nil(mt, err, "error in ListCollectionNames")
+			assert.Equal(mt, got, []string{"encryptedCollection.ecoc"}, "expected ['encryptedCollection.ecoc'], got: %v", got)
+
+			indexSpecs, err := mt.DB.Collection("coll").Indexes().ListSpecifications(context.Background())
+			assert.Nil(mt, err, "error in Indexes().ListSpecifications: %v", err)
+			assert.Equal(mt, len(indexSpecs), 2, "expected two indexes on 'coll', got: %v", indexSpecs)
+			assert.Equal(mt, indexSpecs[1].Name, "__safeContent___1", "expected second index to be '__safeContent___1', got %v", indexSpecs[1].Name)
+		}
+	})
+}