feat: add async public key support (#5473)

Author: CarolYeh910

api/unstable/async_offload.h

@@ -0,0 +1,89 @@
+/*
+* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+*
+* Licensed under the Apache License, Version 2.0 (the "License").
+* You may not use this file except in compliance with the License.
+* A copy of the License is located at
+*
+*  http://aws.amazon.com/apache2.0
+*
+* or in the "license" file accompanying this file. This file is distributed
+* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+* express or implied. See the License for the specific language governing
+* permissions and limitations under the License.
+*/
+
+#pragma once
+
+#include <s2n.h>
+
+/**
+ * @file async_offload.h
+ * 
+ * The following APIs enable applications to offload expensive handshake operations that do not require user input.
+ * This model can be useful to move CPU-heavy operations (e.g. cryptographic calculations) out of the main event loop.
+ */
+
+/**
+ * Opaque struct for the async offloading operation
+ */
+struct s2n_async_offload_op;
+
+/**
+ * The type of operations supported by the async offloading callback. Each type is represented by a different bit.
+ * 
+ * S2N_ASYNC_OFFLOAD_ALLOW_ALL will automatically opt in to all the new types added in the future.
+ */
+typedef enum {
+    S2N_ASYNC_OFFLOAD_PKEY_VERIFY = 0x01,
+    /* Max value: ISO C restricts enumerator values to range of ‘int’ before C2X. */
+    S2N_ASYNC_OFFLOAD_ALLOW_ALL = 0x7FFFFFFF,
+} s2n_async_offload_op_type;
+
+/**
+ * The callback function invoked every time an allowed async operation is encountered during the handshake.
+ * 
+ * To perform an operation asynchronously, the following condiditions must be satisfied:
+ * 1) This op type must be included in the allow list;
+ * 2) Async offloading callback returns success and s2n_async_offload_op_perform() is invoked outside the callback.
+ *
+ * If s2n_async_offload_op_perform() is invoked inside the callback, it is equivalent to the synchronous use case.
+ * 
+ * `op` is owned by s2n-tls and will be freed along with s2n_connection eventually.
+ *
+ * @param conn Connection which triggered the async offloading callback
+ * @param op An opaque object representing the async operation
+ * @param ctx Application data provided to the callback via s2n_config_set_async_offload_callback()
+ */
+typedef int (*s2n_async_offload_cb)(struct s2n_connection *conn, struct s2n_async_offload_op *op, void *ctx);
+
+/**
+ * Sets up the custom async offloading callback and configures the offloaded handshake operations via allow_list.
+ * 
+ * The value of allow_list should be the Bit-OR of all the allowed s2n_async_offload_op_type values.
+ * 
+ * S2N_ASYNC_OFFLOAD_ALLOW_ALL provides the performance benefit of offloading all the supported operations;
+ * ensure your callback can support arbitrary operations. Otherwise, only allow operations that fit your use case.
+ *
+ * @param config Config to set the callback
+ * @param allow_list A bit representation of allowed operations
+ * @param fn The function that should be called for each allowed async operation
+ * @param ctx Optional application data passed to the callback
+ */
+S2N_API extern int s2n_config_set_async_offload_callback(struct s2n_config *config, uint32_t allow_list,
+        s2n_async_offload_cb fn, void *ctx);
+
+/**
+ * Performs the operation triggered by the async offloading callback.
+ * 
+ * To execute operations asynchronously, users should spawn a separate thread to invoke s2n_async_offload_op_perform()
+ * and immediately return S2N_SUCCESS from the callback without waiting for that separate thread to complete.
+ * 
+ * s2n_negotiate() will throw an `S2N_ERR_T_BLOCKED` error if the handshake is pending on the async offloading callback.
+ * Retrying s2n_negotiate() will produce the same result until s2n_async_offload_op_perform() is completed.
+ * 
+ * s2n_async_offload_op_perform() can only be called once for each triggered operation.
+ * 
+ * @param op An opaque object representing the async operation
+ */
+S2N_API extern int s2n_async_offload_op_perform(struct s2n_async_offload_op *op);

bindings/rust/extended/s2n-tls-sys/Cargo.toml

@@ -31,6 +31,7 @@ fips = ["aws-lc-rs/fips"]
 pq = []
 internal = []
 stacktrace = []
+unstable-async_offload = []
 unstable-cert_authorities = []
 unstable-cleanup = []
 unstable-crl = []

tests/unit/s2n_async_offload_cb_test.c

@@ -0,0 +1,226 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+#include "api/s2n.h"
+#include "error/s2n_errno.h"
+#include "s2n_test.h"
+#include "testlib/s2n_testlib.h"
+#include "tls/s2n_async_offload.h"
+#include "tls/s2n_cipher_suites.h"
+#include "tls/s2n_connection.h"
+#include "tls/s2n_security_policies.h"
+#include "utils/s2n_safety.h"
+
+#define S2N_ASYNC_OFFLOAD_OP_NONE 0
+
+struct s2n_async_offload_cb_test {
+    unsigned async_test : 1;
+    int result;
+    int invoked_count;
+    struct s2n_async_offload_op *last_seen_op;
+};
+
+int s2n_async_offload_test_callback(struct s2n_connection *conn, struct s2n_async_offload_op *op, void *ctx)
+{
+    EXPECT_NOT_NULL(op);
+    struct s2n_async_offload_cb_test *data = (struct s2n_async_offload_cb_test *) ctx;
+    data->invoked_count += 1;
+    data->last_seen_op = op;
+
+    if (!data->async_test) {
+        EXPECT_SUCCESS(s2n_async_offload_op_perform(op));
+    }
+    return data->result;
+}
+
+static int s2n_test_handshake_async(struct s2n_connection *server_conn, struct s2n_connection *client_conn,
+        struct s2n_async_offload_cb_test *data)
+{
+    for (size_t retry = 0; retry < 20; retry++) {
+        int ret_val = s2n_negotiate_test_server_and_client(server_conn, client_conn);
+
+        if (ret_val != S2N_SUCCESS && s2n_errno == S2N_ERR_ASYNC_BLOCKED) {
+            /* Handshake remains blocked as long as op_perform() is not invoked. */
+            EXPECT_FAILURE_WITH_ERRNO(s2n_negotiate_test_server_and_client(server_conn, client_conn),
+                    S2N_ERR_ASYNC_BLOCKED);
+
+            EXPECT_SUCCESS(s2n_async_offload_op_perform(data->last_seen_op));
+            /* Each operation can only be performed once. */
+            EXPECT_FAILURE_WITH_ERRNO(s2n_async_offload_op_perform(data->last_seen_op), S2N_ERR_INVALID_STATE);
+        } else {
+            return ret_val;
+        }
+    }
+
+    FAIL_MSG("Async offload operation does not terminate");
+}
+
+int main(int argc, char *argv[])
+{
+    BEGIN_TEST();
+
+    DEFER_CLEANUP(struct s2n_cert_chain_and_key *chain_and_key = NULL, s2n_cert_chain_and_key_ptr_free);
+    EXPECT_SUCCESS(s2n_test_cert_chain_and_key_new(&chain_and_key,
+            S2N_DEFAULT_TEST_CERT_CHAIN, S2N_DEFAULT_TEST_PRIVATE_KEY));
+
+    /* Safety Check */
+    {
+        struct s2n_async_offload_cb_test test_data = { 0 };
+        EXPECT_FAILURE_WITH_ERRNO(
+                s2n_config_set_async_offload_callback(NULL, S2N_ASYNC_OFFLOAD_ALLOW_ALL,
+                        s2n_async_offload_test_callback, &test_data),
+                S2N_ERR_NULL);
+
+        DEFER_CLEANUP(struct s2n_config *test_config = s2n_config_new(), s2n_config_ptr_free);
+        EXPECT_NOT_NULL(test_config);
+        EXPECT_EQUAL(test_config->async_offload_allow_list, S2N_ASYNC_OFFLOAD_OP_NONE);
+        EXPECT_FAILURE_WITH_ERRNO(
+                s2n_config_set_async_offload_callback(test_config, S2N_ASYNC_OFFLOAD_ALLOW_ALL,
+                        NULL, &test_data),
+                S2N_ERR_NULL);
+
+        EXPECT_SUCCESS(s2n_config_set_async_offload_callback(test_config, S2N_ASYNC_OFFLOAD_PKEY_VERIFY,
+                s2n_async_offload_test_callback, &test_data));
+        EXPECT_TRUE(s2n_async_offload_op_is_in_allow_list(test_config, S2N_ASYNC_OFFLOAD_PKEY_VERIFY));
+
+        EXPECT_FAILURE_WITH_ERRNO(s2n_async_offload_op_perform(NULL), S2N_ERR_NULL);
+        struct s2n_async_offload_op test_op = { 0 };
+        EXPECT_FAILURE_WITH_ERRNO(s2n_async_offload_op_perform(&test_op), S2N_ERR_INVALID_STATE);
+    }
+
+    /* clang-format off */
+    struct s2n_async_offload_test_case {
+        bool async_test;
+        s2n_async_offload_op_type allow_list;
+        int cb_return;
+        int cb_invoked;
+        bool client_auth;
+        s2n_error expected_error;
+    } test_cases[] = {
+        /* Default option: no op type is allowed. */
+        {
+            .async_test = false,
+            .allow_list = S2N_ASYNC_OFFLOAD_OP_NONE,
+            .cb_return = S2N_SUCCESS,
+            .cb_invoked = 0,
+            .client_auth = true,
+            .expected_error = S2N_ERR_OK,
+        },
+        /* Test a random value that has not been used by any op type. */
+        /* Changing return value does not fail the handshake because the callback is not invoked. */
+        {
+            .async_test = false,
+            .allow_list = 0x70000000,
+            .cb_return = S2N_FAILURE,
+            .cb_invoked = 0,
+            .client_auth = true,
+            .expected_error = S2N_ERR_OK,
+        },
+        /* Async PKEY_VERIFY allowed. Client auth enabled. A successful handshake performs pkey_verify() twice. */
+        {
+            .async_test = false,
+            .allow_list = S2N_ASYNC_OFFLOAD_PKEY_VERIFY,
+            .cb_return = S2N_SUCCESS,
+            .cb_invoked = 2,
+            .client_auth = true,
+            .expected_error = S2N_ERR_OK,
+        },
+        /* Any op type is allowed. Handshake failed because the callback failed in the first attempt. */
+        {
+            .async_test = false,
+            .allow_list = S2N_ASYNC_OFFLOAD_ALLOW_ALL,
+            .cb_return = S2N_FAILURE,
+            .cb_invoked = 1,
+            .client_auth = true,
+            .expected_error = S2N_ERR_CANCELLED,
+        },
+        /* Client auth is enabled. pkey_verify() is performed by both server side and client side. */
+        {
+            .async_test = true,
+            .allow_list = S2N_ASYNC_OFFLOAD_PKEY_VERIFY,
+            .cb_return = S2N_SUCCESS,
+            .cb_invoked = 2,
+            .client_auth = true,
+            .expected_error = S2N_ERR_ASYNC_BLOCKED,
+        },
+        /* Client auth is not enabled. pkey_verify() is performed only by client side. */
+        {
+            .async_test = true,
+            .allow_list = S2N_ASYNC_OFFLOAD_ALLOW_ALL,
+            .cb_return = S2N_SUCCESS,
+            .cb_invoked = 1,
+            .client_auth = false,
+            .expected_error = S2N_ERR_ASYNC_BLOCKED,
+        },
+    };
+    /* clang-format on */
+
+    /* Test with both TLS 1.2 and TLS 1.3 policies */
+    const char *versions[] = { "20240501", "default_tls13" };
+
+    for (int test_idx = 0; test_idx < s2n_array_len(test_cases); test_idx++) {
+        for (int version_idx = 0; version_idx < s2n_array_len(versions); version_idx++) {
+            DEFER_CLEANUP(struct s2n_config *config = s2n_config_new(), s2n_config_ptr_free);
+            EXPECT_NOT_NULL(config);
+            EXPECT_SUCCESS(s2n_config_add_cert_chain_and_key_to_store(config, chain_and_key));
+            EXPECT_SUCCESS(s2n_config_set_verification_ca_location(config, S2N_DEFAULT_TEST_CERT_CHAIN, NULL));
+            EXPECT_SUCCESS(s2n_config_set_cipher_preferences(config, versions[version_idx]));
+
+            struct s2n_async_offload_test_case test_case = test_cases[test_idx];
+            if (test_case.client_auth) {
+                EXPECT_SUCCESS(s2n_config_set_client_auth_type(config, S2N_CERT_AUTH_REQUIRED));
+            }
+
+            struct s2n_async_offload_cb_test data = { .async_test = test_case.async_test, .result = test_case.cb_return };
+            if (test_case.allow_list != S2N_ASYNC_OFFLOAD_OP_NONE) {
+                EXPECT_SUCCESS(s2n_config_set_async_offload_callback(config, test_case.allow_list,
+                        s2n_async_offload_test_callback, &data));
+            }
+
+            DEFER_CLEANUP(struct s2n_connection *server_conn = s2n_connection_new(S2N_SERVER), s2n_connection_ptr_free);
+            EXPECT_NOT_NULL(server_conn);
+            EXPECT_SUCCESS(s2n_connection_set_config(server_conn, config));
+
+            DEFER_CLEANUP(struct s2n_connection *client_conn = s2n_connection_new(S2N_CLIENT), s2n_connection_ptr_free);
+            EXPECT_NOT_NULL(client_conn);
+            EXPECT_SUCCESS(s2n_connection_set_config(client_conn, config));
+            EXPECT_SUCCESS(s2n_connection_set_blinding(client_conn, S2N_SELF_SERVICE_BLINDING));
+            EXPECT_SUCCESS(s2n_set_server_name(client_conn, "localhost"));
+
+            DEFER_CLEANUP(struct s2n_test_io_pair io_pair = { 0 }, s2n_io_pair_close);
+            EXPECT_SUCCESS(s2n_io_pair_init_non_blocking(&io_pair));
+            EXPECT_SUCCESS(s2n_connection_set_io_pair(client_conn, &io_pair));
+            EXPECT_SUCCESS(s2n_connection_set_io_pair(server_conn, &io_pair));
+
+            s2n_error expected_error = test_case.expected_error;
+            if (test_case.async_test) {
+                EXPECT_SUCCESS(s2n_test_handshake_async(server_conn, client_conn, &data));
+            } else if (expected_error == S2N_ERR_OK) {
+                EXPECT_SUCCESS(s2n_negotiate_test_server_and_client(server_conn, client_conn));
+            } else {
+                EXPECT_FAILURE_WITH_ERRNO(s2n_negotiate_test_server_and_client(server_conn, client_conn), expected_error);
+            }
+            EXPECT_EQUAL(data.invoked_count, test_case.cb_invoked);
+
+            if (s2n_is_tls13_fully_supported() && version_idx == 1) {
+                EXPECT_EQUAL(s2n_connection_get_actual_protocol_version(server_conn), S2N_TLS13);
+            } else {
+                EXPECT_EQUAL(s2n_connection_get_actual_protocol_version(server_conn), S2N_TLS12);
+            }
+        }
+    }
+
+    END_TEST();
+}

tests/unit/s2n_connection_size_test.c

@@ -44,7 +44,7 @@ int main(int argc, char **argv)
     }
 
     /* Carefully consider any increases to this number. */
