Update contrib/restricted/aws/aws-c-io to 0.13.0

9 files changed, 856 insertions, 313 deletions

diff --git a/contrib/restricted/aws/aws-c-io/CMakeLists.darwin.txt b/contrib/restricted/aws/aws-c-io/CMakeLists.darwin.txt
index 981541f14de..702233062b3 100644
--- a/contrib/restricted/aws/aws-c-io/CMakeLists.darwin.txt
+++ b/contrib/restricted/aws/aws-c-io/CMakeLists.darwin.txt
@@ -29,6 +29,7 @@ target_compile_options(restricted-aws-aws-c-io PRIVATE
   -DS2N_MADVISE_SUPPORTED
   -DS2N_SIKE_P434_R3_ASM
   -DS2N___RESTRICT__SUPPORTED
+  -DUSE_S2N
   -Wno-everything
 )
 target_include_directories(restricted-aws-aws-c-io PUBLIC
@@ -49,7 +50,8 @@ target_sources(restricted-aws-aws-c-io PRIVATE
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/io.c
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/message_pool.c
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/pem_utils.c
-  ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/pkcs11.c
+  ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/pkcs11_lib.c
+  ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/pkcs11_tls_op_handler.c
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/pki_utils.c
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/posix/host_resolver.c
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/posix/pipe.c
diff --git a/contrib/restricted/aws/aws-c-io/CMakeLists.linux.txt b/contrib/restricted/aws/aws-c-io/CMakeLists.linux.txt
index 65936476862..ffac39a8507 100644
--- a/contrib/restricted/aws/aws-c-io/CMakeLists.linux.txt
+++ b/contrib/restricted/aws/aws-c-io/CMakeLists.linux.txt
@@ -29,6 +29,7 @@ target_compile_options(restricted-aws-aws-c-io PRIVATE
   -DS2N_MADVISE_SUPPORTED
   -DS2N_SIKE_P434_R3_ASM
   -DS2N___RESTRICT__SUPPORTED
+  -DUSE_S2N
   -Wno-everything
 )
 target_include_directories(restricted-aws-aws-c-io PUBLIC
@@ -49,7 +50,8 @@ target_sources(restricted-aws-aws-c-io PRIVATE
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/io.c
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/message_pool.c
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/pem_utils.c
-  ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/pkcs11.c
+  ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/pkcs11_lib.c
+  ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/pkcs11_tls_op_handler.c
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/pki_utils.c
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/posix/host_resolver.c
   ${CMAKE_SOURCE_DIR}/contrib/restricted/aws/aws-c-io/source/posix/pipe.c
diff --git a/contrib/restricted/aws/aws-c-io/include/aws/io/tls_channel_handler.h b/contrib/restricted/aws/aws-c-io/include/aws/io/tls_channel_handler.h
index 1c19c59b580..5e27c9b2849 100644
--- a/contrib/restricted/aws/aws-c-io/include/aws/io/tls_channel_handler.h
+++ b/contrib/restricted/aws/aws-c-io/include/aws/io/tls_channel_handler.h
@@ -42,6 +42,39 @@ enum aws_tls_cipher_pref {
     AWS_IO_TLS_CIPHER_PREF_END_RANGE = 0xFFFF
 };
 
+/**
+ * The hash algorithm of a TLS private key operation. Any custom private key operation handlers are expected to perform
+ * operations on the input TLS data using the correct hash algorithm or fail the operation.
+ */
+enum aws_tls_hash_algorithm {
+    AWS_TLS_HASH_UNKNOWN,
+    AWS_TLS_HASH_SHA1,
+    AWS_TLS_HASH_SHA224,
+    AWS_TLS_HASH_SHA256,
+    AWS_TLS_HASH_SHA384,
+    AWS_TLS_HASH_SHA512,
+};
+
+/**
+ * The signature of a TLS private key operation. Any custom private key operation handlers are expected to perform
+ * operations on the input TLS data using the correct signature algorithm or fail the operation.
+ */
+enum aws_tls_signature_algorithm {
+    AWS_TLS_SIGNATURE_UNKNOWN,
+    AWS_TLS_SIGNATURE_RSA,
+    AWS_TLS_SIGNATURE_ECDSA,
+};
+
+/**
+ * The TLS private key operation that needs to be performed by a custom private key operation handler when making
+ * a connection using mutual TLS.
+ */
+enum aws_tls_key_operation_type {
+    AWS_TLS_KEY_OPERATION_UNKNOWN,
+    AWS_TLS_KEY_OPERATION_SIGN,
+    AWS_TLS_KEY_OPERATION_DECRYPT,
+};
+
 struct aws_tls_ctx {
     struct aws_allocator *alloc;
     void *impl;
@@ -102,6 +135,13 @@ struct aws_tls_connection_options {
     uint32_t timeout_ms;
 };
 
+/**
+ * A struct containing all of the data needed for a private key operation when
+ * making a mutual TLS connection. This struct contains the data that needs
+ * to be operated on, like performing a sign operation or a decrypt operation.
+ */
+struct aws_tls_key_operation;
+
 struct aws_tls_ctx_options {
     struct aws_allocator *allocator;
 
@@ -201,17 +241,13 @@ struct aws_tls_ctx_options {
     void *ctx_options_extension;
 
     /**
-     * Set if using PKCS#11 for private key operations.
-     * See aws_tls_ctx_pkcs11_options for more details.
+     * Set if using custom private key operations.
+     * See aws_custom_key_op_handler for more details
+     *
+     * Note: Custom key operations (and PKCS#11 integration) hasn't been tested with TLS 1.3, so don't use
+     * cipher preferences that allow TLS 1.3. If this is set, we will always use non TLS 1.3 preferences.
      */
-    struct {
-        struct aws_pkcs11_lib *lib;                  /* required */
-        struct aws_string *user_pin;                 /* NULL if token uses "protected authentication path" */
-        struct aws_string *token_label;              /* optional */
-        struct aws_string *private_key_object_label; /* optional */
-        uint64_t slot_id;                            /* optional */
-        bool has_slot_id;
-    } pkcs11;
+    struct aws_custom_key_op_handler *custom_key_op_handler;
 };
 
 struct aws_tls_negotiated_protocol_message {
@@ -304,6 +340,81 @@ AWS_IO_API int aws_tls_ctx_options_init_client_mtls(
     const struct aws_byte_cursor *pkey);
 
 /**
+ * vtable for aws_custom_key_op_handler.
+ */
+struct aws_custom_key_op_handler_vtable {
+    /**
+     * Called when the a TLS handshake has an operation it needs the custom key operation handler to perform.
+     * NOTE: You must call aws_tls_key_operation_complete() or aws_tls_key_operation_complete_with_error()
+     * otherwise the TLS handshake will stall the TLS connection indefinitely and leak memory.
+     */
+    void (*on_key_operation)(struct aws_custom_key_op_handler *key_op_handler, struct aws_tls_key_operation *operation);
+};
+
+/**
+ * The custom key operation that is used when performing a mutual TLS handshake. This can
+ * be extended to provide custom private key operations, like PKCS11 or similar.
+ */
+struct aws_custom_key_op_handler {
+    /**
+     * A void* intended to be populated with a reference to whatever class is extending this class. For example,
+     * if you have extended aws_custom_key_op_handler with a custom struct, you would put a pointer to this struct
+     * to *impl so you can retrieve it back in the vtable functions.
+     */
+    void *impl;
+
+    /**
+     * A vtable containing all of the functions the aws_custom_key_op_handler implements. Is intended to be extended.
+     * NOTE: Use "aws_custom_key_op_handler_<func>" to access vtable functions.
+     */
+    const struct aws_custom_key_op_handler_vtable *vtable;
+
+    /**
+     * A reference count for handling memory usage.
+     * Use aws_custom_key_op_handler_acquire and aws_custom_key_op_handler_release to increase/decrease count.
+     */
+    struct aws_ref_count ref_count;
+};
+
+/**
+ * Increases the reference count for the passed-in aws_custom_key_op_handler and returns it.
+ */
+AWS_IO_API struct aws_custom_key_op_handler *aws_custom_key_op_handler_acquire(
+    struct aws_custom_key_op_handler *key_op_handler);
+
+/**
+ * Decreases the reference count for the passed-in aws_custom_key_op_handler and returns NULL.
+ */
+AWS_IO_API struct aws_custom_key_op_handler *aws_custom_key_op_handler_release(
+    struct aws_custom_key_op_handler *key_op_handler);
+
+/**
+ * Calls the on_key_operation vtable function. See aws_custom_key_op_handler_vtable for function details.
+ */
+AWS_IO_API void aws_custom_key_op_handler_perform_operation(
+    struct aws_custom_key_op_handler *key_op_handler,
+    struct aws_tls_key_operation *operation);
+
+/**
+ * Initializes options for use with mutual TLS in client mode,
+ * where private key operations are handled by custom code.
+ *
+ * Note: cert_file_contents will be copied into a new buffer after this
+ * function is called, so you do not need to keep that data alive
+ * after calling this function.
+ *
+ * @param options               aws_tls_ctx_options to be initialized.
+ * @param allocator             Allocator to use.
+ * @param custom                Options for custom key operations.
+ * @param cert_file_contents    The contents of a certificate file.
+ */
+AWS_IO_API int aws_tls_ctx_options_init_client_mtls_with_custom_key_operations(
+    struct aws_tls_ctx_options *options,
+    struct aws_allocator *allocator,
+    struct aws_custom_key_op_handler *custom,
+    const struct aws_byte_cursor *cert_file_contents);
+
+/**
  * This struct exists as a graceful way to pass many arguments when
  * calling init-with-pkcs11 functions on aws_tls_ctx_options (this also makes
  * it easy to introduce optional arguments in the future).
@@ -706,6 +817,65 @@ AWS_IO_API struct aws_byte_buf aws_tls_handler_protocol(struct aws_channel_handl
  */
 AWS_IO_API struct aws_byte_buf aws_tls_handler_server_name(struct aws_channel_handler *handler);
 
+/**************************** TLS KEY OPERATION *******************************/
+
+/* Note: Currently this assumes the user knows what key is being used for key/cert pairs
+         but s2n supports multiple cert/key pairs. This functionality is not used in the
+         CRT currently, but in the future, we may need to implement this */
+
+/**
+ * Complete a successful TLS private key operation by providing its output.
+ * The output is copied into the TLS connection.
+ * The operation is freed by this call.
+ *
+ * You MUST call this or aws_tls_key_operation_complete_with_error().
+ * Failure to do so will stall the TLS connection indefinitely and leak memory.
+ */
+AWS_IO_API
+void aws_tls_key_operation_complete(struct aws_tls_key_operation *operation, struct aws_byte_cursor output);
+
+/**
+ * Complete an failed TLS private key operation.
+ * The TLS connection will fail.
+ * The operation is freed by this call.
+ *
+ * You MUST call this or aws_tls_key_operation_complete().
+ * Failure to do so will stall the TLS connection indefinitely and leak memory.
+ */
+AWS_IO_API
+void aws_tls_key_operation_complete_with_error(struct aws_tls_key_operation *operation, int error_code);
+
+/**
+ * Returns the input data that needs to be operated on by the custom key operation.
+ */
+AWS_IO_API
+struct aws_byte_cursor aws_tls_key_operation_get_input(const struct aws_tls_key_operation *operation);
+
+/**
+ * Returns the type of operation that needs to be performed by the custom key operation.
+ * If the implementation cannot perform the operation,
+ * use aws_tls_key_operation_complete_with_error() to preventing stalling the TLS connection.
+ */
+AWS_IO_API
+enum aws_tls_key_operation_type aws_tls_key_operation_get_type(const struct aws_tls_key_operation *operation);
+
+/**
+ * Returns the algorithm the operation is expected to be operated with.
+ * If the implementation does not support the signature algorithm,
+ * use aws_tls_key_operation_complete_with_error() to preventing stalling the TLS connection.
+ */
+AWS_IO_API
+enum aws_tls_signature_algorithm aws_tls_key_operation_get_signature_algorithm(
+    const struct aws_tls_key_operation *operation);
+
+/**
+ * Returns the algorithm the operation digest is signed with.
+ * If the implementation does not support the digest algorithm,
+ * use aws_tls_key_operation_complete_with_error() to preventing stalling the TLS connection.
+ */
+AWS_IO_API
+enum aws_tls_hash_algorithm aws_tls_key_operation_get_digest_algorithm(const struct aws_tls_key_operation *operation);
+
 /********************************* Misc TLS related *********************************/
 
 /*
@@ -718,6 +888,24 @@ AWS_IO_API int aws_channel_setup_client_tls(
     struct aws_channel_slot *right_of_slot,
     struct aws_tls_connection_options *tls_options);
 
+/**
+ * Given enum, return string like: AWS_TLS_HASH_SHA256 -> "SHA256"
+ */
+AWS_IO_API
+const char *aws_tls_hash_algorithm_str(enum aws_tls_hash_algorithm hash);
+
+/**
+ * Given enum, return string like: AWS_TLS_SIGNATURE_RSA -> "RSA"
+ */
+AWS_IO_API
+const char *aws_tls_signature_algorithm_str(enum aws_tls_signature_algorithm signature);
+
+/**
+ * Given enum, return string like: AWS_TLS_SIGNATURE_RSA -> "RSA"
+ */
+AWS_IO_API
+const char *aws_tls_key_operation_type_str(enum aws_tls_key_operation_type operation_type);
+
 AWS_EXTERN_C_END
 
 #endif /* AWS_IO_TLS_CHANNEL_HANDLER_H */
diff --git a/contrib/restricted/aws/aws-c-io/source/pkcs11.c b/contrib/restricted/aws/aws-c-io/source/pkcs11_lib.c
index 943f153fd52..8047d118c79 100644
--- a/contrib/restricted/aws/aws-c-io/source/pkcs11.c
+++ b/contrib/restricted/aws/aws-c-io/source/pkcs11_lib.c
@@ -25,7 +25,7 @@
  * See https://tools.ietf.org/html/rfc3447#page-43
  * (Notes to help understand what's going on here with DER encoding)
  * 0x30 nn - Sequence of tags, nn bytes, including hash, nn = mm+jj+4 (PKCS11 DigestInfo)
- *   0x30 mm - Subsequence of tags, mm bytes (ii+4) (PKCS11 
+ *   0x30 mm - Subsequence of tags, mm bytes (ii+4) (PKCS11
  *     0x06 ii - OID encoding, ii bytes, see X.680 - this identifies the hash algorithm
  *     0x05 00 - NULL
  *   0x04 jj - OCTET, nn = mm + jj + 4
@@ -38,29 +38,6 @@ static const uint8_t SHA512_PREFIX_TO_RSA_SIG[] = { 0x30, 0x51, 0x30, 0x0d, 0x06
 static const uint8_t SHA224_PREFIX_TO_RSA_SIG[] = { 0x30, 0x2d, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x04, 0x05, 0x00, 0x04, 0x1c };
 /* clang-format on */
 
-const char *aws_tls_hash_algorithm_str(enum aws_tls_hash_algorithm hash) {
-    /* clang-format off */
-    switch (hash) {
-        case (AWS_TLS_HASH_SHA1): return "SHA1";
-        case (AWS_TLS_HASH_SHA224): return "SHA224";
-        case (AWS_TLS_HASH_SHA256): return "SHA256";
-        case (AWS_TLS_HASH_SHA384): return "SHA384";
-        case (AWS_TLS_HASH_SHA512): return "SHA512";
-        default: return "<UNKNOWN HASH ALGORITHM>";
-    }
-    /* clang-format on */
-}
-
-const char *aws_tls_signature_algorithm_str(enum aws_tls_signature_algorithm signature) {
-    /* clang-format off */
-    switch (signature) {
-        case (AWS_TLS_SIGNATURE_RSA): return "RSA";
-        case (AWS_TLS_SIGNATURE_ECDSA): return "ECDSA";
-        default: return "<UNKNOWN SIGNATURE ALGORITHM>";
-    }
-    /* clang-format on */
-}
-
 /* Return c-string for PKCS#11 CKR_* contants. */
 const char *aws_pkcs11_ckr_str(CK_RV rv) {
     /* clang-format off */
diff --git a/contrib/restricted/aws/aws-c-io/source/pkcs11_private.h b/contrib/restricted/aws/aws-c-io/source/pkcs11_private.h
index 8eeca4604fa..314c739199d 100644
--- a/contrib/restricted/aws/aws-c-io/source/pkcs11_private.h
+++ b/contrib/restricted/aws/aws-c-io/source/pkcs11_private.h
@@ -5,7 +5,7 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0.
  */
-#include <aws/io/io.h>
+#include <aws/io/tls_channel_handler.h>
 
 /* These defines must exist before the official PKCS#11 headers are included */
 #define CK_PTR *
@@ -28,23 +28,9 @@
  */
 
 struct aws_pkcs11_lib;
+struct aws_pkcs11_tls_key_handler;
 struct aws_string;
 
-enum aws_tls_hash_algorithm {
-    AWS_TLS_HASH_UNKNOWN = -1,
-    AWS_TLS_HASH_SHA1,
-    AWS_TLS_HASH_SHA224,
-    AWS_TLS_HASH_SHA256,
-    AWS_TLS_HASH_SHA384,
-    AWS_TLS_HASH_SHA512,
-};
-
-enum aws_tls_signature_algorithm {
-    AWS_TLS_SIGNATURE_UNKNOWN = -1,
-    AWS_TLS_SIGNATURE_RSA,
-    AWS_TLS_SIGNATURE_ECDSA,
-};
-
 AWS_EXTERN_C_BEGIN
 
 /**
@@ -152,16 +138,20 @@ AWS_IO_API
 int aws_pkcs11_asn1_enc_ubigint(struct aws_byte_buf *const buffer, struct aws_byte_cursor bigint);
 
 /**
- * Given enum, return string like: AWS_TLS_HASH_SHA256 -> "SHA256"
- */
-AWS_IO_API
-const char *aws_tls_hash_algorithm_str(enum aws_tls_hash_algorithm hash);
-
-/**
- * Given enum, return string like: AWS_TLS_SIGNATURE_RSA -> "RSA"
+ * Creates a new PKCS11 TLS operation handler with an associated aws_custom_key_op_handler
+ * with a reference count set to 1.
+ *
+ * The PKCS11 TLS operation handler will automatically be destroyed when the reference count reaches zero
+ * on the aws_custom_key_op_handler.
  */
 AWS_IO_API
-const char *aws_tls_signature_algorithm_str(enum aws_tls_signature_algorithm signature);
+struct aws_custom_key_op_handler *aws_pkcs11_tls_op_handler_new(
+    struct aws_allocator *allocator,
+    struct aws_pkcs11_lib *pkcs11_lib,
+    const struct aws_byte_cursor *user_pin,
+    const struct aws_byte_cursor *match_token_label,
+    const struct aws_byte_cursor *match_private_key_label,
+    const uint64_t *match_slot_id);
 
 AWS_EXTERN_C_END
 #endif /* AWS_IO_PKCS11_PRIVATE_H */
diff --git a/contrib/restricted/aws/aws-c-io/source/pkcs11_tls_op_handler.c b/contrib/restricted/aws/aws-c-io/source/pkcs11_tls_op_handler.c
new file mode 100644
index 00000000000..6d155cbd9e5
--- /dev/null
+++ b/contrib/restricted/aws/aws-c-io/source/pkcs11_tls_op_handler.c
@@ -0,0 +1,221 @@
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+#include <aws/io/pkcs11.h>
+
+#include "pkcs11_private.h"
+
+#include <aws/common/mutex.h>
+#include <aws/common/string.h>
+#include <aws/io/logging.h>
+
+struct aws_pkcs11_tls_op_handler {
+    /* The custom key operation handler needed for the callbacks */
+    struct aws_custom_key_op_handler base;
+
+    struct aws_allocator *alloc;
+    struct aws_pkcs11_lib *lib;
+
+    /* Use a single PKCS#11 session for all TLS connections on an aws_tls_ctx.
+     * We do this because PKCS#11 tokens may only support a
+     * limited number of sessions (PKCS11-UG-v2.40 section 2.6.7).
+     * If this one shared session turns out to be a severe bottleneck,
+     * we could look into other setups (ex: put session on its own thread,
+     * 1 session per event-loop, 1 session per connection, etc).
+     *
+     * The lock must be held while performing session operations.
+     * Otherwise, it would not be safe for multiple threads to share a
+     * session (PKCS11-UG-v2.40 section 2.6.7). The lock isn't needed for
+     * setup and teardown though, since we ensure nothing parallel is going
+     * on at these times */
+    struct aws_mutex session_lock;
+    CK_SESSION_HANDLE session_handle;
+    CK_OBJECT_HANDLE private_key_handle;
+    CK_KEY_TYPE private_key_type;
+};
+
+static void s_aws_custom_key_op_handler_destroy(struct aws_custom_key_op_handler *key_op_handler) {
+
+    struct aws_pkcs11_tls_op_handler *handler = (struct aws_pkcs11_tls_op_handler *)key_op_handler->impl;
+
+    if (handler->session_handle != 0) {
+        aws_pkcs11_lib_close_session(handler->lib, handler->session_handle);
+    }
+    aws_mutex_clean_up(&handler->session_lock);
+    aws_pkcs11_lib_release(handler->lib);
+
+    aws_mem_release(handler->alloc, handler);
+}
+
+/**
+ * Performs the PKCS11 TLS private key operation. This is called automatically when performing a mutual TLS handshake.
+ */
+void s_aws_pkcs11_tls_op_handler_do_operation(
+    struct aws_custom_key_op_handler *handler,
+    struct aws_tls_key_operation *operation) {
+
+    struct aws_pkcs11_tls_op_handler *pkcs11_handler = (struct aws_pkcs11_tls_op_handler *)handler->impl;
+    struct aws_byte_buf output_buf; /* initialized later */
+    AWS_ZERO_STRUCT(output_buf);
+
+    /*********** BEGIN CRITICAL SECTION ***********/
+    aws_mutex_lock(&pkcs11_handler->session_lock);
+    bool success_while_locked = false;
+
+    switch (aws_tls_key_operation_get_type(operation)) {
+        case AWS_TLS_KEY_OPERATION_DECRYPT:
+            if (aws_pkcs11_lib_decrypt(
+                    pkcs11_handler->lib,
+                    pkcs11_handler->session_handle,
+                    pkcs11_handler->private_key_handle,
+                    pkcs11_handler->private_key_type,
+                    aws_tls_key_operation_get_input(operation),
+                    pkcs11_handler->alloc,
+                    &output_buf)) {
+
+                goto unlock;
+            }
+            break;
+
+        case AWS_TLS_KEY_OPERATION_SIGN:
+            if (aws_pkcs11_lib_sign(
+                    pkcs11_handler->lib,
+                    pkcs11_handler->session_handle,
+                    pkcs11_handler->private_key_handle,
+                    pkcs11_handler->private_key_type,
+                    aws_tls_key_operation_get_input(operation),
+                    pkcs11_handler->alloc,
+                    aws_tls_key_operation_get_digest_algorithm(operation),
+                    aws_tls_key_operation_get_signature_algorithm(operation),
+                    &output_buf)) {
+
+                goto unlock;
+            }
+            break;
+
+        default:
+            AWS_LOGF_ERROR(
+                AWS_LS_IO_PKCS11,
+                "PKCS11 Handler %p: Unknown TLS key operation with value of %u",
+                (void *)handler,
+                aws_tls_key_operation_get_type(operation));
+            aws_raise_error(AWS_ERROR_INVALID_STATE);
+            goto unlock;
+    }
+
+    success_while_locked = true;
+unlock:
+    aws_mutex_unlock(&pkcs11_handler->session_lock);
+    /*********** END CRITICAL SECTION ***********/
+
+    if (success_while_locked) {
+        aws_tls_key_operation_complete(operation, aws_byte_cursor_from_buf(&output_buf));
+    } else {
+        aws_tls_key_operation_complete_with_error(operation, aws_last_error());
+    }
+
+    aws_byte_buf_clean_up(&output_buf);
+}
+
+static struct aws_custom_key_op_handler_vtable s_aws_custom_key_op_handler_vtable = {
+    .on_key_operation = s_aws_pkcs11_tls_op_handler_do_operation,
+};
+
+struct aws_custom_key_op_handler *aws_pkcs11_tls_op_handler_new(
+    struct aws_allocator *allocator,
+    struct aws_pkcs11_lib *pkcs11_lib,
+    const struct aws_byte_cursor *user_pin,
+    const struct aws_byte_cursor *match_token_label,
+    const struct aws_byte_cursor *match_private_key_label,
+    const uint64_t *match_slot_id) {
+
+    bool success = false;
+
+    struct aws_pkcs11_tls_op_handler *pkcs11_handler =
+        aws_mem_calloc(allocator, 1, sizeof(struct aws_pkcs11_tls_op_handler));
+
+    // Optional data
+    struct aws_string *pkcs_user_pin = NULL;
+    struct aws_string *pkcs_token_label = NULL;
+    struct aws_string *pkcs_private_key_object_label = NULL;
+
+    aws_ref_count_init(
+        &pkcs11_handler->base.ref_count,
+        &pkcs11_handler->base,
+        (aws_simple_completion_callback *)s_aws_custom_key_op_handler_destroy);
+
+    pkcs11_handler->base.impl = (void *)pkcs11_handler;
+    pkcs11_handler->base.vtable = &s_aws_custom_key_op_handler_vtable;
+
+    pkcs11_handler->alloc = allocator;
+
+    /* pkcs11_lib is required */
+    if (pkcs11_lib == NULL) {
+        aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
+        AWS_LOGF_ERROR(AWS_LS_IO_PKCS11, "PKCS11 Handler %p new: PKCS11 library is null", (void *)pkcs11_handler);
+        goto done;
+    }
+    pkcs11_handler->lib = aws_pkcs11_lib_acquire(pkcs11_lib); /* cannot fail */
+    aws_mutex_init(&pkcs11_handler->session_lock);
+
+    /* user_pin is optional */
+    if (user_pin->ptr != NULL) {
+        pkcs_user_pin = aws_string_new_from_cursor(allocator, user_pin);
+    }
+
+    /* token_label is optional */
+    if (match_token_label->ptr != NULL) {
+        pkcs_token_label = aws_string_new_from_cursor(allocator, match_token_label);
+    }
+
+    /* private_key_object_label is optional */
+    if (match_private_key_label->ptr != NULL) {
+        pkcs_private_key_object_label = aws_string_new_from_cursor(allocator, match_private_key_label);
+    }
+
+    CK_SLOT_ID slot_id;
+    if (aws_pkcs11_lib_find_slot_with_token(pkcs11_handler->lib, match_slot_id, pkcs_token_label, &slot_id /*out*/)) {
+        goto done;
+    }
+
+    if (aws_pkcs11_lib_open_session(pkcs11_handler->lib, slot_id, &pkcs11_handler->session_handle)) {
+        goto done;
+    }
+
+    if (pkcs_user_pin != NULL) {
+        if (aws_pkcs11_lib_login_user(pkcs11_handler->lib, pkcs11_handler->session_handle, pkcs_user_pin)) {
+            goto done;
+        }
+    }
+
+    if (aws_pkcs11_lib_find_private_key(
+            pkcs11_handler->lib,
+            pkcs11_handler->session_handle,
+            pkcs_private_key_object_label,
+            &pkcs11_handler->private_key_handle /*out*/,
+            &pkcs11_handler->private_key_type /*out*/)) {
+        goto done;
+    }
+    success = true;
+
+done:
+
+    /* CLEANUP */
+    if (pkcs_user_pin != NULL) {
+        aws_string_destroy_secure(pkcs_user_pin);
+    }
+    if (pkcs_token_label != NULL) {
+        aws_string_destroy(pkcs_token_label);
+    }
+    if (pkcs_private_key_object_label != NULL) {
+        aws_string_destroy(pkcs_private_key_object_label);
+    }
+
+    if (success) {
+        return &pkcs11_handler->base;
+    } else {
+        aws_custom_key_op_handler_release(&pkcs11_handler->base);
+        return NULL;
+    }
+}
diff --git a/contrib/restricted/aws/aws-c-io/source/retry_strategy.c b/contrib/restricted/aws/aws-c-io/source/retry_strategy.c
index ff1288d142b..5c92799d283 100644
--- a/contrib/restricted/aws/aws-c-io/source/retry_strategy.c
+++ b/contrib/restricted/aws/aws-c-io/source/retry_strategy.c
@@ -5,14 +5,16 @@
 #include <aws/io/retry_strategy.h>
 
 void aws_retry_strategy_acquire(struct aws_retry_strategy *retry_strategy) {
-    aws_atomic_fetch_add_explicit(&retry_strategy->ref_count, 1, aws_memory_order_relaxed);
+    size_t old_value = aws_atomic_fetch_add_explicit(&retry_strategy->ref_count, 1, aws_memory_order_relaxed);
+    AWS_ASSERT(old_value > 0 && "aws_retry_strategy refcount had been zero, it's invalid to use it again.");
+    (void)old_value;
 }
 
 void aws_retry_strategy_release(struct aws_retry_strategy *retry_strategy) {
     if (retry_strategy) {
-        size_t ref_count = aws_atomic_fetch_sub_explicit(&retry_strategy->ref_count, 1, aws_memory_order_seq_cst);
-
-        if (ref_count == 1) {
+        size_t old_value = aws_atomic_fetch_sub_explicit(&retry_strategy->ref_count, 1, aws_memory_order_seq_cst);
+        AWS_ASSERT(old_value > 0 && "aws_retry_strategy refcount has gone negative");
+        if (old_value == 1) {
             retry_strategy->vtable->destroy(retry_strategy);
         }
     }
@@ -50,7 +52,9 @@ int aws_retry_token_record_success(struct aws_retry_token *token) {
 }
 
 void aws_retry_token_acquire(struct aws_retry_token *token) {
-    aws_atomic_fetch_add_explicit(&token->ref_count, 1u, aws_memory_order_relaxed);
+    size_t old_value = aws_atomic_fetch_add_explicit(&token->ref_count, 1u, aws_memory_order_relaxed);
+    AWS_ASSERT(old_value > 0 && "aws_retry_token refcount had been zero, it's invalid to use it again.");
+    (void)old_value;
 }
 
 void aws_retry_token_release(struct aws_retry_token *token) {
@@ -58,9 +62,9 @@ void aws_retry_token_release(struct aws_retry_token *token) {
         AWS_PRECONDITION(token->retry_strategy);
         AWS_PRECONDITION(token->retry_strategy->vtable->release_token);
 
-        size_t prev_count = aws_atomic_fetch_sub_explicit(&token->ref_count, 1u, aws_memory_order_seq_cst);
-
-        if (prev_count == 1u) {
+        size_t old_value = aws_atomic_fetch_sub_explicit(&token->ref_count, 1u, aws_memory_order_seq_cst);
+        AWS_ASSERT(old_value > 0 && "aws_retry_token refcount has gone negative");
+        if (old_value == 1u) {
             token->retry_strategy->vtable->release_token(token);
         }
     }
diff --git a/contrib/restricted/aws/aws-c-io/source/s2n/s2n_tls_channel_handler.c b/contrib/restricted/aws/aws-c-io/source/s2n/s2n_tls_channel_handler.c
index 5c3561154da..2bc9e5aeed4 100644
--- a/contrib/restricted/aws/aws-c-io/source/s2n/s2n_tls_channel_handler.c
+++ b/contrib/restricted/aws/aws-c-io/source/s2n/s2n_tls_channel_handler.c
@@ -11,13 +11,10 @@
 #include <aws/io/event_loop.h>
 #include <aws/io/file_utils.h>
 #include <aws/io/logging.h>
-#include <aws/io/pkcs11.h>
 #include <aws/io/private/pki_utils.h>
 #include <aws/io/private/tls_channel_handler_shared.h>
 #include <aws/io/statistics.h>
 
-#include "../pkcs11_private.h"
-
 #include <aws/common/encoding.h>
 #include <aws/common/string.h>
 #include <aws/common/task_scheduler.h>
@@ -67,7 +64,6 @@ struct s2n_handler {
         NEGOTIATION_SUCCEEDED,
     } state;
     struct s2n_delayed_shutdown_task delayed_shutdown_task;
-    struct aws_channel_task async_pkey_task;
 };
 
 struct s2n_ctx {
@@ -77,25 +73,29 @@ struct s2n_ctx {
     /* Only used in special circumstances (ex: have cert but no key, because key is in PKCS#11) */
     struct s2n_cert_chain_and_key *custom_cert_chain_and_key;
 
-    /* Use a single PKCS#11 session for all TLS connections on this s2n_ctx.
-     * We do this because PKCS#11 tokens may only support a
-     * limited number of sessions (PKCS11-UG-v2.40 section 2.6.7).
-     * If this one shared session turns out to be a severe bottleneck,
-     * we could look into other setups (ex: put session on its own thread,
-     * 1 session per event-loop, 1 session per connection, etc).
+    /**
+     * Custom key operations to perform when a private key operation is required in the TLS handshake.
+     * Only will be used if non-NULL, otherwise this is ignored and the standard private key operations
+     * are performed instead.
+     * NOTE: PKCS11 also is done via this custom_key_handler.
      *
-     * The lock must be held while performing session operations.
-     * Otherwise, it would not be safe for multiple threads to share a
-     * session (PKCS11-UG-v2.40 section 2.6.7). The lock isn't needed for
-     * setup and teardown though, since we ensure nothing parallel is going
-     * on at these times */
-    struct {
-        struct aws_pkcs11_lib *lib;
-        struct aws_mutex session_lock;
-        CK_SESSION_HANDLE session_handle;
-        CK_OBJECT_HANDLE private_key_handle;
-        CK_KEY_TYPE private_key_type;
-    } pkcs11;
+     * See aws_custom_key_op_handler in tls_channel_handler.h for more details.
+     */
+    struct aws_custom_key_op_handler *custom_key_handler;
+};
+
+struct aws_tls_key_operation {
+    struct aws_allocator *alloc;
+    struct s2n_async_pkey_op *s2n_op;
+    struct s2n_handler *s2n_handler;
+    enum aws_tls_key_operation_type operation_type;
+    enum aws_tls_signature_algorithm signature_algorithm;
+    enum aws_tls_hash_algorithm digest_algorithm;
+    struct aws_byte_buf input_data;
+    struct aws_channel_task completion_task;
+    int completion_error_code;
+
+    struct aws_atomic_var complete_count;
 };
 
 AWS_STATIC_STRING_FROM_LITERAL(s_debian_path, "/etc/ssl/certs");
@@ -642,59 +642,169 @@ static enum aws_tls_hash_algorithm s_s2n_to_aws_hash_algorithm(s2n_tls_hash_algo
     }
 }
 
-/* This task performs the PKCS#11 private key operations.
- * This task is scheduled because the s2n async private key operation is not allowed to complete synchronously */
-static void s_s2n_pkcs11_async_pkey_task(
+static void s_tls_key_operation_destroy(struct aws_tls_key_operation *operation) {
+    if (operation->s2n_op) {
+        s2n_async_pkey_op_free(operation->s2n_op);
+    }
+    if (operation->s2n_handler) {
+        aws_channel_release_hold(operation->s2n_handler->slot->channel);
+    }
+    aws_byte_buf_clean_up(&operation->input_data);
+    aws_mem_release(operation->alloc, operation);
+}
+
+/* This task finishes a private key operation on the event-loop thread.
+ * If the operation was successful, TLS negotiation is resumed.
+ * If the operation failed, the channel is shut down */
+static void s_tls_key_operation_completion_task(
     struct aws_channel_task *channel_task,
     void *arg,
     enum aws_task_status status) {
 
-    struct s2n_handler *s2n_handler = AWS_CONTAINER_OF(channel_task, struct s2n_handler, async_pkey_task);
+    (void)channel_task;
+    struct aws_tls_key_operation *operation = arg;
+    struct s2n_handler *s2n_handler = operation->s2n_handler;
     struct aws_channel_handler *handler = &s2n_handler->handler;
-    struct s2n_async_pkey_op *op = arg;
-    bool success = false;
-
-    uint8_t *input_data = NULL;     /* allocated later */
-    struct aws_byte_buf output_buf; /* initialized later */
-    AWS_ZERO_STRUCT(output_buf);
 
     /* if things started failing since this task was scheduled, just clean up and bail out */
     if (status != AWS_TASK_STATUS_RUN_READY || s2n_handler->state != NEGOTIATION_ONGOING) {
         goto clean_up;
     }
 
-    AWS_LOGF_TRACE(AWS_LS_IO_TLS, "id=%p: Running PKCS#11 async pkey task", (void *)handler);
+    if (operation->completion_error_code == 0) {
+        if (s2n_async_pkey_op_apply(operation->s2n_op, s2n_handler->connection)) {
+            AWS_LOGF_ERROR(AWS_LS_IO_TLS, "id=%p: Failed applying s2n async pkey op", (void *)handler);
+            operation->completion_error_code = AWS_ERROR_INVALID_STATE;
+        }
+    }
+
+    if (operation->completion_error_code == 0) {
+        s_drive_negotiation(handler);
+    } else {
+        aws_channel_shutdown(s2n_handler->slot->channel, operation->completion_error_code);
+    }
+
+clean_up:
+    s_tls_key_operation_destroy(operation);
+}
+
+/* Common implementation for aws_tls_key_operation_complete() and aws_tls_key_operation_complete_with_error()
+ * This is called exactly once. Schedules a task to actually finish things up on the event-loop thread. */
+static void s_tls_key_operation_complete_common(
+    struct aws_tls_key_operation *operation,
+    int error_code,
+    const struct aws_byte_cursor *output) {
+
+    AWS_ASSERT((error_code != 0) ^ (output != NULL)); /* error_code XOR output must be set */
 
-    /* We check all s2n_async_pkey_op functions for success,
-     * but they shouldn't fail if they're called correctly.
-     * Even if the output is bad, the failure will happen later in s2n_negotiate() */
+    /* Ensure this can only be called once and exactly once. */
+    size_t complete_count = aws_atomic_fetch_add(&operation->complete_count, 1);
+    AWS_FATAL_ASSERT(complete_count == 0 && "TLS key operation marked complete multiple times");
 
+    struct s2n_handler *s2n_handler = operation->s2n_handler;
+    struct aws_channel_handler *handler = &s2n_handler->handler;
+
+    if (output != NULL) {
+        /* Immediately pass output through to s2n_op. */
+        if (s2n_async_pkey_op_set_output(operation->s2n_op, output->ptr, output->len)) {
+            AWS_LOGF_ERROR(AWS_LS_IO_TLS, "id=%p: Failed setting output on s2n async pkey op", (void *)handler);
+            error_code = AWS_ERROR_INVALID_STATE;
+            goto done;
+        }
+    }
+
+done:
+    operation->completion_error_code = error_code;
+
+    /* Schedule a task to finish the operation.
+     * We schedule a task because the user might
+     * have completed the operation asynchronously,
+     * but we need to be on the event-loop thread to
+     * resume TLS negotiation. */
+    aws_channel_task_init(
+        &operation->completion_task,
+        s_tls_key_operation_completion_task,
+        operation,
+        "tls_key_operation_completion_task");
+    aws_channel_schedule_task_now(s2n_handler->slot->channel, &operation->completion_task);
+}
+
+void aws_tls_key_operation_complete(struct aws_tls_key_operation *operation, struct aws_byte_cursor output) {
+    if (operation == NULL) {
+        AWS_LOGF_ERROR(AWS_LS_IO_TLS, "Operation complete: operation is null and therefore cannot be set to complete!");
+        return;
+    }
+
+    AWS_LOGF_DEBUG(
+        AWS_LS_IO_TLS,
+        "id=%p: TLS key operation complete with %zu bytes of output data",
+        (void *)operation->s2n_handler,
+        output.len);
+    s_tls_key_operation_complete_common(operation, 0, &output);
+}
+
+void aws_tls_key_operation_complete_with_error(struct aws_tls_key_operation *operation, int error_code) {
+    if (operation == NULL) {
+        AWS_LOGF_ERROR(
+            AWS_LS_IO_TLS, "Operation complete with error: operation is null and therefore cannot be set to complete!");
+        return;
+    }
+
+    if (error_code == 0) {
+        error_code = AWS_ERROR_UNKNOWN;
+        AWS_LOGF_ERROR(
+            AWS_LS_IO_TLS,
+            "id=%p: TLS key operation completed with error, but no error-code set. Using %s",
+            (void *)operation->s2n_handler,
+            aws_error_name(error_code));
+    }
+
+    AWS_LOGF_ERROR(
+        AWS_LS_IO_TLS,
+        "id=%p: TLS key operation complete with error %s",
+        (void *)operation->s2n_handler,
+        aws_error_name(error_code));
+
+    s_tls_key_operation_complete_common(operation, error_code, NULL);
+}
+
+static struct aws_tls_key_operation *s_tls_key_operation_new(
+    struct aws_channel_handler *handler,
+    struct s2n_async_pkey_op *s2n_op) {
+
+    struct s2n_handler *s2n_handler = handler->impl;
+
+    struct aws_tls_key_operation *operation = aws_mem_calloc(handler->alloc, 1, sizeof(struct aws_tls_key_operation));
+    operation->alloc = handler->alloc;
+
+    /* Copy input data */
     uint32_t input_size = 0;
-    if (s2n_async_pkey_op_get_input_size(op, &input_size)) {
+    if (s2n_async_pkey_op_get_input_size(s2n_op, &input_size)) {
         AWS_LOGF_ERROR(AWS_LS_IO_TLS, "id=%p: Failed querying s2n async pkey op size", (void *)handler);
         aws_raise_error(AWS_ERROR_INVALID_STATE);
         goto error;
     }
 
-    input_data = aws_mem_acquire(handler->alloc, input_size);
-    if (s2n_async_pkey_op_get_input(op, input_data, input_size)) {
+    aws_byte_buf_init(&operation->input_data, operation->alloc, input_size); /* cannot fail */
+    if (s2n_async_pkey_op_get_input(s2n_op, operation->input_data.buffer, input_size)) {
         AWS_LOGF_ERROR(AWS_LS_IO_TLS, "id=%p: Failed querying s2n async pkey input", (void *)handler);
         aws_raise_error(AWS_ERROR_INVALID_STATE);
         goto error;
     }
-    struct aws_byte_cursor input_cursor = aws_byte_cursor_from_array(input_data, input_size);
+    operation->input_data.len = input_size;
 
-    s2n_async_pkey_op_type op_type = 0;
-    if (s2n_async_pkey_op_get_op_type(op, &op_type)) {
+    /* Get operation type */
+    s2n_async_pkey_op_type s2n_op_type = 0;
+    if (s2n_async_pkey_op_get_op_type(s2n_op, &s2n_op_type)) {
         AWS_LOGF_ERROR(AWS_LS_IO_TLS, "id=%p: Failed querying s2n async pkey op type", (void *)handler);
         aws_raise_error(AWS_ERROR_INVALID_STATE);
         goto error;
     }
 
-    /* Gather additional information if this is a SIGN operation */
-    enum aws_tls_signature_algorithm aws_sign_alg = 0;
-    enum aws_tls_hash_algorithm aws_digest_alg = 0;
-    if (op_type == S2N_ASYNC_SIGN) {
+    if (s2n_op_type == S2N_ASYNC_SIGN) {
+        operation->operation_type = AWS_TLS_KEY_OPERATION_SIGN;
+
+        /* Gather additional information if this is a SIGN operation */
         s2n_tls_signature_algorithm s2n_sign_alg = 0;
         if (s2n_connection_get_selected_client_cert_signature_algorithm(s2n_handler->connection, &s2n_sign_alg)) {
             AWS_LOGF_ERROR(AWS_LS_IO_TLS, "id=%p: Failed getting s2n client cert signature algorithm", (void *)handler);
@@ -702,8 +812,8 @@ static void s_s2n_pkcs11_async_pkey_task(
             goto error;
         }
 
-        aws_sign_alg = s_s2n_to_aws_signature_algorithm(s2n_sign_alg);
-        if (aws_sign_alg == AWS_TLS_SIGNATURE_UNKNOWN) {
+        operation->signature_algorithm = s_s2n_to_aws_signature_algorithm(s2n_sign_alg);
+        if (operation->signature_algorithm == AWS_TLS_SIGNATURE_UNKNOWN) {
             AWS_LOGF_ERROR(
                 AWS_LS_IO_TLS,
                 "id=%p: Cannot sign with s2n_tls_signature_algorithm=%d. Algorithm currently unsupported",
@@ -720,8 +830,8 @@ static void s_s2n_pkcs11_async_pkey_task(
             goto error;
         }
 
-        aws_digest_alg = s_s2n_to_aws_hash_algorithm(s2n_digest_alg);
-        if (aws_digest_alg == AWS_TLS_HASH_UNKNOWN) {
+        operation->digest_algorithm = s_s2n_to_aws_hash_algorithm(s2n_digest_alg);
+        if (operation->digest_algorithm == AWS_TLS_HASH_UNKNOWN) {
             AWS_LOGF_ERROR(
                 AWS_LS_IO_TLS,
                 "id=%p: Cannot sign digest created with s2n_tls_hash_algorithm=%d. Algorithm currently unsupported",
@@ -730,113 +840,75 @@ static void s_s2n_pkcs11_async_pkey_task(
             aws_raise_error(AWS_IO_TLS_DIGEST_ALGORITHM_UNSUPPORTED);
             goto error;
         }
-    }
-
-    /*********** BEGIN CRITICAL SECTION ***********/
-    aws_mutex_lock(&s2n_handler->s2n_ctx->pkcs11.session_lock);
-    bool success_while_locked = false;
-
-    switch (op_type) {
-        case S2N_ASYNC_DECRYPT:
-            if (aws_pkcs11_lib_decrypt(
-                    s2n_handler->s2n_ctx->pkcs11.lib,
-                    s2n_handler->s2n_ctx->pkcs11.session_handle,
-                    s2n_handler->s2n_ctx->pkcs11.private_key_handle,
-                    s2n_handler->s2n_ctx->pkcs11.private_key_type,
-                    input_cursor,
-                    handler->alloc,
-                    &output_buf)) {
-
-                AWS_LOGF_ERROR(
-                    AWS_LS_IO_TLS,
-                    "id=%p: PKCS#11 decrypt failed, error %s",
-                    (void *)handler,
-                    aws_error_name(aws_last_error()));
-                goto unlock;
-            }
-            break;
-
-        case S2N_ASYNC_SIGN:
-            if (aws_pkcs11_lib_sign(
-                    s2n_handler->s2n_ctx->pkcs11.lib,
-                    s2n_handler->s2n_ctx->pkcs11.session_handle,
-                    s2n_handler->s2n_ctx->pkcs11.private_key_handle,
-                    s2n_handler->s2n_ctx->pkcs11.private_key_type,
-                    input_cursor,
-                    handler->alloc,
-                    aws_digest_alg,
-                    aws_sign_alg,
-                    &output_buf)) {
-
-                AWS_LOGF_ERROR(
-                    AWS_LS_IO_TLS,
-                    "id=%p: PKCS#11 sign failed, error %s",
-                    (void *)handler,
-                    aws_error_name(aws_last_error()));
-                goto unlock;
-            }
-            break;
-
-        default:
-            AWS_LOGF_ERROR(AWS_LS_IO_TLS, "id=%p: Unknown s2n_async_pkey_op_type:%d", (void *)handler, (int)op_type);
-            aws_raise_error(AWS_ERROR_INVALID_STATE);
-            goto unlock;
-    }
-
-    success_while_locked = true;
-unlock:
-    aws_mutex_unlock(&s2n_handler->s2n_ctx->pkcs11.session_lock);
-    /*********** END CRITICAL SECTION ***********/
-
-    if (!success_while_locked) {
-        goto error;
-    }
 
-    AWS_LOGF_TRACE(
-        AWS_LS_IO_TLS, "id=%p: PKCS#11 operation complete. output-size:%zu", (void *)handler, output_buf.len);
+    } else if (s2n_op_type == S2N_ASYNC_DECRYPT) {
+        operation->operation_type = AWS_TLS_KEY_OPERATION_DECRYPT;
 
-    if (s2n_async_pkey_op_set_output(op, output_buf.buffer, output_buf.len)) {
-        AWS_LOGF_ERROR(AWS_LS_IO_TLS, "id=%p: Failed setting output on s2n async pkey op", (void *)handler);
+    } else {
+        AWS_LOGF_ERROR(AWS_LS_IO_TLS, "id=%p: Unknown s2n async pkey op type:%d", (void *)handler, (int)s2n_op_type);
         aws_raise_error(AWS_ERROR_INVALID_STATE);
         goto error;
     }
 
-    if (s2n_async_pkey_op_apply(op, s2n_handler->connection)) {
-        AWS_LOGF_ERROR(AWS_LS_IO_TLS, "id=%p: Failed applying s2n async pkey op", (void *)handler);
-        aws_raise_error(AWS_ERROR_INVALID_STATE);
-        goto error;
-    }
+    /* Keep channel alive until operation completes */
+    operation->s2n_handler = s2n_handler;
+    aws_channel_acquire_hold(s2n_handler->slot->channel);
+
+    /* Set this to zero so we can track how many times complete has been called */
+    aws_atomic_init_int(&operation->complete_count, 0);
 
-    /* Success! */
-    success = true;
-    goto clean_up;
+    /* Set this last. We don't want to take ownership of s2n_op until we know setup was 100% successful */
+    operation->s2n_op = s2n_op;
 
+    return operation;
 error:
-    aws_channel_shutdown(s2n_handler->slot->channel, aws_last_error());
+    s_tls_key_operation_destroy(operation);
+    return NULL;
+}
 
-clean_up:
-    s2n_async_pkey_op_free(op);
-    aws_mem_release(handler->alloc, input_data);
-    aws_byte_buf_clean_up(&output_buf);
+struct aws_byte_cursor aws_tls_key_operation_get_input(const struct aws_tls_key_operation *operation) {
+    return aws_byte_cursor_from_buf(&operation->input_data);
+}
 
-    if (success) {
-        s_drive_negotiation(handler);
-    }
+enum aws_tls_key_operation_type aws_tls_key_operation_get_type(const struct aws_tls_key_operation *operation) {
+    return operation->operation_type;
 }
 
-static int s_s2n_pkcs11_async_pkey_callback(struct s2n_connection *conn, struct s2n_async_pkey_op *op) {
+enum aws_tls_signature_algorithm aws_tls_key_operation_get_signature_algorithm(
+    const struct aws_tls_key_operation *operation) {
+    return operation->signature_algorithm;
+}
+
+enum aws_tls_hash_algorithm aws_tls_key_operation_get_digest_algorithm(const struct aws_tls_key_operation *operation) {
+    return operation->digest_algorithm;
+}
+
+static int s_s2n_async_pkey_callback(struct s2n_connection *conn, struct s2n_async_pkey_op *s2n_op) {
     struct s2n_handler *s2n_handler = s2n_connection_get_ctx(conn);
     struct aws_channel_handler *handler = &s2n_handler->handler;
 
     AWS_ASSERT(conn == s2n_handler->connection);
     (void)conn;
 
-    /* Schedule a task to do the work.
-     * s2n can't deal with the async private key operation completing synchronously, so we can't just do it now */
-    AWS_LOGF_TRACE(AWS_LS_IO_TLS, "id=%p: async pkey callback received, scheduling PKCS#11 task", (void *)handler);
+    AWS_LOGF_TRACE(AWS_LS_IO_TLS, "id=%p: s2n async pkey callback received", (void *)handler);
+
+    /* Create the AWS wrapper around s2n_async_pkey_op */
+    struct aws_tls_key_operation *operation = s_tls_key_operation_new(handler, s2n_op);
+    if (operation == NULL) {
+        s2n_async_pkey_op_free(s2n_op);
+        return S2N_FAILURE;
+    }
+
+    AWS_LOGF_DEBUG(
+        AWS_LS_IO_TLS,
+        "id=%p: Begin TLS key operation. type=%s input_data.len=%zu signature=%s digest=%s",
+        (void *)operation,
+        aws_tls_key_operation_type_str(operation->operation_type),
+        operation->input_data.len,
+        aws_tls_signature_algorithm_str(operation->signature_algorithm),
+        aws_tls_hash_algorithm_str(operation->digest_algorithm));
 
-    aws_channel_task_init(&s2n_handler->async_pkey_task, s_s2n_pkcs11_async_pkey_task, op, "s2n_pkcs11_async_pkey_op");
-    aws_channel_schedule_task_now(s2n_handler->slot->channel, &s2n_handler->async_pkey_task);
+    aws_custom_key_op_handler_perform_operation(s2n_handler->s2n_ctx->custom_key_handler, operation);
 
     return S2N_SUCCESS;
 }
@@ -1205,16 +1277,12 @@ struct aws_channel_handler *aws_tls_server_handler_new(
 
 static void s_s2n_ctx_destroy(struct s2n_ctx *s2n_ctx) {
     if (s2n_ctx != NULL) {
-        if (s2n_ctx->pkcs11.session_handle != 0) {
-            aws_pkcs11_lib_close_session(s2n_ctx->pkcs11.lib, s2n_ctx->pkcs11.session_handle);
-        }
-        aws_mutex_clean_up(&s2n_ctx->pkcs11.session_lock);
-        aws_pkcs11_lib_release(s2n_ctx->pkcs11.lib);
         s2n_config_free(s2n_ctx->s2n_config);
 
         if (s2n_ctx->custom_cert_chain_and_key) {
             s2n_cert_chain_and_key_free(s2n_ctx->custom_cert_chain_and_key);
         }
+        s2n_ctx->custom_key_handler = aws_custom_key_op_handler_release(s2n_ctx->custom_key_handler);
 
         aws_mem_release(s2n_ctx->ctx.alloc, s2n_ctx);
     }
@@ -1240,44 +1308,6 @@ static int s2n_monotonic_clock_time_nanoseconds(void *context, uint64_t *time_in
     return 0;
 }
 
-static int s_tls_ctx_pkcs11_setup(struct s2n_ctx *s2n_ctx, const struct aws_tls_ctx_options *options) {
-    /* PKCS#11 options were already sanitized (ie: check for required args) in tls_channel_handler.c */
-
-    /* anything initialized in this function is cleaned up during s_s2n_ctx_destroy()
-     * so don't worry about cleaning up unless it's some tmp heap allocation */
-
-    s2n_ctx->pkcs11.lib = aws_pkcs11_lib_acquire(options->pkcs11.lib); /* cannot fail */
-    aws_mutex_init(&s2n_ctx->pkcs11.session_lock);
-
-    CK_SLOT_ID slot_id = 0;
-    if (aws_pkcs11_lib_find_slot_with_token(
-            s2n_ctx->pkcs11.lib,
-            options->pkcs11.has_slot_id ? &options->pkcs11.slot_id : NULL,
-            options->pkcs11.token_label,
-            &slot_id /*out*/)) {
-        return AWS_OP_ERR;
-    }
-
-    if (aws_pkcs11_lib_open_session(s2n_ctx->pkcs11.lib, slot_id, &s2n_ctx->pkcs11.session_handle)) {
-        return AWS_OP_ERR;
-    }
-
-    if (aws_pkcs11_lib_login_user(s2n_ctx->pkcs11.lib, s2n_ctx->pkcs11.session_handle, options->pkcs11.user_pin)) {
-        return AWS_OP_ERR;
-    }
-
-    if (aws_pkcs11_lib_find_private_key(
-            s2n_ctx->pkcs11.lib,
-            s2n_ctx->pkcs11.session_handle,
-            options->pkcs11.private_key_object_label,
-            &s2n_ctx->pkcs11.private_key_handle /*out*/,
-            &s2n_ctx->pkcs11.private_key_type /*out*/)) {
-        return AWS_OP_ERR;
-    }
-
-    return AWS_OP_SUCCESS;
-}
-
 static void s_log_and_raise_s2n_errno(const char *msg) {
     AWS_LOGF_ERROR(
         AWS_LS_IO_TLS, "%s: %s (%s)", msg, s2n_strerror(s2n_errno, "EN"), s2n_strerror_debug(s2n_errno, "EN"));
@@ -1322,7 +1352,7 @@ static struct aws_tls_ctx *s_tls_ctx_new(
         goto cleanup_s2n_config;
     }
 
-    if (options->pkcs11.lib != NULL) {
+    if (options->custom_key_op_handler != NULL) {
         /* PKCS#11 integration hasn't been tested with TLS 1.3, so don't use cipher preferences that allow 1.3 */
         switch (options->minimum_tls_version) {
             case AWS_IO_SSLv3:
@@ -1414,14 +1444,12 @@ static struct aws_tls_ctx *s_tls_ctx_new(
             s_log_and_raise_s2n_errno("ctx: Failed to add certificate and private key");
             goto cleanup_s2n_config;
         }
-    } else if (options->pkcs11.lib != NULL) {
-        AWS_LOGF_DEBUG(AWS_LS_IO_TLS, "ctx: PKCS#11 has been set, setting it up now.");
-        if (s_tls_ctx_pkcs11_setup(s2n_ctx, options)) {
-            goto cleanup_s2n_config;
-        }
+    } else if (options->custom_key_op_handler != NULL) {
+
+        s2n_ctx->custom_key_handler = aws_custom_key_op_handler_acquire(options->custom_key_op_handler);
 
-        /* set callback so that we can do private key operations through PKCS#11 */
-        if (s2n_config_set_async_pkey_callback(s2n_ctx->s2n_config, s_s2n_pkcs11_async_pkey_callback)) {
+        /* set callback so that we can do custom private key operations */
+        if (s2n_config_set_async_pkey_callback(s2n_ctx->s2n_config, s_s2n_async_pkey_callback)) {
             s_log_and_raise_s2n_errno("ctx: failed to set private key callback");
             goto cleanup_s2n_config;
         }
diff --git a/contrib/restricted/aws/aws-c-io/source/tls_channel_handler.c b/contrib/restricted/aws/aws-c-io/source/tls_channel_handler.c
index 65acdf191de..27218bbb61f 100644
--- a/contrib/restricted/aws/aws-c-io/source/tls_channel_handler.c
+++ b/contrib/restricted/aws/aws-c-io/source/tls_channel_handler.c
@@ -13,6 +13,8 @@
 
 #define AWS_DEFAULT_TLS_TIMEOUT_MS 10000
 
+#include "./pkcs11_private.h"
+
 #include <aws/common/string.h>
 
 void aws_tls_ctx_options_init_default_client(struct aws_tls_ctx_options *options, struct aws_allocator *allocator) {
@@ -40,11 +42,7 @@ void aws_tls_ctx_options_clean_up(struct aws_tls_ctx_options *options) {
 #endif
 
     aws_string_destroy(options->alpn_list);
-
-    aws_pkcs11_lib_release(options->pkcs11.lib);
-    aws_string_destroy_secure(options->pkcs11.user_pin);
-    aws_string_destroy(options->pkcs11.token_label);
-    aws_string_destroy(options->pkcs11.private_key_object_label);
+    aws_custom_key_op_handler_release(options->custom_key_op_handler);
 
     AWS_ZERO_STRUCT(*options);
 }
@@ -134,85 +132,131 @@ error:
 #endif
 }
 
-int aws_tls_ctx_options_init_client_mtls_with_pkcs11(
+int aws_tls_ctx_options_init_client_mtls_with_custom_key_operations(
     struct aws_tls_ctx_options *options,
     struct aws_allocator *allocator,
-    const struct aws_tls_ctx_pkcs11_options *pkcs11_options) {
+    struct aws_custom_key_op_handler *custom,
+    const struct aws_byte_cursor *cert_file_contents) {
 
-#if defined(_WIN32) || defined(__APPLE__)
+#if !USE_S2N
+    (void)options;
     (void)allocator;
-    (void)pkcs11_options;
+    (void)custom;
+    (void)cert_file_contents;
     AWS_ZERO_STRUCT(*options);
-    AWS_LOGF_ERROR(AWS_LS_IO_TLS, "static: This platform does not currently support TLS with PKCS#11.");
-    return aws_raise_error(AWS_ERROR_PLATFORM_NOT_SUPPORTED);
+    AWS_LOGF_ERROR(
+        AWS_LS_IO_TLS, "static: This platform does not currently support TLS with custom private key operations.");
+    return aws_raise_error(AWS_ERROR_UNIMPLEMENTED);
 #else
 
     aws_tls_ctx_options_init_default_client(options, allocator);
 
-    /* pkcs11_lib is required */
-    if (pkcs11_options->pkcs11_lib == NULL) {
-        AWS_LOGF_ERROR(AWS_LS_IO_TLS, "static: A PKCS#11 library must be specified.");
-        aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
+    /* on_key_operation is required */
+    AWS_ASSERT(custom != NULL);
+    AWS_ASSERT(custom->vtable != NULL);
+    AWS_ASSERT(custom->vtable->on_key_operation != NULL);
+
+    /* Hold a reference to the custom key operation handler so it cannot be destroyed */
+    options->custom_key_op_handler = aws_custom_key_op_handler_acquire((struct aws_custom_key_op_handler *)custom);
+
+    /* Copy the certificate data from the cursor */
+    AWS_ASSERT(cert_file_contents != NULL);
+    aws_byte_buf_init_copy_from_cursor(&options->certificate, allocator, *cert_file_contents);
+
+    /* Make sure the certificate is set and valid */
+    if (aws_sanitize_pem(&options->certificate, allocator)) {
+        AWS_LOGF_ERROR(AWS_LS_IO_TLS, "static: Invalid certificate. File must contain PEM encoded data");
         goto error;
     }
-    options->pkcs11.lib = aws_pkcs11_lib_acquire(pkcs11_options->pkcs11_lib); /* cannot fail */
 
-    /* user_pin is optional */
-    if (pkcs11_options->user_pin.ptr != NULL) {
-        options->pkcs11.user_pin = aws_string_new_from_cursor(allocator, &pkcs11_options->user_pin);
-    }
+    return AWS_OP_SUCCESS;
 
-    /* slot_id is optional */
-    if (pkcs11_options->slot_id != NULL) {
-        options->pkcs11.slot_id = *pkcs11_options->slot_id;
-        options->pkcs11.has_slot_id = true;
-    }
+error:
+    aws_tls_ctx_options_clean_up(options);
+    return AWS_OP_ERR;
 
-    /* token_label is optional */
-    if (pkcs11_options->token_label.ptr != NULL) {
-        options->pkcs11.token_label = aws_string_new_from_cursor(allocator, &pkcs11_options->token_label);
-    }
+#endif /* PLATFORM-SUPPORTS-CUSTOM-KEY-OPERATIONS */
+}
 
-    /* private_key_object_label is optional */
-    if (pkcs11_options->private_key_object_label.ptr != NULL) {
-        options->pkcs11.private_key_object_label =
-            aws_string_new_from_cursor(allocator, &pkcs11_options->private_key_object_label);
+int aws_tls_ctx_options_init_client_mtls_with_pkcs11(
+    struct aws_tls_ctx_options *options,
+    struct aws_allocator *allocator,
+    const struct aws_tls_ctx_pkcs11_options *pkcs11_options) {
+
+#if defined(USE_S2N)
+
+    struct aws_custom_key_op_handler *pkcs11_handler = aws_pkcs11_tls_op_handler_new(
+        allocator,
+        pkcs11_options->pkcs11_lib,
+        &pkcs11_options->user_pin,
+        &pkcs11_options->token_label,
+        &pkcs11_options->private_key_object_label,
+        pkcs11_options->slot_id);
+
+    struct aws_byte_buf tmp_cert_buf;
+    AWS_ZERO_STRUCT(tmp_cert_buf);
+    bool success = false;
+    int custom_key_result = AWS_OP_ERR;
+
+    if (pkcs11_handler == NULL) {
+        aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
+        goto finish;
     }
 
-    /* certificate required, but there are multiple ways to pass it in */
-    if ((pkcs11_options->cert_file_path.ptr != NULL) && (pkcs11_options->cert_file_contents.ptr != NULL)) {
-        AWS_LOGF_ERROR(
-            AWS_LS_IO_TLS, "static: Both certificate filepath and contents are specified. Only one may be set.");
+    if ((pkcs11_options->cert_file_contents.ptr != NULL) && (pkcs11_options->cert_file_path.ptr != NULL)) {
+        AWS_LOGF_ERROR(AWS_LS_IO_TLS, "static: Cannot use certificate AND certificate file path, only one can be set");
         aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
-        goto error;
-    } else if (pkcs11_options->cert_file_path.ptr != NULL) {
+        goto finish;
+    } else if (pkcs11_options->cert_file_contents.ptr != NULL) {
+        custom_key_result = aws_tls_ctx_options_init_client_mtls_with_custom_key_operations(
+            options, allocator, pkcs11_handler, &pkcs11_options->cert_file_contents);
+        success = true;
+    } else {
         struct aws_string *tmp_string = aws_string_new_from_cursor(allocator, &pkcs11_options->cert_file_path);
-        int op = aws_byte_buf_init_from_file(&options->certificate, allocator, aws_string_c_str(tmp_string));
+        int op = aws_byte_buf_init_from_file(&tmp_cert_buf, allocator, aws_string_c_str(tmp_string));
         aws_string_destroy(tmp_string);
+
         if (op != AWS_OP_SUCCESS) {
-            goto error;
+            goto finish;
         }
-    } else if (pkcs11_options->cert_file_contents.ptr != NULL) {
-        if (aws_byte_buf_init_copy_from_cursor(&options->certificate, allocator, pkcs11_options->cert_file_contents)) {
-            goto error;
-        }
-    } else {
-        AWS_LOGF_ERROR(AWS_LS_IO_TLS, "static: A certificate must be specified.");
-        aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
-        goto error;
+
+        struct aws_byte_cursor tmp_cursor = aws_byte_cursor_from_buf(&tmp_cert_buf);
+        custom_key_result = aws_tls_ctx_options_init_client_mtls_with_custom_key_operations(
+            options, allocator, pkcs11_handler, &tmp_cursor);
+        success = true;
     }
 
-    if (aws_sanitize_pem(&options->certificate, allocator)) {
-        AWS_LOGF_ERROR(AWS_LS_IO_TLS, "static: Invalid certificate. File must contain PEM encoded data");
-        goto error;
+finish:
+
+    if (pkcs11_handler != NULL) {
+        /**
+         * Calling aws_tls_ctx_options_init_client_mtls_with_custom_key_operations will have this options
+         * hold a reference to the custom key operations, but creating the TLS operations handler using
+         * aws_pkcs11_tls_op_handler_set_certificate_data adds a reference too, so we need to release
+         * this reference so the only thing (currently) holding a reference is the TLS options itself and
+         * not this function.
+         */
+        aws_custom_key_op_handler_release(pkcs11_handler);
     }
+    if (success == false) {
+        aws_tls_ctx_options_clean_up(options);
+    }
+    aws_byte_buf_clean_up(&tmp_cert_buf);
 
-    /* Success! */
-    return AWS_OP_SUCCESS;
+    if (success) {
+        return custom_key_result;
+    } else {
+        return AWS_OP_ERR;
+    }
+
+#else /* Platform does not support S2N */
+
+    (void)allocator;
+    (void)pkcs11_options;
+    AWS_ZERO_STRUCT(*options);
+    AWS_LOGF_ERROR(AWS_LS_IO_TLS, "static: This platform does not currently support TLS with PKCS#11.");
+    return aws_raise_error(AWS_ERROR_PLATFORM_NOT_SUPPORTED);
 
-error:
-    aws_tls_ctx_options_clean_up(options);
-    return AWS_OP_ERR;
 #endif /* PLATFORM-SUPPORTS-PKCS11-TLS */
 }
 
@@ -739,3 +783,90 @@ void aws_tls_ctx_release(struct aws_tls_ctx *ctx) {
         aws_ref_count_release(&ctx->ref_count);
     }
 }
+
+const char *aws_tls_hash_algorithm_str(enum aws_tls_hash_algorithm hash) {
+    /* clang-format off */
+    switch (hash) {
+        case (AWS_TLS_HASH_SHA1): return "SHA1";
+        case (AWS_TLS_HASH_SHA224): return "SHA224";
+        case (AWS_TLS_HASH_SHA256): return "SHA256";
+        case (AWS_TLS_HASH_SHA384): return "SHA384";
+        case (AWS_TLS_HASH_SHA512): return "SHA512";
+        default: return "<UNKNOWN HASH ALGORITHM>";
+    }
+    /* clang-format on */
+}
+
+const char *aws_tls_signature_algorithm_str(enum aws_tls_signature_algorithm signature) {
+    /* clang-format off */
+    switch (signature) {
+        case (AWS_TLS_SIGNATURE_RSA): return "RSA";
+        case (AWS_TLS_SIGNATURE_ECDSA): return "ECDSA";
+        default: return "<UNKNOWN SIGNATURE ALGORITHM>";
+    }
+    /* clang-format on */
+}
+
+const char *aws_tls_key_operation_type_str(enum aws_tls_key_operation_type operation_type) {
+    /* clang-format off */
+    switch (operation_type) {
+        case (AWS_TLS_KEY_OPERATION_SIGN): return "SIGN";
+        case (AWS_TLS_KEY_OPERATION_DECRYPT): return "DECRYPT";
+        default: return "<UNKNOWN OPERATION TYPE>";
+    }
+    /* clang-format on */
+}
+
+#if !USE_S2N
+void aws_tls_key_operation_complete(struct aws_tls_key_operation *operation, struct aws_byte_cursor output) {
+    (void)operation;
+    (void)output;
+}
+
+void aws_tls_key_operation_complete_with_error(struct aws_tls_key_operation *operation, int error_code) {
+    (void)operation;
+    (void)error_code;
+}
+
+struct aws_byte_cursor aws_tls_key_operation_get_input(const struct aws_tls_key_operation *operation) {
+    (void)operation;
+    return aws_byte_cursor_from_array(NULL, 0);
+}
+
+enum aws_tls_key_operation_type aws_tls_key_operation_get_type(const struct aws_tls_key_operation *operation) {
+    (void)operation;
+    return AWS_TLS_KEY_OPERATION_UNKNOWN;
+}
+
+enum aws_tls_signature_algorithm aws_tls_key_operation_get_signature_algorithm(
+    const struct aws_tls_key_operation *operation) {
+    (void)operation;
+    return AWS_TLS_SIGNATURE_UNKNOWN;
+}
+
+enum aws_tls_hash_algorithm aws_tls_key_operation_get_digest_algorithm(const struct aws_tls_key_operation *operation) {
+    (void)operation;
+    return AWS_TLS_HASH_UNKNOWN;
+}
+
+#endif
+
+struct aws_custom_key_op_handler *aws_custom_key_op_handler_acquire(struct aws_custom_key_op_handler *key_op_handler) {
+    if (key_op_handler != NULL) {
+        aws_ref_count_acquire(&key_op_handler->ref_count);
+    }
+    return key_op_handler;
+}
+
+struct aws_custom_key_op_handler *aws_custom_key_op_handler_release(struct aws_custom_key_op_handler *key_op_handler) {
+    if (key_op_handler != NULL) {
+        aws_ref_count_release(&key_op_handler->ref_count);
+    }
+    return NULL;
+}
+
+void aws_custom_key_op_handler_perform_operation(
+    struct aws_custom_key_op_handler *key_op_handler,
+    struct aws_tls_key_operation *operation) {
+    key_op_handler->vtable->on_key_operation(key_op_handler, operation);
+}