feat: Add key update to ktls feature (#5484)

maddeleine · lrstewart · web-flow · commit 30f40f2345a8 · 2025-09-24T20:15:54.000Z
Co-authored-by: Lindsay Stewart &lt;stewart.r.lindsay@gmail.com&gt;
diff --git a/codebuild/spec/buildspec_ktls_keyupdate.yml b/codebuild/spec/buildspec_ktls_keyupdate.yml
@@ -25,6 +25,7 @@ env:
     NIX_CACHE_BUCKET: "s3://s2n-tls-nixcachebucket-x86-64?region=us-west-2"
     NIX_INSTALLER: "https://nixos.org/nix/install"
     S2N_KTLS_TESTING_EXPECTED: 1
+    S2N_KTLS_KEYUPDATE_TESTING_EXPECTED: 1
 phases:
   pre_build:
     commands:
diff --git a/error/s2n_errno.c b/error/s2n_errno.c
@@ -309,6 +309,7 @@ static const char *no_such_error = "Internal s2n error";
     ERR_ENTRY(S2N_ERR_KTLS_RENEG, "kTLS does not support secure renegotiation") \
     ERR_ENTRY(S2N_ERR_KTLS_KEYUPDATE, "Received KeyUpdate from peer, but kernel does not support updating tls keys") \
     ERR_ENTRY(S2N_ERR_KTLS_KEY_LIMIT, "Reached key encryption limit, but kernel does not support updating tls keys") \
+    ERR_ENTRY(S2N_ERR_KTLS_SOCKOPT, "A call to sockopt failed when attempting to update the keys in the kernel") \
     ERR_ENTRY(S2N_ERR_UNEXPECTED_CERT_REQUEST, "Client forbids mutual authentication, but server requested a cert") \
     ERR_ENTRY(S2N_ERR_MISSING_CERT_REQUEST, "Client requires mutual authentication, but server did not request a cert") \
     ERR_ENTRY(S2N_ERR_MISSING_CLIENT_CERT, "Server requires client certificate") \
diff --git a/error/s2n_errno.h b/error/s2n_errno.h
@@ -328,6 +328,7 @@ typedef enum {
     S2N_ERR_KTLS_ENABLE,
     S2N_ERR_KTLS_BAD_CMSG,
     S2N_ERR_KTLS_RENEG,
+    S2N_ERR_KTLS_SOCKOPT,
     S2N_ERR_ATOMIC,
     S2N_ERR_KTLS_KEY_LIMIT,
     S2N_ERR_SECURITY_POLICY_INCOMPATIBLE_CERT,
diff --git a/tests/unit/s2n_key_update_test.c b/tests/unit/s2n_key_update_test.c
@@ -148,7 +148,7 @@ int main(int argc, char **argv)
             EXPECT_SUCCESS(s2n_connection_free(conn));
         };
 
-        /* Key update messages not allowed with ktls */
+        /* Key update messages not allowed with ktls by default */
         {
             DEFER_CLEANUP(struct s2n_stuffer input, s2n_stuffer_free);
             EXPECT_SUCCESS(s2n_stuffer_growable_alloc(&input, 0));
@@ -175,8 +175,10 @@ int main(int argc, char **argv)
             conn->ktls_send_enabled = true;
             conn->ktls_recv_enabled = false;
             EXPECT_SUCCESS(s2n_stuffer_write_uint8(&input, S2N_KEY_UPDATE_REQUESTED));
-            EXPECT_FAILURE_WITH_ERRNO(s2n_key_update_recv(conn, &input), S2N_ERR_KTLS_KEYUPDATE);
+            EXPECT_SUCCESS(s2n_key_update_recv(conn, &input));
             EXPECT_EQUAL(s2n_stuffer_data_available(&input), 0);
+            s2n_blocked_status blocked = S2N_NOT_BLOCKED;
+            EXPECT_FAILURE_WITH_ERRNO(s2n_key_update_send(conn, &blocked), S2N_ERR_KTLS_KEYUPDATE);
 
             /* Succeeds if only sending with ktls and no update requested:
              * No kernel key update would be required.
@@ -527,7 +529,7 @@ int main(int argc, char **argv)
             /* Fails if KeyUpdate required */
             EXPECT_FAILURE_WITH_ERRNO(
                     s2n_key_update_send(conn, &blocked),
-                    S2N_ERR_KTLS_KEY_LIMIT);
+                    S2N_ERR_KTLS_KEYUPDATE);
             EXPECT_TRUE(s2n_atomic_flag_test(&conn->key_update_pending));
         };
     };
