feat: add dynamic buffer capabilities (#3472)

api/s2n.h

@@ -1639,6 +1639,20 @@ extern int s2n_connection_prefer_throughput(struct s2n_connection *conn);
 S2N_API
 extern int s2n_connection_prefer_low_latency(struct s2n_connection *conn);
 
+/**
+ * Configure the connection to free IO buffers when they are not currently in use.
+ *
+ * This configuration can be used to minimize connection memory footprint size, at the cost
+ * of more calls to alloc and free. Some of these costs can be mitigated by configuring s2n-tls
+ * to use an allocator that includes thread-local caches or lock-free allocation patterns.
+ *
+ * @param conn The connection object being update
+ * @param enabled Set to `true` if dynamic buffers are enabled; `false` if disabled
+ * @returns S2N_SUCCESS on success. S2N_FAILURE on failure
+ */
+S2N_API
+extern int s2n_connection_set_dynamic_buffers(struct s2n_connection *conn, bool enabled);
+
 /**
  * Provides a smooth transition from s2n_connection_prefer_low_latency() to s2n_connection_prefer_throughput().
  *

stuffer/s2n_stuffer.c

@@ -105,12 +105,20 @@ int s2n_stuffer_growable_alloc(struct s2n_stuffer *stuffer, const uint32_t size)
 int s2n_stuffer_free(struct s2n_stuffer *stuffer)
 {
     POSIX_PRECONDITION(s2n_stuffer_validate(stuffer));
-    if (stuffer != NULL) {
-        if (stuffer->alloced) {
-            POSIX_GUARD(s2n_free(&stuffer->blob));
-        }
-        *stuffer = (struct s2n_stuffer) {0};
+    if (stuffer->alloced) {
+        POSIX_GUARD(s2n_free(&stuffer->blob));
     }
+    *stuffer = (struct s2n_stuffer) {0};
+    return S2N_SUCCESS;
+}
+
+int s2n_stuffer_free_without_wipe(struct s2n_stuffer *stuffer)
+{
+    POSIX_PRECONDITION(s2n_stuffer_validate(stuffer));
+    if (stuffer->alloced) {
+        POSIX_GUARD(s2n_free_without_wipe(&stuffer->blob));
+    }
+    *stuffer = (struct s2n_stuffer) {0};
     return S2N_SUCCESS;
 }
 

stuffer/s2n_stuffer.h

@@ -68,6 +68,13 @@ extern int s2n_stuffer_init(struct s2n_stuffer *stuffer, struct s2n_blob *in);
 extern int s2n_stuffer_alloc(struct s2n_stuffer *stuffer, const uint32_t size);
 extern int s2n_stuffer_growable_alloc(struct s2n_stuffer *stuffer, const uint32_t size);
 extern int s2n_stuffer_free(struct s2n_stuffer *stuffer);
+/**
+ * Frees the stuffer without zeroizing the contained data.
+ *
+ * This should only be used in scenarios where the data is encrypted or has been
+ * cleared with `s2n_stuffer_erase_and_read`. In most cases, prefer `s2n_stuffer_free`.
+ */
+extern int s2n_stuffer_free_without_wipe(struct s2n_stuffer *stuffer);
 extern int s2n_stuffer_resize(struct s2n_stuffer *stuffer, const uint32_t size);
 extern int s2n_stuffer_resize_if_empty(struct s2n_stuffer *stuffer, const uint32_t size);
 extern int s2n_stuffer_rewind_read(struct s2n_stuffer *stuffer, const uint32_t size);

tests/unit/s2n_self_talk_io_mem_test.c

@@ -219,6 +219,82 @@ int main(int argc, char **argv)
         EXPECT_SUCCESS(s2n_connection_free(server_conn));
     }
 
