encrypt retry fix

mongodb · Dec 11, 2024 · afd7d11 · afd7d11
1 parent 46213b6
commit afd7d11
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/src/mongocrypt-ctx-rewrap-many-datakey.c b/src/mongocrypt-ctx-rewrap-many-datakey.c
@@ -136,6 +136,19 @@ static mongocrypt_kms_ctx_t *_next_kms_ctx_encrypt(mongocrypt_ctx_t *ctx) {
     mongocrypt_ctx_t *dkctx = NULL;
 
     BSON_ASSERT_PARAM(ctx);
+    _mongocrypt_ctx_rmd_datakey_t *it = rmdctx->datakeys;
+
+    /* Check if any need retry */
+    {
+        _mongocrypt_ctx_rmd_datakey_t *it = rmdctx->datakeys;
+        while (it != NULL) {
+            _mongocrypt_ctx_datakey_t *dkctx = (_mongocrypt_ctx_datakey_t *)it->dkctx;
+            if (dkctx->kms.should_retry) {
+                return &dkctx->kms;
+            }
+            it = it->next;
+        }
+    }
 
     /* No more datakey contexts requiring KMS. */
     if (!rmdctx->datakeys_iter) {
@@ -166,6 +179,7 @@ static bool _add_new_datakey(mongocrypt_ctx_t *ctx, key_returned_t *key) {
     datakey->dkctx = mongocrypt_ctx_new(ctx->crypt);
     datakey->next = rmdctx->datakeys;
     datakey->doc = key->doc;
+    // add KMS
     rmdctx->datakeys = datakey; /* Ownership transfer. */
 
     /* Set new provider and master key (rewrapManyDataKeyOpts). */