diff --git a/tests/unit/s2n_ktls_test.c b/tests/unit/s2n_ktls_test.c
@@ -639,5 +639,76 @@ int main(int argc, char **argv)
         EXPECT_TRUE(config->ktls_tls13_enabled);
     };
 
+    /* Test: s2n_ktls_check_estimated_record_limit sets key_update_pending flag */
+    {
+        /* Create a cipher suite with an artificially lowered encryption limit */
+        const size_t test_encryption_limit = 3;
+        struct s2n_record_algorithm test_record_alg = *s2n_tls13_aes_128_gcm_sha256.record_alg;
+        test_record_alg.encryption_limit = test_encryption_limit;
+        struct s2n_cipher_suite test_cipher_suite = s2n_tls13_aes_128_gcm_sha256;
+        test_cipher_suite.record_alg = &test_record_alg;
+
+        DEFER_CLEANUP(struct s2n_connection *conn = s2n_connection_new(S2N_CLIENT),
+                s2n_connection_ptr_free);
+        EXPECT_NOT_NULL(conn);
+        DEFER_CLEANUP(struct s2n_config *config = s2n_config_new(), s2n_config_ptr_free);
+        EXPECT_SUCCESS(s2n_config_ktls_enable_unsafe_tls13(config));
+        EXPECT_SUCCESS(s2n_connection_set_config(conn, config));
+
+        EXPECT_NOT_NULL(conn->secure);
+        conn->secure->cipher_suite = &test_cipher_suite;
+        conn->actual_protocol_version = S2N_TLS13;
+
+        /* Test: 0 records sent so far */
+        {
+            /* Requested bytes exactly hit encryption limit */
+            size_t bytes_requested = S2N_TLS_MAXIMUM_FRAGMENT_LENGTH * test_encryption_limit;
+            EXPECT_OK(s2n_ktls_check_estimated_record_limit(conn, bytes_requested));
+            EXPECT_FALSE(s2n_atomic_flag_test(&conn->key_update_pending));
+
+            /* Requested bytes go over the encryption limit */
+            bytes_requested = (S2N_TLS_MAXIMUM_FRAGMENT_LENGTH * test_encryption_limit) + 1;
+            EXPECT_ERROR_WITH_ERRNO(s2n_ktls_check_estimated_record_limit(conn, bytes_requested), S2N_ERR_INVALID_ARGUMENT);
+            EXPECT_FALSE(s2n_atomic_flag_test(&conn->key_update_pending));
+        }
+
+        /* Test: 2 records sent so far */
+        {
+            EXPECT_OK(s2n_ktls_set_estimated_sequence_number(conn, S2N_TLS_MAXIMUM_FRAGMENT_LENGTH * 2));
+
+            /* Requested bytes exactly hit encryption limit */
+            size_t bytes_requested = S2N_TLS_MAXIMUM_FRAGMENT_LENGTH;
+            EXPECT_OK(s2n_ktls_check_estimated_record_limit(conn, bytes_requested));
+            EXPECT_FALSE(s2n_atomic_flag_test(&conn->key_update_pending));
+
+            /* Requested bytes go over the encryption limit */
+            bytes_requested = S2N_TLS_MAXIMUM_FRAGMENT_LENGTH + 1;
+            EXPECT_OK(s2n_ktls_check_estimated_record_limit(conn, bytes_requested));
+            EXPECT_TRUE(s2n_atomic_flag_test(&conn->key_update_pending));
+            s2n_atomic_flag_clear(&conn->key_update_pending);
+
+            /* Requested bytes go over the encryption limit twice */
+            bytes_requested = (S2N_TLS_MAXIMUM_FRAGMENT_LENGTH * test_encryption_limit) + 1;
+            EXPECT_ERROR_WITH_ERRNO(s2n_ktls_check_estimated_record_limit(conn, bytes_requested), S2N_ERR_INVALID_ARGUMENT);
+            EXPECT_FALSE(s2n_atomic_flag_test(&conn->key_update_pending));
+        }
+
+        /* Test: 3 records sent so far (exactly at encryption limit) */
+        {
+            EXPECT_OK(s2n_ktls_set_estimated_sequence_number(conn, S2N_TLS_MAXIMUM_FRAGMENT_LENGTH * 3));
+
+            /* Requested bytes go over the encryption limit */
+            size_t bytes_requested = 1;
+            EXPECT_OK(s2n_ktls_check_estimated_record_limit(conn, bytes_requested));
+            EXPECT_TRUE(s2n_atomic_flag_test(&conn->key_update_pending));
+            s2n_atomic_flag_clear(&conn->key_update_pending);
+
+            /* Requested bytes go over the encryption limit twice */
+            bytes_requested = (S2N_TLS_MAXIMUM_FRAGMENT_LENGTH * test_encryption_limit) + 1;
+            EXPECT_ERROR_WITH_ERRNO(s2n_ktls_check_estimated_record_limit(conn, bytes_requested), S2N_ERR_INVALID_ARGUMENT);
+            EXPECT_FALSE(s2n_atomic_flag_test(&conn->key_update_pending));
+        }
+    }
+
     END_TEST();
 }
