update NeedKms logic

Author: qingyang-hu

Taskfile.yml

@@ -142,7 +142,7 @@ tasks:
     - go test -exec "env PKG_CONFIG_PATH=${PKG_CONFIG_PATH} LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" ${BUILD_TAGS} -v -timeout {{.TEST_TIMEOUT}}s ./internal/integration -run TestClientSideEncryptionProse/kms_tls_tests >> test.suite
 
   evg-test-retry-kms-requests:
-    - go test -exec "env PKG_CONFIG_PATH=${PKG_CONFIG_PATH} LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" ${BUILD_TAGS} -v -timeout {{.TEST_TIMEOUT}}s ./internal/integration -run TestClientSideEncryptionProse/kms_retry_tests >> test.suite
+    - go test -exec "env PKG_CONFIG_PATH=${PKG_CONFIG_PATH} LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" ${BUILD_TAGS} -v -timeout 300s ./internal/integration -run TestClientSideEncryptionProse/kms_retry_tests >> test.suite
 
   evg-test-load-balancers:
     # Load balancer should be tested with all unified tests as well as tests in the following

etc/install-libmongocrypt.sh

@@ -3,7 +3,7 @@
 # This script installs libmongocrypt into an "install" directory.
 set -eux
 
-LIBMONGOCRYPT_TAG="1.11.0"
+LIBMONGOCRYPT_TAG="1.12.0"
 
 # Install libmongocrypt based on OS.
 if [ "Windows_NT" = "${OS:-}" ]; then

mongo/client_encryption.go

@@ -188,12 +188,14 @@ func (ce *ClientEncryption) CreateDataKey(
 	}
 
 	// create data key document
+	fmt.Println("CreateDataKey")
 	dataKeyDoc, err := ce.crypt.CreateDataKey(ctx, kmsProvider, co)
 	if err != nil {
 		return bson.Binary{}, err
 	}
 
 	// insert key into key vault
+	fmt.Println("InsertOne")
 	_, err = ce.keyVaultColl.InsertOne(ctx, dataKeyDoc)
 	if err != nil {
 		return bson.Binary{}, err

x/mongo/driver/crypt.go

@@ -260,6 +260,7 @@ func (c *crypt) executeStateMachine(ctx context.Context, cryptCtx *mongocrypt.Co
 	var err error
 	for {
 		state := cryptCtx.State()
+		fmt.Println("state", state)
 		switch state {
 		case mongocrypt.NeedMongoCollInfo:
 			err = c.collectionInfo(ctx, cryptCtx, db)
@@ -341,6 +342,7 @@ func (c *crypt) retrieveKeys(ctx context.Context, cryptCtx *mongocrypt.Context)
 }
 
 func (c *crypt) decryptKeys(cryptCtx *mongocrypt.Context) error {
+	c.mongoCrypt.EnableRetry()
 	for {
 		kmsCtx := cryptCtx.NextKmsContext()
 		if kmsCtx == nil {
@@ -376,8 +378,10 @@ func (c *crypt) decryptKey(kmsCtx *mongocrypt.KmsContext) error {
 	if tlsCfg == nil {
 		tlsCfg = &tls.Config{MinVersion: tls.VersionTLS12}
 	}
+	fmt.Println("dial", addr)
 	conn, err := tls.Dial("tcp", addr, tlsCfg)
 	if err != nil {
+		fmt.Println("dial error", err)
 		return err
 	}
 	defer func() {
@@ -388,6 +392,7 @@ func (c *crypt) decryptKey(kmsCtx *mongocrypt.KmsContext) error {
 		return err
 	}
 	if _, err = conn.Write(msg); err != nil {
+		fmt.Println("conn write", err)
 		return err
 	}
 
@@ -400,6 +405,11 @@ func (c *crypt) decryptKey(kmsCtx *mongocrypt.KmsContext) error {
 		res := make([]byte, bytesNeeded)
 		bytesRead, err := conn.Read(res)
 		if err != nil && !errors.Is(err, io.EOF) {
+			fail := kmsCtx.Fail()
+			fmt.Println("conn read", err, fail)
+			if fail {
+				continue
+			}
 			return err
 		}
 

x/mongo/driver/mongocrypt/mongocrypt.go

@@ -522,3 +522,8 @@ func (m *MongoCrypt) GetKmsProviders(ctx context.Context) (bsoncore.Document, er
 	}
 	return builder.Build(), nil
 }
+
+// EnableRetry enables retry.
+func (m *MongoCrypt) EnableRetry() {
+	_ = C.mongocrypt_setopt_retry_kms(m.wrapped, true)
+}

x/mongo/driver/mongocrypt/mongocrypt_kms_context.go

@@ -11,6 +11,7 @@ package mongocrypt
 
 // #include <mongocrypt.h>
 import "C"
+import "time"
 
 // KmsContext represents a mongocrypt_kms_ctx_t handle.
 type KmsContext struct {
@@ -41,6 +42,8 @@ func (kc *KmsContext) KMSProvider() string {
 
 // Message returns the message to send to the KMS.
 func (kc *KmsContext) Message() ([]byte, error) {
+	time.Sleep(time.Duration(C.mongocrypt_kms_ctx_usleep(kc.wrapped)) * time.Microsecond)
+
 	msgBinary := newBinary()
 	defer msgBinary.close()
 
@@ -74,3 +77,8 @@ func (kc *KmsContext) createErrorFromStatus() error {
 	C.mongocrypt_kms_ctx_status(kc.wrapped, status)
 	return errorFromStatus(status)
 }
+
+// Fail returns a boolean indicating whether the failed request may be retried.
+func (kc *KmsContext) Fail() bool {
+	return bool(C.mongocrypt_kms_ctx_fail(kc.wrapped))
+}

x/mongo/driver/mongocrypt/mongocrypt_kms_context_not_enabled.go

@@ -37,3 +37,8 @@ func (kc *KmsContext) BytesNeeded() int32 {
 func (kc *KmsContext) FeedResponse([]byte) error {
 	panic(cseNotSupportedMsg)
 }
+
+// Fail returns a boolean indicating whether the failed request may be retried.
+func (kc *KmsContext) Fail() bool {
+	panic(cseNotSupportedMsg)
+}

x/mongo/driver/mongocrypt/mongocrypt_not_enabled.go

@@ -95,3 +95,8 @@ func (m *MongoCrypt) Close() {
 func (m *MongoCrypt) GetKmsProviders(context.Context) (bsoncore.Document, error) {
 	panic(cseNotSupportedMsg)
 }
+
+// EnableRetry enables retry.
+func (m *MongoCrypt) EnableRetry() {
+	panic(cseNotSupportedMsg)
+}