RUBY-3429 Retry failed KMS requests

.evergreen/run-tests.sh

@@ -219,6 +219,7 @@ if test -n "$FLE"; then
   python3 -u .evergreen/csfle/kms_http_server.py --ca_file .evergreen/x509gen/ca.pem --cert_file .evergreen/x509gen/server.pem --port 8002 --require_client_cert &
   python3 -u .evergreen/csfle/kms_kmip_server.py &
   python3 -u .evergreen/csfle/fake_azure.py &
+  python3 -u .evergreen/csfle/kms_failpoint_server.py --port 9003 &
 
   # Obtain temporary AWS credentials
   PYTHON=python3 . .evergreen/csfle/set-temp-creds.sh

gemfiles/standard.rb

@@ -68,6 +68,6 @@ def standard_dependencies
     gem 'ruby-lsp', platforms: :mri
   end
 
-  gem 'libmongocrypt-helper', '~> 1.11.0' if ENV['FLE'] == 'helper'
+  gem 'libmongocrypt-helper', '~> 1.12.0' if ENV['FLE'] == 'helper'
 end
 # rubocop:enable Metrics/AbcSize, Metrics/MethodLength, Metrics/BlockLength

lib/mongo/crypt/binding.rb

@@ -83,7 +83,7 @@ class Binding
       # will cause a `LoadError`.
       #
       # @api private
-      MIN_LIBMONGOCRYPT_VERSION = Gem::Version.new("1.7.0")
+      MIN_LIBMONGOCRYPT_VERSION = Gem::Version.new("1.12.0")
 
       # @!method self.mongocrypt_version(len)
       #   @api private
@@ -1113,6 +1113,62 @@ def self.check_kms_ctx_status(kms_context)
         end
       end
 
+      # @!method self.mongocrypt_kms_ctx_usleep(ctx)
+      #   @api private
+      #
+      #   Indicates how long to sleep before sending KMS request.
+      #
+      #   @param [ FFI::Pointer ] ctx A pointer to a mongocrypt_ctx_t object.
+      #   @return [ int64 ] A 64-bit encoded number of microseconds of how long to sleep.
+      attach_function :mongocrypt_kms_ctx_usleep, [:pointer], :int64
+
+      # Returns number of milliseconds to sleep before sending KMS request
+      # for the given KMS context.
+      #
+      # @param [ Mongo::Crypt::KmsContext ] kms_context KMS Context we are going
+      #   to send KMS request for.
+      # @return [ Integer ] A 64-bit encoded number of microseconds to sleep.
+      def self.kms_ctx_usleep(kms_context)
+        mongocrypt_kms_ctx_usleep(kms_context.kms_ctx_p)
+      end
+
+      # @!method self.mongocrypt_kms_ctx_fail(ctx)
+      #   @api private
+      #
+      # Indicate a network-level failure.
+      #
+      # @param [ FFI::Pointer ] ctx A pointer to a mongocrypt_ctx_t object.
+      # @return [ Boolean ] whether the failed request may be retried.
+      attach_function :mongocrypt_kms_ctx_fail, [:pointer], :bool
+
+      # Check whether the last failed request for the KMS context may be retried.
+      #
+      # @param [ Mongo::Crypt::KmsContext ] kms_context KMS Context
+      # @return [ true, false ] whether the failed request may be retried.
+      def self.kms_ctx_fail(kms_context)
+        mongocrypt_kms_ctx_fail(kms_context.kms_ctx_p)
+      end
+
+      # @!method self.mongocrypt_setopt_retry_kms(crypt, enable)
+      #   @api private
+      #
+      # Enable or disable KMS retry behavior.
+      #
+      # @param [ FFI::Pointer ] crypt A pointer to a mongocrypt_t object
+      # @param [ Boolean ] enable A boolean indicating whether to retry operations.
+      # @return [ Boolean ] indicating success.
+      attach_function :mongocrypt_setopt_retry_kms, [:pointer, :bool], :bool
+
+      # Enable or disable KMS retry behavior.
+      #
+      # @param [ Mongo::Crypt::Handle ] handle
+      # @param [ true, false ] value whether to retry operations.
+      # @return [ true, fale ] true is the option was set, otherwise false.
+      def self.kms_ctx_setopt_retry_kms(handle, value)
+        mongocrypt_setopt_retry_kms(handle.ref, value)
+      end
+
+
       # @!method self.mongocrypt_kms_ctx_done(ctx)
       #   @api private
       #

lib/mongo/crypt/context.rb

@@ -49,7 +49,6 @@ def initialize(mongocrypt_handle, io)
           Binding.mongocrypt_ctx_new(@mongocrypt_handle.ref),
           Binding.method(:mongocrypt_ctx_destroy)
         )