diff --git a/tests/unit/s2n_self_talk_ktls_test.c b/tests/unit/s2n_self_talk_ktls_test.c
@@ -97,6 +97,7 @@ int main(int argc, char **argv)
      * where we think we're testing it.
      */
     const bool ktls_expected = (getenv("S2N_KTLS_TESTING_EXPECTED") != NULL);
+    const bool ktls_keyupdate_expected = (getenv("S2N_KTLS_KEYUPDATE_TESTING_EXPECTED") != NULL);
 
     if (!s2n_ktls_is_supported_on_platform() && !ktls_expected) {
         END_TEST();
@@ -460,22 +461,6 @@ int main(int argc, char **argv)
                 EXPECT_BYTEARRAY_EQUAL(test_data, buffer, read);
             }
         };
-
-        /* Test: receiving KeyUpdate message */
-        {
-            EXPECT_SUCCESS(s2n_connection_set_blinding(reader, S2N_SELF_SERVICE_BLINDING));
-
-            /* Write KeyUpdate message */
-            s2n_atomic_flag_set(&writer->key_update_pending);
-            int written = s2n_send(writer, test_data, sizeof(test_data), &blocked);
-            EXPECT_EQUAL(written, sizeof(test_data));
-
-            /* Read KeyUpdate message */
-            uint8_t buffer[sizeof(test_data)] = { 0 };
-            EXPECT_FAILURE_WITH_ERRNO(
-                    s2n_recv(reader, buffer, sizeof(buffer), &blocked),
-                    S2N_ERR_KTLS_KEYUPDATE);
-        };
     };
 
     /* Test: s2n_shutdown
@@ -698,5 +683,124 @@ int main(int argc, char **argv)
         }
     }
 
+    /* Test: Keyupdates with KTLS */
+    if (ktls_keyupdate_expected) {
+        /* Cipher suite with an artificially lowered encryption limit */
+        const size_t test_encryption_limit = 1;
+        struct s2n_record_algorithm test_record_alg = *s2n_tls13_aes_128_gcm_sha256.record_alg;
+        test_record_alg.encryption_limit = test_encryption_limit;
+        struct s2n_cipher_suite test_cipher_suite = s2n_tls13_aes_128_gcm_sha256;
+        test_cipher_suite.record_alg = &test_record_alg;
+
+        /* Test: Sending key update with KTLS */
+        if (ktls_send_supported) {
+            /* Test: Multiple key updates are allowed as long as they can be sent over multiple
+             * s2n_send calls with ktls. */
+            {
+                DEFER_CLEANUP(struct s2n_connection *client = s2n_connection_new(S2N_CLIENT),
+                        s2n_connection_ptr_free);
+                EXPECT_NOT_NULL(client);
+                EXPECT_SUCCESS(s2n_connection_set_config(client, config));
+                EXPECT_SUCCESS(s2n_connection_set_cipher_preferences(client, "default_tls13"));
+                EXPECT_SUCCESS(s2n_connection_set_blinding(client, S2N_SELF_SERVICE_BLINDING));
+
+                DEFER_CLEANUP(struct s2n_connection *server = s2n_connection_new(S2N_SERVER),
+                        s2n_connection_ptr_free);
+                EXPECT_NOT_NULL(server);
+                EXPECT_SUCCESS(s2n_connection_set_config(server, config));
+                EXPECT_SUCCESS(s2n_connection_set_cipher_preferences(server, "default_tls13"));
+                EXPECT_SUCCESS(s2n_connection_set_blinding(server, S2N_SELF_SERVICE_BLINDING));
+
+                DEFER_CLEANUP(struct s2n_test_io_pair io_pair = { 0 }, s2n_io_pair_close);
+                EXPECT_OK(s2n_new_inet_socket_pair(&io_pair));
+                EXPECT_OK(s2n_setup_connections(server, client, &io_pair));
+
+                EXPECT_SUCCESS(s2n_connection_ktls_enable_send(server));
+
+                /* Reset server cipher suite to trigger a key update after sending one record */
+                EXPECT_NOT_NULL(server->secure);
+                EXPECT_EQUAL(server->secure->cipher_suite, &s2n_tls13_aes_128_gcm_sha256);
+                server->secure->cipher_suite = &test_cipher_suite;
+
+                /* This will require a keyupdate mid-send, which is not allowed with ktls */
+                uint8_t exceeds_record_limit[S2N_TLS_MAXIMUM_FRAGMENT_LENGTH * 2] = { 0 };
+                s2n_blocked_status blocked = S2N_NOT_BLOCKED;
+                EXPECT_FAILURE_WITH_ERRNO(s2n_send(server, exceeds_record_limit, sizeof(exceeds_record_limit),
+                                                  &blocked),
+                        S2N_ERR_INVALID_ARGUMENT);
+
+                uint8_t large_test_data[S2N_TLS_MAXIMUM_FRAGMENT_LENGTH] = { "Hello there" };
+                /* Each send call will include one key update */
+                for (int i = 0; i < 10; i++) {
+                    int written = s2n_send(server, large_test_data, sizeof(large_test_data), &blocked);
+                    EXPECT_EQUAL(written, sizeof(large_test_data));
+                    EXPECT_EQUAL(blocked, S2N_NOT_BLOCKED);
+                }
+                EXPECT_EQUAL(server->recv_key_updated, 0);
+                EXPECT_EQUAL(server->send_key_updated, 9);
+
+                /* We only get one record per s2n_recv call, so we call it ten times */
+                for (size_t i = 0; i < 10; i++) {
+                    uint8_t buffer[sizeof(large_test_data)] = { 1 };
+                    int read = s2n_recv(client, buffer, sizeof(buffer), &blocked);
+                    EXPECT_EQUAL(read, sizeof(large_test_data));
+                    EXPECT_BYTEARRAY_EQUAL(large_test_data, buffer, read);
+                }
+
+                EXPECT_EQUAL(client->recv_key_updated, 9);
+                EXPECT_EQUAL(client->send_key_updated, 0);
+            }
+        };
+
+        /* Test: Receiving key update with KTLS */
+        if (ktls_recv_supported) {
+            DEFER_CLEANUP(struct s2n_connection *client = s2n_connection_new(S2N_CLIENT),
+                    s2n_connection_ptr_free);
+            EXPECT_NOT_NULL(client);
+            EXPECT_SUCCESS(s2n_connection_set_config(client, config));
+            EXPECT_SUCCESS(s2n_connection_set_cipher_preferences(client, "default_tls13"));
+            EXPECT_SUCCESS(s2n_connection_set_blinding(client, S2N_SELF_SERVICE_BLINDING));
+
+            DEFER_CLEANUP(struct s2n_connection *server = s2n_connection_new(S2N_SERVER),
+                    s2n_connection_ptr_free);
+            EXPECT_NOT_NULL(server);
+            EXPECT_SUCCESS(s2n_connection_set_config(server, config));
+            EXPECT_SUCCESS(s2n_connection_set_cipher_preferences(server, "default_tls13"));
+            EXPECT_SUCCESS(s2n_connection_set_blinding(server, S2N_SELF_SERVICE_BLINDING));
+
+            DEFER_CLEANUP(struct s2n_test_io_pair io_pair = { 0 }, s2n_io_pair_close);
+            EXPECT_OK(s2n_new_inet_socket_pair(&io_pair));
+            EXPECT_OK(s2n_setup_connections(server, client, &io_pair));
+
+            EXPECT_SUCCESS(s2n_connection_ktls_enable_recv(client));
+
+            /* Reset server cipher suite to trigger a key update after sending one record */
+            EXPECT_NOT_NULL(server->secure);
+            EXPECT_EQUAL(server->secure->cipher_suite, &s2n_tls13_aes_128_gcm_sha256);
+            server->secure->cipher_suite = &test_cipher_suite;
+
+            uint8_t large_test_data[S2N_DEFAULT_FRAGMENT_LENGTH * 10] = { "Hello there" };
+            s2n_blocked_status blocked = S2N_NOT_BLOCKED;
+            int written = s2n_send(server, large_test_data, sizeof(large_test_data), &blocked);
+            EXPECT_EQUAL(written, sizeof(large_test_data));
+            EXPECT_EQUAL(blocked, S2N_NOT_BLOCKED);
+
+            /* Sent 10 records and the encryption limit is 1 record so the send key will be updated
+             * 9 times. */
+            EXPECT_EQUAL(server->recv_key_updated, 0);
+            EXPECT_EQUAL(server->send_key_updated, 9);
+
+            /* We only get one record per s2n_recv call, so we call it ten times */
+            for (size_t i = 0; i < 10; i++) {
+                uint8_t buffer[sizeof(large_test_data) / 10] = { 1 };
+                int read = s2n_recv(client, buffer, sizeof(buffer), &blocked);
+                EXPECT_EQUAL(read, sizeof(large_test_data) / 10);
+                EXPECT_BYTEARRAY_EQUAL(buffer, large_test_data + (i * S2N_DEFAULT_FRAGMENT_LENGTH), read);
+            }
+
+            EXPECT_EQUAL(client->recv_key_updated, 9);
+            EXPECT_EQUAL(client->send_key_updated, 0);
+        };
+    }
     END_TEST();
 }
