aws-crt 0.1.4 → 0.1.5

data/aws-crt-ffi/crt/s2n/tests/unit/s2n_fork_generation_number_test.c

@@ -0,0 +1,335 @@
+/*
+ * 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.
+ */
+
+/* For clone() */
+#define _GNU_SOURCE
+
+#include "s2n_test.h"
+#include "utils/s2n_fork_detection.h"
+
+#include <pthread.h>
+#include <sched.h>
+#include <stdio.h>
+#include <sys/wait.h>
+
+#define NUMBER_OF_FGN_TEST_CASES 4
+#define MAX_NUMBER_OF_TEST_THREADS 2
+#define FORK_LEVEL_FOR_TESTS 2
+/* Before calling s2n_get_fork_generation_number() set the argument to this
+ * value to avoid any unlucky collisions
+ */
+#define UNEXPECTED_RETURNED_FGN 0xFF
+
+#define CLONE_TEST_NO 0
+#define CLONE_TEST_YES 1
+#define CLONE_TEST_DETERMINE_AT_RUNTIME 2
+
+struct fgn_test_case {
+    const char *test_case_label;
+    int (*test_case_cb)(struct fgn_test_case *test_case);
+    int test_case_must_pass_clone_test;
+};
+
+static void s2n_verify_child_exit_status(pid_t proc_pid)
+{
+    int status = 0;
+    EXPECT_EQUAL(waitpid(proc_pid, &status, __WALL), proc_pid);
+
+    /* Check that child exited with EXIT_SUCCESS. If not, this indicates
+     * that an error was encountered in the unit tests executed in that
+     * child process.
+     */
+    EXPECT_NOT_EQUAL(WIFEXITED(status), 0);
+    EXPECT_EQUAL(WEXITSTATUS(status), EXIT_SUCCESS);
+}
+
+static void * s2n_unit_test_thread_get_fgn(void *expected_fork_generation_number)
+{
+    uint64_t return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+    EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+    EXPECT_EQUAL(return_fork_generation_number, *(uint64_t *) expected_fork_generation_number);
+
+    return NULL;
+}
+
+static int s2n_unit_test_thread(uint64_t expected_fork_generation_number)
+{
+    pthread_t threads[MAX_NUMBER_OF_TEST_THREADS];
+
+    for (size_t thread_index = 0; thread_index < MAX_NUMBER_OF_TEST_THREADS; thread_index++) {
+        EXPECT_EQUAL(pthread_create(&threads[thread_index], NULL, &s2n_unit_test_thread_get_fgn, (void *) &expected_fork_generation_number), 0);
+    }
+
+    /* Wait for all threads to finish */
+    for (size_t thread_index = 0; thread_index < MAX_NUMBER_OF_TEST_THREADS; thread_index++) {
+       pthread_join(threads[thread_index], NULL);
+    }
+
+    return S2N_SUCCESS;
+}
+
+static int s2n_unit_test_fork(uint64_t parent_process_fgn, int fork_level)
+{
+    pid_t proc_pid = fork();
+    EXPECT_TRUE(proc_pid >= 0);
+
+    fork_level = fork_level - 1;
+
+    if (proc_pid == 0) {
+        /* In child */
+        uint64_t return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+        EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+        EXPECT_EQUAL(return_fork_generation_number, parent_process_fgn + 1);
+
+        /* Verify stability */
+        return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+        EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+        EXPECT_EQUAL(return_fork_generation_number, parent_process_fgn + 1);
+
+        /* Verify in threads */
+        EXPECT_EQUAL(s2n_unit_test_thread(return_fork_generation_number), S2N_SUCCESS);
+
+        if (fork_level > 0) {
+            /* Fork again and verify fork generation number */
+            EXPECT_EQUAL(s2n_unit_test_fork(parent_process_fgn + 1, fork_level), S2N_SUCCESS);
+        }
+
+        /* Exit code EXIT_SUCCESS means that tests in this process finished
+         * successfully. Any errors would have exited the process with an
+         * exit code != EXIT_SUCCESS. We verify this in the parent process.
+         */
+        exit(EXIT_SUCCESS);
+    }
+    else {
+        s2n_verify_child_exit_status(proc_pid);
+
+        /* Verify stability */
+        uint64_t return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+        EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+        EXPECT_EQUAL(return_fork_generation_number, parent_process_fgn);
+    }
+
+    return S2N_SUCCESS;
+}
+
+/* Similar test to unit_test_fork() but verify in threads first */
+static int s2n_unit_test_fork_check_threads_first(uint64_t parent_process_fgn)
+{
+    pid_t proc_pid = fork();
+    EXPECT_TRUE(proc_pid >= 0);
+
+    if (proc_pid == 0) {
+        /* In child. Verify threads first. */
+        EXPECT_EQUAL(s2n_unit_test_thread(parent_process_fgn + 1), S2N_SUCCESS);
+
+        /* Then in the thread spawned when forking */
+        uint64_t return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+        EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+        EXPECT_EQUAL(return_fork_generation_number, parent_process_fgn + 1);
+
+        /* Verify stability */
+        return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+        EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+        EXPECT_EQUAL(return_fork_generation_number, parent_process_fgn + 1);
+
+        /* Exit code EXIT_SUCCESS means that tests in this process finished
+         * successfully. Any errors would have exited the process with an
+         * exit code != EXIT_SUCCESS. We verify this in the parent process.
+         */
+        exit(EXIT_SUCCESS);
+    }
+    else {
+        s2n_verify_child_exit_status(proc_pid);
+
+        /* Verify stability */
+        uint64_t return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+        EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+        EXPECT_EQUAL(return_fork_generation_number, parent_process_fgn);
+    }
+
+    return S2N_SUCCESS;
+}
+
+static int s2n_unit_test_clone_child_process(void *parent_process_fgn)
+{
+    /* In child */
+    uint64_t local_parent_process_fgn = *(uint64_t *) parent_process_fgn;
+    uint64_t return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+    EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+    EXPECT_EQUAL(return_fork_generation_number, local_parent_process_fgn + 1);
+
+    /* Verify stability */
+    return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+    EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+    EXPECT_EQUAL(return_fork_generation_number, local_parent_process_fgn + 1);
+
+    /* Verify in threads */
+    EXPECT_EQUAL(s2n_unit_test_thread(return_fork_generation_number), S2N_SUCCESS);
+
+    /* This translates to the exit code for this child process */
+    return EXIT_SUCCESS;
+}
+
+#define PROCESS_CHILD_STACK_SIZE (1024 * 1024) /* Suggested by clone() man page... */
+static int s2n_unit_test_clone(uint64_t parent_process_fgn)
+{
+#if defined(S2N_CLONE_SUPPORTED)
+    /* Verify stability */
+    uint64_t return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+    EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+    EXPECT_EQUAL(return_fork_generation_number, parent_process_fgn);
+
+    /* Use stack memory for this... We don't exit unit_test_clone() before this
+     * memory has served its purpose.
+     * Why? Using dynamically allocated memory causes Valgrind to squat on the
+     * allocated memory when the child process exists.
+     */
+    char process_child_stack[PROCESS_CHILD_STACK_SIZE];
+    EXPECT_NOT_NULL(process_child_stack);
+
+    int proc_pid = clone(s2n_unit_test_clone_child_process, (void *) (process_child_stack + PROCESS_CHILD_STACK_SIZE), 0, (void *) &return_fork_generation_number);
+    EXPECT_NOT_EQUAL(proc_pid, -1);
+
+    s2n_verify_child_exit_status(proc_pid);
+
+    /* Verify stability */
+    return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+    EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+    EXPECT_EQUAL(return_fork_generation_number, parent_process_fgn);
+#endif
+
+    return S2N_SUCCESS;
+}
+
+static int s2n_unit_tests_common(struct fgn_test_case *test_case)
+{
+    uint64_t return_fork_generation_number = 0;
+
+    EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+    EXPECT_EQUAL(return_fork_generation_number, 0);
+
+    /* Should be idempotent if no fork event occurred */
+    return_fork_generation_number = UNEXPECTED_RETURNED_FGN;
+    EXPECT_OK(s2n_get_fork_generation_number(&return_fork_generation_number));
+    EXPECT_EQUAL(return_fork_generation_number, 0);
+
+    /* Should be idempotent in threaded environment as well */
+    EXPECT_EQUAL(s2n_unit_test_thread(return_fork_generation_number), S2N_SUCCESS);
+
+    /* Cached FGN should increment if a fork event occurs */
+    EXPECT_EQUAL(s2n_unit_test_fork(return_fork_generation_number, FORK_LEVEL_FOR_TESTS), S2N_SUCCESS);
+    EXPECT_EQUAL(s2n_unit_test_fork_check_threads_first(return_fork_generation_number), S2N_SUCCESS);
+
+    /* Some fork detection mechanisms can also detect forks through clone() */
+    if (test_case->test_case_must_pass_clone_test == CLONE_TEST_YES) {
+        EXPECT_EQUAL((s2n_is_madv_wipeonfork_supported() == true) || (s2n_is_map_inherit_zero_supported() == true), true);
+        EXPECT_EQUAL(s2n_unit_test_clone(return_fork_generation_number), S2N_SUCCESS);
+    }
+    else if (test_case->test_case_must_pass_clone_test == CLONE_TEST_DETERMINE_AT_RUNTIME) {
+        if ((s2n_is_madv_wipeonfork_supported() == true) ||
+            (s2n_is_map_inherit_zero_supported() == true)) {
+            EXPECT_EQUAL(s2n_unit_test_clone(return_fork_generation_number), S2N_SUCCESS);
+        }
+    }
+
+    return S2N_SUCCESS;
+}
+
+static int s2n_test_case_default_cb(struct fgn_test_case *test_case)
+{
+    EXPECT_EQUAL(s2n_unit_tests_common(test_case), S2N_SUCCESS);
+
+    return S2N_SUCCESS;
+}
+
+static int s2n_test_case_pthread_atfork_cb(struct fgn_test_case *test_case)
+{
+    POSIX_GUARD_RESULT(s2n_ignore_wipeonfork_and_inherit_zero_for_testing());
+    EXPECT_EQUAL(s2n_unit_tests_common(test_case), S2N_SUCCESS);
+
+    return S2N_SUCCESS;
+}
+
+static int s2n_test_case_madv_wipeonfork_cb(struct fgn_test_case *test_case)
+{
+    if (s2n_is_madv_wipeonfork_supported() == false) {
+        TEST_DEBUG_PRINT("s2n_fork_generation_number_test.c test case not supported. Skipping.\nTest case: %s\n", test_case->test_case_label);
+        return S2N_SUCCESS;
+    }
+
+    POSIX_GUARD_RESULT(s2n_ignore_pthread_atfork_for_testing());
+    EXPECT_EQUAL(s2n_unit_tests_common(test_case), S2N_SUCCESS);
+
+    return S2N_SUCCESS;
+}
+
+static int s2n_test_case_map_inherit_zero_cb(struct fgn_test_case *test_case)
+{
+    if (s2n_is_map_inherit_zero_supported() == false) {
+        TEST_DEBUG_PRINT("s2n_fork_generation_number_test.c test case not supported. Skipping.\nTest case: %s\n", test_case->test_case_label);
+        return S2N_SUCCESS;
+    }
+
+    POSIX_GUARD_RESULT(s2n_ignore_pthread_atfork_for_testing());
+    EXPECT_EQUAL(s2n_unit_tests_common(test_case), S2N_SUCCESS);
+
+    return S2N_SUCCESS;
+}
+
+struct fgn_test_case fgn_test_cases[NUMBER_OF_FGN_TEST_CASES] = {
+    {"Default fork detect mechanisms.", s2n_test_case_default_cb, CLONE_TEST_DETERMINE_AT_RUNTIME},
+    {"Only pthread_atfork fork detection mechanism.", s2n_test_case_pthread_atfork_cb, CLONE_TEST_NO},
+    {"Only madv_wipeonfork fork detection mechanism.", s2n_test_case_madv_wipeonfork_cb, CLONE_TEST_YES},
+    {"Only map_inheret_zero fork detection mechanism.", s2n_test_case_map_inherit_zero_cb, CLONE_TEST_YES}
+};
+
+int main(int argc, char **argv)
+{
+    BEGIN_TEST();
+
+    EXPECT_TRUE(s2n_array_len(fgn_test_cases) == NUMBER_OF_FGN_TEST_CASES);
+
+    /* Create NUMBER_OF_FGN_TEST_CASES number of child processes that run each
+     * test case.
+     *
+     * Fork detection is lazily initialised on first invocation of
+     * s2n_get_fork_generation_number(). Hence, it is important that childs are
+     * created before calling into the fork detection code.
+     */
+    pid_t proc_pids[NUMBER_OF_FGN_TEST_CASES] = {0};
+
+    for (size_t i = 0; i < NUMBER_OF_FGN_TEST_CASES; i++) {
+
+        proc_pids[i] = fork();
+        EXPECT_TRUE(proc_pids[i] >= 0);
+
+        if (proc_pids[i] == 0) {
+            /* In child */
+            EXPECT_EQUAL(fgn_test_cases[i].test_case_cb(&fgn_test_cases[i]), S2N_SUCCESS);
+
+            /* Exit code EXIT_SUCCESS means that tests in this process finished
+             * successfully. Any errors would have exited the process with an
+             * exit code != EXIT_SUCCESS. We verify this in the parent process.
+             * Also prevents child from creating more childs.
+             */
+            exit(EXIT_SUCCESS);
+        }
+        else {
+            s2n_verify_child_exit_status(proc_pids[i]);
+        }
+    }
+
+    END_TEST();
+}

