Fix s2n_shutdown + failed recv bug (aws#4350)

lrstewart · Jan 18, 2024 · d946c2d · d946c2d
1 parent 7f84701
commit d946c2d
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 11 deletions.
diff --git a/bindings/rust/s2n-tls-tokio/tests/shutdown.rs b/bindings/rust/s2n-tls-tokio/tests/shutdown.rs
@@ -190,12 +190,8 @@ async fn shutdown_with_blinding() -> Result<(), Box<dyn std::error::Error>> {
     // Shutdown MUST NOT complete faster than minimal blinding time.
     assert!(time_elapsed > common::MIN_BLINDING_SECS);
 
-    // TODO: While the server SHOULD successfully shutdown, there is currently
-    // a C bug preventing it from doing so: https://github.com/aws/s2n-tls/pull/4350
-    let io_error = result.unwrap_err();
-    let error: error::Error = io_error.try_into()?;
-    assert!(error.kind() == error::ErrorType::IOError);
-    assert!(error.name() == "S2N_ERR_IO");
+    // Server MUST eventually successfully shutdown
+    assert!(result.is_ok());
 
     // Shutdown MUST have sent the close_notify message needed by the peer
     // to also shutdown successfully.

diff --git a/tests/unit/s2n_shutdown_test.c b/tests/unit/s2n_shutdown_test.c
@@ -356,6 +356,61 @@ int main(int argc, char **argv)
         EXPECT_SUCCESS(s2n_shutdown(conn, &blocked));
     };
 
+    /* Test: previous failed partial reads do not affect reading close_notify */
+    {
+        DEFER_CLEANUP(struct s2n_connection *conn = s2n_connection_new(S2N_SERVER),
+                s2n_connection_ptr_free);
+        EXPECT_NOT_NULL(conn);
+        EXPECT_OK(s2n_skip_handshake(conn));
+        EXPECT_SUCCESS(s2n_connection_set_blinding(conn, S2N_SELF_SERVICE_BLINDING));
+
+        /* Set the version so that a record header with the wrong version will
+         * be rejected as invalid.
+         */
+        conn->actual_protocol_version_established = true;
+        conn->actual_protocol_version = S2N_TLS13;
+
+        s2n_blocked_status blocked = S2N_NOT_BLOCKED;
+        DEFER_CLEANUP(struct s2n_stuffer input = { 0 }, s2n_stuffer_free);
+        EXPECT_SUCCESS(s2n_stuffer_growable_alloc(&input, 0));
+        DEFER_CLEANUP(struct s2n_stuffer output = { 0 }, s2n_stuffer_free);
+        EXPECT_SUCCESS(s2n_stuffer_growable_alloc(&output, 0));
+        EXPECT_SUCCESS(s2n_connection_set_io_stuffers(&input, &output, conn));
+
+        /* Receive a malformed record.
+         * We want reading this record to leave our IO in a bad state.
+         */
+        uint8_t header_bytes[] = {
+            /* record type */
+            TLS_HANDSHAKE,
+            /* bad protocol version */
+            0,
+            0,
+            /* zero length */
+            0,
+            0,
+        };
+        EXPECT_SUCCESS(s2n_stuffer_write_bytes(&input, header_bytes, sizeof(header_bytes)));
+        uint8_t recv_buffer[1] = { 0 };
+        EXPECT_FAILURE_WITH_ERRNO(
+                s2n_recv(conn, recv_buffer, sizeof(recv_buffer), &blocked),
+                S2N_ERR_BAD_MESSAGE);
+        EXPECT_TRUE(s2n_stuffer_space_remaining(&conn->header_in) < sizeof(header_bytes));
+
+        /* Clear the blinding delay so that we can call s2n_shutdown */
+        EXPECT_TRUE(conn->delay > 0);
+        conn->delay = 0;
+
+        /* Make the valid close_notify available */
+        EXPECT_SUCCESS(s2n_stuffer_write_bytes(&input, alert_record_header, sizeof(alert_record_header)));
+        EXPECT_SUCCESS(s2n_stuffer_write_bytes(&input, close_notify_alert, sizeof(close_notify_alert)));
+
+        /* Successfully shutdown.
+         * The initial bad call to s2n_recv should not affect shutdown.
+         */
+        EXPECT_SUCCESS(s2n_shutdown(conn, &blocked));
+    };
+
     /* Test: s2n_shutdown with aggressive socket close */
     {
         DEFER_CLEANUP(struct s2n_connection *server_conn = s2n_connection_new(S2N_SERVER),

diff --git a/tls/s2n_shutdown.c b/tls/s2n_shutdown.c
@@ -121,16 +121,18 @@ int s2n_shutdown(struct s2n_connection *conn, s2n_blocked_status *blocked)
     int isSSLv2 = false;
     *blocked = S2N_BLOCKED_ON_READ;
     while (!s2n_atomic_flag_test(&conn->close_notify_received)) {
+        /* Reset IO. Make sure we do this before attempting to read a record in
+         * case a previous failed read left IO in a bad state.
+         */
+        POSIX_GUARD(s2n_stuffer_wipe(&conn->header_in));
+        POSIX_GUARD(s2n_stuffer_wipe(&conn->in));
+        conn->in_status = ENCRYPTED;
+
         POSIX_GUARD(s2n_read_full_record(conn, &record_type, &isSSLv2));
         POSIX_ENSURE(!isSSLv2, S2N_ERR_BAD_MESSAGE);
         if (record_type == TLS_ALERT) {
             POSIX_GUARD(s2n_process_alert_fragment(conn));
         }
-
-        /* Wipe and keep trying */
-        POSIX_GUARD(s2n_stuffer_wipe(&conn->header_in));
-        POSIX_GUARD(s2n_stuffer_wipe(&conn->in));
-        conn->in_status = ENCRYPTED;
     }
 
     *blocked = S2N_NOT_BLOCKED;