diff --git a/tls/s2n_key_update.c b/tls/s2n_key_update.c
@@ -18,6 +18,7 @@
 #include "crypto/s2n_sequence.h"
 #include "error/s2n_errno.h"
 #include "tls/s2n_connection.h"
+#include "tls/s2n_ktls.h"
 #include "tls/s2n_record.h"
 #include "tls/s2n_tls.h"
 #include "tls/s2n_tls13_handshake.h"
@@ -39,14 +40,13 @@ S2N_RESULT s2n_set_key_update_request_for_testing(s2n_peer_key_update request)
 int s2n_key_update_recv(struct s2n_connection *conn, struct s2n_stuffer *request)
 {
     POSIX_ENSURE_REF(conn);
+    POSIX_ENSURE_REF(conn->config);
     POSIX_ENSURE(conn->actual_protocol_version >= S2N_TLS13, S2N_ERR_BAD_MESSAGE);
     POSIX_ENSURE(!s2n_connection_is_quic_enabled(conn), S2N_ERR_BAD_MESSAGE);
-    POSIX_ENSURE(!conn->ktls_recv_enabled, S2N_ERR_KTLS_KEYUPDATE);
 
     uint8_t key_update_request = 0;
     POSIX_GUARD(s2n_stuffer_read_uint8(request, &key_update_request));
     if (key_update_request == S2N_KEY_UPDATE_REQUESTED) {
-        POSIX_ENSURE(!conn->ktls_send_enabled, S2N_ERR_KTLS_KEYUPDATE);
         s2n_atomic_flag_set(&conn->key_update_pending);
     } else {
         POSIX_ENSURE(key_update_request == S2N_KEY_UPDATE_NOT_REQUESTED, S2N_ERR_BAD_MESSAGE);
@@ -59,6 +59,11 @@ int s2n_key_update_recv(struct s2n_connection *conn, struct s2n_stuffer *request
         POSIX_GUARD(s2n_update_application_traffic_keys(conn, S2N_CLIENT, RECEIVING));
     }
 
+    /* We need to update the socket with the new keys in the case of ktls */
+    if (conn->ktls_recv_enabled) {
+        POSIX_GUARD_RESULT(s2n_ktls_key_update_process(conn));
+    }
+
     return S2N_SUCCESS;
 }
 
@@ -73,7 +78,7 @@ int s2n_key_update_send(struct s2n_connection *conn, s2n_blocked_status *blocked
     POSIX_GUARD(s2n_check_record_limit(conn, &sequence_number));
 
     if (s2n_atomic_flag_test(&conn->key_update_pending)) {
-        POSIX_ENSURE(!conn->ktls_send_enabled, S2N_ERR_KTLS_KEY_LIMIT);
+        POSIX_ENSURE(!conn->ktls_send_enabled, S2N_ERR_KTLS_KEYUPDATE);
 
         /* Flush any buffered records to ensure an empty output buffer.
          *
diff --git a/tls/s2n_key_update.h b/tls/s2n_key_update.h
@@ -27,3 +27,4 @@ typedef enum {
 
 int s2n_key_update_recv(struct s2n_connection *conn, struct s2n_stuffer *request);
 int s2n_key_update_send(struct s2n_connection *conn, s2n_blocked_status *blocked);
+int s2n_key_update_write(struct s2n_blob *out);
diff --git a/tls/s2n_ktls.c b/tls/s2n_ktls.c
diff --git a/tls/s2n_ktls.h b/tls/s2n_ktls.h
diff --git a/tls/s2n_ktls_io.c b/tls/s2n_ktls_io.c

Original file line number	Diff line number	Diff line change
`@@ -27,3 +27,4 @@ typedef enum {`
`27`	`27`
`28`	`28`	`int s2n_key_update_recv(struct s2n_connection conn, struct s2n_stuffer request);`
`29`	`29`	`int s2n_key_update_send(struct s2n_connection conn, s2n_blocked_status blocked);`
	`30`	`+int s2n_key_update_write(struct s2n_blob *out);`