+    /* Test that dynamic buffers work correctly */
+    {
+        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));
+
+        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));
+
+        /* Enable the dynamic buffers setting */
+        EXPECT_SUCCESS(s2n_connection_set_dynamic_buffers(client_conn, true));
+        EXPECT_SUCCESS(s2n_connection_set_dynamic_buffers(server_conn, true));
+
+        /* Configure the connection to use stuffers instead of fds.
+         * This will let us block the send.
+         */
+        DEFER_CLEANUP(struct s2n_stuffer client_in = { 0 }, s2n_stuffer_free);
+        DEFER_CLEANUP(struct s2n_stuffer client_out = { 0 }, s2n_stuffer_free);
+        EXPECT_SUCCESS(s2n_stuffer_growable_alloc(&client_in, 0));
+        EXPECT_SUCCESS(s2n_stuffer_growable_alloc(&client_out, 0));
+        EXPECT_SUCCESS(s2n_connection_set_io_stuffers(&client_in, &client_out, client_conn));
+        EXPECT_SUCCESS(s2n_connection_set_io_stuffers(&client_out, &client_in, server_conn));
+
+        /* Do handshake */
+        EXPECT_SUCCESS(s2n_negotiate_test_server_and_client(server_conn, client_conn));
+
+        /* all IO buffers should be empty after the handshake */
+        EXPECT_EQUAL(client_conn->in.blob.size, 0);
+        EXPECT_EQUAL(client_conn->out.blob.size, 0);
+        EXPECT_EQUAL(server_conn->in.blob.size, 0);
+        EXPECT_EQUAL(server_conn->out.blob.size, 0);
+
+        /* block the server from sending */
+        EXPECT_SUCCESS(s2n_stuffer_free(&client_in));
+
+        s2n_blocked_status blocked = 0;
+        int send_status = S2N_SUCCESS;
+
+        /* choose a large enough payload to send a full record, 4k is a common page size so we'll go with that */
+        uint8_t buf[4096] = { 42 };
+
+        send_status = s2n_send(server_conn, &buf, s2n_array_len(buf), &blocked);
+
+        /* the first send call should block */
+        EXPECT_FAILURE(send_status);
+        EXPECT_EQUAL(blocked, S2N_BLOCKED_ON_WRITE);
+
+        /* the `out` buffer should not be freed until it's completely flushed to the socket */
+        EXPECT_NOT_EQUAL(server_conn->out.blob.size, 0);
+
+        /* unblock the send call by letting the stuffer grow */
+        EXPECT_SUCCESS(s2n_stuffer_growable_alloc(&client_in, 0));
+
+        send_status = s2n_send(server_conn, &buf, s2n_array_len(buf), &blocked);
+        EXPECT_SUCCESS(send_status);
+
+        /* the entire payload should have been sent */
+        EXPECT_EQUAL(send_status, s2n_array_len(buf));
+
+        /* make sure the `out` buffer was freed after sending */
+        EXPECT_EQUAL(server_conn->out.blob.size, 0);
+
+        /* Receive half of the payload on the first call */
+        EXPECT_EQUAL(s2n_recv(client_conn, &buf, s2n_array_len(buf) / 2, &blocked), s2n_array_len(buf) / 2);
+
+        /* the `in` buffer should not be freed until it's completely flushed to the application */
+        EXPECT_NOT_EQUAL(client_conn->in.blob.size, 0);
+
+        /* Receive the second half of the payload on the second call */
+        EXPECT_EQUAL(s2n_recv(client_conn, &buf, s2n_array_len(buf) / 2, &blocked), s2n_array_len(buf) / 2);
+
+        /* at this point the application has received the full message and the `in` buffer should be freed */
+        EXPECT_EQUAL(client_conn->in.blob.size, 0);
+    }
+
     EXPECT_SUCCESS(s2n_config_free(config));
     EXPECT_SUCCESS(s2n_cert_chain_and_key_free(chain_and_key));
     END_TEST();

tls/s2n_connection.c

@@ -1208,6 +1208,13 @@ int s2n_connection_prefer_low_latency(struct s2n_connection *conn)
     return S2N_SUCCESS;
 }
 
