diff --git a/src/mongocrypt-ctx-rewrap-many-datakey.c b/src/mongocrypt-ctx-rewrap-many-datakey.c index 32fc6f793..a9deb3067 100644 --- 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). */