CDRIVER-5517 retry KMS encrypt requests on transient errors

.evergreen/config_generator/components/funcs/run_mock_kms_servers.py

@@ -47,6 +47,7 @@ class RunMockKMSServers(Function):
                 python -u kms_http_server.py --ca_file ../x509gen/ca.pem --cert_file ../x509gen/expired.pem --port 9000 &
                 python -u kms_http_server.py --ca_file ../x509gen/ca.pem --cert_file ../x509gen/wrong-host.pem --port 9001 &
                 python -u kms_http_server.py --ca_file ../x509gen/ca.pem --cert_file ../x509gen/server.pem --require_client_cert --port 9002 &
+                python -u kms_failpoint_server.py --port 9003 &
                 python -u kms_kmip_server.py &
                 deactivate
                 echo "Starting mock KMS TLS servers... done."

.evergreen/generated_configs/functions.yml

@@ -384,6 +384,7 @@ functions:
             python -u kms_http_server.py --ca_file ../x509gen/ca.pem --cert_file ../x509gen/expired.pem --port 9000 &
             python -u kms_http_server.py --ca_file ../x509gen/ca.pem --cert_file ../x509gen/wrong-host.pem --port 9001 &
             python -u kms_http_server.py --ca_file ../x509gen/ca.pem --cert_file ../x509gen/server.pem --require_client_cert --port 9002 &
+            python -u kms_failpoint_server.py --port 9003 &
             python -u kms_kmip_server.py &
             deactivate
             echo "Starting mock KMS TLS servers... done."

.evergreen/scripts/compile-libmongocrypt.sh

@@ -9,11 +9,13 @@ compile_libmongocrypt() {
   # libmongocrypt's kms-message in `src/kms-message`. Run
   # `.evergreen/scripts/kms-divergence-check.sh` to ensure that there is no
   # divergence in the copied files.
-  # TODO: once 1.12.0 is released replace the following with:
+
+  # TODO: once 1.12.0 is released (containing MONGOCRYPT-599) replace the following with:
   # git clone -q --depth=1 https://github.com/mongodb/libmongocrypt --branch 1.12.0 || return
   {
     git clone -q https://github.com/mongodb/libmongocrypt || return
-    git -C libmongocrypt checkout bca8e7dc1ecb7b1c039132e07de5e0db2703c701
+    # Check out commit containing MONGOCRYPT-599
+    git -C libmongocrypt checkout 7aeaec4ae1369c7d3c5b3aea6f1da35c5e9478b0
   }
 
   declare -a crypt_cmake_flags=(

src/libmongoc/src/mongoc/mongoc-crypt.c

@@ -542,8 +542,9 @@ _state_need_kms (_state_machine_t *state_machine, bson_error_t *error)
    mongocrypt_binary_t *http_reply = NULL;
    const char *endpoint;
    const int32_t sockettimeout = MONGOC_DEFAULT_SOCKETTIMEOUTMS;
-   kms_ctx = mongocrypt_ctx_next_kms_ctx (state_machine->ctx);
-   while (kms_ctx) {
+   int64_t sleep_usec = 0;
+
+   while ((kms_ctx = mongocrypt_ctx_next_kms_ctx (state_machine->ctx))) {
       mongoc_iovec_t iov;
       const mongoc_ssl_opt_t *ssl_opt;
       const char *provider;
@@ -576,6 +577,11 @@ _state_need_kms (_state_machine_t *state_machine, bson_error_t *error)
          goto fail;
       }
 
+      sleep_usec = mongocrypt_kms_ctx_usleep (kms_ctx);
+      if (sleep_usec > 0) {
+         _mongoc_usleep (sleep_usec);
+      }
+
       mongoc_stream_destroy (tls_stream);
       tls_stream = _get_stream (endpoint, sockettimeout, ssl_opt, error);
 #ifdef MONGOC_ENABLE_SSL_SECURE_CHANNEL
@@ -585,14 +591,31 @@ _state_need_kms (_state_machine_t *state_machine, bson_error_t *error)
       }
 #endif
       if (!tls_stream) {
-         goto fail;
+         if (mongocrypt_kms_ctx_fail (kms_ctx)) {
+            continue;
+         } else {
+            /* TLS errors are set in _get_stream */
+            goto fail;
+         }
       }
 
       iov.iov_base = (char *) mongocrypt_binary_data (http_req);
       iov.iov_len = mongocrypt_binary_len (http_req);
 
       if (!_mongoc_stream_writev_full (tls_stream, &iov, 1, sockettimeout, error)) {
-         goto fail;
+         if (mongocrypt_kms_ctx_fail (kms_ctx)) {
+            continue;
+         } else {
+            bson_error_t kms_error;
+            BSON_ASSERT (!_kms_ctx_check_error (kms_ctx, &kms_error, true));
+            bson_set_error (error,
+                            MONGOC_ERROR_STREAM,
+                            MONGOC_ERROR_STREAM_SOCKET,
+                            "%s. Failed to write to KMS stream: %s",
+                            kms_error.message,
+                            endpoint);
+            goto fail;
+         }
       }
 
       /* Read and feed reply. */
@@ -608,17 +631,21 @@ _state_need_kms (_state_machine_t *state_machine, bson_error_t *error)
          }
 
          read_ret = mongoc_stream_read (tls_stream, buf, bytes_needed, 1 /* min_bytes. */, sockettimeout);
-         if (read_ret == -1) {
-            bson_set_error (
-               error, MONGOC_ERROR_STREAM, MONGOC_ERROR_STREAM_SOCKET, "failed to read from KMS stream: %d", errno);
-            goto fail;
+         if (read_ret <= 0) {
+            if (mongocrypt_kms_ctx_fail (kms_ctx)) {
+               break; // Stop reading reply.
+            } else {
+               bson_error_t kms_error;
+               BSON_ASSERT (!_kms_ctx_check_error (kms_ctx, &kms_error, true));
+               bson_set_error (error,
+                               MONGOC_ERROR_STREAM,
+                               MONGOC_ERROR_STREAM_SOCKET,
+                               "%s. Failed to read from KMS stream to: %s",
+                               kms_error.message,
+                               endpoint);
+               goto fail;
+            }
          }
-
-         if (read_ret == 0) {
-            bson_set_error (error, MONGOC_ERROR_STREAM, MONGOC_ERROR_STREAM_SOCKET, "unexpected EOF from KMS stream");
-            goto fail;
-         }
-
          mongocrypt_binary_destroy (http_reply);
 
          BSON_ASSERT (bson_in_range_signed (uint32_t, read_ret));
@@ -628,7 +655,6 @@ _state_need_kms (_state_machine_t *state_machine, bson_error_t *error)
             goto fail;
          }
       }
