From 51c49fd1c7ea25345f95c5e7ccde4d67e4b1f8cc Mon Sep 17 00:00:00 2001
From: maddeleine <59030281+maddeleine@users.noreply.github.com>
Date: Mon, 8 Jan 2024 10:30:31 -0800
Subject: [PATCH 01/13] feat: Publishes mdbook to Github Pages (#4343)

---
 .github/workflows/usage_guide.yml | 73 +++++++++++++++++++++++++++++++
 README.md                         |  4 +-
 2 files changed, 75 insertions(+), 2 deletions(-)
 create mode 100644 .github/workflows/usage_guide.yml

diff --git a/.github/workflows/usage_guide.yml b/.github/workflows/usage_guide.yml
new file mode 100644
index 00000000000..3f6ced8887c
--- /dev/null
+++ b/.github/workflows/usage_guide.yml
@@ -0,0 +1,73 @@
+name: Publish Usage Guide
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+env:
+  CDN: https://d3fqnyekunr9xg.cloudfront.net
+
+# By default dependabot only receives read permissions. Explicitly give it write
+# permissions which is needed by the ouzi-dev/commit-status-updater task.
+#
+# Updating status is relatively safe (doesnt modify source code) and caution
+# should be taken before adding more permissions.
+permissions:
+  contents: write
+  statuses: write
+
+jobs:
+  build-deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout s2n-tls repo
+        uses: actions/checkout@v4
+
+      - uses: dtolnay/rust-toolchain@stable
+
+      - name: Set override
+        run: rustup override set stable
+
+      - uses: camshaft/install@v1
+        with:
+          crate: mdbook
+
+      - name: Build book
+        run: |
+          cd docs/usage-guide
+          mdbook build
+      
+      - name: Deploy documentation to gh-pages
+        uses: JamesIves/github-pages-deploy-action@v4.5.0
+        if: github.event_name == 'push'
+        with:
+          destination_dir: usage-guide
+          publish_dir: docs/usage-guide/book
+  
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v4.0.1
+        if: github.event_name == 'push' || github.repository == github.event.pull_request.head.repo.full_name
+        with:
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          aws-region: us-west-1
+      
+      - name: Upload to S3
+        if: github.event_name == 'push' || github.repository == github.event.pull_request.head.repo.full_name
+        id: s3
+        run: |
+          TARGET="${{ github.sha }}/book"
+          aws s3 sync docs/usage-guide/book "s3://s2n-tls-ci-artifacts/$TARGET" --acl private --follow-symlinks
+          URL="$CDN/$TARGET/index.html"
+          echo "URL=$URL" >> $GITHUB_OUTPUT
+      
+      - name: Output mdbook url 
+        uses: ouzi-dev/commit-status-updater@v2.0.1
+        if: github.event_name == 'push' || github.repository == github.event.pull_request.head.repo.full_name
+        with:
+          name: "book / url"
+          status: "success"
+          url: "${{ steps.s3.outputs.URL }}"
diff --git a/README.md b/README.md
index 19e91e33d70..75ae99a5cc8 100644
--- a/README.md
+++ b/README.md
@@ -44,7 +44,7 @@ If you have any questions about submitting PRs, s2n-tls API usage, or something
 
 ## Documentation
 
-s2n-tls uses [Doxygen](https://doxygen.nl/index.html) to document its public API. The latest s2n-tls documentation can be found on [GitHub pages](https://aws.github.io/s2n-tls/doxygen/). The [Usage Guide](docs/usage-guide/) explains how different TLS features can be configured and used.
+s2n-tls uses [Doxygen](https://doxygen.nl/index.html) to document its public API. The latest s2n-tls documentation can be found on [GitHub pages](https://aws.github.io/s2n-tls/doxygen/). The [Usage Guide](https://aws.github.io/s2n-tls/usage-guide/) explains how different TLS features can be configured and used.
 
 Documentation for older versions or branches of s2n-tls can be generated locally. To generate the documentation, install doxygen and run `doxygen docs/doxygen/Doxyfile`. The doxygen documentation can now be found at `docs/doxygen/output/html/index.html`.
 
@@ -77,7 +77,7 @@ int bytes_written;
 bytes_written = s2n_send(conn, "Hello World", sizeof("Hello World"), &blocked);
 ```
 
-For details on building the s2n-tls library and how to use s2n-tls in an application you are developing, see the [usage guide][Usage Guide](docs/usage-guide).
+For details on building the s2n-tls library and how to use s2n-tls in an application you are developing, see the [Usage Guide](https://aws.github.io/s2n-tls/usage-guide).
 
 ## s2n-tls features
 

From d606b53c25fb61ce6d81fd4ab68ee5687bf0ae3e Mon Sep 17 00:00:00 2001
From: maddeleine <59030281+maddeleine@users.noreply.github.com>
Date: Mon, 8 Jan 2024 14:17:36 -0800
Subject: [PATCH 02/13] bug: Fixes mdbook action (#4345)

---
 .github/workflows/usage_guide.yml   | 4 ++--
 docs/usage-guide/topics/ch01-api.md | 4 ----
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/usage_guide.yml b/.github/workflows/usage_guide.yml
index 3f6ced8887c..bbdd404562b 100644
--- a/.github/workflows/usage_guide.yml
+++ b/.github/workflows/usage_guide.yml
@@ -44,8 +44,8 @@ jobs:
         uses: JamesIves/github-pages-deploy-action@v4.5.0
         if: github.event_name == 'push'
         with:
-          destination_dir: usage-guide
-          publish_dir: docs/usage-guide/book
+          target-folder: usage-guide
+          folder: docs/usage-guide/book
   
       - name: Configure AWS credentials
         uses: aws-actions/configure-aws-credentials@v4.0.1
diff --git a/docs/usage-guide/topics/ch01-api.md b/docs/usage-guide/topics/ch01-api.md
index 4d3c019d0b9..3f5bd55b9c8 100644
--- a/docs/usage-guide/topics/ch01-api.md
+++ b/docs/usage-guide/topics/ch01-api.md
@@ -14,10 +14,6 @@ The [VERSIONING.rst](https://github.com/aws/s2n-tls/blob/main/VERSIONING.rst) do
 
 s2n-tls uses [Doxygen](https://doxygen.nl/index.html) to document its public API. The latest s2n-tls documentation can be found on [GitHub pages](https://aws.github.io/s2n-tls/doxygen/).
 
-Documentation for older versions or branches of s2n-tls can be generated locally. To generate the documentation, install doxygen and run `doxygen docs/doxygen/Doxyfile`. The doxygen documentation can now be found at `docs/doxygen/output/html/index.html`.
-
-Doxygen installation instructions are available at the [Doxygen](https://doxygen.nl/download.html) webpage.
-
 The doxygen documentation should be used in conjunction with this guide.
 
 ## Examples

From e5e7b019aeb4e49fb5c241f4fd84cdc6ac0e4e17 Mon Sep 17 00:00:00 2001
From: Lindsay Stewart <slindsay@amazon.com>
Date: Tue, 9 Jan 2024 10:46:39 -0800
Subject: [PATCH 03/13] ktls: improve messaging around freed handshakes (#4346)

---
 api/unstable/ktls.h | 6 ++++--
 tls/s2n_ktls.c      | 6 ++++++
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/api/unstable/ktls.h b/api/unstable/ktls.h
index ce2453f502d..ee395f5f006 100644
--- a/api/unstable/ktls.h
+++ b/api/unstable/ktls.h
@@ -44,7 +44,8 @@
  * Enables sending using kTLS on a given connection.
  *
  * See above for the limitations on when kTLS can be enabled. Additionally,
- * s2n_connection_ktls_enable_send must be called after the handshake completes.
+ * s2n_connection_ktls_enable_send must be called after the handshake completes
+ * but before the handshake is freed with s2n_connection_free_handshake.
  * It may be called after some application data is sent and received without kTLS,
  * but there must be no pending application data that requires flushing. If these
  * requirements are not met, enabling kTLS will fail with an error.
@@ -74,7 +75,8 @@ S2N_API int s2n_connection_ktls_enable_send(struct s2n_connection *conn);
  * Enables receiving using kTLS on a given connection.
  *
  * See above for the limitations on when kTLS can be enabled. Additionally,
- * s2n_connection_ktls_enable_recv must be called after the handshake completes.
+ * s2n_connection_ktls_enable_recv must be called after the handshake completes
+ * but before the handshake is freed with s2n_connection_free_handshake.
  * It may be called after some application data is sent and received without kTLS,
  * but there must be no buffered application data that requires draining. If these
  * requirements are not met, enabling kTLS will fail with an error.
diff --git a/tls/s2n_ktls.c b/tls/s2n_ktls.c
index 8731791e0ed..ed3e2024d26 100644
--- a/tls/s2n_ktls.c
+++ b/tls/s2n_ktls.c
@@ -55,6 +55,12 @@ static S2N_RESULT s2n_ktls_validate(struct s2n_connection *conn, s2n_ktls_mode k
     /* kTLS enable should only be called once the handshake has completed. */
     RESULT_ENSURE(is_handshake_complete(conn), S2N_ERR_HANDSHAKE_NOT_COMPLETE);
 
+    /* kTLS uses the prf_space to recalculate the keys, but the prf_space may be
+     * freed by s2n_connection_free_handshake to reduce the connection size.
+     * Explicitly check for prf_space here to avoid a confusing S2N_ERR_NULL later.
+     */
+    RESULT_ENSURE(conn->prf_space, S2N_ERR_INVALID_STATE);
+
     /* For now, only allow TlS1.3 if explicitly enabled.
      *
      * TLS1.3 is potentially more dangerous to enable than TLS1.2, since the kernel

From edebdae2ebfff08bba94d25f908514f383d791c5 Mon Sep 17 00:00:00 2001
From: Cameron Bytheway <bythewc@amazon.com>
Date: Tue, 9 Jan 2024 13:10:41 -0700
Subject: [PATCH 04/13] fix(bindings): remove optional cmake dependency (#4347)

---
 bindings/rust/s2n-tls-sys/build.rs            | 79 +------------------
 .../rust/s2n-tls-sys/templates/Cargo.template |  5 +-
 2 files changed, 4 insertions(+), 80 deletions(-)

diff --git a/bindings/rust/s2n-tls-sys/build.rs b/bindings/rust/s2n-tls-sys/build.rs
index b7ae351986d..0ee83cfe68c 100644
--- a/bindings/rust/s2n-tls-sys/build.rs
+++ b/bindings/rust/s2n-tls-sys/build.rs
@@ -8,17 +8,6 @@ fn main() {
     if external.is_enabled() {
         external.link();
     } else {
-        #[cfg(feature = "cmake")]
-        {
-            // branch on a runtime value so we don't get unused code warnings
-            if option_env("CARGO_FEATURE_CMAKE").is_some() {
-                build_cmake();
-            } else {
-                build_vendored();
-            }
-        }
-
-        #[cfg(not(feature = "cmake"))]
         build_vendored();
     }
 }
@@ -93,26 +82,7 @@ fn build_vendored() {
 
     let mut build = builder(&libcrypto);
 
-    // TODO: update rust bindings to handle no pq-crypto dir
-
-    let pq = option_env("CARGO_FEATURE_PQ").is_some();
-
-    // TODO each pq section needs to be built separately since it
-    //      has its own relative include paths
-    assert!(!pq, "pq builds are not currently supported without cmake");
-
-    build.files(include!("./files.rs").iter().copied().filter(|file| {
-        // the pq entry file is still needed
-        if *file == "lib/pq-crypto/s2n_pq.c" {
-            return true;
-        }
-
-        if file.starts_with("lib/pq-crypto/") {
-            return pq;
-        }
-
-        true
-    }));
+    build.files(include!("./files.rs"));
 
     if env("PROFILE") == "release" {
         // fortify source is only available in release mode
@@ -127,10 +97,6 @@ fn build_vendored() {
         }
     }
 
-    if !pq {
-        build.define("S2N_NO_PQ", "1");
-    }
-
     let out_dir = PathBuf::from(env("OUT_DIR"));
 
     let features = FeatureDetector::new(&out_dir, &libcrypto);
@@ -214,49 +180,6 @@ fn builder(libcrypto: &Libcrypto) -> cc::Build {
     build
 }
 
-#[cfg(feature = "cmake")]
-fn build_cmake() {
-    let mut config = cmake::Config::new("lib");
-
-    let libcrypto = Libcrypto::default();
-
-    config
-        .register_dep(&format!("aws_lc_{}", libcrypto.version))
-        .configure_arg("-DBUILD_TESTING=off");
-
-    if option_env("CARGO_FEATURE_PQ").is_none() {
-        config.configure_arg("-DS2N_NO_PQ=on");
-    }
-
-    let dst = config.build();
-
-    let lib = search(dst.join("lib64"))
-        .or_else(|| search(dst.join("lib")))
-        .or_else(|| search(dst.join("build").join("lib")))
-        .expect("could not build libs2n");
-
-    // link the built artifact
-    if lib.join("libs2n.a").exists() {
-        println!("cargo:rustc-link-lib=static=s2n");
-    } else {
-        println!("cargo:rustc-link-lib=s2n");
-    }
-
-    println!("cargo:include={}", dst.join("include").display());
-
-    // tell rust we're linking with libcrypto
-    println!("cargo:rustc-link-lib={}", libcrypto.link);
-
-    fn search(path: PathBuf) -> Option<PathBuf> {
-        if path.exists() {
-            println!("cargo:rustc-link-search={}", path.display());
-            Some(path)
-        } else {
-            None
-        }
-    }
-}
-
 #[derive(PartialEq, Eq, PartialOrd, Ord)]
 struct Libcrypto {
     version: String,
diff --git a/bindings/rust/s2n-tls-sys/templates/Cargo.template b/bindings/rust/s2n-tls-sys/templates/Cargo.template
index bdaf60b8360..326593a6958 100644
--- a/bindings/rust/s2n-tls-sys/templates/Cargo.template
+++ b/bindings/rust/s2n-tls-sys/templates/Cargo.template
@@ -24,8 +24,10 @@ include = [
 
 [features]
 default = []
+# preserve the cmake feature in case any consumers had it enabled before
+cmake = []
 quic = []
-pq = ["cmake"] # the pq build logic is complicated so just use cmake instead
+pq = []
 internal = []
 stacktrace = []
 <TOKEN_REPLACED_WITH_UNSTABLE_FEATURES>
@@ -38,7 +40,6 @@ libc = "0.2"
 
 [build-dependencies]
 cc = { version = "1.0", features = ["parallel"] }
-cmake = { version = "0.1", optional = true }
 
 [dev-dependencies]
 jobserver = "=0.1.26" # newer versions require rust 1.66, see https://github.com/aws/s2n-tls/issues/4241

From 35c9f1825a735737088e50d06fc70a4e85b0ead0 Mon Sep 17 00:00:00 2001
From: James Mayclin <maycj@amazon.com>
Date: Wed, 10 Jan 2024 14:56:02 -0800
Subject: [PATCH 05/13] fix: stack-use-after-scope variable ordering (#4355)

Co-authored-by: Lindsay Stewart <stewart.r.lindsay@gmail.com>
---
 tls/s2n_ktls_io.c | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tls/s2n_ktls_io.c b/tls/s2n_ktls_io.c
index a1fd655bc9c..34a1de6f02d 100644
--- a/tls/s2n_ktls_io.c
+++ b/tls/s2n_ktls_io.c
@@ -423,8 +423,9 @@ ssize_t s2n_ktls_sendv_with_offset(struct s2n_connection *conn, const struct iov
     POSIX_GUARD_RESULT(s2n_sendv_with_offset_total_size(bufs, count_in, offs_in, &total_bytes));
     POSIX_GUARD_RESULT(s2n_ktls_check_estimated_record_limit(conn, total_bytes));
 
-    DEFER_CLEANUP(struct s2n_blob new_bufs = { 0 }, s2n_free_or_wipe);
+    /* The order of new_bufs and new_bufs_mem matters. See https://github.com/aws/s2n-tls/issues/4354 */
     uint8_t new_bufs_mem[S2N_MAX_STACK_IOVECS_MEM] = { 0 };
+    DEFER_CLEANUP(struct s2n_blob new_bufs = { 0 }, s2n_free_or_wipe);
     POSIX_GUARD(s2n_blob_init(&new_bufs, new_bufs_mem, sizeof(new_bufs_mem)));
     if (offs > 0) {
         POSIX_GUARD_RESULT(s2n_ktls_update_bufs_with_offset(&bufs, &count, offs, &new_bufs));

From 7c471bb8d7a924c849bfbef4562432931b3f0247 Mon Sep 17 00:00:00 2001
From: James Mayclin <maycj@amazon.com>
Date: Thu, 11 Jan 2024 18:56:59 -0800
Subject: [PATCH 06/13] ci: cmake asan buildspec (#4048)

---
 codebuild/README.md               | 41 ++++++++++++++++++++
 codebuild/spec/buildspec_asan.yml | 63 +++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+)
 create mode 100644 codebuild/README.md
 create mode 100644 codebuild/spec/buildspec_asan.yml

diff --git a/codebuild/README.md b/codebuild/README.md
new file mode 100644
index 00000000000..313cfdaa4d8
--- /dev/null
+++ b/codebuild/README.md
@@ -0,0 +1,41 @@
+# Docker Image Structure
+The codebuild specifications are run on a custom docker images that have the test dependencies installed. The docker image structure is described below.
+
+### libcrypto
+Various libcryptos are installed to `/usr/local/$LIBCRYPTO` directories. For example
+```
+# non-exhaustive list
+/usr/local/openssl-1.0.2/lib/libcrypto.a
+/usr/local/openssl-1.0.2/lib/libcrypto.so
+/usr/local/openssl-1.0.2/lib/libcrypto.so.1.0.0
+/usr/local/openssl-1.0.2/lib/pkgconfig/libcrypto.pc
+/usr/local/openssl-3.0/lib64/libcrypto.a
+/usr/local/openssl-3.0/lib64/libcrypto.so.3
+/usr/local/openssl-3.0/lib64/libcrypto.so
+/usr/local/openssl-3.0/lib64/pkgconfig/libcrypto.pc
+/usr/local/boringssl/lib/libcrypto.so
+/usr/local/awslc/lib/libcrypto.a
+/usr/local/awslc/lib/libcrypto.so
+```
+
+Packages installed from the `apt` package manager can generally be found in `/usr/lib`. For example, our 32 bit build uses the 32 bit `i386` libcrypto, and it's artifacts are located at
+```
+/usr/lib/i386-linux-gnu/libcrypto.a
+/usr/lib/i386-linux-gnu/libcrypto.so.3
+/usr/lib/i386-linux-gnu/libcrypto.so
+/usr/lib/i386-linux-gnu/pkgconfig/libcrypto.pc
+```
+
+When the docker image is available locally, the structure can be easily examined by attaching an interactive terminal to the container with the following command
+```
+docker run --entrypoint /bin/bash -it --privileged <image id>
+```
+
+Then the `find` command can be used to look at the various artifacts that are available.
+```
+sudo find / -name libcrypto* # list all libcrypto artifacts
+```
+or
+```
+sudo find / -name clang* # find all clang binaries
+```
\ No newline at end of file
diff --git a/codebuild/spec/buildspec_asan.yml b/codebuild/spec/buildspec_asan.yml
new file mode 100644
index 00000000000..d6ea0e4cb0e
--- /dev/null
+++ b/codebuild/spec/buildspec_asan.yml
@@ -0,0 +1,63 @@
+---
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use
+# this file except in compliance with the License. A copy of the License is
+# located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied. See the License for the specific language governing permissions and
+# limitations under the License.
+version: 0.2
+
+# This buildspec runs on an Ubuntu22 image. That configuration is a property of
+# the codebuild job itself.
+
+# Codebuild's matrix jobs have non-differentiated names so use batch-list
+# instead.
+batch:
+  build-list:
+    # awslc is the happy path libcrypto for s2n-tls
+    - identifier: awslc
+      env:
+        compute-type: BUILD_GENERAL1_LARGE
+        variables:
+          S2N_LIBCRYPTO: awslc
+    # s2n-tls takes different code paths for ossl3, so make sure we run asan on
+    # it. See pr 4033 for a historical motivating example.
+    - identifier: openssl_3_0
+      env:
+        compute-type: BUILD_GENERAL1_LARGE
+        variables:
+          S2N_LIBCRYPTO: openssl-3.0
+    # openssl 1.1.1 is a widely deployed version of openssl.
+    - identifier: openssl_1_1_1
+      env:
+        compute-type: BUILD_GENERAL1_LARGE
+        variables:
+          S2N_LIBCRYPTO: openssl-1.1.1
+    # openssl 1.0.2 is the default distributed on AL2, and AL2 is still widely
+    # deployed
+    - identifier: openssl_1_0_2
+      env:
+        compute-type: BUILD_GENERAL1_LARGE
+        variables:
+          S2N_LIBCRYPTO: openssl-1.0.2
+
+phases:
+  build:
+    on-failure: ABORT
+    commands:
+      - |
+        cmake . -Bbuild \
+          -DCMAKE_C_COMPILER=/usr/bin/clang \
+          -DCMAKE_PREFIX_PATH=/usr/local/$S2N_LIBCRYPTO \
+          -DASAN=ON
+      - cmake --build ./build -- -j $(nproc)
+  post_build:
+    on-failure: ABORT
+    commands:
+      - CTEST_OUTPUT_ON_FAILURE=1 CTEST_PARALLEL_LEVEL=$(nproc) make -C build test

From d573ce238fe8f1c5c2f5cf7f8c65997c236155fd Mon Sep 17 00:00:00 2001
From: Lindsay Stewart <slindsay@amazon.com>
Date: Wed, 17 Jan 2024 15:33:38 -0800
Subject: [PATCH 07/13] bindings: clean up blinding tests (#4356)

---
 .../rust/s2n-tls-tokio/tests/handshake.rs     |  28 ++---
 bindings/rust/s2n-tls-tokio/tests/shutdown.rs | 108 ++++++++++--------
 2 files changed, 68 insertions(+), 68 deletions(-)

diff --git a/bindings/rust/s2n-tls-tokio/tests/handshake.rs b/bindings/rust/s2n-tls-tokio/tests/handshake.rs
index f37c2232085..84669f746c6 100644
--- a/bindings/rust/s2n-tls-tokio/tests/handshake.rs
+++ b/bindings/rust/s2n-tls-tokio/tests/handshake.rs
@@ -163,26 +163,18 @@ async fn handshake_error_with_blinding() -> Result<(), Box<dyn std::error::Error
 
     let client = TlsConnector::new(client_config.clone());
     let server = TlsAcceptor::new(server_config.clone());
+    let (server_stream, client_stream) = common::get_streams().await?;
+
+    let time_start = time::Instant::now();
+    let result = common::run_negotiate(&client, client_stream, &server, server_stream).await;
+    let time_elapsed = time_start.elapsed();
 
     // Handshake MUST NOT finish faster than minimal blinding time.
-    let (server_stream, client_stream) = common::get_streams().await?;
-    let timeout = time::timeout(
-        common::MIN_BLINDING_SECS,
-        common::run_negotiate(&client, client_stream, &server, server_stream),
-    )
-    .await;
-    assert!(timeout.is_err());
-
-    // Handshake MUST eventually gracefully close after blinding
-    let (server_stream, client_stream) = common::get_streams().await?;
-    let timeout = time::timeout(
-        common::MAX_BLINDING_SECS.mul_f32(1.1),
-        common::run_negotiate(&client, client_stream, &server, server_stream),
-    )
-    .await;
-    let result = timeout?;
-    assert!(result.is_err());
-    assert_eq!(result.unwrap_err().kind(), ErrorType::ProtocolError);
+    assert!(time_elapsed > common::MIN_BLINDING_SECS);
+
+    // Handshake MUST eventually gracefully fail after blinding
+    let error = result.unwrap_err();
+    assert_eq!(error.kind(), ErrorType::ProtocolError);
 
     Ok(())
 }
diff --git a/bindings/rust/s2n-tls-tokio/tests/shutdown.rs b/bindings/rust/s2n-tls-tokio/tests/shutdown.rs
index d5e59fc5f52..9147651789c 100644
--- a/bindings/rust/s2n-tls-tokio/tests/shutdown.rs
+++ b/bindings/rust/s2n-tls-tokio/tests/shutdown.rs
@@ -6,11 +6,15 @@ use s2n_tls_tokio::{TlsAcceptor, TlsConnector, TlsStream};
 use std::{convert::TryFrom, sync::Arc};
 use tokio::{
     io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
-    join, time,
+    time,
 };
 
 pub mod common;
 
+// An arbitrary but very long timeout.
+// No valid single IO operation should take anywhere near 10 minutes.
+pub const LONG_TIMEOUT: time::Duration = time::Duration::from_secs(600);
+
 async fn read_until_shutdown<S: AsyncRead + AsyncWrite + Unpin>(
     stream: &mut TlsStream<S>,
 ) -> Result<(), std::io::Error> {
@@ -157,7 +161,15 @@ async fn shutdown_with_blinding() -> Result<(), Box<dyn std::error::Error>> {
     let (mut client, mut server) =
         common::run_negotiate(&client, client_stream, &server, server_stream).await?;
 
-    // Trigger a blinded error for the server.
+    // Attempt to shutdown the client. This will eventually fail because the
+    // server has not written the close_notify message yet, but it will at least
+    // write the close_notify message that the server needs.
+    // Because time is mocked for testing, this does not actually take LONG_TIMEOUT.
+    // TODO: replace this with a half-close once the bindings support half-close.
+    let timeout = time::timeout(LONG_TIMEOUT, client.shutdown()).await;
+    assert!(timeout.is_err());
+
+    // Setup a bad record for the next read
     overrides.next_read(Some(Box::new(|_, _, buf| {
         // Parsing the header is one of the blinded operations
         // in s2n_recv, so provide a malformed header.
@@ -165,26 +177,29 @@ async fn shutdown_with_blinding() -> Result<(), Box<dyn std::error::Error>> {
         buf.put_slice(&zeroed_header);
         Ok(()).into()
     })));
+
+    // Trigger the blinded error
     let mut received = [0; 1];
     let result = server.read_exact(&mut received).await;
     assert!(result.is_err());
 
+    let time_start = time::Instant::now();
+    let result = server.shutdown().await;
+    let time_elapsed = time_start.elapsed();
+
     // Shutdown MUST NOT complete faster than minimal blinding time.
-    let (timeout, _) = join!(
-        time::timeout(common::MIN_BLINDING_SECS, server.shutdown()),
-        time::timeout(common::MIN_BLINDING_SECS, read_until_shutdown(&mut client)),
-    );
-    assert!(timeout.is_err());
+    assert!(time_elapsed > common::MIN_BLINDING_SECS);
 
-    // Shutdown MUST eventually complete after blinding.
-    //
-    // We check for completion, but not for success. At the moment, the
-    // call to s2n_shutdown will fail due to issues in the underlying C library.
-    let (timeout, _) = join!(
-        time::timeout(common::MAX_BLINDING_SECS, server.shutdown()),
-        time::timeout(common::MAX_BLINDING_SECS, read_until_shutdown(&mut client)),
-    );
-    assert!(timeout.is_ok());
+    // 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");
+
+    // Shutdown MUST have sent the close_notify message needed by the peer
+    // to also shutdown successfully.
+    client.shutdown().await?;
 
     Ok(())
 }
@@ -204,33 +219,31 @@ async fn shutdown_with_blinding_bad_close_record() -> Result<(), Box<dyn std::er
     let (mut client, mut server) =
         common::run_negotiate(&client, client_stream, &server, server_stream).await?;
 
-    // Turn the closure alert to a bad message
+    // Setup a bad record for the next read
     overrides.next_read(Some(Box::new(|_, _, buf| {
         // Parsing the header is one of the blinded operations
-        // in s2n_recv, so provide a malformed header.
+        // in s2n_shutdown, so provide a malformed header.
         let zeroed_header = [23, 0, 0, 0, 0];
         buf.put_slice(&zeroed_header);
         Ok(()).into()
     })));
 
+    let time_start = time::Instant::now();
+    let result = server.shutdown().await;
+    let time_elapsed = time_start.elapsed();
+
     // Shutdown MUST NOT complete faster than minimal blinding time.
-    let (timeout, _) = join!(
-        time::timeout(common::MIN_BLINDING_SECS, server.shutdown()),
-        time::timeout(common::MIN_BLINDING_SECS, read_until_shutdown(&mut client)),
-    );
-    assert!(timeout.is_err());
+    assert!(time_elapsed > common::MIN_BLINDING_SECS);
 
-    // Shutdown MUST eventually complete after blinding.
-    //
-    // We check for completion, but not for success. At the moment, the
-    // call to s2n_shutdown will fail due to issues in the underlying C library.
-    let (timeout, _) = join!(
-        time::timeout(common::MAX_BLINDING_SECS, server.shutdown()),
-        time::timeout(common::MAX_BLINDING_SECS, read_until_shutdown(&mut client)),
-    );
-    // timeout should be OK, but shutdown should return an error because of
-    // the bad record.
-    assert!(matches!(timeout, Ok(Err(_))));
+    // Shutdown MUST eventually complete with the correct error after blinding.
+    let io_error = result.unwrap_err();
+    let error: error::Error = io_error.try_into()?;
+    assert!(error.kind() == error::ErrorType::ProtocolError);
+    assert!(error.name() == "S2N_ERR_BAD_MESSAGE");
+
+    // Shutdown MUST have sent the close_notify message needed by the peer
+    // to also shutdown successfully.
+    client.shutdown().await?;
 
     Ok(())
 }
@@ -247,10 +260,10 @@ async fn shutdown_with_poll_blinding() -> Result<(), Box<dyn std::error::Error>>
     let (server_stream, client_stream) = common::get_streams().await?;
     let server_stream = common::TestStream::new(server_stream);
     let overrides = server_stream.overrides();
-    let (mut client, mut server) =
+    let (_, mut server) =
         common::run_negotiate(&client, client_stream, &server, server_stream).await?;
 
-    // Trigger a blinded error for the server.
+    // Setup a bad record for the next read
     overrides.next_read(Some(Box::new(|_, _, buf| {
         // Parsing the header is one of the blinded operations
         // in s2n_recv, so provide a malformed header.
@@ -258,26 +271,21 @@ async fn shutdown_with_poll_blinding() -> Result<(), Box<dyn std::error::Error>>
         buf.put_slice(&zeroed_header);
         Ok(()).into()
     })));
+
+    // Trigger the blinded error
     let mut received = [0; 1];
     let result = server.read_exact(&mut received).await;
     assert!(result.is_err());
 
+    let time_start = time::Instant::now();
+    let result = server.apply_blinding().await;
+    let time_elapsed = time_start.elapsed();
+
     // poll_blinding MUST NOT complete faster than minimal blinding time.
-    let (timeout, _) = join!(
-        time::timeout(common::MIN_BLINDING_SECS, server.apply_blinding()),
-        time::timeout(common::MIN_BLINDING_SECS, read_until_shutdown(&mut client)),
-    );
-    assert!(timeout.is_err());
+    assert!(time_elapsed > common::MIN_BLINDING_SECS);
 
-    // Shutdown MUST eventually complete after blinding.
-    //
-    // We check for completion, but not for success. At the moment, the
-    // call to s2n_shutdown will fail due to issues in the underlying C library.
-    let (timeout, _) = join!(
-        time::timeout(common::MAX_BLINDING_SECS, server.apply_blinding()),
-        time::timeout(common::MAX_BLINDING_SECS, read_until_shutdown(&mut client)),
-    );
-    assert!(timeout.is_ok());
+    // poll_blinding MUST eventually complete
+    assert!(result.is_ok());
 
     Ok(())
 }

From 028d3159d7e5358f4eb0318a3662bc189da646c7 Mon Sep 17 00:00:00 2001
From: Lindsay Stewart <slindsay@amazon.com>
Date: Wed, 17 Jan 2024 17:05:28 -0800
Subject: [PATCH 08/13] Move client hello parsing out of unstable (#4359)

---
 api/s2n.h                                    | 29 ++++++++++++++++++++
 api/unstable/fingerprint.h                   | 26 ------------------
 docs/usage-guide/topics/ch10-client-hello.md | 18 +++++++++++-
 3 files changed, 46 insertions(+), 27 deletions(-)

diff --git a/api/s2n.h b/api/s2n.h
index c77b97a8c9f..84e3cf01fb7 100644
--- a/api/s2n.h
+++ b/api/s2n.h
@@ -1393,6 +1393,35 @@ struct s2n_client_hello;
  */
 S2N_API extern struct s2n_client_hello *s2n_connection_get_client_hello(struct s2n_connection *conn);
 
+/**
+ * Creates an s2n_client_hello from bytes representing a ClientHello message.
+ *
+ * The input bytes should include the message header (message type and length),
+ * but not the record header.
+ *
+ * Unlike s2n_connection_get_client_hello, the s2n_client_hello returned by this
+ * method is owned by the application and must be freed with s2n_client_hello_free.
+ *
+ * This method does not support SSLv2 ClientHellos.
+ *
+ * @param bytes The raw bytes representing the ClientHello.
+ * @param size The size of raw_message.
+ * @returns A new s2n_client_hello on success, or NULL on failure.
+ */
+S2N_API extern struct s2n_client_hello *s2n_client_hello_parse_message(const uint8_t *bytes, uint32_t size);
+
+/**
+ * Frees an s2n_client_hello structure.
+ *
+ * This method should be called to free s2n_client_hellos returned by
+ * s2n_client_hello_parse_message. It will error if passed an s2n_client_hello
+ * returned by s2n_connection_get_client_hello and owned by the connection.
+ *
+ * @param ch The structure to be freed.
+ * @returns S2N_SUCCESS on success, S2N_FAILURE on failure.
+ */
+S2N_API extern int s2n_client_hello_free(struct s2n_client_hello **ch);
+
 /**
  * Function to determine the size of the raw Client Hello buffer. 
  *
diff --git a/api/unstable/fingerprint.h b/api/unstable/fingerprint.h
index 90ecbfb1448..8ad675e8866 100644
--- a/api/unstable/fingerprint.h
+++ b/api/unstable/fingerprint.h
@@ -74,29 +74,3 @@ S2N_API int s2n_client_hello_get_fingerprint_hash(struct s2n_client_hello *ch,
 S2N_API int s2n_client_hello_get_fingerprint_string(struct s2n_client_hello *ch,
         s2n_fingerprint_type type, uint32_t max_size,
         uint8_t *output, uint32_t *output_size);
-
-/**
- * Creates an s2n_client_hello from bytes representing a ClientHello message.
- *
- * Unlike s2n_connection_get_client_hello, the s2n_client_hello returned by this
- * method is owned by the application and must be freed with s2n_client_hello_free.
- *
- * This method does not support SSLv2 ClientHellos.
- *
- * @param bytes The raw bytes representing the ClientHello.
- * @param size The size of raw_message.
- * @returns A new s2n_client_hello on success, or NULL on failure.
- */
-S2N_API struct s2n_client_hello *s2n_client_hello_parse_message(const uint8_t *bytes, uint32_t size);
-
-/**
- * Frees an s2n_client_hello structure.
- *
- * This method should be called to free s2n_client_hellos returned by
- * s2n_client_hello_parse_message. It will error if passed an s2n_client_hello
- * returned by s2n_connection_get_client_hello and owned by the connection.
- *
- * @param ch The structure to be freed.
- * @returns S2N_SUCCESS on success, S2N_FAILURE on failure.
- */
-S2N_API int s2n_client_hello_free(struct s2n_client_hello **ch);
diff --git a/docs/usage-guide/topics/ch10-client-hello.md b/docs/usage-guide/topics/ch10-client-hello.md
index ac90da8cad0..d3b37368ce1 100644
--- a/docs/usage-guide/topics/ch10-client-hello.md
+++ b/docs/usage-guide/topics/ch10-client-hello.md
@@ -1,14 +1,30 @@
 # Examining the Client Hello
 
+## Getting a Client Hello
+
+### From a connection
 s2n-tls stores the received Client Hello and makes it available to the application. Call `s2n_connection_get_client_hello()` to get a pointer to the `s2n_client_hello` struct storing the Client Hello message. A NULL value will be returned if the connection has not yet received the Client Hello. The earliest point in the handshake when this struct is available is during the [Client Hello Callback](#client-hello-callback). The stored Client Hello message will not be available after calling `s2n_connection_free_handshake()`.
 
+### From raw bytes
+s2n-tls can parse a Client Hello from raw bytes. Call `s2n_client_hello_parse_message()`
+with raw bytes representing a Client Hello message (including the message header, but excluding
+the record header). The returned pointer to a `s2n_client_hello` struct behaves
+the same as a pointer returned from `s2n_connection_get_client_hello()`, except
+that the memory is owned by the application and must be freed with `s2n_client_hello_free()`.
+
+## Examining the message
+
 Call `s2n_client_hello_get_raw_message()` to retrieve the complete Client Hello message with the random bytes on it zeroed out.
 
 Call `s2n_client_hello_get_cipher_suites()` to retrieve the list of cipher suites sent by the client.
 
 Call `s2n_client_hello_get_session_id()` to retrieve the session ID sent by the client in the ClientHello message. Note that this value may not be the session ID eventually associated with this particular connection since the session ID can change when the server sends the Server Hello. The official session ID can be retrieved with `s2n_connection_get_session_id()`after the handshake completes.
 
-Call `s2n_client_hello_get_extensions()` to retrieve the entire list of extensions sent in the Client Hello. Calling `s2n_client_hello_get_extension_by_id()` allows you to interrogate the `s2n_client_hello` struct for a specific extension.
+Call `s2n_client_hello_get_extensions()` to retrieve the entire list of extensions sent in the Client Hello. Call `s2n_client_hello_get_extension_by_id()` to retrieve a specific extension. Because `s2n_client_hello_get_extension_by_id()` doesn't distinguish between zero-length extensions and missing extensions,
+`s2n_client_hello_has_extension()` should be used to check for the existence of an extension.
+
+Call `s2n_client_hello_get_supported_groups()` to retrieve the entire list of
+supported groups sent by the client.
 
 ## SSLv2
 s2n-tls will not negotiate SSLv2, but will accept SSLv2 ClientHellos advertising a

From 078d1e9b6f756a29a2fbe285d2510459df275a24 Mon Sep 17 00:00:00 2001
From: Lindsay Stewart <slindsay@amazon.com>
Date: Wed, 17 Jan 2024 18:29:24 -0800
Subject: [PATCH 09/13] ktls: add method to track key updates (#4364)

---
 api/unstable/ktls.h              | 18 ++++++++++++++++++
 tests/unit/s2n_connection_test.c | 28 ++++++++++++++++++++++++++++
 tests/unit/s2n_key_update_test.c | 24 ++++++++++++++++++++++++
 tls/s2n_connection.c             | 13 +++++++++++++
 tls/s2n_connection.h             |  4 ++++
 tls/s2n_tls13_handshake.c        |  8 ++++++++
 6 files changed, 95 insertions(+)

diff --git a/api/unstable/ktls.h b/api/unstable/ktls.h
index ee395f5f006..0d9c28fd588 100644
--- a/api/unstable/ktls.h
+++ b/api/unstable/ktls.h
@@ -102,6 +102,8 @@ S2N_API int s2n_connection_ktls_enable_recv(struct s2n_connection *conn);
  *
  * Enabling TLS1.3 with this method is considered "unsafe" because the kernel
  * currently doesn't support updating encryption keys, which is required in TLS1.3.
+ * s2n_connection_get_key_update_counts can be used to gather metrics on whether
+ * key updates are occurring on your connections before enabling TLS1.3.
  *
  * In order to safely enable TLS1.3, an application must ensure that its peer will
  * not send any KeyUpdate messages. If s2n-tls receives a KeyUpdate message while
@@ -119,6 +121,22 @@ S2N_API int s2n_connection_ktls_enable_recv(struct s2n_connection *conn);
  */
 S2N_API int s2n_config_ktls_enable_unsafe_tls13(struct s2n_config *config);
 
+/**
+ * Reports the number of times sending and receiving keys have been updated.
+ *
+ * This only applies to TLS1.3. Earlier versions do not support key updates.
+ *
+ * @warning s2n-tls only tracks up to UINT8_MAX (255) key updates. If this method
+ * reports 255 updates, then more than 255 updates may have occurred.
+ *
+ * @param conn A pointer to the connection.
+ * @param send_key_updates Number of times the sending key was updated.
+ * @param recv_key_updates Number of times the receiving key was updated.
+ * @returns S2N_SUCCESS if successful, S2N_FAILURE otherwise.
+ */
+S2N_API int s2n_connection_get_key_update_counts(struct s2n_connection *conn,
+        uint8_t *send_key_updates, uint8_t *recv_key_updates);
+
 /**
  * Sends the contents of a file as application data.
  *
diff --git a/tests/unit/s2n_connection_test.c b/tests/unit/s2n_connection_test.c
index 9e331586071..61c03d7ac0c 100644
--- a/tests/unit/s2n_connection_test.c
+++ b/tests/unit/s2n_connection_test.c
@@ -15,6 +15,7 @@
 
 #include "tls/s2n_connection.h"
 
+#include "api/unstable/ktls.h"
 #include "crypto/s2n_hash.h"
 #include "s2n_test.h"
 #include "testlib/s2n_testlib.h"
@@ -878,6 +879,33 @@ int main(int argc, char **argv)
         };
     };
 
+    /* Test s2n_connection_get_key_update_counts */
+    {
+        /* Safety */
+        DEFER_CLEANUP(struct s2n_connection *conn = s2n_connection_new(S2N_CLIENT),
+                s2n_connection_ptr_free);
+        EXPECT_NOT_NULL(conn);
+        uint8_t send_count = 0, recv_count = 0;
+        EXPECT_FAILURE_WITH_ERRNO(
+                s2n_connection_get_key_update_counts(NULL, &send_count, &recv_count),
+                S2N_ERR_NULL);
+        EXPECT_FAILURE_WITH_ERRNO(
+                s2n_connection_get_key_update_counts(conn, NULL, &recv_count),
+                S2N_ERR_NULL);
+        EXPECT_FAILURE_WITH_ERRNO(
+                s2n_connection_get_key_update_counts(conn, &send_count, NULL),
+                S2N_ERR_NULL);
+
+        /* Returns counts */
+        const uint8_t expected_send_count = 10;
+        conn->send_key_updated = expected_send_count;
+        const uint8_t expected_recv_count = 255;
+        conn->recv_key_updated = expected_recv_count;
+        EXPECT_SUCCESS(s2n_connection_get_key_update_counts(conn, &send_count, &recv_count));
+        EXPECT_EQUAL(send_count, expected_send_count);
+        EXPECT_EQUAL(recv_count, expected_recv_count);
+    }
+
     EXPECT_SUCCESS(s2n_cert_chain_and_key_free(ecdsa_chain_and_key));
     EXPECT_SUCCESS(s2n_cert_chain_and_key_free(rsa_chain_and_key));
     END_TEST();
diff --git a/tests/unit/s2n_key_update_test.c b/tests/unit/s2n_key_update_test.c
index 40e629c98d0..23955cdd2a8 100644
--- a/tests/unit/s2n_key_update_test.c
+++ b/tests/unit/s2n_key_update_test.c
@@ -221,6 +221,7 @@ int main(int argc, char **argv)
             EXPECT_SUCCESS(s2n_key_update_recv(server_conn, &input));
             EXPECT_EQUAL(server_conn->secure->client_sequence_number[0], 0);
             EXPECT_FALSE(s2n_atomic_flag_test(&server_conn->key_update_pending));
+            EXPECT_EQUAL(server_conn->recv_key_updated, 1);
 
             EXPECT_SUCCESS(s2n_connection_free(server_conn));
         };
@@ -243,6 +244,7 @@ int main(int argc, char **argv)
             EXPECT_SUCCESS(s2n_key_update_recv(client_conn, &input));
             EXPECT_EQUAL(client_conn->secure->server_sequence_number[0], 0);
             EXPECT_FALSE(s2n_atomic_flag_test(&client_conn->key_update_pending));
+            EXPECT_EQUAL(client_conn->recv_key_updated, 1);
 
             EXPECT_SUCCESS(s2n_connection_free(client_conn));
         };
@@ -294,6 +296,26 @@ int main(int argc, char **argv)
             /* KeyUpdate still pending */
             EXPECT_TRUE(s2n_atomic_flag_test(&conn->key_update_pending));
         };
+
+        /* Receiving a KeyUpdate doesn't overflow the key update count */
+        {
+            DEFER_CLEANUP(struct s2n_stuffer input, s2n_stuffer_free);
+            EXPECT_SUCCESS(s2n_stuffer_growable_alloc(&input, 0));
+
+            DEFER_CLEANUP(struct s2n_connection *conn = s2n_connection_new(S2N_SERVER),
+                    s2n_connection_ptr_free);
+            EXPECT_NOT_NULL(conn);
+            conn->actual_protocol_version = S2N_TLS13;
+            conn->secure->cipher_suite = cipher_suite_with_limit;
+            conn->recv_key_updated = UINT8_MAX;
+
+            conn->secure->client_sequence_number[0] = 1;
+            EXPECT_SUCCESS(s2n_stuffer_write_uint8(&input, S2N_KEY_UPDATE_NOT_REQUESTED));
+            EXPECT_SUCCESS(s2n_key_update_recv(conn, &input));
+            EXPECT_EQUAL(conn->secure->client_sequence_number[0], 0);
+            EXPECT_EQUAL(conn->recv_key_updated, UINT8_MAX);
+            EXPECT_EQUAL(conn->send_key_updated, 0);
+        };
     };
 
     /* s2n_key_update_send */
@@ -319,6 +341,7 @@ int main(int argc, char **argv)
             EXPECT_EQUAL(s2n_atomic_flag_test(&client_conn->key_update_pending), false);
             EXPECT_BYTEARRAY_EQUAL(client_conn->secure->client_sequence_number, zeroed_sequence_number, S2N_TLS_SEQUENCE_NUM_LEN);
             EXPECT_TRUE(s2n_stuffer_data_available(&stuffer) > 0);
+            EXPECT_EQUAL(client_conn->send_key_updated, 1);
 
             EXPECT_SUCCESS(s2n_stuffer_free(&stuffer));
             EXPECT_SUCCESS(s2n_connection_free(client_conn));
@@ -346,6 +369,7 @@ int main(int argc, char **argv)
             EXPECT_EQUAL(s2n_atomic_flag_test(&client_conn->key_update_pending), false);
             EXPECT_BYTEARRAY_EQUAL(client_conn->secure->client_sequence_number, zeroed_sequence_number, S2N_TLS_SEQUENCE_NUM_LEN);
             EXPECT_TRUE(s2n_stuffer_data_available(&stuffer) > 0);
+            EXPECT_EQUAL(client_conn->send_key_updated, 1);
 
             EXPECT_SUCCESS(s2n_stuffer_free(&stuffer));
             EXPECT_SUCCESS(s2n_connection_free(client_conn));
diff --git a/tls/s2n_connection.c b/tls/s2n_connection.c
index 2b601bdb6b5..cabc2ea9f15 100644
--- a/tls/s2n_connection.c
+++ b/tls/s2n_connection.c
@@ -25,6 +25,8 @@
 #include <unistd.h>
 
 #include "api/s2n.h"
+/* Required for s2n_connection_get_key_update_counts */
+#include "api/unstable/ktls.h"
 #include "crypto/s2n_certificate.h"
 #include "crypto/s2n_cipher.h"
 #include "crypto/s2n_crypto.h"
@@ -1689,3 +1691,14 @@ S2N_RESULT s2n_connection_get_sequence_number(struct s2n_connection *conn,
 
     return S2N_RESULT_OK;
 }
+
+int s2n_connection_get_key_update_counts(struct s2n_connection *conn,
+        uint8_t *send_key_updates, uint8_t *recv_key_updates)
+{
+    POSIX_ENSURE_REF(conn);
+    POSIX_ENSURE_REF(send_key_updates);
+    POSIX_ENSURE_REF(recv_key_updates);
+    *send_key_updates = conn->send_key_updated;
+    *recv_key_updates = conn->recv_key_updated;
+    return S2N_SUCCESS;
+}
diff --git a/tls/s2n_connection.h b/tls/s2n_connection.h
index 529ef4bab80..1f8c735f127 100644
--- a/tls/s2n_connection.h
+++ b/tls/s2n_connection.h
@@ -383,6 +383,10 @@ struct s2n_connection {
      * The writer clears it after a KeyUpdate is sent.
      */
     s2n_atomic_flag key_update_pending;
+
+    /* Track KeyUpdates for metrics */
+    uint8_t send_key_updated;
+    uint8_t recv_key_updated;
 };
 
 S2N_CLEANUP_RESULT s2n_connection_ptr_free(struct s2n_connection **s2n_connection);
diff --git a/tls/s2n_tls13_handshake.c b/tls/s2n_tls13_handshake.c
index dcbf9017655..a4fb5d49c99 100644
--- a/tls/s2n_tls13_handshake.c
+++ b/tls/s2n_tls13_handshake.c
@@ -182,13 +182,21 @@ int s2n_update_application_traffic_keys(struct s2n_connection *conn, s2n_mode mo
     s2n_tls13_key_blob(app_key, conn->secure->cipher_suite->record_alg->cipher->key_material_size);
 
     /* Derives next generation of traffic key */
+    uint8_t *count = NULL;
     POSIX_GUARD(s2n_tls13_derive_traffic_keys(&keys, &app_secret_update, &app_key, &app_iv));
     if (status == RECEIVING) {
         POSIX_GUARD(conn->secure->cipher_suite->record_alg->cipher->set_decryption_key(old_key, &app_key));
+        count = &conn->recv_key_updated;
     } else {
         POSIX_GUARD(conn->secure->cipher_suite->record_alg->cipher->set_encryption_key(old_key, &app_key));
+        count = &conn->send_key_updated;
     }
 
+    /* Increment the count.
+     * Don't treat overflows as errors-- we only do best-effort reporting.
+     */
+    *count = MIN(UINT8_MAX, *count + 1);
+
     /* According to https://tools.ietf.org/html/rfc8446#section-5.3:
      * Each sequence number is set to zero at the beginning of a connection and
      * whenever the key is changed; the first record transmitted under a particular traffic key

From 7f84701e4efef70f65d7f0208d190f440f5b7a37 Mon Sep 17 00:00:00 2001
From: Alex Weibel <aweibel@amazon.com>
Date: Thu, 18 Jan 2024 10:16:01 -0800
Subject: [PATCH 10/13] Add new PQ TLS Policies (#4327)

---
 tests/unit/s2n_security_rules_test.c |  2 ++
 tls/s2n_cipher_preferences.c         | 40 ++++++++++++++++++++++++++++
 tls/s2n_cipher_preferences.h         |  3 +++
 tls/s2n_kem_preferences.c            | 16 +++++++++++
 tls/s2n_kem_preferences.h            |  1 +
 tls/s2n_security_policies.c          | 36 +++++++++++++++++++++++++
 tls/s2n_security_rules.c             | 14 ++++++++++
 7 files changed, 112 insertions(+)

diff --git a/tests/unit/s2n_security_rules_test.c b/tests/unit/s2n_security_rules_test.c
index 5c69d9f2824..60d304f4699 100644
--- a/tests/unit/s2n_security_rules_test.c
+++ b/tests/unit/s2n_security_rules_test.c
@@ -126,6 +126,7 @@ int main(int argc, char **argv)
         .signature_preferences = &valid_sig_prefs,
         .certificate_signature_preferences = &valid_sig_prefs,
         .ecc_preferences = &valid_ecc_prefs,
+        .kem_preferences = &kem_preferences_null,
         .minimum_protocol_version = VALID_VERSION,
     };
     const struct s2n_security_policy invalid_policy = {
@@ -133,6 +134,7 @@ int main(int argc, char **argv)
         .signature_preferences = &invalid_sig_prefs,
         .certificate_signature_preferences = &invalid_sig_prefs,
         .ecc_preferences = &invalid_ecc_prefs,
+        .kem_preferences = &kem_preferences_null,
         .minimum_protocol_version = EXAMPLE_INVALID_VERSION,
     };
 
diff --git a/tls/s2n_cipher_preferences.c b/tls/s2n_cipher_preferences.c
index cd3406d2a92..3552c72aaa8 100644
--- a/tls/s2n_cipher_preferences.c
+++ b/tls/s2n_cipher_preferences.c
@@ -1891,6 +1891,46 @@ const struct s2n_cipher_preferences cipher_preferences_pq_tls_1_3_2023_06_01 = {
     .allow_chacha20_boosting = false,
 };
 
+struct s2n_cipher_suite *cipher_suites_20231213[] = {
+    &s2n_tls13_aes_128_gcm_sha256,
+    &s2n_tls13_aes_256_gcm_sha384,
+    &s2n_ecdhe_ecdsa_with_aes_128_gcm_sha256,
+    &s2n_ecdhe_rsa_with_aes_128_gcm_sha256,
+    &s2n_ecdhe_ecdsa_with_aes_256_gcm_sha384,
+    &s2n_ecdhe_rsa_with_aes_256_gcm_sha384,
+    &s2n_ecdhe_ecdsa_with_aes_128_cbc_sha256,
+    &s2n_ecdhe_rsa_with_aes_128_cbc_sha256,
+    &s2n_ecdhe_ecdsa_with_aes_256_cbc_sha384,
+    &s2n_ecdhe_rsa_with_aes_256_cbc_sha384,
+    &s2n_rsa_with_aes_128_gcm_sha256,
+    &s2n_rsa_with_aes_256_gcm_sha384,
+    &s2n_rsa_with_aes_128_cbc_sha256,
+    &s2n_rsa_with_aes_256_cbc_sha256,
+};
+
+const struct s2n_cipher_preferences cipher_preferences_20231213 = {
+    .count = s2n_array_len(cipher_suites_20231213),
+    .suites = cipher_suites_20231213,
+};
+
+struct s2n_cipher_suite *cipher_suites_20231214[] = {
+    &s2n_tls13_aes_128_gcm_sha256,
+    &s2n_tls13_aes_256_gcm_sha384,
+    &s2n_ecdhe_ecdsa_with_aes_128_gcm_sha256,
+    &s2n_ecdhe_rsa_with_aes_128_gcm_sha256,
+    &s2n_ecdhe_ecdsa_with_aes_256_gcm_sha384,
+    &s2n_ecdhe_rsa_with_aes_256_gcm_sha384,
+    &s2n_ecdhe_ecdsa_with_aes_128_cbc_sha256,
+    &s2n_ecdhe_rsa_with_aes_128_cbc_sha256,
+    &s2n_ecdhe_ecdsa_with_aes_256_cbc_sha384,
+    &s2n_ecdhe_rsa_with_aes_256_cbc_sha384,
+};
+
+const struct s2n_cipher_preferences cipher_preferences_20231214 = {
+        .count = s2n_array_len(cipher_suites_20231214),
+        .suites = cipher_suites_20231214,
+};
+
 struct s2n_cipher_suite *cipher_suites_kms_fips_tls_1_2_2018_10[] = {
     &s2n_ecdhe_rsa_with_aes_256_gcm_sha384,
     &s2n_ecdhe_rsa_with_aes_128_gcm_sha256,
diff --git a/tls/s2n_cipher_preferences.h b/tls/s2n_cipher_preferences.h
index a9f1830c821..dc68e1fbb8c 100644
--- a/tls/s2n_cipher_preferences.h
+++ b/tls/s2n_cipher_preferences.h
@@ -55,6 +55,9 @@ extern const struct s2n_cipher_preferences cipher_preferences_20210816_gcm;
 extern const struct s2n_cipher_preferences cipher_preferences_20210825;
 extern const struct s2n_cipher_preferences cipher_preferences_20210825_gcm;
 extern const struct s2n_cipher_preferences cipher_preferences_20210831;
+extern const struct s2n_cipher_preferences cipher_preferences_20231213;
+extern const struct s2n_cipher_preferences cipher_preferences_20231214;
+
 extern const struct s2n_cipher_preferences cipher_preferences_default_fips;
 
 extern const struct s2n_cipher_preferences cipher_preferences_test_all;
diff --git a/tls/s2n_kem_preferences.c b/tls/s2n_kem_preferences.c
index b30bb86695b..7a30d34f3ee 100644
--- a/tls/s2n_kem_preferences.c
+++ b/tls/s2n_kem_preferences.c
@@ -34,6 +34,13 @@ const struct s2n_kem_group *pq_kem_groups_r3_2023_06[] = {
     &s2n_x25519_kyber_512_r3,
 };
 
+const struct s2n_kem_group *pq_kem_groups_r3_2023_12[] = {
+    &s2n_secp256r1_kyber_768_r3,
+    &s2n_secp384r1_kyber_768_r3,
+    &s2n_secp521r1_kyber_1024_r3,
+    &s2n_secp256r1_kyber_512_r3,
+};
+
 const struct s2n_kem_preferences kem_preferences_pq_tls_1_0_2021_05 = {
     .kem_count = s2n_array_len(pq_kems_r3_2021_05),
     .kems = pq_kems_r3_2021_05,
@@ -59,6 +66,15 @@ const struct s2n_kem_preferences kem_preferences_pq_tls_1_3_2023_06 = {
     .tls13_pq_hybrid_draft_revision = 5
 };
 
+/* Same as kem_preferences_pq_tls_1_3_2023_06, but without x25519 */
+const struct s2n_kem_preferences kem_preferences_pq_tls_1_3_2023_12 = {
+    .kem_count = 0,
+    .kems = NULL,
+    .tls13_kem_group_count = s2n_array_len(pq_kem_groups_r3_2023_12),
+    .tls13_kem_groups = pq_kem_groups_r3_2023_12,
+    .tls13_pq_hybrid_draft_revision = 5
+};
+
 const struct s2n_kem_preferences kem_preferences_all = {
     .kem_count = s2n_array_len(pq_kems_r3_2021_05),
     .kems = pq_kems_r3_2021_05,
diff --git a/tls/s2n_kem_preferences.h b/tls/s2n_kem_preferences.h
index 59536da4e24..0d10b45a08c 100644
--- a/tls/s2n_kem_preferences.h
+++ b/tls/s2n_kem_preferences.h
@@ -47,6 +47,7 @@ extern const struct s2n_kem_group *pq_kem_groups_r3_2023_06[];
 extern const struct s2n_kem_preferences kem_preferences_pq_tls_1_0_2021_05;
 extern const struct s2n_kem_preferences kem_preferences_pq_tls_1_0_2023_01;
 extern const struct s2n_kem_preferences kem_preferences_pq_tls_1_3_2023_06;
+extern const struct s2n_kem_preferences kem_preferences_pq_tls_1_3_2023_12;
 extern const struct s2n_kem_preferences kem_preferences_all;
 extern const struct s2n_kem_preferences kem_preferences_null;
 
diff --git a/tls/s2n_security_policies.c b/tls/s2n_security_policies.c
index 2b53f6bd9e9..0ac9e960da9 100644
--- a/tls/s2n_security_policies.c
+++ b/tls/s2n_security_policies.c
@@ -731,6 +731,39 @@ const struct s2n_security_policy security_policy_pq_tls_1_2_2023_10_10 = {
     .ecc_preferences = &s2n_ecc_preferences_20200310,
 };
 
+/* General purpose "mostly" FIPS + PQ policy (with the exception of supporting RSA Key Exchange for backwards compatibility). */
+const struct s2n_security_policy security_policy_pq_20231213 = {
+    .minimum_protocol_version = S2N_TLS12,
+    .cipher_preferences = &cipher_preferences_20231213,
+    .kem_preferences = &kem_preferences_pq_tls_1_3_2023_12,
+    .signature_preferences = &s2n_signature_preferences_20230317,
+    .ecc_preferences = &s2n_ecc_preferences_20201021,
+};
+
+/* General purpose FIPS + PQ policy that meets all current FIPS requirements. */
+const struct s2n_security_policy security_policy_pq_20231214 = {
+    .minimum_protocol_version = S2N_TLS12,
+    .cipher_preferences = &cipher_preferences_20231214,
+    .kem_preferences = &kem_preferences_pq_tls_1_3_2023_12,
+    .signature_preferences = &s2n_signature_preferences_20230317,
+    .ecc_preferences = &s2n_ecc_preferences_20201021,
+    .rules = {
+            [S2N_FIPS_140_3] = true,
+    },
+};
+
+/* FIPS + PQ Policy that uses KMS's FIPS cipher preference list and meets all current FIPS requirements. */
+const struct s2n_security_policy security_policy_pq_20231215 = {
+    .minimum_protocol_version = S2N_TLS12,
+    .cipher_preferences = &cipher_preferences_kms_fips_tls_1_2_2021_08,
+    .kem_preferences = &kem_preferences_pq_tls_1_3_2023_12,
+    .signature_preferences = &s2n_signature_preferences_20230317,
+    .ecc_preferences = &s2n_ecc_preferences_20201021,
+    .rules = {
+            [S2N_FIPS_140_3] = true,
+    },
+};
+
 const struct s2n_security_policy security_policy_kms_fips_tls_1_2_2018_10 = {
     .minimum_protocol_version = S2N_TLS12,
     .cipher_preferences = &cipher_preferences_kms_fips_tls_1_2_2018_10,
@@ -1103,6 +1136,9 @@ struct s2n_security_policy_selection security_policy_selection[] = {
     { .version = "PQ-TLS-1-2-2023-10-08", .security_policy = &security_policy_pq_tls_1_2_2023_10_08, .ecc_extension_required = 0, .pq_kem_extension_required = 0 },
     { .version = "PQ-TLS-1-2-2023-10-09", .security_policy = &security_policy_pq_tls_1_2_2023_10_09, .ecc_extension_required = 0, .pq_kem_extension_required = 0 },
     { .version = "PQ-TLS-1-2-2023-10-10", .security_policy = &security_policy_pq_tls_1_2_2023_10_10, .ecc_extension_required = 0, .pq_kem_extension_required = 0 },
+    { .version = "PQ-TLS-1-2-2023-12-13", .security_policy = &security_policy_pq_20231213, .ecc_extension_required = 0, .pq_kem_extension_required = 0 },
+    { .version = "PQ-TLS-1-2-2023-12-14", .security_policy = &security_policy_pq_20231214, .ecc_extension_required = 0, .pq_kem_extension_required = 0 },
+    { .version = "PQ-TLS-1-2-2023-12-15", .security_policy = &security_policy_pq_20231215, .ecc_extension_required = 0, .pq_kem_extension_required = 0 },
     { .version = "20140601", .security_policy = &security_policy_20140601, .ecc_extension_required = 0, .pq_kem_extension_required = 0 },
     { .version = "20141001", .security_policy = &security_policy_20141001, .ecc_extension_required = 0, .pq_kem_extension_required = 0 },
     { .version = "20150202", .security_policy = &security_policy_20150202, .ecc_extension_required = 0, .pq_kem_extension_required = 0 },
diff --git a/tls/s2n_security_rules.c b/tls/s2n_security_rules.c
index 9d85199d51b..ebeefe99430 100644
--- a/tls/s2n_security_rules.c
+++ b/tls/s2n_security_rules.c
@@ -166,6 +166,20 @@ S2N_RESULT s2n_security_rule_validate_policy(const struct s2n_security_rule *rul
                 "curve", curve->name, i + 1));
     }
 
+    const struct s2n_kem_preferences *kem_prefs = policy->kem_preferences;
+    RESULT_ENSURE_REF(kem_prefs);
+    for (size_t i = 0; i < kem_prefs->tls13_kem_group_count; i++) {
+        const struct s2n_kem_group *kem_group = kem_prefs->tls13_kem_groups[i];
+        const struct s2n_ecc_named_curve *curve = kem_group->curve;
+        RESULT_ENSURE_REF(curve);
+        bool is_valid = false;
+        RESULT_ENSURE_REF(rule->validate_curve);
+        RESULT_GUARD(rule->validate_curve(curve, &is_valid));
+        RESULT_GUARD(s2n_security_rule_result_process(result, is_valid,
+                error_msg_format_name, rule->name, policy_name,
+                "curve", curve->name, i + 1));
+    }
+
     bool is_valid = false;
     RESULT_ENSURE_REF(rule->validate_version);
     RESULT_GUARD(rule->validate_version(policy->minimum_protocol_version, &is_valid));

From d946c2de8ecf3090f970fc9d82cb2b67e8193079 Mon Sep 17 00:00:00 2001
From: Lindsay Stewart <slindsay@amazon.com>
Date: Thu, 18 Jan 2024 15:53:14 -0800
Subject: [PATCH 11/13] Fix s2n_shutdown + failed recv bug (#4350)

---
 bindings/rust/s2n-tls-tokio/tests/shutdown.rs |  8 +--
 tests/unit/s2n_shutdown_test.c                | 55 +++++++++++++++++++
 tls/s2n_shutdown.c                            | 12 ++--
 3 files changed, 64 insertions(+), 11 deletions(-)

diff --git a/bindings/rust/s2n-tls-tokio/tests/shutdown.rs b/bindings/rust/s2n-tls-tokio/tests/shutdown.rs
index 9147651789c..29931bd08bf 100644
--- 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
index 61fddacfc9a..e2054a6ef50 100644
--- 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
index 6086322a15e..3ad818d6c00 100644
--- 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;

From 5744711c70e577a7c0a37a60e8cd9c5640cbe1ea Mon Sep 17 00:00:00 2001
From: Lindsay Stewart <slindsay@amazon.com>
Date: Thu, 18 Jan 2024 18:16:14 -0800
Subject: [PATCH 12/13] bindings: fix handling of s2n_shutdown errors (#4358)

---
 bindings/rust/s2n-tls-tokio/src/lib.rs        | 111 ++++++----------
 .../rust/s2n-tls-tokio/tests/common/mod.rs    |   2 +
 .../rust/s2n-tls-tokio/tests/common/stream.rs |  32 ++++-
 bindings/rust/s2n-tls-tokio/tests/shutdown.rs | 123 +++++++++++++++++-
 4 files changed, 195 insertions(+), 73 deletions(-)

diff --git a/bindings/rust/s2n-tls-tokio/src/lib.rs b/bindings/rust/s2n-tls-tokio/src/lib.rs
index 52f31daaa69..3e6d60e687b 100644
--- a/bindings/rust/s2n-tls-tokio/src/lib.rs
+++ b/bindings/rust/s2n-tls-tokio/src/lib.rs
@@ -2,7 +2,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use errno::{set_errno, Errno};
-use pin_project_lite::pin_project;
 use s2n_tls::{
     config::Config,
     connection::{Builder, Connection},
@@ -12,7 +11,7 @@ use s2n_tls::{
 use std::{
     fmt,
     future::Future,
-    io, mem,
+    io,
     os::raw::{c_int, c_void},
     pin::Pin,
     task::{
@@ -139,18 +138,6 @@ where
     }
 }
 
-pin_project! {
-struct BlindingState {
-    #[pin]
-    timer: Sleep,
-
-    // The remembered error if we got into blinding because of
-    // an error, or Ok(()) if we didn't. After returning the error,
-    // this goes back to Ok(()).
-    remembered_error: Result<(), Error>,
-}
-}
-
 pub struct TlsStream<S, C = Connection>
 where
     C: AsRef<Connection> + AsMut<Connection> + Unpin,
@@ -158,7 +145,8 @@ where
 {
     conn: C,
     stream: S,
-    blinding: Option<Pin<Box<BlindingState>>>,
+    blinding: Option<Pin<Box<Sleep>>>,
+    shutdown_error: Option<Error>,
 }
 
 impl<S, C> TlsStream<S, C>
@@ -182,6 +170,7 @@ where
             conn,
             stream,
             blinding: None,
+            shutdown_error: None,
         };
         TlsHandshake {
             tls: &mut tls,
@@ -255,35 +244,6 @@ where
         })
     }
 
-    // Sets the blinding timer to the remaining blinding delay and possibly
-    // remembers an error.
-    //
-    // Returns the error if there was no blinding needed and the error
-    // did not need to be remembered.
-    fn set_blinding_timer(
-        self: Pin<&mut Self>,
-        mut remembered_error: Result<(), Error>,
-    ) -> Result<(), Error> {
-        let tls = self.get_mut();
-
-        if tls.blinding.is_none() {
-            let delay = tls.as_ref().remaining_blinding_delay()?;
-            if !delay.is_zero() {
-                // Sleep operates at the milisecond resolution, so add an extra
-                // millisecond to account for any stray nanoseconds.
-                let safety = Duration::from_millis(1);
-                // Return the error *later*, after the blinding is done
-                let remembered_error = mem::replace(&mut remembered_error, Ok(()));
-                tls.blinding = Some(Box::pin(BlindingState {
-                    timer: sleep(delay.saturating_add(safety)),
-                    remembered_error,
-                }));
-            }
-        }
-
-        remembered_error
-    }
-
     /// Polls the blinding timer, if there is any.
     ///
     /// s2n has a "blinding" functionality - when a bad behavior from the peer
@@ -296,25 +256,24 @@ where
     /// before dropping an s2n connection, you should wait until either
     /// `poll_blinding` or `poll_shutdown` (which calls `poll_blinding`
     /// internally) returns ready.
-    pub fn poll_blinding(
-        mut self: Pin<&mut Self>,
-        ctx: &mut Context<'_>,
-    ) -> Poll<Result<(), Error>> {
-        self.as_mut().set_blinding_timer(Ok(()))?;
-
+    pub fn poll_blinding(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Result<(), Error>> {
         let tls = self.get_mut();
 
-        if let Some(blinding) = &mut tls.blinding {
-            ready!(blinding.as_mut().project().timer.as_mut().poll(ctx));
-
-            // Set blinding to None to ensure the next go can have blinding
-            let mut blinding = tls.blinding.take().unwrap();
+        if tls.blinding.is_none() {
+            let delay = tls.as_ref().remaining_blinding_delay()?;
+            if !delay.is_zero() {
+                // Sleep operates at the milisecond resolution, so add an extra
+                // millisecond to account for any stray nanoseconds.
+                let safety = Duration::from_millis(1);
+                tls.blinding = Some(Box::pin(sleep(delay.saturating_add(safety))));
+            }
+        };
 
-            // If there is an error, return it
-            mem::replace(blinding.as_mut().project().remembered_error, Ok(()))?;
+        if let Some(timer) = tls.blinding.as_mut() {
+            ready!(timer.as_mut().poll(ctx));
+            tls.blinding = None;
         }
 
-        // Otherwise we are OK
         Poll::Ready(Ok(()))
     }
 
@@ -404,19 +363,33 @@ where
     fn poll_shutdown(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<io::Result<()>> {
         ready!(self.as_mut().poll_blinding(ctx))?;
 
-        let status = ready!(self.as_mut().with_io(ctx, |mut context| {
-            context.conn.as_mut().poll_shutdown().map(|r| r.map(|_| ()))
-        }));
+        // s2n_shutdown must not be called again if it errors
+        if self.shutdown_error.is_none() {
+            let result = ready!(self.as_mut().with_io(ctx, |mut context| {
+                context.conn.as_mut().poll_shutdown().map(|r| r.map(|_| ()))
+            }));
+            if let Err(error) = result {
+                self.shutdown_error = Some(error);
+                // s2n_shutdown reading might have triggered blinding again
+                ready!(self.as_mut().poll_blinding(ctx))?;
+            }
+        };
+
+        let tcp_result = ready!(Pin::new(&mut self.as_mut().stream).poll_shutdown(ctx));
 
-        if let Err(e) = status {
-            // In case of an error shutting down, make sure you wait for
-            // the blinding timeout.
-            self.as_mut().set_blinding_timer(Err(e))?;
-            ready!(self.as_mut().poll_blinding(ctx))?;
-            unreachable!("should have returned the error we just put in!");
-        }
+        if let Some(err) = self.shutdown_error.take() {
+            // poll methods shouldn't be called again after returning Ready, but
+            // nothing actually prevents it so poll_shutdown should handle it.
+            // s2n_shutdown can be polled indefinitely after succeeding, but not after failing.
+            // s2n_tls::error::Error isn't cloneable, so we can't just return the same error
+            // if poll_shutdown is called again. Instead, save a different error.
+            let next_error = Error::application("Shutdown called again after error".into());
+            self.shutdown_error = Some(next_error);
 
-        Pin::new(&mut self.as_mut().stream).poll_shutdown(ctx)
+            Ready(Err(io::Error::from(err)))
+        } else {
+            Ready(tcp_result)
+        }
     }
 }
 
diff --git a/bindings/rust/s2n-tls-tokio/tests/common/mod.rs b/bindings/rust/s2n-tls-tokio/tests/common/mod.rs
index b9fbf7832d8..fbcab0e1e06 100644
--- a/bindings/rust/s2n-tls-tokio/tests/common/mod.rs
+++ b/bindings/rust/s2n-tls-tokio/tests/common/mod.rs
@@ -40,6 +40,8 @@ pub static RSA_KEY_PEM: &[u8] = include_bytes!(concat!(
 pub const MIN_BLINDING_SECS: Duration = Duration::from_secs(10);
 pub const MAX_BLINDING_SECS: Duration = Duration::from_secs(30);
 
+pub static TEST_STR: &str = "hello world";
+
 pub async fn get_streams() -> Result<(TcpStream, TcpStream), tokio::io::Error> {
     let localhost = "127.0.0.1".to_owned();
     let listener = TcpListener::bind(format!("{}:0", localhost)).await?;
diff --git a/bindings/rust/s2n-tls-tokio/tests/common/stream.rs b/bindings/rust/s2n-tls-tokio/tests/common/stream.rs
index 6ef86cfce0c..889f2926282 100644
--- a/bindings/rust/s2n-tls-tokio/tests/common/stream.rs
+++ b/bindings/rust/s2n-tls-tokio/tests/common/stream.rs
@@ -14,11 +14,13 @@ use tokio::{
 
 type ReadFn = Box<dyn Fn(Pin<&mut TcpStream>, &mut Context, &mut ReadBuf) -> Poll<io::Result<()>>>;
 type WriteFn = Box<dyn Fn(Pin<&mut TcpStream>, &mut Context, &[u8]) -> Poll<io::Result<usize>>>;
+type ShutdownFn = Box<dyn Fn(Pin<&mut TcpStream>, &mut Context) -> Poll<io::Result<()>>>;
 
 #[derive(Default)]
 struct OverrideMethods {
     next_read: Option<ReadFn>,
     next_write: Option<WriteFn>,
+    next_shutdown: Option<ShutdownFn>,
 }
 
 #[derive(Default)]
@@ -36,6 +38,22 @@ impl Overrides {
             overrides.next_write = input;
         }
     }
+
+    pub fn next_shutdown(&self, input: Option<ShutdownFn>) {
+        if let Ok(mut overrides) = self.0.lock() {
+            overrides.next_shutdown = input;
+        }
+    }
+
+    pub fn is_consumed(&self) -> bool {
+        if let Ok(overrides) = self.0.lock() {
+            overrides.next_read.is_none()
+                && overrides.next_write.is_none()
+                && overrides.next_shutdown.is_none()
+        } else {
+            false
+        }
+    }
 }
 
 unsafe impl Send for Overrides {}
@@ -100,7 +118,17 @@ impl AsyncWrite for TestStream {
         Pin::new(&mut self.stream).poll_flush(ctx)
     }
 
-    fn poll_shutdown(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<io::Result<()>> {
-        Pin::new(&mut self.stream).poll_shutdown(ctx)
+    fn poll_shutdown(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        let s = self.get_mut();
+        let stream = Pin::new(&mut s.stream);
+        let action = match s.overrides.0.lock() {
+            Ok(mut overrides) => overrides.next_shutdown.take(),
+            _ => None,
+        };
+        if let Some(f) = action {
+            (f)(stream, ctx)
+        } else {
+            stream.poll_shutdown(ctx)
+        }
     }
 }
diff --git a/bindings/rust/s2n-tls-tokio/tests/shutdown.rs b/bindings/rust/s2n-tls-tokio/tests/shutdown.rs
index 29931bd08bf..e67b87599dd 100644
--- a/bindings/rust/s2n-tls-tokio/tests/shutdown.rs
+++ b/bindings/rust/s2n-tls-tokio/tests/shutdown.rs
@@ -3,7 +3,12 @@
 
 use s2n_tls::error;
 use s2n_tls_tokio::{TlsAcceptor, TlsConnector, TlsStream};
-use std::{convert::TryFrom, sync::Arc};
+use std::{
+    convert::TryFrom,
+    io,
+    sync::Arc,
+    task::Poll::{Pending, Ready},
+};
 use tokio::{
     io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
     time,
@@ -164,7 +169,11 @@ async fn shutdown_with_blinding() -> Result<(), Box<dyn std::error::Error>> {
     // Attempt to shutdown the client. This will eventually fail because the
     // server has not written the close_notify message yet, but it will at least
     // write the close_notify message that the server needs.
-    // Because time is mocked for testing, this does not actually take LONG_TIMEOUT.
+    //
+    // Because this test begins paused and relies on auto-advancing, this does
+    // not actually require waiting LONG_TIMEOUT. See the tokio `pause()` docs:
+    // https://docs.rs/tokio/latest/tokio/time/fn.pause.html
+    //
     // TODO: replace this with a half-close once the bindings support half-close.
     let timeout = time::timeout(LONG_TIMEOUT, client.shutdown()).await;
     assert!(timeout.is_err());
@@ -285,3 +294,113 @@ async fn shutdown_with_poll_blinding() -> Result<(), Box<dyn std::error::Error>>
 
     Ok(())
 }
+
+#[tokio::test(start_paused = true)]
+async fn shutdown_with_tcp_error() -> Result<(), Box<dyn std::error::Error>> {
+    let client = TlsConnector::new(common::client_config()?.build()?);
+    let server = TlsAcceptor::new(common::server_config()?.build()?);
+
+    let (server_stream, client_stream) = common::get_streams().await?;
+    let server_stream = common::TestStream::new(server_stream);
+    let overrides = server_stream.overrides();
+
+    let (mut client, mut server) =
+        common::run_negotiate(&client, client_stream, &server, server_stream).await?;
+
+    // Attempt to shutdown the client. This will eventually fail because the
+    // server has not written the close_notify message yet, but it will at least
+    // write the close_notify message that the server needs.
+    //
+    // Because this test begins paused and relies on auto-advancing, this does
+    // not actually require waiting LONG_TIMEOUT. See the tokio `pause()` docs:
+    // https://docs.rs/tokio/latest/tokio/time/fn.pause.html
+    //
+    // TODO: replace this with a half-close once the bindings support half-close.
+    _ = time::timeout(time::Duration::from_secs(600), client.shutdown()).await;
+
+    // The underlying stream should return a unique error on shutdown
+    overrides.next_shutdown(Some(Box::new(|_, _| {
+        Ready(Err(io::Error::new(io::ErrorKind::Other, common::TEST_STR)))
+    })));
+
+    // Shutdown should complete with the correct error from the underlying stream
+    let result = server.shutdown().await;
+    let error = result.unwrap_err().into_inner().unwrap();
+    assert!(error.to_string() == common::TEST_STR);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn shutdown_with_tls_error_and_tcp_error() -> Result<(), Box<dyn std::error::Error>> {
+    let client = TlsConnector::new(common::client_config()?.build()?);
+    let server = TlsAcceptor::new(common::server_config()?.build()?);
+
+    let (server_stream, client_stream) = common::get_streams().await?;
+    let server_stream = common::TestStream::new(server_stream);
+    let overrides = server_stream.overrides();
+
+    let (_, mut server) =
+        common::run_negotiate(&client, client_stream, &server, server_stream).await?;
+
+    // Both s2n_shutdown and the underlying stream should error on shutdown
+    overrides.next_read(Some(Box::new(|_, _, _| {
+        Ready(Err(io::Error::from(io::ErrorKind::Other)))
+    })));
+    overrides.next_shutdown(Some(Box::new(|_, _| {
+        Ready(Err(io::Error::new(io::ErrorKind::Other, common::TEST_STR)))
+    })));
+
+    // Shutdown should complete with the correct error from s2n_shutdown
+    let result = server.shutdown().await;
+    let io_error = result.unwrap_err();
+    let error: error::Error = io_error.try_into()?;
+    // Any non-blocking read error is translated as "IOError"
+    assert!(error.kind() == error::ErrorType::IOError);
+
+    // Even if s2n_shutdown fails, we need to close the underlying stream.
+    // Make sure we called our mock shutdown, consuming it.
+    assert!(overrides.is_consumed());
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn shutdown_with_tls_error_and_tcp_delay() -> Result<(), Box<dyn std::error::Error>> {
+    let client = TlsConnector::new(common::client_config()?.build()?);
+    let server = TlsAcceptor::new(common::server_config()?.build()?);
+
+    let (server_stream, client_stream) = common::get_streams().await?;
+    let server_stream = common::TestStream::new(server_stream);
+    let overrides = server_stream.overrides();
+
+    let (_, mut server) =
+        common::run_negotiate(&client, client_stream, &server, server_stream).await?;
+
+    // We want s2n_shutdown to fail on read in order to ensure that it is only
+    // called once on failure.
+    // If s2n_shutdown were called again, the second call would hang waiting
+    // for nonexistent input from the peer.
+    overrides.next_read(Some(Box::new(|_, _, _| {
+        Ready(Err(io::Error::from(io::ErrorKind::Other)))
+    })));
+
+    // The underlying stream should initially return Pending, delaying shutdown
+    overrides.next_shutdown(Some(Box::new(|_, ctx| {
+        ctx.waker().wake_by_ref();
+        Pending
+    })));
+
+    // Shutdown should complete with the correct error from s2n_shutdown
+    let result = server.shutdown().await;
+    let io_error = result.unwrap_err();
+    let error: error::Error = io_error.try_into()?;
+    // Any non-blocking read error is translated as "IOError"
+    assert!(error.kind() == error::ErrorType::IOError);
+
+    // Even if s2n_shutdown fails, we need to close the underlying stream.
+    // Make sure we at least called our mock shutdown, consuming it.
+    assert!(overrides.is_consumed());
+
+    Ok(())
+}

From 670cb43534ee6515033be53369a2af774f3083b0 Mon Sep 17 00:00:00 2001
From: Sam Clark <3758302+goatgoose@users.noreply.github.com>
Date: Fri, 19 Jan 2024 12:20:49 -0500
Subject: [PATCH 13/13] Fix initialization errors in unit tests (#4370)

---
 .github/s2n_osx.sh              |  9 ++++++++
 tests/testlib/s2n_mem_testlib.c | 39 +++++++++++++++++++++++----------
 tests/unit/s2n_mem_test.c       | 15 ++++++++-----
 utils/s2n_mem.c                 | 33 +++++++++++++++++++++++-----
 utils/s2n_mem.h                 |  5 +++++
 5 files changed, 79 insertions(+), 22 deletions(-)

diff --git a/.github/s2n_osx.sh b/.github/s2n_osx.sh
index 4a9210d7a0d..d166fba9224 100755
--- a/.github/s2n_osx.sh
+++ b/.github/s2n_osx.sh
@@ -28,3 +28,12 @@ cmake . -Bbuild -GNinja \
 
 cmake --build ./build -j $(nproc)
 time CTEST_PARALLEL_LEVEL=$(nproc) ninja -C build test
+
+# Build shared library
+cmake . -Bbuild -GNinja \
+-DCMAKE_BUILD_TYPE=Debug \
+-DCMAKE_PREFIX_PATH=${OPENSSL_1_1_1_INSTALL_DIR} .. \
+-DBUILD_SHARED_LIBS=ON
+
+cmake --build ./build -j $(nproc)
+time CTEST_PARALLEL_LEVEL=$(nproc) ninja -C build test
diff --git a/tests/testlib/s2n_mem_testlib.c b/tests/testlib/s2n_mem_testlib.c
index d25ce76cac2..80ed49bba01 100644
--- a/tests/testlib/s2n_mem_testlib.c
+++ b/tests/testlib/s2n_mem_testlib.c
@@ -16,9 +16,8 @@
 #include "testlib/s2n_mem_testlib.h"
 
 #include "stuffer/s2n_stuffer.h"
-
-/* Required to override memory callbacks at runtime */
-#include "utils/s2n_mem.c"
+#include "utils/s2n_mem.h"
+#include "utils/s2n_safety.h"
 
 struct s2n_mem_test_cb_ctx {
     struct s2n_stuffer mallocs;
@@ -32,25 +31,43 @@ static int s2n_mem_test_free_cb(void *ptr, uint32_t size);
 
 static S2N_RESULT s2n_mem_test_set_callbacks()
 {
-    if (s2n_mem_malloc_cb != s2n_mem_test_malloc_cb) {
-        s2n_mem_malloc_cb_backup = s2n_mem_malloc_cb;
-        s2n_mem_malloc_cb = s2n_mem_test_malloc_cb;
+    s2n_mem_init_callback mem_init_cb = NULL;
+    s2n_mem_cleanup_callback mem_cleanup_cb = NULL;
+    s2n_mem_malloc_callback mem_malloc_cb = NULL;
+    s2n_mem_free_callback mem_free_cb = NULL;
+    RESULT_GUARD(s2n_mem_get_callbacks(&mem_init_cb, &mem_cleanup_cb, &mem_malloc_cb, &mem_free_cb));
+
+    if (mem_malloc_cb != s2n_mem_test_malloc_cb) {
+        s2n_mem_malloc_cb_backup = mem_malloc_cb;
+        mem_malloc_cb = s2n_mem_test_malloc_cb;
     }
-    if (s2n_mem_free_cb != s2n_mem_test_free_cb) {
-        s2n_mem_free_cb_backup = s2n_mem_free_cb;
-        s2n_mem_free_cb = s2n_mem_test_free_cb;
+    if (mem_free_cb != s2n_mem_test_free_cb) {
+        s2n_mem_free_cb_backup = mem_free_cb;
+        mem_free_cb = s2n_mem_test_free_cb;
     }
+
+    RESULT_GUARD(s2n_mem_override_callbacks(mem_init_cb, mem_cleanup_cb, mem_malloc_cb, mem_free_cb));
+
     return S2N_RESULT_OK;
 }
 
 static S2N_RESULT s2n_mem_test_unset_callbacks()
 {
+    s2n_mem_init_callback mem_init_cb = NULL;
+    s2n_mem_cleanup_callback mem_cleanup_cb = NULL;
+    s2n_mem_malloc_callback mem_malloc_cb = NULL;
+    s2n_mem_free_callback mem_free_cb = NULL;
+    RESULT_GUARD(s2n_mem_get_callbacks(&mem_init_cb, &mem_cleanup_cb, &mem_malloc_cb, &mem_free_cb));
+
     if (s2n_mem_malloc_cb_backup != NULL) {
-        s2n_mem_malloc_cb = s2n_mem_malloc_cb_backup;
+        mem_malloc_cb = s2n_mem_malloc_cb_backup;
     }
     if (s2n_mem_free_cb_backup != NULL) {
-        s2n_mem_free_cb = s2n_mem_free_cb_backup;
+        mem_free_cb = s2n_mem_free_cb_backup;
     }
+
+    RESULT_GUARD(s2n_mem_override_callbacks(mem_init_cb, mem_cleanup_cb, mem_malloc_cb, mem_free_cb));
+
     return S2N_RESULT_OK;
 }
 
diff --git a/tests/unit/s2n_mem_test.c b/tests/unit/s2n_mem_test.c
index 0c522bc0adf..4146f615bbc 100644
--- a/tests/unit/s2n_mem_test.c
+++ b/tests/unit/s2n_mem_test.c
@@ -13,12 +13,11 @@
  * permissions and limitations under the License.
  */
 
+#include "utils/s2n_mem.h"
+
 #include "s2n_test.h"
 #include "utils/s2n_safety.h"
 
-/* Required to override memory callbacks at runtime */
-#include "utils/s2n_mem.c"
-
 int s2n_strict_mem_free_cb(void *ptr, uint32_t size)
 {
     POSIX_ENSURE_REF(ptr);
@@ -69,10 +68,14 @@ int main(int argc, char **argv)
          */
         {
             /* Save real free callback */
-            s2n_mem_free_callback saved_free_cb = s2n_mem_free_cb;
+            s2n_mem_init_callback mem_init_cb = NULL;
+            s2n_mem_cleanup_callback mem_cleanup_cb = NULL;
+            s2n_mem_malloc_callback mem_malloc_cb = NULL;
+            s2n_mem_free_callback mem_free_cb = NULL;
+            EXPECT_OK(s2n_mem_get_callbacks(&mem_init_cb, &mem_cleanup_cb, &mem_malloc_cb, &mem_free_cb));
 
             /* Set callback that won't accepts NULLs / zeros */
-            s2n_mem_free_cb = s2n_strict_mem_free_cb;
+            EXPECT_OK(s2n_mem_override_callbacks(mem_init_cb, mem_cleanup_cb, mem_malloc_cb, s2n_strict_mem_free_cb));
 
             /* No-op for empty blob */
             struct s2n_blob blob = { 0 };
@@ -81,7 +84,7 @@ int main(int argc, char **argv)
             EXPECT_SUCCESS(s2n_free_or_wipe(&blob));
 
             /* Restore real free callback */
-            s2n_mem_free_cb = saved_free_cb;
+            EXPECT_OK(s2n_mem_override_callbacks(mem_init_cb, mem_cleanup_cb, mem_malloc_cb, mem_free_cb));
         };
     };
 
diff --git a/utils/s2n_mem.c b/utils/s2n_mem.c
index c5fa06c3a08..d8c79ddee77 100644
--- a/utils/s2n_mem.c
+++ b/utils/s2n_mem.c
@@ -135,18 +135,41 @@ int s2n_mem_set_callbacks(s2n_mem_init_callback mem_init_callback, s2n_mem_clean
         s2n_mem_malloc_callback mem_malloc_callback, s2n_mem_free_callback mem_free_callback)
 {
     POSIX_ENSURE(!initialized, S2N_ERR_INITIALIZED);
+    POSIX_GUARD_RESULT(s2n_mem_override_callbacks(mem_init_callback, mem_cleanup_callback,
+            mem_malloc_callback, mem_free_callback));
+    return S2N_SUCCESS;
+}
 
-    POSIX_ENSURE_REF(mem_init_callback);
-    POSIX_ENSURE_REF(mem_cleanup_callback);
-    POSIX_ENSURE_REF(mem_malloc_callback);
-    POSIX_ENSURE_REF(mem_free_callback);
+S2N_RESULT s2n_mem_override_callbacks(s2n_mem_init_callback mem_init_callback, s2n_mem_cleanup_callback mem_cleanup_callback,
+        s2n_mem_malloc_callback mem_malloc_callback, s2n_mem_free_callback mem_free_callback)
+{
+    RESULT_ENSURE_REF(mem_init_callback);
+    RESULT_ENSURE_REF(mem_cleanup_callback);
+    RESULT_ENSURE_REF(mem_malloc_callback);
+    RESULT_ENSURE_REF(mem_free_callback);
 
     s2n_mem_init_cb = mem_init_callback;
     s2n_mem_cleanup_cb = mem_cleanup_callback;
     s2n_mem_malloc_cb = mem_malloc_callback;
     s2n_mem_free_cb = mem_free_callback;
 
-    return S2N_SUCCESS;
+    return S2N_RESULT_OK;
+}
+
+S2N_RESULT s2n_mem_get_callbacks(s2n_mem_init_callback *mem_init_callback, s2n_mem_cleanup_callback *mem_cleanup_callback,
+        s2n_mem_malloc_callback *mem_malloc_callback, s2n_mem_free_callback *mem_free_callback)
+{
+    RESULT_ENSURE_REF(mem_init_callback);
+    RESULT_ENSURE_REF(mem_cleanup_callback);
+    RESULT_ENSURE_REF(mem_malloc_callback);
+    RESULT_ENSURE_REF(mem_free_callback);
+
+    *mem_init_callback = s2n_mem_init_cb;
+    *mem_cleanup_callback = s2n_mem_cleanup_cb;
+    *mem_malloc_callback = s2n_mem_malloc_cb;
+    *mem_free_callback = s2n_mem_free_cb;
+
+    return S2N_RESULT_OK;
 }
 
 int s2n_alloc(struct s2n_blob *b, uint32_t size)
diff --git a/utils/s2n_mem.h b/utils/s2n_mem.h
index a66b90ed5bc..4b87ab39059 100644
--- a/utils/s2n_mem.h
+++ b/utils/s2n_mem.h
@@ -44,3 +44,8 @@ int s2n_dup(struct s2n_blob *from, struct s2n_blob *to);
  * Prefer s2n_free. Only use this method if completely necessary.
  */
 int s2n_free_or_wipe(struct s2n_blob *b);
+
+S2N_RESULT s2n_mem_override_callbacks(s2n_mem_init_callback mem_init_callback, s2n_mem_cleanup_callback mem_cleanup_callback,
+        s2n_mem_malloc_callback mem_malloc_callback, s2n_mem_free_callback mem_free_callback);
+S2N_RESULT s2n_mem_get_callbacks(s2n_mem_init_callback *mem_init_callback, s2n_mem_cleanup_callback *mem_cleanup_callback,
+        s2n_mem_malloc_callback *mem_malloc_callback, s2n_mem_free_callback *mem_free_callback);