GODRIVER-2149 Add private API to allow overriding CSFLE functionality (#789)

Author: benjirewis

mongo/change_stream.go

@@ -94,7 +94,7 @@ type changeStreamConfig struct {
 	streamType     StreamType
 	collectionName string
 	databaseName   string
-	crypt          *driver.Crypt
+	crypt          driver.Crypt
 }
 
 func newChangeStream(ctx context.Context, config changeStreamConfig, pipeline interface{},

mongo/client.go

@@ -68,7 +68,7 @@ type Client struct {
 	keyVaultClientFLE *Client
 	keyVaultCollFLE   *Collection
 	mongocryptdFLE    *mcryptClient
-	cryptFLE          *driver.Crypt
+	cryptFLE          driver.Crypt
 	metadataClientFLE *Client
 	internalClientFLE *Client
 }
@@ -634,6 +634,8 @@ func (c *Client) configure(opts *options.ClientOptions) error {
 		if err := c.configureAutoEncryption(opts); err != nil {
 			return err
 		}
+	} else {
+		c.cryptFLE = opts.Crypt
 	}
 
 	// OCSP cache

mongo/client_encryption.go

@@ -22,7 +22,7 @@ import (
 
 // ClientEncryption is used to create data keys and explicitly encrypt and decrypt BSON values.
 type ClientEncryption struct {
-	crypt          *driver.Crypt
+	crypt          driver.Crypt
 	keyVaultClient *Client
 	keyVaultColl   *Collection
 }

mongo/integration/client_side_encryption_test.go

@@ -22,6 +22,8 @@ import (
 	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
 	"go.mongodb.org/mongo-driver/mongo/options"
 	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
+	"go.mongodb.org/mongo-driver/x/mongo/driver"
+	mcopts "go.mongodb.org/mongo-driver/x/mongo/driver/mongocrypt/options"
 )
 
 // createDataKeyAndEncrypt creates a data key with the alternate name @keyName.
@@ -235,3 +237,146 @@ func TestClientSideEncryptionWithExplicitSessions(t *testing.T) {
 		assert.NotEqual(mt, lsid, session.ID(), "expected different lsid, but got %v", lsid)
 	})
 }
+
+// customCrypt is a test implementation of the driver.Crypt interface. It keeps track of the number of times its
+// methods have been called.
+type customCrypt struct {
+	numEncryptCalls              int
+	numDecryptCalls              int
+	numCreateDataKeyCalls        int
+	numEncryptExplicitCalls      int
+	numDecryptExplicitCalls      int
+	numCloseCalls                int
+	numBypassAutoEncryptionCalls int
+}
+
+var (
+	_     driver.Crypt = (*customCrypt)(nil)
+	mySSN              = "123456789"
+)
+
+// Encrypt encrypts the given command.
+func (c *customCrypt) Encrypt(_ context.Context, _ string, cmd bsoncore.Document) (bsoncore.Document, error) {
+	c.numEncryptCalls++
+	elems, err := cmd.Elements()
+	if err != nil {
+		return nil, err
+	}
+
+	encryptedCmd := bsoncore.NewDocumentBuilder()
+	for _, elem := range elems {
+		// "encrypt" ssn element as "hidden"
+		if elem.Key() == "ssn" {
+			encryptedCmd = encryptedCmd.AppendString("ssn", "hidden")
+		} else {
+			encryptedCmd = encryptedCmd.AppendValue(elem.Key(), elem.Value())
+		}
+	}
+	return encryptedCmd.Build(), nil
+}
+
+// Decrypt decrypts the given command response.
+func (c *customCrypt) Decrypt(_ context.Context, cmdResponse bsoncore.Document) (bsoncore.Document, error) {
+	c.numDecryptCalls++
+	elems, err := cmdResponse.Elements()
+	if err != nil {
+		return nil, err
+	}
+
+	decryptedCmdResponse := bsoncore.NewDocumentBuilder()
+	for _, elem := range elems {
+		// "decrypt" ssn element as mySSN
+		if elem.Key() == "ssn" {
+			decryptedCmdResponse = decryptedCmdResponse.AppendString("ssn", mySSN)
+		} else {
+			decryptedCmdResponse = decryptedCmdResponse.AppendValue(elem.Key(), elem.Value())
+		}
+	}
+	return decryptedCmdResponse.Build(), nil
+}
+
+// CreateDataKey implements the driver.Crypt interface.
+func (c *customCrypt) CreateDataKey(_ context.Context, _ string, _ *mcopts.DataKeyOptions) (bsoncore.Document, error) {
+	c.numCreateDataKeyCalls++
+	return nil, nil
+}
+
+// EncryptExplicit implements the driver.Crypt interface.
+func (c *customCrypt) EncryptExplicit(_ context.Context, _ bsoncore.Value, _ *mcopts.ExplicitEncryptionOptions) (byte, []byte, error) {
+	c.numEncryptExplicitCalls++
+	return 0, nil, nil
+}
+
+// DecryptExplicit implements the driver.Crypt interface.
+func (c *customCrypt) DecryptExplicit(_ context.Context, _ byte, _ []byte) (bsoncore.Value, error) {
+	c.numDecryptExplicitCalls++
+	return bsoncore.Value{}, nil
+}
+
+// Close implements the driver.Crypt interface.
+func (c *customCrypt) Close() {
+	c.numCloseCalls++
+}
+
+// BypassAutoEncryption implements the driver.Crypt interface.
+func (c *customCrypt) BypassAutoEncryption() bool {
+	c.numBypassAutoEncryptionCalls++
+	return false
+}
+
+func TestClientSideEncryptionCustomCrypt(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().MinServerVersion("4.2").Enterprise(true).CreateClient(false))
+	defer mt.Close()
+
+	kmsProvidersMap := map[string]map[string]interface{}{
+		"local": {"key": localMasterKey},
+	}
+
+	mt.Run("auto encryption and decryption", func(mt *mtest.T) {
+		aeOpts := options.AutoEncryption().
+			SetKmsProviders(kmsProvidersMap).
+			SetKeyVaultNamespace("keyvault.datakeys")
+		clientOpts := options.Client().
+			ApplyURI(mtest.ClusterURI()).
+			SetAutoEncryptionOptions(aeOpts)
+		cc := &customCrypt{}
+		clientOpts.Crypt = cc
+		testutil.AddTestServerAPIVersion(clientOpts)
+
+		client, err := mongo.Connect(mtest.Background, clientOpts)
+		defer client.Disconnect(mtest.Background)
+		assert.Nil(mt, err, "Connect error: %v", err)
+
+		coll := client.Database("db").Collection("coll")
+		defer func() { _ = coll.Drop(mtest.Background) }()
+
+		doc := bson.D{{"foo", "bar"}, {"ssn", mySSN}}
+		_, err = coll.InsertOne(mtest.Background, doc)
+		assert.Nil(mt, err, "InsertOne error: %v", err)
+
+		res := coll.FindOne(mtest.Background, bson.D{{"foo", "bar"}})
+		assert.Nil(mt, res.Err(), "FindOne error: %v", err)
+
+		rawRes, err := res.DecodeBytes()
+		assert.Nil(mt, err, "DecodeBytes error: %v", err)
+		ssn, ok := rawRes.Lookup("ssn").StringValueOK()
+		assert.True(mt, ok, "expected 'ssn' value to be type string, got %T", ssn)
+		assert.Equal(mt, ssn, mySSN, "expected 'ssn' value %q, got %q", mySSN, ssn)
+
+		// Assert customCrypt methods are called the correct number of times.
+		assert.Equal(mt, cc.numEncryptCalls, 1,
+			"expected 1 call to Encrypt, got %v", cc.numEncryptCalls)
+		assert.Equal(mt, cc.numDecryptCalls, 1,
+			"expected 1 call to Decrypt, got %v", cc.numDecryptCalls)
+		assert.Equal(mt, cc.numCreateDataKeyCalls, 0,
+			"expected 0 calls to CreateDataKey, got %v", cc.numCreateDataKeyCalls)
+		assert.Equal(mt, cc.numEncryptExplicitCalls, 0,
+			"expected 0 calls to EncryptExplicit, got %v", cc.numEncryptExplicitCalls)
+		assert.Equal(mt, cc.numDecryptExplicitCalls, 0,
+			"expected 0 calls to DecryptExplicit, got %v", cc.numDecryptExplicitCalls)
+		assert.Equal(mt, cc.numCloseCalls, 0,
+			"expected 0 calls to Close, got %v", cc.numCloseCalls)
+		assert.Equal(mt, cc.numBypassAutoEncryptionCalls, 2,
+			"expected 2 calls to BypassAutoEncryption, got %v", cc.numBypassAutoEncryptionCalls)
+	})
+}