-      kms_ctx = mongocrypt_ctx_next_kms_ctx (state_machine->ctx);
    }
    /* When NULL is returned by mongocrypt_ctx_next_kms_ctx, this can either be
     * an error or end-of-list. */
@@ -1366,6 +1392,7 @@ _mongoc_crypt_new (const bson_t *kms_providers,
    crypt = bson_malloc0 (sizeof (*crypt));
    crypt->kmsid_to_tlsopts = mcd_mapof_kmsid_to_tlsopts_new ();
    crypt->handle = mongocrypt_new ();
+   mongocrypt_setopt_retry_kms (crypt->handle, true);
 
    // Stash away a copy of the user's kmsProviders in case we need to lazily
    // load credentials.

src/libmongoc/tests/test-mongoc-client-side-encryption.c

@@ -31,6 +31,7 @@
 #include "mongoc/mongoc-client-side-encryption-private.h"
 
 #include "mongoc/mongoc-uri.h"
+#include "mongoc/mongoc-http-private.h"
 
 static void
 _before_test (json_test_ctx_t *ctx, const bson_t *test)
@@ -2596,7 +2597,7 @@ test_kms_tls_cert_wrong_host (void *unused)
    mongoc_client_destroy (client);
 }
 
-typedef enum { NO_CLIENT_CERT, WITH_TLS, INVALID_HOSTNAME, EXPIRED, WITH_NAMES } tls_test_ce_t;
+typedef enum { NO_CLIENT_CERT, WITH_TLS, INVALID_HOSTNAME, EXPIRED, WITH_NAMES, RETRY } tls_test_ce_t;
 
 static mongoc_client_encryption_t *
 _tls_test_make_client_encryption (mongoc_client_t *keyvault_client, tls_test_ce_t test_ce)
