RUST-1894 Retry KMS requests on transient errors (#37)

Author: isabelatkinson

mongocrypt-sys/src/bindings.rs

@@ -1,4 +1,4 @@
-/* automatically generated by rust-bindgen 0.63.0 */
+/* automatically generated by rust-bindgen 0.64.0 */
 
 extern "C" {
     #[doc = " Returns the version string for libmongocrypt.\n\n @param[out] len  An optional length of the returned string. May be NULL.\n @returns a NULL terminated version string for libmongocrypt."]
@@ -158,6 +158,10 @@ extern "C" {
         log_ctx: *mut ::std::os::raw::c_void,
     ) -> bool;
 }
+extern "C" {
+    #[doc = " Enable or disable KMS retry behavior.\n\n @param[in] crypt The @ref mongocrypt_t object.\n @param[in] enable A boolean indicating whether to retry operations.\n @pre @ref mongocrypt_init has not been called on @p crypt.\n @returns A boolean indicating success. If false, an error status is set.\n Retrieve it with @ref mongocrypt_ctx_status"]
+    pub fn mongocrypt_setopt_retry_kms(crypt: *mut mongocrypt_t, enable: bool) -> bool;
+}
 extern "C" {
     #[doc = " Configure an AWS KMS provider on the @ref mongocrypt_t object.\n\n This has been superseded by the more flexible:\n @ref mongocrypt_setopt_kms_providers\n\n @param[in] crypt The @ref mongocrypt_t object.\n @param[in] aws_access_key_id The AWS access key ID used to generate KMS\n messages.\n @param[in] aws_access_key_id_len The string length (in bytes) of @p\n aws_access_key_id. Pass -1 to determine the string length with strlen (must\n be NULL terminated).\n @param[in] aws_secret_access_key The AWS secret access key used to generate\n KMS messages.\n @param[in] aws_secret_access_key_len The string length (in bytes) of @p\n aws_secret_access_key. Pass -1 to determine the string length with strlen\n (must be NULL terminated).\n @pre @ref mongocrypt_init has not been called on @p crypt.\n @returns A boolean indicating success. If false, an error status is set.\n Retrieve it with @ref mongocrypt_ctx_status"]
     pub fn mongocrypt_setopt_kms_provider_aws(
@@ -238,7 +242,7 @@ extern "C" {
     ) -> *const ::std::os::raw::c_char;
 }
 extern "C" {
-    #[doc = " @brief Obtain a 64-bit constant encoding the version of the loaded\n crypt_shared library, if available.\n\n @param[in] crypt The mongocrypt_t object after a successul call to\n mongocrypt_init.\n\n @return A 64-bit encoded version number, with the version encoded as four\n sixteen-bit integers, or zero if no crypt_shared library was loaded.\n\n The version is encoded as four 16-bit numbers, from high to low:\n\n - Major version\n - Minor version\n - Revision\n - Reserved\n\n For example, version 6.2.1 would be encoded as: 0x0006'0002'0001'0000"]
+    #[doc = " @brief Obtain a 64-bit constant encoding the version of the loaded\n crypt_shared library, if available.\n\n @param[in] crypt The mongocrypt_t object after a successful call to\n mongocrypt_init.\n\n @return A 64-bit encoded version number, with the version encoded as four\n sixteen-bit integers, or zero if no crypt_shared library was loaded.\n\n The version is encoded as four 16-bit numbers, from high to low:\n\n - Major version\n - Minor version\n - Revision\n - Reserved\n\n For example, version 6.2.1 would be encoded as: 0x0006'0002'0001'0000"]
     pub fn mongocrypt_crypt_shared_lib_version(crypt: *const mongocrypt_t) -> u64;
 }
 #[repr(C)]
@@ -411,7 +415,7 @@ pub struct _mongocrypt_kms_ctx_t {
 #[doc = " Manages a single KMS HTTP request/response."]
 pub type mongocrypt_kms_ctx_t = _mongocrypt_kms_ctx_t;
 extern "C" {
-    #[doc = " Get the next KMS handle.\n\n Multiple KMS handles may be retrieved at once. Drivers may do this to fan\n out multiple concurrent KMS HTTP requests. Feeding multiple KMS requests\n is thread-safe.\n\n If KMS handles are being handled synchronously, the driver can reuse the same\n TLS socket to send HTTP requests and receive responses.\n\n @param[in] ctx A @ref mongocrypt_ctx_t.\n @returns a new @ref mongocrypt_kms_ctx_t or NULL."]
+    #[doc = " Get the next KMS handle.\n\n Multiple KMS handles may be retrieved at once. Drivers may do this to fan\n out multiple concurrent KMS HTTP requests. Feeding multiple KMS requests\n is thread-safe.\n\n If KMS handles are being handled synchronously, the driver can reuse the same\n TLS socket to send HTTP requests and receive responses.\n\n The returned KMS handle does not outlive `ctx`.\n\n @param[in] ctx A @ref mongocrypt_ctx_t.\n @returns a new @ref mongocrypt_kms_ctx_t or NULL."]
     pub fn mongocrypt_ctx_next_kms_ctx(ctx: *mut mongocrypt_ctx_t) -> *mut mongocrypt_kms_ctx_t;
 }
 extern "C" {
@@ -432,13 +436,21 @@ extern "C" {
     #[doc = " Indicates how many bytes to feed into @ref mongocrypt_kms_ctx_feed.\n\n @param[in] kms The @ref mongocrypt_kms_ctx_t.\n @returns The number of requested bytes."]
     pub fn mongocrypt_kms_ctx_bytes_needed(kms: *mut mongocrypt_kms_ctx_t) -> u32;
 }
+extern "C" {
+    #[doc = " Indicates how long to sleep before sending this request.\n\n @param[in] kms The @ref mongocrypt_kms_ctx_t.\n @returns How long to sleep in microseconds."]
+    pub fn mongocrypt_kms_ctx_usleep(kms: *mut mongocrypt_kms_ctx_t) -> i64;
+}
 extern "C" {
     #[doc = " Feed bytes from the HTTP response.\n\n Feeding more bytes than what has been returned in @ref\n mongocrypt_kms_ctx_bytes_needed is an error.\n\n @param[in] kms The @ref mongocrypt_kms_ctx_t.\n @param[in] bytes The bytes to feed. The viewed data is copied. It is valid to\n destroy @p bytes with @ref mongocrypt_binary_destroy immediately after.\n @returns A boolean indicating success. If false, an error status is set.\n Retrieve it with @ref mongocrypt_kms_ctx_status"]
     pub fn mongocrypt_kms_ctx_feed(
         kms: *mut mongocrypt_kms_ctx_t,
         bytes: *mut mongocrypt_binary_t,
     ) -> bool;
 }
+extern "C" {
+    #[doc = " Indicate a network-level failure.\n\n @param[in] kms The @ref mongocrypt_kms_ctx_t.\n @return A boolean indicating whether the failed request may be retried."]
+    pub fn mongocrypt_kms_ctx_fail(kms: *mut mongocrypt_kms_ctx_t) -> bool;
+}
 extern "C" {
     #[doc = " Get the status associated with a @ref mongocrypt_kms_ctx_t object.\n\n @param[in] kms The @ref mongocrypt_kms_ctx_t object.\n @param[out] status Receives the status.\n\n @returns A boolean indicating success. If false, an error status is set."]
     pub fn mongocrypt_kms_ctx_status(
@@ -475,7 +487,7 @@ extern "C" {
     #[doc = " Destroy and free all memory associated with a @ref mongocrypt_ctx_t.\n\n @param[in] ctx A @ref mongocrypt_ctx_t."]
     pub fn mongocrypt_ctx_destroy(ctx: *mut mongocrypt_ctx_t);
 }
-#[doc = " An crypto AES-256-CBC encrypt or decrypt function.\n\n Note, @p in is already padded. Encrypt with padding disabled.\n @param[in] ctx An optional context object that may have been set when hooks\n were enabled.\n @param[in] key An encryption key (32 bytes for AES_256).\n @param[in] iv An initialization vector (16 bytes for AES_256);\n @param[in] in The input.\n @param[out] out A preallocated byte array for the output. See @ref\n mongocrypt_binary_data.\n @param[out] bytes_written Set this to the number of bytes written to @p out.\n @param[out] status An optional status to pass error messages. See @ref\n mongocrypt_status_set.\n @returns A boolean indicating success. If returning false, set @p status\n with a message indiciating the error using @ref mongocrypt_status_set."]
+#[doc = " An crypto AES-256-CBC encrypt or decrypt function.\n\n Note, @p in is already padded. Encrypt with padding disabled.\n @param[in] ctx An optional context object that may have been set when hooks\n were enabled.\n @param[in] key An encryption key (32 bytes for AES_256).\n @param[in] iv An initialization vector (16 bytes for AES_256);\n @param[in] in The input.\n @param[out] out A preallocated byte array for the output. See @ref\n mongocrypt_binary_data.\n @param[out] bytes_written Set this to the number of bytes written to @p out.\n @param[out] status An optional status to pass error messages. See @ref\n mongocrypt_status_set.\n @returns A boolean indicating success. If returning false, set @p status\n with a message indicating the error using @ref mongocrypt_status_set."]
 pub type mongocrypt_crypto_fn = ::std::option::Option<
     unsafe extern "C" fn(
         ctx: *mut ::std::os::raw::c_void,
@@ -487,7 +499,7 @@ pub type mongocrypt_crypto_fn = ::std::option::Option<
         status: *mut mongocrypt_status_t,
     ) -> bool,
 >;
-#[doc = " A crypto signature or HMAC function.\n\n Currently used in callbacks for HMAC SHA-512, HMAC SHA-256, and RSA SHA-256\n signature.\n\n @param[in] ctx An optional context object that may have been set when hooks\n were enabled.\n @param[in] key An encryption key (32 bytes for HMAC_SHA512).\n @param[in] in The input.\n @param[out] out A preallocated byte array for the output. See @ref\n mongocrypt_binary_data.\n @param[out] status An optional status to pass error messages. See @ref\n mongocrypt_status_set.\n @returns A boolean indicating success. If returning false, set @p status\n with a message indiciating the error using @ref mongocrypt_status_set."]
+#[doc = " A crypto signature or HMAC function.\n\n Currently used in callbacks for HMAC SHA-512, HMAC SHA-256, and RSA SHA-256\n signature.\n\n @param[in] ctx An optional context object that may have been set when hooks\n were enabled.\n @param[in] key An encryption key (32 bytes for HMAC_SHA512).\n @param[in] in The input.\n @param[out] out A preallocated byte array for the output. See @ref\n mongocrypt_binary_data.\n @param[out] status An optional status to pass error messages. See @ref\n mongocrypt_status_set.\n @returns A boolean indicating success. If returning false, set @p status\n with a message indicating the error using @ref mongocrypt_status_set."]
 pub type mongocrypt_hmac_fn = ::std::option::Option<
     unsafe extern "C" fn(
         ctx: *mut ::std::os::raw::c_void,
@@ -497,7 +509,7 @@ pub type mongocrypt_hmac_fn = ::std::option::Option<
         status: *mut mongocrypt_status_t,
     ) -> bool,
 >;
-#[doc = " A crypto hash (SHA-256) function.\n\n @param[in] ctx An optional context object that may have been set when hooks\n were enabled.\n @param[in] in The input.\n @param[out] out A preallocated byte array for the output. See @ref\n mongocrypt_binary_data.\n @param[out] status An optional status to pass error messages. See @ref\n mongocrypt_status_set.\n @returns A boolean indicating success. If returning false, set @p status\n with a message indiciating the error using @ref mongocrypt_status_set."]
+#[doc = " A crypto hash (SHA-256) function.\n\n @param[in] ctx An optional context object that may have been set when hooks\n were enabled.\n @param[in] in The input.\n @param[out] out A preallocated byte array for the output. See @ref\n mongocrypt_binary_data.\n @param[out] status An optional status to pass error messages. See @ref\n mongocrypt_status_set.\n @returns A boolean indicating success. If returning false, set @p status\n with a message indicating the error using @ref mongocrypt_status_set."]
 pub type mongocrypt_hash_fn = ::std::option::Option<
     unsafe extern "C" fn(
         ctx: *mut ::std::os::raw::c_void,
@@ -506,7 +518,7 @@ pub type mongocrypt_hash_fn = ::std::option::Option<
         status: *mut mongocrypt_status_t,
     ) -> bool,
 >;
-#[doc = " A crypto secure random function.\n\n @param[in] ctx An optional context object that may have been set when hooks\n were enabled.\n @param[out] out A preallocated byte array for the output. See @ref\n mongocrypt_binary_data.\n @param[in] count The number of random bytes requested.\n @param[out] status An optional status to pass error messages. See @ref\n mongocrypt_status_set.\n @returns A boolean indicating success. If returning false, set @p status\n with a message indiciating the error using @ref mongocrypt_status_set."]
+#[doc = " A crypto secure random function.\n\n @param[in] ctx An optional context object that may have been set when hooks\n were enabled.\n @param[out] out A preallocated byte array for the output. See @ref\n mongocrypt_binary_data.\n @param[in] count The number of random bytes requested.\n @param[out] status An optional status to pass error messages. See @ref\n mongocrypt_status_set.\n @returns A boolean indicating success. If returning false, set @p status\n with a message indicating the error using @ref mongocrypt_status_set."]
 pub type mongocrypt_random_fn = ::std::option::Option<
     unsafe extern "C" fn(
         ctx: *mut ::std::os::raw::c_void,
@@ -583,9 +595,16 @@ extern "C" {
     ) -> bool;
 }
 extern "C" {
-    #[doc = " Set options for explicit encryption with the \"range\" algorithm.\n\n @p opts is a BSON document of the form:\n {\n    \"min\": Optional<BSON value>,\n    \"max\": Optional<BSON value>,\n    \"sparsity\": Int64,\n    \"precision\": Optional<Int32>,\n    \"trimFactor\": Optional<Int32>\n }\n\n @param[in] ctx The @ref mongocrypt_ctx_t object.\n @param[in] opts BSON.\n @pre @p ctx has not been initialized.\n @returns A boolean indicating success. If false, an error status is set.\n Retrieve it with @ref mongocrypt_ctx_status"]
+    #[doc = " Set options for explicit encryption with the \"range\" algorithm.\n\n @p opts is a BSON document of the form:\n {\n    \"min\": Optional<BSON value>,\n    \"max\": Optional<BSON value>,\n    \"sparsity\": Optional<Int64>,\n    \"precision\": Optional<Int32>,\n    \"trimFactor\": Optional<Int32>\n }\n\n @param[in] ctx The @ref mongocrypt_ctx_t object.\n @param[in] opts BSON.\n @pre @p ctx has not been initialized.\n @returns A boolean indicating success. If false, an error status is set.\n Retrieve it with @ref mongocrypt_ctx_status"]
     pub fn mongocrypt_ctx_setopt_algorithm_range(
         ctx: *mut mongocrypt_ctx_t,
         opts: *mut mongocrypt_binary_t,
     ) -> bool;
 }
+extern "C" {
+    #[doc = " Set the expiration time for the data encryption key cache. Defaults to 60 seconds if not set.\n\n @param[in] ctx The @ref mongocrypt_ctx_t object.\n @param[in] cache_expiration_ms The cache expiration time in milliseconds. If zero, the cache\n never expires."]
+    pub fn mongocrypt_setopt_key_expiration(
+        crypt: *mut mongocrypt_t,
+        cache_expiration_ms: u64,
+    ) -> bool;
+}

mongocrypt/src/ctx.rs

@@ -702,6 +702,16 @@ impl<'scope> KmsCtx<'scope> {
         }
     }
 
+    /// How many microseconds to sleep before sending a request.
+    pub fn sleep_micros(&self) -> i64 {
+        unsafe { sys::mongocrypt_kms_ctx_usleep(self.inner) }
+    }
+
+    /// Whether a failed request can be retried.
+    pub fn retry_failure(&self) -> bool {
+        unsafe { sys::mongocrypt_kms_ctx_fail(self.inner) }
+    }
+
     /// Indicates how many bytes to feed into `feed`.
     pub fn bytes_needed(&self) -> u32 {
         unsafe { sys::mongocrypt_kms_ctx_bytes_needed(self.inner) }
@@ -784,7 +794,7 @@ impl KmsProvider {
     pub fn local() -> Self {
         Self {
             provider_type: KmsProviderType::Local,
-            name: None
+            name: None,
         }
     }
 
@@ -830,7 +840,8 @@ impl KmsProvider {
             KmsProviderType::Local => "local",
             KmsProviderType::Kmip => "kmip",
             KmsProviderType::Other(ref other) => other,
-        }.to_string();
+        }
+        .to_string();
         if let Some(ref name) = self.name {
             full_name.push(':');
             full_name.push_str(name);
@@ -842,9 +853,7 @@ impl KmsProvider {
     /// type followed by an optional ":" and name, e.g. "aws" or "aws:name".
     pub fn from_string(name: &str) -> Self {
         let (provider_type, name) = match name.split_once(':') {
-            Some((provider_type, name)) => {
-                (provider_type, Some(name.to_string()))
-            }
+            Some((provider_type, name)) => (provider_type, Some(name.to_string())),
             None => (name, None),
         };
         let provider_type = match provider_type {
@@ -857,7 +866,7 @@ impl KmsProvider {
         };
         Self {
             provider_type,
-            name
+            name,
         }
     }
 }

mongocrypt/src/lib.rs

@@ -248,6 +248,17 @@ impl CryptBuilder {
         Ok(self)
     }
 
+    /// Enable or disable KMS retry behavior.
+    pub fn retry_kms(self, enable: bool) -> Result<Self> {
+        unsafe {
+            let ok = sys::mongocrypt_setopt_retry_kms(*self.inner.borrow(), enable);
+            if !ok {
+                return Err(self.status().as_error())
+            }
+        }
+        Ok(self)
+    }
+
     pub fn build(mut self) -> Result<Crypt> {
         let _guard = CRYPT_LOCK.lock().unwrap();