-    const uint16_t max_connection_size = 4350;
+    const uint16_t max_connection_size = 4360;
     const uint16_t min_connection_size = max_connection_size * 0.9;
 
     size_t connection_size = sizeof(struct s2n_connection);

tests/unit/s2n_self_talk_offload_signing_test.c

@@ -60,7 +60,7 @@ static S2N_RESULT s2n_test_pkey_sign(const struct s2n_pkey *pkey,
     return S2N_RESULT_OK;
 }
 
-static S2N_RESULT s2n_async_pkey_sign(struct s2n_cert_chain_and_key *complete_chain)
+static S2N_RESULT s2n_async_pkey_sign_test(struct s2n_cert_chain_and_key *complete_chain)
 {
     RESULT_ENSURE_REF(pkey_op);
     RESULT_ENSURE_REF(pkey_op_conn);
@@ -134,7 +134,7 @@ static S2N_RESULT s2n_do_test_handshake(struct s2n_config *config, struct s2n_ce
 
     while (s2n_negotiate_test_server_and_client(server_conn, client_conn) != S2N_SUCCESS) {
         EXPECT_EQUAL(s2n_errno, S2N_ERR_ASYNC_BLOCKED);
-        RESULT_GUARD(s2n_async_pkey_sign(complete_chain));
+        RESULT_GUARD(s2n_async_pkey_sign_test(complete_chain));
     }
 
     RESULT_ENSURE_EQ(client_conn->actual_protocol_version, expected_protocol_version);