@@ -2619,7 +2620,9 @@ _tls_test_make_client_encryption (mongoc_client_t *keyvault_client, tls_test_ce_
    char *ca_file = test_framework_getenv_required ("MONGOC_TEST_CSFLE_TLS_CA_FILE");
    char *certificate_key_file = test_framework_getenv_required ("MONGOC_TEST_CSFLE_TLS_CERTIFICATE_KEY_FILE");
 
-   if (test_ce == WITH_TLS) {
+   if (test_ce == WITH_TLS || test_ce == RETRY) {
+      const char *port = test_ce == RETRY ? "9003" : "9002";
+
       kms_providers = tmp_bson ("{'aws': {'accessKeyId': '%s', 'secretAccessKey': '%s' }}",
                                 mongoc_test_aws_access_key_id,
                                 mongoc_test_aws_secret_access_key);
@@ -2629,19 +2632,21 @@ _tls_test_make_client_encryption (mongoc_client_t *keyvault_client, tls_test_ce_
       bson_concat (kms_providers,
                    tmp_bson ("{'azure': {'tenantId': '%s', 'clientId': '%s', "
                              "'clientSecret': '%s', "
-                             "'identityPlatformEndpoint': '127.0.0.1:9002' }}",
+                             "'identityPlatformEndpoint': '127.0.0.1:%s' }}",
                              mongoc_test_azure_tenant_id,
                              mongoc_test_azure_client_id,
-                             mongoc_test_azure_client_secret));
+                             mongoc_test_azure_client_secret,
+                             port));
       bson_concat (
          tls_opts,
          tmp_bson ("{'azure': {'tlsCaFile': '%s', 'tlsCertificateKeyFile': '%s' }}", ca_file, certificate_key_file));
 
       bson_concat (kms_providers,
                    tmp_bson ("{'gcp': { 'email': '%s', 'privateKey': '%s', "
-                             "'endpoint': '127.0.0.1:9002' }}",
+                             "'endpoint': '127.0.0.1:%s' }}",
                              mongoc_test_gcp_email,
-                             mongoc_test_gcp_privatekey));
+                             mongoc_test_gcp_privatekey,
+                             port));
       bson_concat (
          tls_opts,
          tmp_bson ("{'gcp': {'tlsCaFile': '%s', 'tlsCertificateKeyFile': '%s' }}", ca_file, certificate_key_file));
@@ -3228,6 +3233,60 @@ test_kms_tls_options_extra_rejected (void *unused)
    mongoc_client_destroy (keyvault_client);
 }
 
+static const char *failpoint_server_ip = "127.0.0.1";
+static const int failpoint_server_port = 9003;
+
+static void
+reset_failpoints (mongoc_ssl_opt_t *ssl_opts)
+{
+   mongoc_http_request_t req;
+   mongoc_http_response_t res;
+   bool r;
+   bson_error_t error = {0};
+
+   _mongoc_http_request_init (&req);
+   _mongoc_http_response_init (&res);
+
+   req.method = "POST";
+   req.host = failpoint_server_ip;
+   req.port = failpoint_server_port;
+   req.path = "/reset";
+
+   r = _mongoc_http_send (&req, 10000, true, ssl_opts, &res, &error);
+   ASSERT_OR_PRINT (r, error);
+   _mongoc_http_response_cleanup (&res);
+}
+
+static void
+set_retry_failpoint (mongoc_ssl_opt_t *ssl_opts, bool network, uint32_t count)
+{
+   mongoc_http_request_t req;
+   mongoc_http_response_t res;
+   bool r;
+   bson_error_t error = {0};
+
+   _mongoc_http_request_init (&req);
+   _mongoc_http_response_init (&res);
+
+   req.method = "POST";
+   req.host = failpoint_server_ip;
+   req.port = failpoint_server_port;
+   if (network) {
+      req.path = "/set_failpoint/network";
+   } else {
+      req.path = "/set_failpoint/http";
+   }
+   req.extra_headers = "Content-Type: application/json\r\n";
+   char count_json[25];
+   sprintf (count_json, "{\"count\": %" PRIu32 "}", count);
+   req.body = count_json;
+   req.body_len = strlen (count_json);
+
+   r = _mongoc_http_send (&req, 10000, true, ssl_opts, &res, &error);
+   ASSERT_OR_PRINT (r, error);
+   _mongoc_http_response_cleanup (&res);
+}
+
 /* ee_fixture is a fixture for the Explicit Encryption prose test. */
 typedef struct {
    bson_value_t key1ID;
@@ -6232,6 +6291,88 @@ test_range_explicit_encryption_applies_defaults (void *unused)
    mongoc_client_destroy (keyVaultClient);
 }
 