mongo/options/clientoptions.go

@@ -138,6 +138,13 @@ type ClientOptions struct {
 	// release.
 	AuthenticateToAnything *bool
 
+	// Crypt specifies a custom driver.Crypt to be used to encrypt and decrypt documents. The default is no
+	// encryption.
+	//
+	// Deprecated: This option is for internal use only and should not be set (see GODRIVER-2149). It may be
+	// changed or removed in any release.
+	Crypt driver.Crypt
+
 	// Deployment specifies a custom deployment to use for the new Client.
 	//
 	// Deprecated: This option is for internal use only and should not be set. It may be changed or removed in any
@@ -828,6 +835,9 @@ func MergeClientOptions(opts ...*ClientOptions) *ClientOptions {
 		if opt.ConnectTimeout != nil {
 			c.ConnectTimeout = opt.ConnectTimeout
 		}
+		if opt.Crypt != nil {
+			c.Crypt = opt.Crypt
+		}
 		if opt.HeartbeatInterval != nil {
 			c.HeartbeatInterval = opt.HeartbeatInterval
 		}

x/mongo/driver/batch_cursor.go

@@ -32,7 +32,7 @@ type BatchCursor struct {
 	firstBatch           bool
 	cmdMonitor           *event.CommandMonitor
 	postBatchResumeToken bsoncore.Document
-	crypt                *Crypt
+	crypt                Crypt
 	serverAPI            *ServerAPIOptions
 
 	// legacy server (< 3.2) fields
@@ -130,7 +130,7 @@ type CursorOptions struct {
 	MaxTimeMS      int64
 	Limit          int32
 	CommandMonitor *event.CommandMonitor
-	Crypt          *Crypt
+	Crypt          Crypt
 	ServerAPI      *ServerAPIOptions
 }
 

x/mongo/driver/crypt.go

@@ -43,24 +43,46 @@ type CryptOptions struct {
 	BypassAutoEncryption bool
 }
 
-// Crypt consumes the libmongocrypt.MongoCrypt type to iterate the mongocrypt state machine and perform encryption
+// Crypt is an interface implemented by types that can encrypt and decrypt instances of
+// bsoncore.Document.
+//
+// Users should rely on the driver's crypt type (used by default) for encryption and decryption
+// unless they are perfectly confident in another implementation of Crypt.
+type Crypt interface {
+	// Encrypt encrypts the given command.
+	Encrypt(ctx context.Context, db string, cmd bsoncore.Document) (bsoncore.Document, error)
+	// Decrypt decrypts the given command response.
+	Decrypt(ctx context.Context, cmdResponse bsoncore.Document) (bsoncore.Document, error)
+	// CreateDataKey creates a data key using the given KMS provider and options.
+	CreateDataKey(ctx context.Context, kmsProvider string, opts *options.DataKeyOptions) (bsoncore.Document, error)
+	// EncryptExplicit encrypts the given value with the given options.
+	EncryptExplicit(ctx context.Context, val bsoncore.Value, opts *options.ExplicitEncryptionOptions) (byte, []byte, error)
+	// DecryptExplicit decrypts the given encrypted value.
+	DecryptExplicit(ctx context.Context, subtype byte, data []byte) (bsoncore.Value, error)
+	// Close cleans up any resources associated with the Crypt instance.
+	Close()
+	// BypassAutoEncryption returns true if auto-encryption should be bypassed.
+	BypassAutoEncryption() bool
+}
+
+// crypt consumes the libmongocrypt.MongoCrypt type to iterate the mongocrypt state machine and perform encryption
 // and decryption.
-type Crypt struct {
+type crypt struct {
 	mongoCrypt *mongocrypt.MongoCrypt
 	collInfoFn CollectionInfoFn
 	keyFn      KeyRetrieverFn
 	markFn     MarkCommandFn
 
-	BypassAutoEncryption bool
+	bypassAutoEncryption bool
 }
 
 // NewCrypt creates a new Crypt instance configured with the given AutoEncryptionOptions.
-func NewCrypt(opts *CryptOptions) (*Crypt, error) {
-	c := &Crypt{
+func NewCrypt(opts *CryptOptions) (Crypt, error) {
+	c := &crypt{
 		collInfoFn:           opts.CollInfoFn,
 		keyFn:                opts.KeyFn,
 		markFn:               opts.MarkFn,
-		BypassAutoEncryption: opts.BypassAutoEncryption,
+		bypassAutoEncryption: opts.BypassAutoEncryption,
 	}
 
 	mongocryptOpts := options.MongoCrypt().SetKmsProviders(opts.KmsProviders).SetLocalSchemaMap(opts.SchemaMap)
@@ -74,8 +96,8 @@ func NewCrypt(opts *CryptOptions) (*Crypt, error) {
 }
 
 // Encrypt encrypts the given command.
-func (c *Crypt) Encrypt(ctx context.Context, db string, cmd bsoncore.Document) (bsoncore.Document, error) {
-	if c.BypassAutoEncryption {
+func (c *crypt) Encrypt(ctx context.Context, db string, cmd bsoncore.Document) (bsoncore.Document, error) {
+	if c.bypassAutoEncryption {
 		return cmd, nil
 	}
 
@@ -89,7 +111,7 @@ func (c *Crypt) Encrypt(ctx context.Context, db string, cmd bsoncore.Document) (
 }
 
 // Decrypt decrypts the given command response.
-func (c *Crypt) Decrypt(ctx context.Context, cmdResponse bsoncore.Document) (bsoncore.Document, error) {
+func (c *crypt) Decrypt(ctx context.Context, cmdResponse bsoncore.Document) (bsoncore.Document, error) {
 	cryptCtx, err := c.mongoCrypt.CreateDecryptionContext(cmdResponse)
 	if err != nil {
 		return nil, err
@@ -100,7 +122,7 @@ func (c *Crypt) Decrypt(ctx context.Context, cmdResponse bsoncore.Document) (bso
 }
 
 // CreateDataKey creates a data key using the given KMS provider and options.
-func (c *Crypt) CreateDataKey(ctx context.Context, kmsProvider string, opts *options.DataKeyOptions) (bsoncore.Document, error) {
+func (c *crypt) CreateDataKey(ctx context.Context, kmsProvider string, opts *options.DataKeyOptions) (bsoncore.Document, error) {
 	cryptCtx, err := c.mongoCrypt.CreateDataKeyContext(kmsProvider, opts)
 	if err != nil {
 		return nil, err
@@ -111,7 +133,7 @@ func (c *Crypt) CreateDataKey(ctx context.Context, kmsProvider string, opts *opt
 }
 
 // EncryptExplicit encrypts the given value with the given options.
-func (c *Crypt) EncryptExplicit(ctx context.Context, val bsoncore.Value, opts *options.ExplicitEncryptionOptions) (byte, []byte, error) {
+func (c *crypt) EncryptExplicit(ctx context.Context, val bsoncore.Value, opts *options.ExplicitEncryptionOptions) (byte, []byte, error) {
 	idx, doc := bsoncore.AppendDocumentStart(nil)
 	doc = bsoncore.AppendValueElement(doc, "v", val)
 	doc, _ = bsoncore.AppendDocumentEnd(doc, idx)
@@ -132,7 +154,7 @@ func (c *Crypt) EncryptExplicit(ctx context.Context, val bsoncore.Value, opts *o
 }
 
 // DecryptExplicit decrypts the given encrypted value.
-func (c *Crypt) DecryptExplicit(ctx context.Context, subtype byte, data []byte) (bsoncore.Value, error) {
+func (c *crypt) DecryptExplicit(ctx context.Context, subtype byte, data []byte) (bsoncore.Value, error) {
 	idx, doc := bsoncore.AppendDocumentStart(nil)
 	doc = bsoncore.AppendBinaryElement(doc, "v", subtype, data)
 	doc, _ = bsoncore.AppendDocumentEnd(doc, idx)
@@ -152,11 +174,15 @@ func (c *Crypt) DecryptExplicit(ctx context.Context, subtype byte, data []byte)
 }
 
 // Close cleans up any resources associated with the Crypt instance.
-func (c *Crypt) Close() {
+func (c *crypt) Close() {
 	c.mongoCrypt.Close()
 }
 
-func (c *Crypt) executeStateMachine(ctx context.Context, cryptCtx *mongocrypt.Context, db string) (bsoncore.Document, error) {
+func (c *crypt) BypassAutoEncryption() bool {
+	return c.bypassAutoEncryption
+}
+
+func (c *crypt) executeStateMachine(ctx context.Context, cryptCtx *mongocrypt.Context, db string) (bsoncore.Document, error) {
 	var err error
 	for {
 		state := cryptCtx.State()
@@ -180,7 +206,7 @@ func (c *Crypt) executeStateMachine(ctx context.Context, cryptCtx *mongocrypt.Co
 	}
 }
 
-func (c *Crypt) collectionInfo(ctx context.Context, cryptCtx *mongocrypt.Context, db string) error {
+func (c *crypt) collectionInfo(ctx context.Context, cryptCtx *mongocrypt.Context, db string) error {
 	op, err := cryptCtx.NextOperation()
 	if err != nil {
 		return err
@@ -199,7 +225,7 @@ func (c *Crypt) collectionInfo(ctx context.Context, cryptCtx *mongocrypt.Context
 	return cryptCtx.CompleteOperation()
 }
 
-func (c *Crypt) markCommand(ctx context.Context, cryptCtx *mongocrypt.Context, db string) error {
+func (c *crypt) markCommand(ctx context.Context, cryptCtx *mongocrypt.Context, db string) error {
 	op, err := cryptCtx.NextOperation()
 	if err != nil {
 		return err
@@ -216,7 +242,7 @@ func (c *Crypt) markCommand(ctx context.Context, cryptCtx *mongocrypt.Context, d
 	return cryptCtx.CompleteOperation()
 }
 
-func (c *Crypt) retrieveKeys(ctx context.Context, cryptCtx *mongocrypt.Context) error {
+func (c *crypt) retrieveKeys(ctx context.Context, cryptCtx *mongocrypt.Context) error {
 	op, err := cryptCtx.NextOperation()
 	if err != nil {
 		return err
@@ -236,7 +262,7 @@ func (c *Crypt) retrieveKeys(ctx context.Context, cryptCtx *mongocrypt.Context)
 	return cryptCtx.CompleteOperation()
 }
 
-func (c *Crypt) decryptKeys(ctx context.Context, cryptCtx *mongocrypt.Context) error {
+func (c *crypt) decryptKeys(ctx context.Context, cryptCtx *mongocrypt.Context) error {
 	for {
 		kmsCtx := cryptCtx.NextKmsContext()
 		if kmsCtx == nil {
@@ -251,7 +277,7 @@ func (c *Crypt) decryptKeys(ctx context.Context, cryptCtx *mongocrypt.Context) e
 	return cryptCtx.FinishKmsContexts()
 }
 
-func (c *Crypt) decryptKey(ctx context.Context, kmsCtx *mongocrypt.KmsContext) error {
+func (c *crypt) decryptKey(ctx context.Context, kmsCtx *mongocrypt.KmsContext) error {
 	host, err := kmsCtx.HostName()
 	if err != nil {
 		return err