updates

Author: qingyang-hu

internal/integration/unified/client_entity.go

@@ -7,11 +7,9 @@
 package unified
 
 import (
-	"bytes"
 	"context"
-	"encoding/base64"
+	"crypto/tls"
 	"fmt"
-	"os"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -35,26 +33,11 @@ import (
 // exceed the default truncation length.
 const defaultMaxDocumentLen = 10_000
 
-var (
-	// Security-sensitive commands that should be ignored in command monitoring by default.
-	securitySensitiveCommands = []string{
-		"authenticate", "saslStart", "saslContinue", "getnonce",
-		"createUser", "updateUser", "copydbgetnonce", "copydbsaslstart", "copydb",
-	}
-
-	awsAccessKeyID         = os.Getenv("FLE_AWS_KEY")
-	awsSecretAccessKey     = os.Getenv("FLE_AWS_SECRET")
-	awsTempAccessKeyID     = os.Getenv("CSFLE_AWS_TEMP_ACCESS_KEY_ID")
-	awsTempSecretAccessKey = os.Getenv("CSFLE_AWS_TEMP_SECRET_ACCESS_KEY")
-	awsTempSessionToken    = os.Getenv("CSFLE_AWS_TEMP_SESSION_TOKEN")
-	azureTenantID          = os.Getenv("FLE_AZURE_TENANTID")
-	azureClientID          = os.Getenv("FLE_AZURE_CLIENTID")
-	azureClientSecret      = os.Getenv("FLE_AZURE_CLIENTSECRET")
-	gcpEmail               = os.Getenv("FLE_GCP_EMAIL")
-	gcpPrivateKey          = os.Getenv("FLE_GCP_PRIVATEKEY")
-
-	placeholderDoc = bsoncore.NewDocumentBuilder().AppendInt32("$$placeholder", 1).Build()
-)
+// Security-sensitive commands that should be ignored in command monitoring by default.
+var securitySensitiveCommands = []string{
+	"authenticate", "saslStart", "saslContinue", "getnonce",
+	"createUser", "updateUser", "copydbgetnonce", "copydbsaslstart", "copydb",
+}
 
 // clientEntity is a wrapper for a mongo.Client object that also holds additional information required during test
 // execution.
@@ -295,128 +278,52 @@ func createAutoEncryptionOptions(opts bson.Raw) (*options.AutoEncryptionOptions,
 			if err != nil {
 				return nil, err
 			}
-			retrieveProviderData := func(doc bson.Raw, key, defaultVal string) (any, error) {
-				e := doc.Lookup(key)
-				if e.IsZero() {
-					return nil, nil
+			for _, elem := range elems {
+				key := elem.Key()
+				opt := elem.Value().Document()
+				provider, err := getKmsProvider(key, opt)
+				if err != nil {
+					return nil, err
 				}
-				switch e.Type {
-				case bson.TypeString:
-					return e.StringValue(), nil
-				case bson.TypeEmbeddedDocument:
-					if bytes.Equal(e.Document(), placeholderDoc) {
-						return defaultVal, nil
-					}
+				if provider == nil {
+					continue
 				}
-				return nil, fmt.Errorf("unexpected %s in kms provider: %v", key, e)
-			}
-			for _, elem := range elems {
-				provider := make(map[string]any)
-				providerT := elem.Key()
-				providerOpt := elem.Value().Document()
-				switch providerT {
-				case "aws":
-					accessKeyID := awsAccessKeyID
-					secretAccessKey := awsSecretAccessKey
-
-					// replace with temporary access, if sessionToken placeholder exists
-					v, err := retrieveProviderData(providerOpt, "sessionToken", "$$placeholder")
-					if err != nil {
-						return nil, err
-					}
-					if v == "$$placeholder" {
-						provider["sessionToken"] = awsTempSessionToken
-						accessKeyID = awsTempAccessKeyID
-						secretAccessKey = awsTempSecretAccessKey
-					} else if v != nil {
-						provider["sessionToken"] = v
-					}
-
-					for _, e := range []struct {
-						key        string
-						defaultVal string
-					}{
-						{"accessKeyId", accessKeyID},
-						{"secretAccessKey", secretAccessKey},
-					} {
-						v, err = retrieveProviderData(providerOpt, e.key, e.defaultVal)
-						if err != nil {
-							return nil, err
-						}
-						if v != nil {
-							provider[e.key] = v
-						}
-					}
-				case "azure":
-					for _, e := range []struct {
-						key        string
-						defaultVal string
-					}{
-						{"tenantId", azureTenantID},
-						{"clientId", azureClientID},
-						{"clientSecret", azureClientSecret},
-					} {
-						v, err := retrieveProviderData(providerOpt, e.key, e.defaultVal)
-						if err != nil {
-							return nil, err
-						}
-						if v != nil {
-							provider[e.key] = v
-						}
-					}
-				case "gcp":
-					for _, e := range []struct {
-						key        string
-						defaultVal string
-					}{
-						{"email", gcpEmail},
-						{"privateKey", gcpPrivateKey},
-					} {
-						v, err := retrieveProviderData(providerOpt, e.key, e.defaultVal)
-						if err != nil {
-							return nil, err
-						}
-						if v != nil {
-							provider[e.key] = v
-						}
-					}
-				case "kmip":
-					v, err := retrieveProviderData(providerOpt, "endpoint", "localhost:5698")
-					if err != nil {
-						return nil, err
-					}
-					if v != nil {
-						provider["endpoint"] = v
-					}
-				case "local", "local:name2":
-					str := providerOpt.Lookup("key").StringValue()
-					key, err := base64.StdEncoding.DecodeString(str)
+				providers[key] = provider
+				if key == "kmip" && tlsClientCertificateKeyFile != "" && tlsCAFile != "" {
+					cfg, err := options.BuildTLSConfig(map[string]any{
+						"tlsCertificateKeyFile": tlsClientCertificateKeyFile,
+						"tlsCAFile":             tlsCAFile,
+					})
 					if err != nil {
-						return nil, err
+						return nil, fmt.Errorf("error constructing tls config: %w", err)
 					}
-					provider["key"] = key
-				default:
-					return nil, fmt.Errorf("unrecognized KMS provider: %v", provider)
-				}
-				if len(provider) > 0 {
-					providers[providerT] = provider
+					aeo.SetTLSConfig(map[string]*tls.Config{
+						"kmip": cfg,
+					})
 				}
 			}
 			aeo.SetKmsProviders(providers)
 		case "schemaMap":
 			var schemaMap map[string]any
 			err := bson.Unmarshal(opt.Document(), &schemaMap)
 			if err != nil {
-				return nil, err
+				return nil, fmt.Errorf("error creating schema map: %v", err)
 			}
 			aeo.SetSchemaMap(schemaMap)
 		case "keyVaultNamespace":
 			kvnsFound = true
 			aeo.SetKeyVaultNamespace(opt.StringValue())
-		case "bypassQueryAnalysis":
-			aeo.SetBypassQueryAnalysis(opt.Boolean())
 		case "bypassAutoEncryption":
 			aeo.SetBypassAutoEncryption(opt.Boolean())
+		case "encryptedFieldsMap":
+			var encryptedFieldsMap map[string]any
+			err := bson.Unmarshal(opt.Document(), &encryptedFieldsMap)
+			if err != nil {
+				return nil, fmt.Errorf("error creating encryptedFieldsMap: %v", err)
+			}
+			aeo.SetEncryptedFieldsMap(encryptedFieldsMap)
+		case "bypassQueryAnalysis":
+			aeo.SetBypassQueryAnalysis(opt.Boolean())
 		default:
 			return nil, fmt.Errorf("unrecognized option: %v", name)
 		}