+static void
+_test_retry_with_masterkey (const char *provider, bson_t *masterkey)
+{
+   mongoc_client_t *keyvault_client = test_framework_new_default_client ();
+   mongoc_client_encryption_t *client_encryption = _tls_test_make_client_encryption (keyvault_client, RETRY);
+   bson_error_t error = {0};
+   bson_value_t keyid;
+   mongoc_client_encryption_datakey_opts_t *dkopts;
+   char *ca_file = test_framework_getenv_required ("MONGOC_TEST_CSFLE_TLS_CA_FILE");
+   char *pem_file = test_framework_getenv_required ("MONGOC_TEST_CSFLE_TLS_CERTIFICATE_KEY_FILE");
+   mongoc_ssl_opt_t ssl_opts = {.ca_file = ca_file, .pem_file = pem_file};
+   bool res;
+
+   bson_value_t to_encrypt = {.value_type = BSON_TYPE_INT32, .value.v_int32 = 123};
+   bson_value_t encrypted_field = {0};
+   mongoc_client_encryption_encrypt_opts_t *encrypt_opts = mongoc_client_encryption_encrypt_opts_new ();
+   mongoc_client_encryption_encrypt_opts_set_algorithm (encrypt_opts,
+                                                        MONGOC_AEAD_AES_256_CBC_HMAC_SHA_512_DETERMINISTIC);
+
+   reset_failpoints (&ssl_opts);
+
+   // Case 1: createDataKey and encrypt with TCP retry
+   dkopts = mongoc_client_encryption_datakey_opts_new ();
+   mongoc_client_encryption_datakey_opts_set_masterkey (dkopts, masterkey);
+   set_retry_failpoint (&ssl_opts, true, 1);
+   res = mongoc_client_encryption_create_datakey (client_encryption, provider, dkopts, &keyid, &error);
+   ASSERT_OR_PRINT (res, error);
+
+   set_retry_failpoint (&ssl_opts, true, 1);
+   mongoc_client_encryption_encrypt_opts_set_keyid (encrypt_opts, &keyid);
+   res = mongoc_client_encryption_encrypt (client_encryption, &to_encrypt, encrypt_opts, &encrypted_field, &error);
+   ASSERT_OR_PRINT (res, error);
+   bson_value_destroy (&keyid);
+   bson_value_destroy (&encrypted_field);
+   mongoc_client_encryption_datakey_opts_destroy (dkopts);
+
+   // Case 2: createDataKey and encrypt with HTTP retry
+   dkopts = mongoc_client_encryption_datakey_opts_new ();
+   mongoc_client_encryption_datakey_opts_set_masterkey (dkopts, masterkey);
+   set_retry_failpoint (&ssl_opts, false, 1);
+   res = mongoc_client_encryption_create_datakey (client_encryption, provider, dkopts, &keyid, &error);
+   ASSERT_OR_PRINT (res, error);
+
+   set_retry_failpoint (&ssl_opts, false, 1);
+   mongoc_client_encryption_encrypt_opts_set_keyid (encrypt_opts, &keyid);
+   res = mongoc_client_encryption_encrypt (client_encryption, &to_encrypt, encrypt_opts, &encrypted_field, &error);
+   ASSERT_OR_PRINT (res, error);
+   bson_value_destroy (&keyid);
+   bson_value_destroy (&encrypted_field);
+   mongoc_client_encryption_datakey_opts_destroy (dkopts);
+
+   // Case 3: createDataKey fails after too many retries
+   dkopts = mongoc_client_encryption_datakey_opts_new ();
+   mongoc_client_encryption_datakey_opts_set_masterkey (dkopts, masterkey);
+   set_retry_failpoint (&ssl_opts, true, 4);
+   res = mongoc_client_encryption_create_datakey (client_encryption, provider, dkopts, &keyid, &error);
+   ASSERT_ERROR_CONTAINS (error, MONGOC_ERROR_STREAM, MONGOC_ERROR_STREAM_SOCKET, "KMS request failed after");
+
+   bson_value_destroy (&keyid);
+   mongoc_client_encryption_datakey_opts_destroy (dkopts);
+
+   bson_free (ca_file);
+   bson_free (pem_file);
+   mongoc_client_encryption_encrypt_opts_destroy (encrypt_opts);
+   mongoc_client_encryption_destroy (client_encryption);
+   mongoc_client_destroy (keyvault_client);
+}
+
+/* Prose test 23: KMS Retry Tests */
+static void
+test_kms_retry (void *unused)
+{
+   bson_t *aws_masterkey = tmp_bson (BSON_STR ({"region" : "r", "key" : "k", "endpoint" : "127.0.0.1:9003"}));
+   bson_t *azure_masterkey = tmp_bson (BSON_STR ({"keyVaultEndpoint" : "127.0.0.1:9003", "keyName" : "foo"}));
+   bson_t *gcp_masterkey = tmp_bson (BSON_STR (
+      {"projectId" : "foo", "location" : "bar", "keyRing" : "baz", "keyName" : "qux", "endpoint" : "127.0.0.1:9003"}));
+
+   _test_retry_with_masterkey ("aws", aws_masterkey);
+   _test_retry_with_masterkey ("azure", azure_masterkey);
+   _test_retry_with_masterkey ("gcp", gcp_masterkey);
+}
+
 void
 test_client_side_encryption_install (TestSuite *suite)
 {
@@ -6412,6 +6553,12 @@ test_client_side_encryption_install (TestSuite *suite)
                       NULL,
                       NULL,
                       test_framework_skip_if_no_client_side_encryption);
+   TestSuite_AddFull (suite,
+                      "/client_side_encryption/kms_retry",
+                      test_kms_retry,
+                      NULL,
+                      NULL,
+                      test_framework_skip_if_no_client_side_encryption);
 
    TestSuite_AddFull (suite,
                       "/client_side_encryption/explicit_encryption/case1",