feat: add s2n_peek_buffered (#4490)

api/s2n.h

@@ -1842,6 +1842,7 @@ S2N_API extern int s2n_connection_prefer_low_latency(struct s2n_connection *conn
  *    until it reports S2N_SUCCESS.
  *
  * 4. s2n_peek reports available decrypted data. It does not report any data
+ *    buffered by this feature. However, s2n_peek_buffered does report data
  *    buffered by this feature.
  *
  * 5. s2n_connection_release_buffers will not release the input buffer if it
@@ -1853,24 +1854,44 @@ S2N_API extern int s2n_connection_prefer_low_latency(struct s2n_connection *conn
  * If you stop calling s2n_recv before it reports S2N_ERR_T_BLOCKED, some of those
  * records may remain in s2n-tls's read buffer. If you read part of a record,
  * s2n_peek will report the remainder of that record as available. But if you don't
- * read any of a record, it remains encrypted and is not reported by s2n_peek.
- * And because the data is buffered in s2n-tls instead of in the file descriptor,
- * another call to `poll` will NOT report any more data available. Your application
- * may hang waiting for more data.
+ * read any of a record, it remains encrypted and is not reported by s2n_peek, but
+ * is still reported by s2n_peek_buffered. And because the data is buffered in s2n-tls
+ * instead of in the file descriptor, another call to `poll` will NOT report any
+ * more data available. Your application may hang waiting for more data.
  *
  * @warning This feature cannot be enabled for a connection that will enable kTLS for receiving.
  *
  * @warning This feature may work with blocking IO, if used carefully. Your blocking
- * IO must support partial reads (so MSG_WAITALL cannot be used). You will need
- * to know how much data will eventually be available rather than relying on
- * S2N_ERR_T_BLOCKED as noted in #3 above.
+ * IO must support partial reads (so MSG_WAITALL cannot be used). You will either
+ * need to know exactly how much data your peer is sending, or will need to use
+ * `s2n_peek` and `s2n_peek_buffered` rather than relying on S2N_ERR_T_BLOCKED
+ * as noted in #3 above.
  *
  * @param conn The connection object being updated
  * @param enabled Set to `true` to enable, `false` to disable.
  * @returns S2N_SUCCESS on success. S2N_FAILURE on failure
  */
 S2N_API extern int s2n_connection_set_recv_buffering(struct s2n_connection *conn, bool enabled);
 
+/**
+ * Reports how many bytes of unprocessed TLS records are buffered due to the optimization
+ * enabled by `s2n_connection_set_recv_buffering`.
+ *
+ * `s2n_peek_buffered` is not a replacement for `s2n_peek`.
+ * While `s2n_peek` reports application data that is ready for the application
+ * to read with no additional processing, `s2n_peek_buffered` reports raw TLS
+ * records that still need to be parsed and likely decrypted. Those records may
+ * contain application data, but they may also only contain TLS control messages.
+ *
+ * If an application needs to determine whether there is any data left to handle
+ * (for example, before calling `poll` to wait on the read file descriptor) then
+ * that application must check both `s2n_peek` and `s2n_peek_buffered`.
+ *
+ * @param conn A pointer to the s2n_connection object
+ * @returns The number of buffered encrypted bytes
+ */
+S2N_API extern uint32_t s2n_peek_buffered(struct s2n_connection *conn);
+
 /**
  * Configure the connection to free IO buffers when they are not currently in use.
  *

tests/unit/s2n_recv_buffering_test.c

@@ -65,6 +65,13 @@ int main(int argc, char **argv)
     EXPECT_SUCCESS(s2n_config_set_cipher_preferences(config, "default_tls13"));
     EXPECT_SUCCESS(s2n_config_disable_x509_verification(config));
 
+    DEFER_CLEANUP(struct s2n_config *multi_config = s2n_config_new(),
+            s2n_config_ptr_free);
+    EXPECT_SUCCESS(s2n_config_add_cert_chain_and_key_to_store(multi_config, chain_and_key));
+    EXPECT_SUCCESS(s2n_config_set_cipher_preferences(multi_config, "default_tls13"));
+    EXPECT_SUCCESS(s2n_config_disable_x509_verification(multi_config));
+    EXPECT_SUCCESS(s2n_config_set_recv_multi_record(multi_config, true));
+
     /* Test: Read a single record */
     uint32_t test_record_size_val = 0;
     {
@@ -217,6 +224,38 @@ int main(int argc, char **argv)
         EXPECT_EQUAL(server->buffer_in.blob.allocated, 0);
     }
 
+    /* Test: Read multiple small records with "multi_record" enabled */
+    {
+        DEFER_CLEANUP(struct s2n_connection *client = s2n_connection_new(S2N_CLIENT),
+                s2n_connection_ptr_free);
+        EXPECT_SUCCESS(s2n_connection_set_config(client, multi_config));
+
+        DEFER_CLEANUP(struct s2n_connection *server = s2n_connection_new(S2N_SERVER),
+                s2n_connection_ptr_free);
+        EXPECT_SUCCESS(s2n_connection_set_config(server, multi_config));
+
+        EXPECT_SUCCESS(s2n_connection_set_recv_buffering(client, true));
+        EXPECT_SUCCESS(s2n_connection_set_recv_buffering(server, true));
+
+        DEFER_CLEANUP(struct s2n_test_io_stuffer_pair io_pair = { 0 }, s2n_io_stuffer_pair_free);
+        EXPECT_OK(s2n_io_stuffer_pair_init(&io_pair));
+        EXPECT_OK(s2n_connections_set_io_stuffer_pair(client, server, &io_pair));
+        EXPECT_SUCCESS(s2n_negotiate_test_server_and_client(server, client));
+
+        struct s2n_recv_wrapper counter = { 0 };
+        EXPECT_OK(s2n_connection_set_counting_read(server, &counter));
+
+        s2n_blocked_status blocked = S2N_NOT_BLOCKED;
+        for (size_t i = 0; i < sizeof(test_data); i++) {
+            EXPECT_EQUAL(s2n_send(client, test_data + i, 1, &blocked), 1);
+        }
+
+        uint8_t buffer[sizeof(test_data)] = { 0 };
+        EXPECT_EQUAL(s2n_recv(server, buffer, sizeof(buffer), &blocked), sizeof(test_data));
+        EXPECT_BYTEARRAY_EQUAL(buffer, test_data, sizeof(test_data));
+        EXPECT_EQUAL(counter.count, 1);
+    }
+
     /* Test: Read the rest of a partial record */
     for (size_t i = 0; i < test_record_size; i++) {
         DEFER_CLEANUP(struct s2n_connection *client = s2n_connection_new(S2N_CLIENT),
@@ -514,5 +553,67 @@ int main(int argc, char **argv)
         EXPECT_BYTEARRAY_EQUAL(buffer, test_data, sizeof(test_data));
     }
 
+    /* Test: s2n_peek_buffered */
+    {
+        EXPECT_EQUAL(s2n_peek_buffered(NULL), 0);
+
+        DEFER_CLEANUP(struct s2n_connection *client = s2n_connection_new(S2N_CLIENT),
+                s2n_connection_ptr_free);
+        EXPECT_SUCCESS(s2n_connection_set_config(client, multi_config));
+
+        DEFER_CLEANUP(struct s2n_connection *server = s2n_connection_new(S2N_SERVER),
+                s2n_connection_ptr_free);
+        EXPECT_SUCCESS(s2n_connection_set_config(server, multi_config));
+
+        struct {
+            uint32_t read_size;
+            uint32_t expect_available;
+            uint32_t expect_buffered;
+        } test_cases[] = {
+            {
+                    .read_size = 1,
+                    .expect_available = sizeof(test_data) - 1,
+                    .expect_buffered = test_record_size,
+            },
+            {
+                    .read_size = sizeof(test_data) - 1,
+                    .expect_available = 1,
+                    .expect_buffered = test_record_size,
+            },
+            {
+                    .read_size = sizeof(test_data),
+                    .expect_available = 0,
+                    .expect_buffered = test_record_size,
+            },
+            {
+                    .read_size = sizeof(test_data) + 1,
+                    .expect_available = sizeof(test_data) - 1,
+                    .expect_buffered = 0,
+            },
+        };
+        for (size_t i = 0; i < s2n_array_len(test_cases); i++) {
+            DEFER_CLEANUP(struct s2n_test_io_stuffer_pair io_pair = { 0 }, s2n_io_stuffer_pair_free);
+            EXPECT_OK(s2n_io_stuffer_pair_init(&io_pair));
+            EXPECT_OK(s2n_connections_set_io_stuffer_pair(client, server, &io_pair));
+            EXPECT_SUCCESS(s2n_negotiate_test_server_and_client(server, client));
+            EXPECT_SUCCESS(s2n_stuffer_wipe(&io_pair.server_in));
+
+            s2n_blocked_status blocked = S2N_NOT_BLOCKED;
+            uint8_t buffer[sizeof(test_data) * 2] = { 0 };
+
+            EXPECT_EQUAL(s2n_send(client, test_data, sizeof(test_data), &blocked), sizeof(test_data));
+            EXPECT_EQUAL(s2n_send(client, test_data, sizeof(test_data), &blocked), sizeof(test_data));
+
+            uint32_t read_size = test_cases[i].read_size;
+            EXPECT_SUCCESS(s2n_connection_set_recv_buffering(server, true));
+            EXPECT_EQUAL(s2n_recv(server, buffer, read_size, &blocked), read_size);
+            EXPECT_EQUAL(s2n_peek_buffered(server), test_cases[i].expect_buffered);
+            EXPECT_EQUAL(s2n_peek(server), test_cases[i].expect_available);
+
+            EXPECT_SUCCESS(s2n_connection_wipe(client));
+            EXPECT_SUCCESS(s2n_connection_wipe(server));
+        }
+    }
+
     END_TEST();
 }

tls/s2n_recv.c

@@ -317,3 +317,11 @@ uint32_t s2n_peek(struct s2n_connection *conn)
 
     return s2n_stuffer_data_available(&conn->in);
 }
+
+uint32_t s2n_peek_buffered(struct s2n_connection *conn)
+{
+    if (conn == NULL) {
+        return 0;
+    }
+    return s2n_stuffer_data_available(&conn->buffer_in);
+}