data/aws-crt-ffi/crt/s2n/tests/unit/s2n_mem_usage_test.c

@@ -47,7 +47,7 @@
 #ifdef __FreeBSD__
 #define MEM_PER_CONNECTION 57
 #else
-#define MEM_PER_CONNECTION 48
+#define MEM_PER_CONNECTION 49
 #endif
 
 /* This is the maximum memory per connection including 4KB of slack */

data/aws-crt-ffi/crt/s2n/tests/unit/s2n_self_talk_client_hello_cb_test.c

@@ -24,6 +24,7 @@
 
 #include "api/s2n.h"
 #include "tls/s2n_connection.h"
+#include "tls/s2n_internal.h"
 
 struct client_hello_context {
     int invoked;
@@ -38,6 +39,7 @@ struct client_hello_context {
      * this flag tests the previous behavior from blocking callbacks
      */
     int legacy_rc_for_server_name_used;
+    bool mark_done;
 };
 
 int mock_client(struct s2n_test_io_pair *io_pair, int expect_failure, int expect_server_name_used)
@@ -181,9 +183,27 @@ int client_hello_fail_handshake(struct s2n_connection *conn, void *ctx)
     /* Return negative value to terminate the handshake */
     return -1;
 
-} 
+}
+
+int s2n_client_hello_poll_cb(struct s2n_connection *conn, void *ctx)
+{
+    struct client_hello_context *client_hello_ctx;
+    if (ctx == NULL) {
+        return -1;
+    }
+    client_hello_ctx = ctx;
+    /* Increment counter to ensure that callback was invoked */
+    client_hello_ctx->invoked++;
+
+    if (client_hello_ctx->mark_done) {
+        EXPECT_SUCCESS(s2n_client_hello_cb_done(conn));
+        return S2N_SUCCESS;
+    }
+
+    return S2N_SUCCESS;
+}
 
