aws-crt 0.1.5 → 0.1.6

data/aws-crt-ffi/crt/s2n/tests/litani/test/run

@@ -14,6 +14,8 @@
 # permissions and limitations under the License.
 
 
+import argparse
+import importlib
 import logging
 import os
 import pathlib
@@ -22,6 +24,25 @@ import subprocess
 import sys
 import uuid
 
+DESCRIPTION = "Execute e2e and unit tests for Litani"
+
+
+def get_args():
+    pars = argparse.ArgumentParser(description=DESCRIPTION)
+    for arg in [{
+            "flags": ["--output-dir"],
+            "help": "output dir for test results",
+            "default": pathlib.Path(__file__).resolve().parent / "output",
+            "type": pathlib.Path
+        }, {
+            "flags": ["--fast"],
+            "help": "run fast tests only",
+            "action": "store_true"
+        }]:
+        flags = arg.pop("flags")
+        pars.add_argument(*flags, **arg)
+    return pars.parse_args()
+
 
 def run_cmd(cmd):
     try:
@@ -32,6 +53,14 @@ def run_cmd(cmd):
         sys.exit(1)
 
 
+def is_slow_test(module_file):
+    try:
+        return importlib.import_module(str(module_file.stem)).SLOW
+    except AttributeError:
+        logging.error("Variable SLOW is missing from: %s", module_file.name)
+        sys.exit(1)
+
+
 def litani_add(litani, counter, *args, **kwargs):
     cmd = [litani, "add-job"]
     for arg in args:
@@ -53,17 +82,28 @@ def collapse(string):
     return re.sub(r"\s+", " ", string)
 
 
-def add_e2e_tests(litani, test_dir, root_dir, counter):
+def add_e2e_tests(litani, test_dir, counter, output_dir, fast):
     e2e_test_dir = test_dir / "e2e"
     # 4 jobs per test (init, add-jobs, run-build, check-run)
     # skip __init__.py and __pycache__
     counter["total"] += (len(os.listdir(e2e_test_dir / "tests")) - 2) * 4
-
+    sys.path.insert(1, str(e2e_test_dir / "tests"))
     for test_file in (e2e_test_dir / "tests").iterdir():
         if test_file.name in ["__init__.py", "__pycache__"]:
             continue
 
-        run_dir = e2e_test_dir / "output" / str(uuid.uuid4())
+        add_transform_jobs = False
+        with open(test_file) as handle:
+            for line in handle:
+                if line.strip().startswith("def transform_jobs("):
+                    add_transform_jobs = True
+                    break
+        if fast and is_slow_test(test_file):
+            continue
+
+        run_dir = output_dir / "e2e_outputs" / str(uuid.uuid4())
+
+        timeout=10 if fast else 0
 
         litani_add(
             litani, counter,
@@ -77,8 +117,10 @@ def add_e2e_tests(litani, test_dir, root_dir, counter):
             ci_stage="test",
             description=f"{test_file.stem}: init",
             outputs=run_dir / ".litani_cache_dir",
-            cwd=run_dir)
+            cwd=run_dir,
+            timeout=timeout)
 
