feat: add dynamic buffer capabilities

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 after they are no longer used.
+ *
+ * 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

@@ -102,18 +102,32 @@ int s2n_stuffer_growable_alloc(struct s2n_stuffer *stuffer, const uint32_t size)
     return S2N_SUCCESS;
 }
 
-int s2n_stuffer_free(struct s2n_stuffer *stuffer)
+static int s2n_stuffer_free_impl(struct s2n_stuffer *stuffer, bool zeroed)
 {
     POSIX_PRECONDITION(s2n_stuffer_validate(stuffer));
     if (stuffer != NULL) {
         if (stuffer->alloced) {
-            POSIX_GUARD(s2n_free(&stuffer->blob));
+            if (zeroed) {
+                POSIX_GUARD(s2n_free(&stuffer->blob));
+            } else {
+                POSIX_GUARD(s2n_free_non_zeroed(&stuffer->blob));
+            }
         }
         *stuffer = (struct s2n_stuffer) {0};
     }
     return S2N_SUCCESS;
 }
 
+int s2n_stuffer_free(struct s2n_stuffer *stuffer)
+{
+    return s2n_stuffer_free_impl(stuffer, true);
+}
+
+int s2n_stuffer_free_non_zeroed(struct s2n_stuffer *stuffer)
+{
+    return s2n_stuffer_free_impl(stuffer, false);
+}
+
 int s2n_stuffer_resize(struct s2n_stuffer *stuffer, const uint32_t size)
 {
     POSIX_PRECONDITION(s2n_stuffer_validate(stuffer));

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_non_zeroed(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

@@ -17,6 +17,7 @@
 
 #include "testlib/s2n_testlib.h"
 
+#include <sys/param.h>
 #include <sys/wait.h>
 #include <unistd.h>
 #include <time.h>
@@ -219,6 +220,68 @@ int main(int argc, char **argv)
         EXPECT_SUCCESS(s2n_connection_free(server_conn));
     }
 
+    /* Test that dynamic buffers work correctly */
+    {
+        struct s2n_connection *client_conn = s2n_connection_new(S2N_CLIENT);
+        EXPECT_NOT_NULL(client_conn);
+        EXPECT_SUCCESS(s2n_connection_set_config(client_conn, config));
+
+        struct s2n_connection *server_conn = s2n_connection_new(S2N_SERVER);
+        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));
+
+        /* Create nonblocking pipes */
+        struct s2n_test_io_pair io_pair;
+        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));
+
+        /* Do handshake */
+        EXPECT_SUCCESS(s2n_negotiate_test_server_and_client(server_conn, client_conn));
+
+        s2n_blocked_status blocked = 0;
+        uint8_t buf[4096] = { 42 };
+
+        size_t sent_amount = 0;
+
+        while (sent_amount < s2n_array_len(buf)) {
+            int send_status = s2n_send(server_conn, &buf, s2n_array_len(buf) - sent_amount, &blocked);
+            EXPECT_SUCCESS(send_status);
+            sent_amount += s2n_array_len(buf);
+        }
+
+        /* make sure the `out` buffer was freed after sending */
+        EXPECT_EQUAL(server_conn->out.blob.size, 0);
+
+        size_t recv_amount = 0;
+        size_t recv_size = 1;
+        while (true) {
+            int recv_status = s2n_recv(client_conn, &buf, recv_size, &blocked);
+            EXPECT_SUCCESS(recv_status);
+            recv_amount += recv_status;
+
+            if (recv_amount >= s2n_array_len(buf)) {
+                break;
+            }
+
+            /* the `in` buffer should not be freed until it's completely flushed to the application */
+            EXPECT_NOT_EQUAL(client_conn->in.blob.size, 0);
+
+            /* increase the receive size exponentially */
+            recv_size = MIN(recv_size * 2, s2n_array_len(buf));
+        }
+
+        /* make sure the `in` buffer was freed after receiving the full message */
+        EXPECT_EQUAL(client_conn->in.blob.size, 0);
+
+        EXPECT_SUCCESS(s2n_connection_free(client_conn));
+        EXPECT_SUCCESS(s2n_connection_free(server_conn));
+    }
+
     EXPECT_SUCCESS(s2n_config_free(config));
     EXPECT_SUCCESS(s2n_cert_chain_and_key_free(chain_and_key));
     END_TEST();

tls/s2n_connection.c

@@ -1292,6 +1292,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);

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;
 

tls/s2n_recv.c

@@ -206,8 +206,19 @@ 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));
+
+    /* 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 */
+        POSIX_GUARD(s2n_stuffer_free_non_zeroed(&conn->in));
+
+        /* reset the stuffer to its initial state */
+        POSIX_GUARD(s2n_stuffer_growable_alloc(&conn->in, 0));
+    }
+
     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,19 @@ 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));
+
+    /* 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 */
+        POSIX_GUARD(s2n_stuffer_free_non_zeroed(&conn->out));
+
+        /* reset the stuffer to its initial state */
+        POSIX_GUARD(s2n_stuffer_growable_alloc(&conn->out, 0));
+    }
+
     conn->send_in_use = false;
     return result;
 }

utils/s2n_mem.c

@@ -273,13 +273,13 @@ int s2n_mem_cleanup(void)
     return S2N_SUCCESS;
 }
 
-int s2n_free(struct s2n_blob *b)
+static int s2n_free_impl(struct s2n_blob *b, bool zeroed)
 {
     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);
+    int zero_rc = zeroed ? s2n_blob_zero(b) : S2N_SUCCESS;
 
     POSIX_ENSURE(initialized, S2N_ERR_NOT_INITIALIZED);
     POSIX_ENSURE(s2n_blob_is_growable(b), S2N_ERR_FREE_STATIC_BLOB);
@@ -293,6 +293,16 @@ int s2n_free(struct s2n_blob *b)
     return S2N_SUCCESS;
 }
 
+int s2n_free(struct s2n_blob *b)
+{
+    return s2n_free_impl(b, true);
+}
+
+int s2n_free_non_zeroed(struct s2n_blob *b)
+{
+    return s2n_free_impl(b, false);
+}
+
 int s2n_blob_zeroize_free(struct s2n_blob *b) {
     POSIX_ENSURE(initialized, S2N_ERR_NOT_INITIALIZED);
     POSIX_ENSURE_REF(b);

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_non_zeroed(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);