-int s2n_negotiate_nonblocking_ch_cb(struct s2n_connection *conn, 
+int s2n_negotiate_nonblocking_ch_cb(struct s2n_connection *conn,
     struct client_hello_context *ch_ctx, bool server_name_used)
 {
     s2n_blocked_status blocked;
@@ -191,7 +211,7 @@ int s2n_negotiate_nonblocking_ch_cb(struct s2n_connection *conn,
     /* negotiate handshake, we should pause after the nonblocking callback is invoked */
     EXPECT_FAILURE_WITH_ERRNO(s2n_negotiate(conn, &blocked), S2N_ERR_ASYNC_BLOCKED);
     EXPECT_EQUAL(blocked, S2N_BLOCKED_ON_APPLICATION_INPUT);
-    
+
     /* verify client hello cb has been invoked */
     EXPECT_EQUAL(ch_ctx->invoked, 1);
 
@@ -201,7 +221,7 @@ int s2n_negotiate_nonblocking_ch_cb(struct s2n_connection *conn,
     }
     /* unless explicitly unblocked we should stay paused */
     EXPECT_FAILURE_WITH_ERRNO(s2n_negotiate(conn, &blocked), S2N_ERR_ASYNC_BLOCKED);
-    EXPECT_EQUAL(blocked, S2N_BLOCKED_ON_APPLICATION_INPUT);    
+    EXPECT_EQUAL(blocked, S2N_BLOCKED_ON_APPLICATION_INPUT);
 
     /* mark the client hello cb complete */
     EXPECT_SUCCESS(s2n_client_hello_cb_done(conn));
@@ -211,11 +231,39 @@ int s2n_negotiate_nonblocking_ch_cb(struct s2n_connection *conn,
     return s2n_negotiate(conn, &blocked);
 }
 
+int s2n_negotiate_nonblocking_poll(struct s2n_connection *conn,
+    struct client_hello_context *ch_ctx)
+{
+    EXPECT_NOT_NULL(conn);
+    EXPECT_NOT_NULL(ch_ctx);
+    int invoked = 0;
+    s2n_blocked_status blocked = S2N_NOT_BLOCKED;
+
+    EXPECT_EQUAL(ch_ctx->invoked, 0);
+
+    do {
+        /* negotiate handshake, we should pause after the nonblocking callback is invoked */
+        EXPECT_FAILURE_WITH_ERRNO(s2n_negotiate(conn, &blocked), S2N_ERR_ASYNC_BLOCKED);
+        EXPECT_EQUAL(blocked, S2N_BLOCKED_ON_APPLICATION_INPUT);
+        invoked++;
+        EXPECT_EQUAL(ch_ctx->invoked, invoked);
+    } while(invoked < 10);
+    EXPECT_EQUAL(ch_ctx->invoked, invoked);
+
+    ch_ctx->mark_done = true;
+
+    /* Expect the callback to complete after 2nd iteration */
+    EXPECT_SUCCESS(s2n_negotiate(conn, &blocked));
+    EXPECT_EQUAL(ch_ctx->invoked, invoked + 1);
+
+    return S2N_SUCCESS;
+}
+
 int s2n_negotiate_blocking_ch_cb(struct s2n_connection *conn, struct client_hello_context *ch_ctx)
 {
     s2n_blocked_status blocked;
     EXPECT_NOT_NULL(conn);
-    
+
     int rc = s2n_negotiate(conn, &blocked);
     /* verify client hello cb has been invoked */
     EXPECT_EQUAL(ch_ctx->invoked, 1);
@@ -329,7 +377,7 @@ int run_test_config_swap_ch_cb(s2n_client_hello_cb_mode cb_mode,
 
     EXPECT_SUCCESS(start_client_conn(&io_pair, &pid, 0 , 1));
     EXPECT_SUCCESS(init_server_conn(&conn, &io_pair, config));
- 
+
     /* do the handshake */
     if ( cb_mode == S2N_CLIENT_HELLO_CB_NONBLOCKING && !ch_ctx->mark_done_during_callback) {
         /* swap the config and mark server_name_used in the async context */
@@ -340,7 +388,7 @@ int run_test_config_swap_ch_cb(s2n_client_hello_cb_mode cb_mode,
             */
         EXPECT_SUCCESS(s2n_negotiate_blocking_ch_cb(conn, ch_ctx));
     }
- 
+
     /* Server name and error are as expected with null connection */
     EXPECT_NULL(s2n_get_server_name(NULL));
     EXPECT_EQUAL(s2n_errno, S2N_ERR_NULL);
@@ -349,7 +397,7 @@ int run_test_config_swap_ch_cb(s2n_client_hello_cb_mode cb_mode,
     EXPECT_STRING_EQUAL(s2n_get_application_protocol(conn), protocols[0]);
 
     EXPECT_SUCCESS(server_recv(conn));
- 
+
     EXPECT_SUCCESS(test_case_clean(conn, pid, config, &io_pair, ch_ctx));
     EXPECT_SUCCESS(s2n_config_free(swap_config));
     return S2N_SUCCESS;
@@ -372,7 +420,7 @@ int run_test_no_config_swap_ch_cb(s2n_client_hello_cb_mode cb_mode,
     EXPECT_SUCCESS(s2n_config_set_client_hello_cb_mode(config, cb_mode));
     EXPECT_SUCCESS(start_client_conn(&io_pair, &pid, 0 , 0));
     EXPECT_SUCCESS(init_server_conn(&conn, &io_pair, config));
- 
+
     /* do the handshake */
     if ( cb_mode == S2N_CLIENT_HELLO_CB_NONBLOCKING ) {
         /* swap the config and mark server_name_used in the async context */
@@ -387,7 +435,7 @@ int run_test_no_config_swap_ch_cb(s2n_client_hello_cb_mode cb_mode,
     EXPECT_EQUAL(s2n_errno, S2N_ERR_NULL);
 
     EXPECT_SUCCESS(server_recv(conn));
- 
+
     EXPECT_SUCCESS(test_case_clean(conn, pid, config, &io_pair, ch_ctx));
     return S2N_SUCCESS;
 }
@@ -425,7 +473,7 @@ int run_test_reject_handshake_ch_cb(s2n_client_hello_cb_mode cb_mode,
 
     /* Ensure that callback was invoked */
     EXPECT_EQUAL(ch_ctx->invoked, 1);
- 
+
     /* shutdown to flush alert, expext failure as client doesn't send close notify */
     EXPECT_FAILURE(s2n_shutdown(conn, &blocked));
     EXPECT_SUCCESS(s2n_connection_free(conn));
@@ -434,6 +482,37 @@ int run_test_reject_handshake_ch_cb(s2n_client_hello_cb_mode cb_mode,
     return S2N_SUCCESS;
 }
 
+int run_test_poll_ch_cb(s2n_client_hello_cb_mode cb_mode,
+    struct s2n_cert_chain_and_key *chain_and_key,
+    struct client_hello_context *ch_ctx)
+{
+    struct s2n_test_io_pair io_pair = { 0 };
+    struct s2n_config *config = s2n_config_new();
+    EXPECT_NOT_NULL(config);
+    struct s2n_connection *conn;
+    pid_t pid = 0;
+
+    EXPECT_SUCCESS(s2n_config_add_cert_chain_and_key_to_store(config, chain_and_key));
+
+    /* Setup ClientHello callback */
+    EXPECT_SUCCESS(s2n_config_set_client_hello_cb(config, s2n_client_hello_poll_cb, ch_ctx));
+    EXPECT_SUCCESS(s2n_config_set_client_hello_cb_mode(config, cb_mode));
+
+    /* Enable callback polling mode */
+    EXPECT_SUCCESS(s2n_config_client_hello_cb_enable_poll(config));
+
+    EXPECT_SUCCESS(start_client_conn(&io_pair, &pid, 0 , 0));
+    EXPECT_SUCCESS(init_server_conn(&conn, &io_pair, config));
+
+    /* negotiate and make assertions */
+    EXPECT_SUCCESS(s2n_negotiate_nonblocking_poll(conn, ch_ctx));
+
+    EXPECT_SUCCESS(server_recv(conn));
+
+    EXPECT_SUCCESS(test_case_clean(conn, pid, config, &io_pair, ch_ctx));
+    return S2N_SUCCESS;
+}
+
 int main(int argc, char **argv)
 {
     struct client_hello_context client_hello_ctx = {0};
@@ -487,6 +566,9 @@ int main(int argc, char **argv)
     EXPECT_SUCCESS(run_test_reject_handshake_ch_cb(S2N_CLIENT_HELLO_CB_NONBLOCKING,
        chain_and_key, &client_hello_ctx));
 
+    EXPECT_SUCCESS(run_test_poll_ch_cb(S2N_CLIENT_HELLO_CB_NONBLOCKING,
+       chain_and_key, &client_hello_ctx));
+
     EXPECT_SUCCESS(s2n_cert_chain_and_key_free(chain_and_key));
     free(cert_chain_pem);
     free(private_key_pem);