-
         @encryption_io = io
         @cached_azure_token = nil
       end
@@ -90,35 +89,13 @@ def run_state_machine(timeout_holder)
           when :done
             return nil
           when :need_mongo_keys
-            filter = Binding.ctx_mongo_op(self)
-
-            @encryption_io.find_keys(filter, timeout_ms: timeout_ms).each do |key|
-              mongocrypt_feed(key) if key
-            end
-
-            mongocrypt_done
+            provide_keys(timeout_ms)
           when :need_mongo_collinfo
-            filter = Binding.ctx_mongo_op(self)
-
-            result = @encryption_io.collection_info(@db_name, filter, timeout_ms: timeout_ms)
-            mongocrypt_feed(result) if result
-
-            mongocrypt_done
+            provide_collection_info(timeout_ms)
           when :need_mongo_markings
-            cmd = Binding.ctx_mongo_op(self)
-
-            result = @encryption_io.mark_command(cmd, timeout_ms: timeout_ms)
-            mongocrypt_feed(result)
-
-            mongocrypt_done
+            provide_markings(timeout_ms)
           when :need_kms
-            while kms_context = Binding.ctx_next_kms_ctx(self) do
-              provider = Binding.kms_ctx_get_kms_provider(kms_context)
-              tls_options = @mongocrypt_handle.kms_tls_options(provider)
-              @encryption_io.feed_kms(kms_context, tls_options)
-            end
-
-            Binding.ctx_kms_done(self)
+            feed_kms
           when :need_kms_credentials
             Binding.ctx_provide_kms_providers(
               self,
@@ -134,6 +111,57 @@ def run_state_machine(timeout_holder)
 
       private
 
+      def provide_markings(timeout_ms)
+        cmd = Binding.ctx_mongo_op(self)
+
+        result = @encryption_io.mark_command(cmd, timeout_ms: timeout_ms)
+        mongocrypt_feed(result)
+
+        mongocrypt_done
+      end
+
+      def provide_collection_info(timeout_ms)
+        filter = Binding.ctx_mongo_op(self)
+
+        result = @encryption_io.collection_info(@db_name, filter, timeout_ms: timeout_ms)
+        mongocrypt_feed(result) if result
+
+        mongocrypt_done
+      end
+
+      def provide_keys(timeout_ms)
+        filter = Binding.ctx_mongo_op(self)
+
+        @encryption_io.find_keys(filter, timeout_ms: timeout_ms).each do |key|
+          mongocrypt_feed(key) if key
+        end
+
+        mongocrypt_done
+      end
+
+      def feed_kms
+        while (kms_context = Binding.ctx_next_kms_ctx(self)) do
+          begin
+            delay = Binding.kms_ctx_usleep(kms_context)
+            sleep(delay / 1_000_000) unless delay.nil?
+            provider = Binding.kms_ctx_get_kms_provider(kms_context)
+            tls_options = @mongocrypt_handle.kms_tls_options(provider)
+            @encryption_io.feed_kms(kms_context, tls_options)
+          rescue Error::KmsError => e
+            if e.network_error?
+              if Binding.kms_ctx_fail(kms_context)
+                next
+              else
+                raise
+              end
+            else
+              raise
+            end
+          end
+        end
+        Binding.ctx_kms_done(self)
+      end
+
       # Indicate that state machine is done feeding I/O responses back to libmongocrypt
       def mongocrypt_done
         Binding.mongocrypt_ctx_mongo_done(ctx_p)

lib/mongo/crypt/encryption_io.rb

@@ -363,8 +363,10 @@ def with_ssl_socket(endpoint, tls_options, timeout_ms: nil)
           tls_options.merge(socket_options)
         )
         yield(mongo_socket.socket)
+      rescue EOFError => e
+        raise Error::KmsError.new("Error when connecting to KMS provider: #{e.class}: #{e.message}", network_error: true)
       rescue => e
-        raise Error::KmsError, "Error when connecting to KMS provider: #{e.class}: #{e.message}"
+        raise Error::KmsError.new("Error when connecting to KMS provider: #{e.class}: #{e.message}")
       ensure
         mongo_socket&.close
       end

lib/mongo/crypt/handle.rb

@@ -71,7 +71,7 @@ def initialize(kms_providers, kms_tls_options, options={})
           Binding.mongocrypt_new,
           Binding.method(:mongocrypt_destroy)
         )
-
+        Binding.kms_ctx_setopt_retry_kms(self, true)
         @kms_providers = kms_providers
         @kms_tls_options =  kms_tls_options
 

lib/mongo/error/kms_error.rb

@@ -20,6 +20,15 @@ class Error
 
     # A KMS-related error during client-side encryption.
     class KmsError < CryptError