+        run_build_input = str(uuid.uuid4())
         litani_add(
             litani, counter,
             command=collapse(f"""
@@ -91,8 +133,28 @@ def add_e2e_tests(litani, test_dir, root_dir, counter):
             ci_stage="test",
             description=f"{test_file.stem}: add jobs",
             inputs=run_dir / ".litani_cache_dir",
+            phony_outputs=run_build_input,
             outputs=f"{run_dir}/output/jobs",
-            cwd=run_dir)
+            cwd=run_dir,
+            timeout=timeout)
+
+        if add_transform_jobs:
+            add_jobs_output = run_build_input
+            run_build_input = str(uuid.uuid4())
+            litani_add(
+                litani, counter,
+                command=collapse(f"""
+                    {e2e_test_dir / 'run'}
+                    --test-file {test_file}
+                    --litani {litani}
+                    --run-dir {run_dir}
+                    --operation transform-jobs"""),
+                pipeline=f"End-to-end: {test_file.stem}",
+                ci_stage="test",
+                description=f"{test_file.stem}: transform jobs",
+                inputs=add_jobs_output,
+                phony_outputs=run_build_input,
+                cwd=run_dir)
 
         litani_add(
             litani, counter,
@@ -105,9 +167,10 @@ def add_e2e_tests(litani, test_dir, root_dir, counter):
             pipeline=f"End-to-end: {test_file.stem}",
             ci_stage="test",
             description=f"{test_file.stem}: run build",
-            inputs=f"{run_dir}/output/jobs",
+            inputs=run_build_input,
             outputs=f"{run_dir}/output/run.json",
-            cwd=run_dir)
+            cwd=run_dir,
+            timeout=timeout)
 
         litani_add(
             litani, counter,
@@ -121,7 +184,8 @@ def add_e2e_tests(litani, test_dir, root_dir, counter):
             ci_stage="report",
             description=f"{test_file.stem}: check run",
             inputs=f"{run_dir}/output/run.json",
-            cwd=run_dir)
+            cwd=run_dir,
+            timeout=timeout)
 
 
 def add_unit_tests(litani, test_dir, root_dir, counter):
@@ -143,16 +207,21 @@ def print_counter(counter):
 
 
 def main():
-    logging.basicConfig(format="run-tests: %(message)s")
+    args = get_args()
+    logging.basicConfig(format="\nrun-tests: %(message)s")
     test_dir = pathlib.Path(__file__).resolve().parent
     root = test_dir.parent
     litani = root / "litani"
 
+    output_dir = args.output_dir.resolve()
+    output_dir.mkdir(exist_ok=True, parents=True)
+    os.chdir(output_dir)
+
     run_cmd([
         litani, "init",
         "--project", "Litani Test Suite",
-        "--output-prefix", test_dir / "output",
-        "--output-symlink", test_dir / "output" / "latest"])
+        "--output-prefix", ".",
+        "--output-symlink", "latest"])
 
     counter = {
         "added": 0,
@@ -160,7 +229,8 @@ def main():
     }
 
     add_unit_tests(litani, test_dir, root, counter)
-    add_e2e_tests(litani, test_dir, root, counter)
+    add_e2e_tests(
+        litani, test_dir, counter, output_dir, args.fast)
     print()
 
     run_cmd([litani, "run-build"])

data/aws-crt-ffi/crt/s2n/tests/s2n_test.h

@@ -39,40 +39,65 @@ int test_count;
 
 #define EXPECT_SUCCESS_WITHOUT_COUNT( function_call )  EXPECT_NOT_EQUAL_WITHOUT_COUNT( (function_call) ,  -1 )
 
-/**
- * This is a very basic, but functional unit testing framework. All testing should
- * happen in main() and start with a BEGIN_TEST() and end with an END_TEST();
+#define END_TEST_PRINT()                                                            \
+    if (isatty(fileno(stdout))) {                                                   \
+        if (test_count) {                                                           \
+            fprintf(stdout, "\033[32;1mPASSED\033[0m %10d tests\n", test_count );   \
+        }                                                                           \
+        else {                                                                      \
+            fprintf(stdout, "\033[33;1mSKIPPED\033[0m       ALL tests\n" );         \
+        }                                                                           \
+    }                                                                               \
+    else {                                                                          \
+        if (test_count) {                                                           \
+            fprintf(stdout, "PASSED %10d tests\n", test_count );                    \
+        }                                                                           \
+        else {                                                                      \
+            fprintf(stdout, "SKIPPED       ALL tests\n" );                          \
+        }                                                                           \
+    }
+
+/* Macros similar to BEGIN_TEST() and END_TEST() but for tests where s2n should
+ * not initialise at the start of the test. Useful for tests that e.g spawn a
+ * number of independent childs at the start of a unit test and where you want
+ * each child to have its own independently initialised s2n.
+ *
+ * BEGIN_TEST() prints unit test information to stdout. But this often gets
+ * buffered by the kernel and will then be flushed in each child spawned. The
+ * result is a number of repeated messages being send to stdout and, in turn,
+ * appear in the logs. At the moment, we think this is better than risking not
+ * having any printing at all.
  */
-#define BEGIN_TEST()                                           \
-  do {                                                         \
-    test_count = 0;                                            \
-    EXPECT_SUCCESS_WITHOUT_COUNT(s2n_in_unit_test_set(true));  \
-    S2N_TEST_OPTIONALLY_ENABLE_FIPS_MODE();                    \
-    EXPECT_SUCCESS_WITHOUT_COUNT(s2n_init());                  \
-    fprintf(stdout, "Running %-50s ... ", __FILE__);           \
-  } while(0)
-
-#define END_TEST()   do { \
-                        EXPECT_SUCCESS_WITHOUT_COUNT(s2n_in_unit_test_set(false));      \
-                        EXPECT_SUCCESS_WITHOUT_COUNT(s2n_cleanup());       \
-                        if (isatty(fileno(stdout))) { \
-                            if (test_count) { \
-                                fprintf(stdout, "\033[32;1mPASSED\033[0m %10d tests\n", test_count ); \
-                            }\
-                            else {\
-                                fprintf(stdout, "\033[33;1mSKIPPED\033[0m       ALL tests\n" ); \
-                            }\
-                       } \
-                       else { \
-                            if (test_count) { \
-                                fprintf(stdout, "PASSED %10d tests\n", test_count ); \
-                            }\
-                            else {\
-                                fprintf(stdout, "SKIPPED       ALL tests\n" ); \
-                            }\
-                       } \
-                       return 0;\
-                    } while(0)
+#define BEGIN_TEST_NO_INIT()                                        \
+    do {                                                            \
+        test_count = 0;                                             \
+        fprintf(stdout, "Running %-50s ... ", __FILE__);            \
+        EXPECT_SUCCESS_WITHOUT_COUNT(s2n_in_unit_test_set(true));   \
+        S2N_TEST_OPTIONALLY_ENABLE_FIPS_MODE();                     \
+    } while(0)
+
+#define END_TEST_NO_INIT()                                          \
+    do {                                                            \
+        EXPECT_SUCCESS_WITHOUT_COUNT(s2n_in_unit_test_set(false));  \
+        END_TEST_PRINT()                                            \
+        return 0;                                                   \
+    } while(0)
+
+/* This is a very basic, but functional unit testing framework. All testing
+ * should happen in main() and start with a BEGIN_TEST() and end with an
+ * END_TEST().
+ */
+#define BEGIN_TEST()                                                \
+    do {                                                            \
+        BEGIN_TEST_NO_INIT();                                       \
+        EXPECT_SUCCESS_WITHOUT_COUNT(s2n_init());                   \
+    } while(0)
+
+#define END_TEST()                                                  \
+    do {                                                            \
+        EXPECT_SUCCESS_WITHOUT_COUNT(s2n_cleanup());                \
+        END_TEST_NO_INIT();                                         \
+    } while(0)
 
 #define FAIL()      FAIL_MSG("")
 

data/aws-crt-ffi/crt/s2n/tests/testlib/s2n_key_schedule_testlib.c

@@ -15,14 +15,26 @@
 
 #include "testlib/s2n_testlib.h"
 
+S2N_RESULT s2n_connection_set_test_transcript_hash(struct s2n_connection *conn,
+        message_type_t message_type, const struct s2n_blob *digest)
+{
+    conn->handshake.handshake_type = conn->handshake.handshake_type & NEGOTIATED;
+    while(s2n_conn_get_current_message_type(conn) != message_type) {
+        conn->handshake.message_number++;
+    }
+    RESULT_CHECKED_MEMCPY(conn->handshake.hashes->transcript_hash_digest,
+            digest->data, digest->size);
+    return S2N_RESULT_OK;
+}
+
 S2N_RESULT s2n_connection_set_test_early_secret(struct s2n_connection *conn,
         const struct s2n_blob *early_secret)
 {
     RESULT_ENSURE_REF(conn);
     RESULT_ENSURE_REF(early_secret);
-    RESULT_CHECKED_MEMCPY(conn->secrets.tls13.early_secret,
+    RESULT_CHECKED_MEMCPY(conn->secrets.tls13.extract_secret,
             early_secret->data, early_secret->size);
-    conn->secrets.tls13.secrets_state = S2N_EARLY_SECRET;
+    conn->secrets.tls13.extract_secret_type = S2N_EARLY_SECRET;
     return S2N_RESULT_OK;
 }
 
@@ -31,9 +43,9 @@ S2N_RESULT s2n_connection_set_test_handshake_secret(struct s2n_connection *conn,
 {
     RESULT_ENSURE_REF(conn);
     RESULT_ENSURE_REF(handshake_secret);
-    RESULT_CHECKED_MEMCPY(conn->secrets.tls13.handshake_secret,
+    RESULT_CHECKED_MEMCPY(conn->secrets.tls13.extract_secret,
             handshake_secret->data, handshake_secret->size);
-    conn->secrets.tls13.secrets_state = S2N_HANDSHAKE_SECRET;
+    conn->secrets.tls13.extract_secret_type = S2N_HANDSHAKE_SECRET;
     return S2N_RESULT_OK;
 }
 
@@ -42,8 +54,8 @@ S2N_RESULT s2n_connection_set_test_master_secret(struct s2n_connection *conn,
 {
     RESULT_ENSURE_REF(conn);
     RESULT_ENSURE_REF(master_secret);
-    RESULT_CHECKED_MEMCPY(conn->secrets.tls13.master_secret,
+    RESULT_CHECKED_MEMCPY(conn->secrets.tls13.extract_secret,
             master_secret->data, master_secret->size);
-    conn->secrets.tls13.secrets_state = S2N_MASTER_SECRET;
+    conn->secrets.tls13.extract_secret_type = S2N_MASTER_SECRET;
     return S2N_RESULT_OK;
 }

data/aws-crt-ffi/crt/s2n/tests/testlib/s2n_pq_kat_test_utils.c

@@ -69,8 +69,8 @@ static S2N_RESULT s2n_drbg_generate_for_pq_kat_tests(struct s2n_drbg *drbg, stru
     RESULT_ENSURE(blob->size <= S2N_DRBG_GENERATE_LIMIT, S2N_ERR_DRBG_REQUEST_SIZE);
 
     /* We do NOT mix in additional entropy */
-    RESULT_GUARD_POSIX(s2n_drbg_bits(drbg, blob));
-    RESULT_GUARD_POSIX(s2n_drbg_update(drbg, &zeros));
+    RESULT_GUARD(s2n_drbg_bits(drbg, blob));
+    RESULT_GUARD(s2n_drbg_update(drbg, &zeros));
 
     return S2N_RESULT_OK;
 }
@@ -145,7 +145,7 @@ static int s2n_test_kem_with_kat(const struct s2n_kem *kem, const char *kat_file
          * we use the custom function s2n_drbg_generate_for_pq_kat_tests() defined above to turn off the
          * prediction resistance. */
         POSIX_GUARD(ReadHex(kat_file, kat_entropy_blob.data, SEED_LENGTH, "seed = "));
-        POSIX_GUARD(s2n_drbg_instantiate(&drbg_for_pq_kats, &personalization_string, S2N_AES_256_CTR_NO_DF_PR));
+        POSIX_GUARD_RESULT(s2n_drbg_instantiate(&drbg_for_pq_kats, &personalization_string, S2N_AES_256_CTR_NO_DF_PR));
 
         /* Generate the public/private key pair */
         POSIX_GUARD(kem->generate_keypair(pk, sk));
@@ -172,7 +172,7 @@ static int s2n_test_kem_with_kat(const struct s2n_kem *kem, const char *kat_file
         POSIX_ENSURE_EQ(memcmp(ss_answer, server_shared_secret, kem->shared_secret_key_length ), 0);
 
         /* Wipe the DRBG; it will reseed for each KAT test vector. */
-        POSIX_GUARD(s2n_drbg_wipe(&drbg_for_pq_kats));
+        POSIX_GUARD_RESULT(s2n_drbg_wipe(&drbg_for_pq_kats));
     }
     fclose(kat_file);
     free(ct);

data/aws-crt-ffi/crt/s2n/tests/testlib/s2n_testlib.h

@@ -74,6 +74,8 @@ S2N_RESULT s2n_append_test_psk_with_early_data(struct s2n_connection *conn, uint
 S2N_RESULT s2n_append_test_chosen_psk_with_early_data(struct s2n_connection *conn, uint32_t max_early_data,
         const struct s2n_cipher_suite *cipher_suite);
 
+S2N_RESULT s2n_connection_set_test_transcript_hash(struct s2n_connection *conn,
+        message_type_t message_type, const struct s2n_blob *digest);
 S2N_RESULT s2n_connection_set_test_early_secret(struct s2n_connection *conn, const struct s2n_blob *early_secret);
 S2N_RESULT s2n_connection_set_test_handshake_secret(struct s2n_connection *conn, const struct s2n_blob *handshake_secret);
 S2N_RESULT s2n_connection_set_test_master_secret(struct s2n_connection *conn, const struct s2n_blob *master_secret);

data/aws-crt-ffi/crt/s2n/tests/unit/s2n_client_hello_retry_test.c

@@ -13,12 +13,15 @@
  * permissions and limitations under the License.
  */
 
+#include "s2n.h"
 #include "s2n_test.h"
 
+#include "stuffer/s2n_stuffer.h"
 #include "testlib/s2n_testlib.h"
 
 #include "tls/extensions/s2n_server_supported_versions.h"
 
+#include "tls/extensions/s2n_cookie.h"
 #include "tls/s2n_cipher_suites.h"
 #include "tls/s2n_security_policies.h"
 #include "tls/s2n_tls.h"
@@ -31,6 +34,7 @@
 #include "tls/s2n_server_hello.c"
 
 #include "error/s2n_errno.h"
+#include "utils/s2n_safety.h"
 
 #define HELLO_RETRY_MSG_NO 1
 #define SERVER_HELLO_MSG_NO 5
@@ -448,6 +452,66 @@ int main(int argc, char **argv)
         EXPECT_SUCCESS(s2n_io_pair_close(&io_pair));
     }
 
+    /*
+     * Self-talk test: HRR with cookie extension
+     * We also wipe the connection to ensure that the cookie stuffer is handled correctly when connections are reused.
+     */
+    {
+        DEFER_CLEANUP(struct s2n_connection *server_conn = s2n_connection_new(S2N_SERVER),
+                s2n_connection_ptr_free);
+        DEFER_CLEANUP(struct s2n_connection *client_conn = s2n_connection_new(S2N_CLIENT),
+                s2n_connection_ptr_free);
+
+        DEFER_CLEANUP(struct s2n_config *server_config = s2n_config_new(),
+                s2n_config_ptr_free);
+        DEFER_CLEANUP(struct s2n_config *client_config = s2n_config_new(),
+                s2n_config_ptr_free);
+        DEFER_CLEANUP(struct s2n_cert_chain_and_key *tls13_chain_and_key,
+                s2n_cert_chain_and_key_ptr_free);
+
+        EXPECT_SUCCESS(s2n_test_cert_chain_and_key_new(&tls13_chain_and_key,
+                S2N_ECDSA_P384_PKCS1_CERT_CHAIN, S2N_ECDSA_P384_PKCS1_KEY));
+        EXPECT_SUCCESS(s2n_config_add_cert_chain_and_key_to_store(client_config, tls13_chain_and_key));
+        EXPECT_SUCCESS(s2n_config_add_cert_chain_and_key_to_store(server_config, tls13_chain_and_key));
+        EXPECT_SUCCESS(s2n_config_disable_x509_verification(client_config));
+
+        EXPECT_SUCCESS(s2n_connection_set_config(server_conn, server_config));
+        EXPECT_SUCCESS(s2n_connection_set_config(client_conn, client_config));
+
+        struct s2n_test_io_pair io_pair;
+
+        /* Call the test in a loop to ensure that s2n_connection_wipe is implemented correctly */
+        for (int i = 0; i < 10; i++) {
+            /* ensure call to s2n_connection_wipe are safe */
+            EXPECT_SUCCESS(s2n_connection_wipe(client_conn));
+            EXPECT_SUCCESS(s2n_connection_wipe(server_conn));
+
+             /* Create nonblocking pipes */
+            EXPECT_SUCCESS(s2n_io_pair_init_non_blocking(&io_pair));
+            EXPECT_SUCCESS(s2n_connections_set_io_pair(client_conn, server_conn, &io_pair));
+
+            /* include cookie data as part of HRR */
+            EXPECT_SUCCESS(s2n_stuffer_skip_write(&server_conn->cookie_stuffer, 500));
+            EXPECT_TRUE(s2n_server_cookie_extension.should_send(server_conn));
+
+            /* Force the HRR path */
+            client_conn->security_policy_override = &security_policy_test_tls13_retry;
+
+            /* Negotiate handshake */
+            EXPECT_SUCCESS(s2n_negotiate_test_server_and_client(server_conn, client_conn));
+
+            /* Verify that HRR handshake */
+            EXPECT_TRUE(s2n_is_hello_retry_handshake(server_conn));
+            EXPECT_TRUE(s2n_is_hello_retry_handshake(client_conn));
+
+            /* Verify client received cookie data */
+            EXPECT_TRUE(s2n_stuffer_data_available(&client_conn->cookie_stuffer) > 0);
+
+            EXPECT_SUCCESS(s2n_shutdown_test_server_and_client(server_conn, client_conn));
+            EXPECT_SUCCESS(s2n_io_pair_close(&io_pair));
+        }
+    }
+
     /* Self-Talk test: the client initiates a handshake with an X25519 share.
      * The server, however does not support x25519 and prefers P-256.
      * The server then sends a HelloRetryRequest that requires the
@@ -605,7 +669,7 @@ int main(int argc, char **argv)
 
         EXPECT_NOT_NULL(server_conn = s2n_connection_new(S2N_SERVER));
         EXPECT_NOT_NULL(client_conn = s2n_connection_new(S2N_CLIENT));
-        
+
         /* A Hello Retry Request has been processed */
         EXPECT_SUCCESS(s2n_set_hello_retry_required(client_conn));
         client_conn->secure.cipher_suite = &s2n_tls13_aes_256_gcm_sha384;
@@ -631,4 +695,4 @@ int main(int argc, char **argv)
     EXPECT_SUCCESS(s2n_disable_tls13_in_test());
 
     END_TEST();
-}                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   
+}

data/aws-crt-ffi/crt/s2n/tests/unit/s2n_connection_test.c

@@ -126,7 +126,7 @@ int main(int argc, char **argv)
      */
     {
         /* Carefully consider any increases to this number. */
-        const uint16_t max_connection_size = 9100;
+        const uint16_t max_connection_size = 9050;
         const uint16_t min_connection_size = max_connection_size * 0.75;
 
         size_t connection_size = sizeof(struct s2n_connection);

data/aws-crt-ffi/crt/s2n/tests/unit/s2n_drbg_test.c

@@ -299,7 +299,7 @@ int check_drgb_version(s2n_drbg_mode mode, int (*generator)(void *, uint32_t), i
         POSIX_GUARD(s2n_rand_set_callbacks(nist_fake_entropy_init_cleanup, nist_fake_entropy_init_cleanup, generator, generator));
 
         /* Instantiate the DRBG */
-        POSIX_GUARD(s2n_drbg_instantiate(&nist_drbg, &personalization_string, mode));
+        POSIX_GUARD_RESULT(s2n_drbg_instantiate(&nist_drbg, &personalization_string, mode));
 
         uint8_t nist_v[16];
 
@@ -309,13 +309,13 @@ int check_drgb_version(s2n_drbg_mode mode, int (*generator)(void *, uint32_t), i
         /* Generate 512 bits (FIRST CALL) */
         uint8_t out[64];
         struct s2n_blob generated = {.data = out, .size = 64 };
-        POSIX_GUARD(s2n_drbg_generate(&nist_drbg, &generated));
+        POSIX_GUARD_RESULT(s2n_drbg_generate(&nist_drbg, &generated));
 
         POSIX_GUARD(s2n_stuffer_read_bytes(&reference_values, nist_v, sizeof(nist_v)));
         POSIX_ENSURE_EQ(memcmp(nist_v, nist_drbg.v, sizeof(nist_drbg.v)), 0);
 
         /* Generate another 512 bits (SECOND CALL) */
-        POSIX_GUARD(s2n_drbg_generate(&nist_drbg, &generated));
+        POSIX_GUARD_RESULT(s2n_drbg_generate(&nist_drbg, &generated));
 
         POSIX_GUARD(s2n_stuffer_read_bytes(&reference_values, nist_v, sizeof(nist_v)));
         POSIX_ENSURE_EQ(memcmp(nist_v, nist_drbg.v, sizeof(nist_drbg.v)), 0);
@@ -331,7 +331,7 @@ int check_drgb_version(s2n_drbg_mode mode, int (*generator)(void *, uint32_t), i
             POSIX_BAIL(S2N_ERR_DRBG);
         }
 
-        POSIX_GUARD(s2n_drbg_wipe(&nist_drbg));
+        POSIX_GUARD_RESULT(s2n_drbg_wipe(&nist_drbg));
     }
     return 0;
 }
@@ -346,21 +346,21 @@ int main(int argc, char **argv)
     struct s2n_drbg aes256_pr_drbg = {0};
     struct s2n_blob blob = {.data = data, .size = 64 };
 
-    EXPECT_SUCCESS(s2n_drbg_instantiate(&aes128_drbg, &blob, S2N_AES_128_CTR_NO_DF_PR));
-    EXPECT_SUCCESS(s2n_drbg_instantiate(&aes256_pr_drbg, &blob, S2N_AES_256_CTR_NO_DF_PR));
+    EXPECT_OK(s2n_drbg_instantiate(&aes128_drbg, &blob, S2N_AES_128_CTR_NO_DF_PR));
+    EXPECT_OK(s2n_drbg_instantiate(&aes256_pr_drbg, &blob, S2N_AES_256_CTR_NO_DF_PR));
 
     struct s2n_config *config;
     EXPECT_NOT_NULL(config = s2n_config_new());
 
     /* Use the AES128 DRBG for 32MB of data */
     for (int i = 0; i < 500000; i++) {
-        EXPECT_SUCCESS(s2n_drbg_generate(&aes128_drbg, &blob));
+        EXPECT_OK(s2n_drbg_generate(&aes128_drbg, &blob));
     }
     EXPECT_EQUAL(aes128_drbg.mixes, 500000);
 
     /* Use the AES256 DRBG with prediction resistance for 32MB of data */
     for (int i = 0; i < 500000; i++) {
-        EXPECT_SUCCESS(s2n_drbg_generate(&aes256_pr_drbg, &blob));
+        EXPECT_OK(s2n_drbg_generate(&aes256_pr_drbg, &blob));
     }
     EXPECT_EQUAL(aes256_pr_drbg.mixes, 500000);
 
@@ -368,19 +368,39 @@ int main(int argc, char **argv)
     /* the DRBG state is 128 bytes, test that we can get more than that */
     blob.size = 129;
     for (int i = 0; i < 10; i++) {
-        EXPECT_SUCCESS(s2n_drbg_generate(&aes128_drbg, &blob));
-        EXPECT_SUCCESS(s2n_drbg_generate(&aes256_pr_drbg, &blob));
+        EXPECT_OK(s2n_drbg_generate(&aes128_drbg, &blob));
+        EXPECT_OK(s2n_drbg_generate(&aes256_pr_drbg, &blob));
     }
     EXPECT_EQUAL(aes128_drbg.mixes, 500010);
     EXPECT_EQUAL(aes256_pr_drbg.mixes, 500010);
 
+    /* Check that ignoring prediction resistance works */
+    EXPECT_OK(s2n_ignore_prediction_resistance_for_testing(true));
+    uint64_t aes128_drbg_mixes_start = aes128_drbg.mixes;
+    uint64_t aes256_pr_drbg_mixes_start = aes256_pr_drbg.mixes;
+    for (int i = 0; i < 10; i++) {
+        EXPECT_OK(s2n_drbg_generate(&aes128_drbg, &blob));
+        EXPECT_OK(s2n_drbg_generate(&aes256_pr_drbg, &blob));
+    }
+    EXPECT_EQUAL(aes128_drbg.mixes, aes128_drbg_mixes_start);
+    EXPECT_EQUAL(aes256_pr_drbg.mixes, aes256_pr_drbg_mixes_start);
+
+    /* Check that we can enable prediction resistance again */
+    EXPECT_OK(s2n_ignore_prediction_resistance_for_testing(false));
+    for (int i = 0; i < 10; i++) {
+        EXPECT_OK(s2n_drbg_generate(&aes128_drbg, &blob));
+        EXPECT_OK(s2n_drbg_generate(&aes256_pr_drbg, &blob));
+    }
+    EXPECT_EQUAL(aes128_drbg.mixes, aes128_drbg_mixes_start + 10);
+    EXPECT_EQUAL(aes256_pr_drbg.mixes, aes256_pr_drbg_mixes_start + 10);
+
     /* Generate 31 (= 16 + 15) bytes. Since the DRBG generates 16 bytes at a time,
      * a common error is to incorrectly fill the last (not-aligned) bytes. Sometimes
      * they are left unchanged and sometimes a single byte is copied in. We ensure
      * that the last 15 bytes are not all equal to guard against this. */
     POSIX_CHECKED_MEMSET((void*)data, 0, 31);
     blob.size = 31;
-    EXPECT_SUCCESS(s2n_drbg_generate(&aes128_drbg, &blob));
+    EXPECT_OK(s2n_drbg_generate(&aes128_drbg, &blob));
     bool bytes_are_all_equal = true;
     for (size_t i = 17; i < 31; i++) {
         if (data[16] != data[i]) {
@@ -392,7 +412,7 @@ int main(int argc, char **argv)
 
     POSIX_CHECKED_MEMSET((void*)data, 0, 31);
     blob.size = 31;
-    EXPECT_SUCCESS(s2n_drbg_generate(&aes256_pr_drbg, &blob));
+    EXPECT_OK(s2n_drbg_generate(&aes256_pr_drbg, &blob));
     bytes_are_all_equal = true;
     for (size_t i = 17; i < 31; i++) {
         if (data[16] != data[i]) {
@@ -402,8 +422,8 @@ int main(int argc, char **argv)
     }
     EXPECT_FALSE(bytes_are_all_equal);
 
-    EXPECT_SUCCESS(s2n_drbg_wipe(&aes128_drbg));
-    EXPECT_SUCCESS(s2n_drbg_wipe(&aes256_pr_drbg));
+    EXPECT_OK(s2n_drbg_wipe(&aes128_drbg));
+    EXPECT_OK(s2n_drbg_wipe(&aes256_pr_drbg));
 
     /* Check everything against the NIST AES 128 vectors with prediction resistance */
     EXPECT_SUCCESS(s2n_stuffer_alloc_ro_from_hex_string(&nist_aes128_reference_entropy, nist_aes128_reference_entropy_hex));