+int s2n_connection_set_dynamic_buffers(struct s2n_connection *conn, bool enabled)
+{
+    POSIX_ENSURE_REF(conn);
+    conn->dynamic_buffers = enabled;
+    return S2N_SUCCESS;
+}
+
 int s2n_connection_set_dynamic_record_threshold(struct s2n_connection *conn, uint32_t resize_threshold, uint16_t timeout_threshold)
 {
     POSIX_ENSURE_REF(conn);
@@ -1482,3 +1489,35 @@ int s2n_connection_get_config(struct s2n_connection *conn, struct s2n_config **c
 
     return S2N_SUCCESS;
 }
+
+S2N_RESULT s2n_connection_dynamic_free_out_buffer(struct s2n_connection *conn)
+{
+    RESULT_ENSURE_REF(conn);
+
+    /* free the out buffer if we're in dynamic mode and it's completely flushed */
+    if (conn->dynamic_buffers && s2n_stuffer_is_consumed(&conn->out)) {
+        /* since outgoing buffers are already encrypted, the buffers don't need to be zeroed, which saves some overhead */
+        RESULT_GUARD_POSIX(s2n_stuffer_free_without_wipe(&conn->out));
+
+        /* reset the stuffer to its initial state */
+        RESULT_GUARD_POSIX(s2n_stuffer_growable_alloc(&conn->out, 0));
+    }
+
+    return S2N_RESULT_OK;
+}
+
+S2N_RESULT s2n_connection_dynamic_free_in_buffer(struct s2n_connection *conn)
+{
+    RESULT_ENSURE_REF(conn);
+
+    /* free the `in` buffer if we're in dynamic mode and it's completely flushed */
+    if (conn->dynamic_buffers && s2n_stuffer_is_consumed(&conn->in)) {
+        /* when copying the buffer into the application, we use `s2n_stuffer_erase_and_read`, which already zeroes the memory */
+        RESULT_GUARD_POSIX(s2n_stuffer_free_without_wipe(&conn->in));
+
+        /* reset the stuffer to its initial state */
+        RESULT_GUARD_POSIX(s2n_stuffer_growable_alloc(&conn->in, 0));
+    }
+
+    return S2N_RESULT_OK;
+}

tls/s2n_connection.h

@@ -127,6 +127,10 @@ struct s2n_connection {
      * This allows multiple records to be written with one socket send. */
     unsigned multirecord_send:1;
 
+    /* If enabled, this connection will free each of its IO buffers after all data
+     * has been flushed */
+    unsigned dynamic_buffers:1;
+
     /* The configuration (cert, key .. etc ) */
     struct s2n_config *config;
 
@@ -390,6 +394,10 @@ int s2n_connection_recv_stuffer(struct s2n_stuffer *stuffer, struct s2n_connecti
 
 S2N_RESULT s2n_connection_wipe_all_keyshares(struct s2n_connection *conn);
 
+/* If dynamic buffers are enabled, the IO buffers may be freed if they are completely consumed */
+S2N_RESULT s2n_connection_dynamic_free_in_buffer(struct s2n_connection *conn);
+S2N_RESULT s2n_connection_dynamic_free_out_buffer(struct s2n_connection *conn);
+
 int s2n_connection_get_cipher_preferences(struct s2n_connection *conn, const struct s2n_cipher_preferences **cipher_preferences);
 int s2n_connection_get_security_policy(struct s2n_connection *conn, const struct s2n_security_policy **security_policy);
 int s2n_connection_get_kem_preferences(struct s2n_connection *conn, const struct s2n_kem_preferences **kem_preferences);

tls/s2n_handshake_io.c

@@ -62,7 +62,7 @@ static int s2n_always_fail_recv(struct s2n_connection *conn)
     POSIX_BAIL(S2N_ERR_HANDSHAKE_UNREACHABLE);
 }
 
-/* Client and Server handlers for each message type we support.  
+/* Client and Server handlers for each message type we support.
  * See http://www.iana.org/assignments/tls-parameters/tls-parameters.xhtml#tls-parameters-7 for the list of handshake message types
  */
 static struct s2n_handshake_action state_machine[] = {
@@ -1455,7 +1455,13 @@ int s2n_negotiate(struct s2n_connection *conn, s2n_blocked_status *blocked)
     POSIX_ENSURE_REF(conn);
     POSIX_ENSURE(!conn->negotiate_in_use, S2N_ERR_REENTRANCY);
     conn->negotiate_in_use = true;
+
     int result = s2n_negotiate_impl(conn, blocked);
+
+    /* finish up sending and receiving */
+    POSIX_GUARD_RESULT(s2n_connection_dynamic_free_in_buffer(conn));
+    POSIX_GUARD_RESULT(s2n_connection_dynamic_free_out_buffer(conn));
+
     conn->negotiate_in_use = false;
     return result;
 }

tls/s2n_recv.c

@@ -206,8 +206,13 @@ ssize_t s2n_recv(struct s2n_connection * conn, void *buf, ssize_t size, s2n_bloc
 {
     POSIX_ENSURE(!conn->recv_in_use, S2N_ERR_REENTRANCY);
     conn->recv_in_use = true;
+
     ssize_t result = s2n_recv_impl(conn, buf, size, blocked);
     POSIX_GUARD_RESULT(s2n_early_data_record_bytes(conn, result));
+
+    /* finish the recv call */
+    POSIX_GUARD_RESULT(s2n_connection_dynamic_free_in_buffer(conn));
+
     conn->recv_in_use = false;
     return result;
 }

tls/s2n_send.c

@@ -209,7 +209,7 @@ ssize_t s2n_sendv_with_offset_impl(struct s2n_connection *conn, const struct iov
         }
 
         POSIX_GUARD(s2n_post_handshake_send(conn, blocked));
-    
+
         /* Write and encrypt the record */
         int written_to_record = s2n_record_writev(conn, TLS_APPLICATION_DATA, bufs, count,
                     conn->current_user_data_consumed + offs, to_write);
@@ -247,8 +247,12 @@ ssize_t s2n_sendv_with_offset(struct s2n_connection *conn, const struct iovec *b
 {
     POSIX_ENSURE(!conn->send_in_use, S2N_ERR_REENTRANCY);
     conn->send_in_use = true;
+
     ssize_t result = s2n_sendv_with_offset_impl(conn, bufs, count, offs, blocked);
     POSIX_GUARD_RESULT(s2n_early_data_record_bytes(conn, result));
+
+    POSIX_GUARD_RESULT(s2n_connection_dynamic_free_out_buffer(conn));
+
     conn->send_in_use = false;
     return result;
 }

utils/s2n_mem.c

@@ -218,7 +218,7 @@ int s2n_free_object(uint8_t **p_data, uint32_t size)
     if (*p_data == NULL) {
         return S2N_SUCCESS;
     }
-    
+
     POSIX_ENSURE(initialized, S2N_ERR_NOT_INITIALIZED);
     struct s2n_blob b = {.data = *p_data, .allocated = size, .size = size, .growable = 1};
 
@@ -275,11 +275,16 @@ int s2n_mem_cleanup(void)
 
 int s2n_free(struct s2n_blob *b)
 {
-    POSIX_PRECONDITION(s2n_blob_validate(b));
-
     /* To avoid memory leaks, don't exit the function until the memory
        has been freed */
     int zero_rc = s2n_blob_zero(b);
+    POSIX_GUARD(s2n_free_without_wipe(b));
+    return zero_rc;
+}
+
+int s2n_free_without_wipe(struct s2n_blob *b)
+{
+    POSIX_PRECONDITION(s2n_blob_validate(b));
 
     POSIX_ENSURE(initialized, S2N_ERR_NOT_INITIALIZED);
     POSIX_ENSURE(s2n_blob_is_growable(b), S2N_ERR_FREE_STATIC_BLOB);
@@ -288,8 +293,6 @@ int s2n_free(struct s2n_blob *b)
 
     *b = (struct s2n_blob) {0};
 
-    POSIX_GUARD(zero_rc);
-
     return S2N_SUCCESS;
 }
 

utils/s2n_mem.h

@@ -26,6 +26,7 @@ int s2n_mem_cleanup(void);
 int s2n_alloc(struct s2n_blob *b, uint32_t size);
 int s2n_realloc(struct s2n_blob *b, uint32_t size);
 int s2n_free(struct s2n_blob *b);
+int s2n_free_without_wipe(struct s2n_blob *b);
 int s2n_blob_zeroize_free(struct s2n_blob *b);
 int s2n_free_object(uint8_t **p_data, uint32_t size);
 int s2n_dup(struct s2n_blob *from, struct s2n_blob *to);