+      def initialize(message, code: nil, network_error: nil)
+        @network_error = network_error
+        super(message, code: code)
+      end
+    end
+
+    # @return [ true, false ] whether this error was caused by a network error.
+    def network_error?
+      @network_error == true
     end
   end
 end

lib/mongo/socket/ssl.rb

@@ -312,6 +312,8 @@ def create_context(options)
             context.options = context.options | OpenSSL::SSL::OP_NO_RENEGOTIATION
           end
 
+          context.options |= OpenSSL::SSL::OP_IGNORE_UNEXPECTED_EOF
+
           if context.respond_to?(:renegotiation_cb=)
             # Disable renegotiation for older Ruby versions per the sample code at
             # https://rubydocs.org/d/ruby-2-6-0/classes/OpenSSL/SSL/SSLContext.html

spec/integration/client_side_encryption/kms_retry_prose_spec.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+def simulate_failure(type, times = 1)
+  # Define the URL and the data
+  url = URI.parse("https://localhost:9003/set_failpoint/#{type}")
+  data = { count: times }.to_json
+  # Create a new HTTP object
+  http = Net::HTTP.new(url.host, url.port)
+  # Enable SSL/TLS
+  http.use_ssl = true
+  http.verify_mode = OpenSSL::SSL::VERIFY_NONE
+  # Load the CA certificate
+  http.ca_file = '.evergreen/x509gen/ca.pem'
+  # Create a new POST request
+  request = Net::HTTP::Post.new(url.path, { 'Content-Type' => 'application/json' })
+  request.body = data
+  # Execute the request
+  http.request(request)
+end
+
+describe 'KMS Retry Prose Spec' do
+  require_libmongocrypt
+  require_enterprise
+  min_server_version '4.2'
+
+  include_context 'define shared FLE helpers'
+
+  let(:key_vault_client) do
+    ClientRegistry.instance.new_local_client(SpecConfig.instance.addresses)
+  end
+
+  let(:client_encryption) do
+    Mongo::ClientEncryption.new(
+      key_vault_client,
+      kms_tls_options: {
+        aws: default_kms_tls_options_for_provider,
+        gcp: default_kms_tls_options_for_provider,
+        azure: default_kms_tls_options_for_provider,
+      },
+      key_vault_namespace: key_vault_namespace,
+      # For some reason libmongocrypt ignores custom endpoints for Azure and CGP
+      # kms_providers: aws_kms_providers.merge(azure_kms_providers).merge(gcp_kms_providers)
+      kms_providers: aws_kms_providers
+    )
+  end
+
+  shared_examples 'kms_retry prose spec' do
+    it 'createDataKey and encrypt with TCP retry' do
+      simulate_failure('network')
+      data_key_id = client_encryption.create_data_key(kms_provider, master_key: master_key)
+      simulate_failure('network')
+      expect {
+        client_encryption.encrypt(123, key_id: data_key_id, algorithm: 'AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic')
+      }.not_to raise_error
+    end
+
+    it 'createDataKey and encrypt with HTTP retry' do
+      simulate_failure('http')
+      data_key_id = client_encryption.create_data_key(kms_provider, master_key: master_key)
+      simulate_failure('http')
+      expect {
+        client_encryption.encrypt(123, key_id: data_key_id, algorithm: 'AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic')
+      }.not_to raise_error
+    end
+
+    it 'createDataKey fails after too many retries' do
+      simulate_failure('network', 4)
+      expect {
+        client_encryption.create_data_key(kms_provider, master_key: master_key)
+      }.to raise_error(Mongo::Error::KmsError)
+    end
+  end
+
+  context 'aws' do
+    let(:kms_provider) { 'aws' }
+
+    let(:master_key) do
+      {
+        region: "foo",
+        key: "bar",
+        endpoint: "127.0.0.1:9003",
+      }
+    end
+
+    include_examples 'kms_retry prose spec'
+  end
+
+  # For some reason libmongocrypt ignores custom endpoints for Azure and CGP
+  xcontext 'gcp' do
+    let(:kms_provider) { 'gcp' }
+
+    let(:master_key) do
+      {
+        project_id: "foo",
+        location: "bar",
+        key_ring: "baz",
+        key_name: "qux",
+        endpoint: "127.0.0.1:9003"
+      }
+    end
+
+    include_examples 'kms_retry prose spec'
+  end
+
+  # For some reason libmongocrypt ignores custom endpoints for Azure and CGP
+  xcontext 'azure' do
+    let(:kms_provider) { 'azure' }
+
+    let(:master_key) do
+      {
+        key_vault_endpoint: "127.0.0.1:9003",
+        key_name: "foo",
+      }
+    end
+
+    include_examples 'kms_retry prose spec'
+  end
+end