update NeedKms logic

Author: qingyang-hu

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

internal/integration/client_side_encryption_prose_test.go

@@ -2988,6 +2988,10 @@ func TestClientSideEncryptionProse(t *testing.T) {
 			mt.Skipf("Skipping test as KMS_FAILPOINT_SERVERS_RUNNING is not set")
 		}
 
+		tlsCfg := &tls.Config{
+			InsecureSkipVerify: true,
+		}
+
 		setFailPoint := func(failure string, count int) error {
 			url := fmt.Sprintf("https://localhost:9003/set_failpoint/%s", failure)
 			var payloadBuf bytes.Buffer
@@ -2999,26 +3003,26 @@ func TestClientSideEncryptionProse(t *testing.T) {
 			}
 
 			client := &http.Client{
-				Transport: &http.Transport{
-					TLSClientConfig: &tls.Config{
-						InsecureSkipVerify: true,
-					},
-				},
+				Transport: &http.Transport{TLSClientConfig: tlsCfg},
 			}
-			_, err = client.Do(req)
-			return err
+			res, err := client.Do(req)
+			if err != nil {
+				return err
+			}
+			return res.Body.Close()
 		}
 
 		keyVaultClient, err := mongo.Connect(options.Client().ApplyURI(mtest.ClusterURI()))
 		require.NoError(mt, err, "error on Connect: %v", err)
 
 		ceo := options.ClientEncryption().
 			SetKeyVaultNamespace("keyvault.datakeys").
-			SetKmsProviders(fullKmsProvidersMap)
+			SetKmsProviders(fullKmsProvidersMap).
+			SetTLSConfig(map[string]*tls.Config{"aws": tlsCfg})
 		clientEncryption, err := mongo.NewClientEncryption(keyVaultClient, ceo)
 		require.NoError(mt, err, "error on NewClientEncryption: %v", err)
 
-		err = setFailPoint("http", 1)
+		err = setFailPoint("network", 1)
 		require.NoError(mt, err, "mock server error: %v", err)
 
 		dkOpts := options.DataKey().SetMasterKey(
@@ -3032,7 +3036,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 		keyID, err = clientEncryption.CreateDataKey(context.Background(), "aws", dkOpts)
 		require.NoError(mt, err, "error in CreateDataKey: %v", err)
 
-		err = setFailPoint("http", 1)
+		err = setFailPoint("network", 1)
 		require.NoError(mt, err, "mock server error: %v", err)
 
 		testVal := bson.RawValue{Type: bson.TypeInt32, Value: bsoncore.AppendInt32(nil, 123)}

x/mongo/driver/crypt.go

@@ -9,9 +9,7 @@ package driver
 import (
 	"context"
 	"crypto/tls"
-	"errors"
 	"fmt"
-	"io"
 	"strings"
 	"time"
 
@@ -399,7 +397,10 @@ 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) {
+		if err != nil {
+			if kmsCtx.Fail() {
+				err = nil
+			}
 			return err
 		}
 

x/mongo/driver/mongocrypt/mongocrypt.go

@@ -53,6 +53,7 @@ func NewMongoCrypt(opts *options.MongoCryptOptions) (*MongoCrypt, error) {
 	if wrapped == nil {
 		return nil, errors.New("could not create new mongocrypt object")
 	}
+	C.mongocrypt_setopt_retry_kms(wrapped, true)
 	httpClient := opts.HTTPClient
 	if httpClient == nil {
 		httpClient = httputil.DefaultHTTPClient
@@ -85,7 +86,7 @@ func NewMongoCrypt(opts *options.MongoCryptOptions) (*MongoCrypt, error) {
 	}
 
 	if opts.BypassQueryAnalysis {
-		C.mongocrypt_setopt_bypass_query_analysis(wrapped)
+		C.mongocrypt_setopt_bypass_query_analysis(crypt.wrapped)
 	}
 
 	// If loading the crypt_shared library isn't disabled, set the default library search path "$SYSTEM"

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)
+}