From 8c557c23efd8aa40d63224307fd58cdc9a1ad5c1 Mon Sep 17 00:00:00 2001
From: Daniel McCarney <daniel@binaryparadox.net>
Date: Wed, 15 Nov 2023 17:37:57 -0500
Subject: [PATCH] lib: add try_take! macro, handling AlreadyUsed Options

In several places we represent something that could be consumed as an
`Option<T>`. When we try to use it, we `take()` the option, match the
result, and return `rustls_result::AlreadyUsed` if `take()` returned
`None`.

Since this pattern is becoming more common with the use of more builder
patterns this commit adds a `try_take!` macro that can do this
repetitive work for us.
---
 src/acceptor.rs |  8 +++-----
 src/cipher.rs   | 20 ++++----------------
 src/lib.rs      | 13 +++++++++++++
 3 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/src/acceptor.rs b/src/acceptor.rs
index 09097f53..a559d841 100644
--- a/src/acceptor.rs
+++ b/src/acceptor.rs
@@ -12,7 +12,8 @@ use crate::rslice::{rustls_slice_bytes, rustls_str};
 use crate::server::rustls_server_config;
 use crate::{
     ffi_panic_boundary, free_box, rustls_result, set_boxed_mut_ptr, to_box, to_boxed_mut_ptr,
-    try_arc_from_ptr, try_callback, try_mut_from_ptr, try_ref_from_ptr, Castable, OwnershipBox,
+    try_arc_from_ptr, try_callback, try_mut_from_ptr, try_ref_from_ptr, try_take, Castable,
+    OwnershipBox,
 };
 use rustls_result::NullParameter;
 
@@ -398,10 +399,7 @@ impl rustls_accepted {
     ) -> rustls_result {
         ffi_panic_boundary! {
             let accepted: &mut Option<Accepted> = try_mut_from_ptr!(accepted);
-            let accepted = match accepted.take() {
-                Some(a) => a,
-                None => return rustls_result::AlreadyUsed,
-            };
+            let accepted = try_take!(accepted);
             let config: Arc<ServerConfig> = try_arc_from_ptr!(config);
             match accepted.into_connection(config) {
                 Ok(built) => {
diff --git a/src/cipher.rs b/src/cipher.rs
index 721b55a4..463f0726 100644
--- a/src/cipher.rs
+++ b/src/cipher.rs
@@ -19,10 +19,10 @@ use crate::error::{map_error, rustls_result};
 use crate::rslice::{rustls_slice_bytes, rustls_str};
 use crate::{
     ffi_panic_boundary, free_arc, to_arc_const_ptr, to_boxed_mut_ptr, try_box_from_ptr,
-    try_mut_from_ptr, try_ref_from_ptr, try_slice, Castable, OwnershipArc, OwnershipBox,
+    try_mut_from_ptr, try_ref_from_ptr, try_slice, try_take, Castable, OwnershipArc, OwnershipBox,
     OwnershipRef,
 };
-use rustls_result::{AlreadyUsed, NullParameter};
+use rustls_result::NullParameter;
 
 /// An X.509 certificate, as used in rustls.
 /// Corresponds to `Certificate` in the Rust API.
@@ -566,13 +566,7 @@ impl rustls_allow_any_authenticated_client_builder {
                 Err(_) => return rustls_result::CertificateRevocationListParseError,
             };
 
-            let client_cert_verifier = match client_cert_verifier_builder.take() {
-                None => {
-                    return AlreadyUsed;
-                },
-                Some(x) => x,
-            };
-
+            let client_cert_verifier = try_take!(client_cert_verifier_builder);
             match client_cert_verifier.with_crls(crls_der) {
                 Ok(v) => client_cert_verifier_builder.replace(v),
                 Err(e) => return map_error(rustls::Error::InvalidCertRevocationList(e)),
@@ -710,13 +704,7 @@ impl rustls_allow_any_anonymous_or_authenticated_client_builder {
                 Err(_) => return rustls_result::CertificateRevocationListParseError,
             };
 
-            let client_cert_verifier = match client_cert_verifier_builder.take() {
-                None => {
-                    return AlreadyUsed;
-                },
-                Some(x) => x,
-            };
-
+            let client_cert_verifier = try_take!(client_cert_verifier_builder);
             match client_cert_verifier.with_crls(crls_der) {
                 Ok(v) => client_cert_verifier_builder.replace(v),
                 Err(e) => return map_error(rustls::Error::InvalidCertRevocationList(e)),
diff --git a/src/lib.rs b/src/lib.rs
index 6486113a..04a7e1df 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -643,6 +643,19 @@ macro_rules! try_callback {
 
 pub(crate) use try_callback;
 
+macro_rules! try_take {
+    ( $var:ident ) => {
+        match $var.take() {
+            None => {
+                return $crate::rustls_result::AlreadyUsed;
+            }
+            Some(x) => x,
+        }
+    };
+}
+
+pub(crate) use try_take;
+
 /// Returns a static string containing the rustls-ffi version as well as the
 /// rustls version. The string is alive for the lifetime of the program and does
 /// not need to be freed.