diff --git a/codebuild/bin/grep_simple_mistakes.sh b/codebuild/bin/grep_simple_mistakes.sh index 2d10b3576ab..1dbe9fe9208 100755 --- a/codebuild/bin/grep_simple_mistakes.sh +++ b/codebuild/bin/grep_simple_mistakes.sh @@ -160,6 +160,22 @@ for file in $S2N_FILES_ASSERT_DOUBLE_SEMICOLON; do fi done +############################################# +# Assert that we don't call malloc directly outside of utils/s2n_mem.c or tests. +# All other allocations must use s2n_alloc or s2n_realloc. +############################################# +S2N_FILES_ASSERT_MALLOC=$(find "$PWD" -type f -name "s2n*.c" \ + -not -path "*/utils/s2n_mem.c" \ + -not -path "*/tests/*" \ + -not -path "*/bin/*") +for file in $S2N_FILES_ASSERT_MALLOC; do + RESULT_MALLOC=`grep -n "= malloc(" $file` + if [ "${#RESULT_MALLOC}" != "0" ]; then + FAILED=1 + printf "\e[1;34mFound malloc in $file:\e[0m\n$RESULT_MALLOC\n\n" + fi +done + ############################################# ## Assert that there are no new uses of S2N_ERROR_IF # TODO add crypto, tls (see https://github.com/aws/s2n-tls/issues/2635) diff --git a/error/s2n_errno.c b/error/s2n_errno.c index b055670f78e..eee4b136cd9 100644 --- a/error/s2n_errno.c +++ b/error/s2n_errno.c @@ -299,6 +299,7 @@ static const char *no_such_error = "Internal s2n error"; ERR_ENTRY(S2N_ERR_KTLS_ENABLE_CRYPTO, "An error occurred when attempting to enable kTLS on socket.") \ ERR_ENTRY(S2N_ERR_KTLS_BAD_CMSG, "Error handling cmsghdr.") \ ERR_ENTRY(S2N_ERR_ATOMIC, "Atomic operations in this environment would require locking") \ + ERR_ENTRY(S2N_ERR_TEST_ASSERTION, "Test assertion failed") \ /* clang-format on */ #define ERR_STR_CASE(ERR, str) \ diff --git a/error/s2n_errno.h b/error/s2n_errno.h index a6991a298df..7de875d457c 100644 --- a/error/s2n_errno.h +++ b/error/s2n_errno.h @@ -229,6 +229,7 @@ typedef enum { S2N_ERR_LIBCRYPTO_VERSION_NUMBER_MISMATCH, S2N_ERR_LIBCRYPTO_VERSION_NAME_MISMATCH, S2N_ERR_OSSL_PROVIDER, + S2N_ERR_TEST_ASSERTION, S2N_ERR_T_INTERNAL_END, /* S2N_ERR_T_USAGE */ diff --git a/tests/testlib/s2n_mem_testlib.c b/tests/testlib/s2n_mem_testlib.c new file mode 100644 index 00000000000..d25ce76cac2 --- /dev/null +++ b/tests/testlib/s2n_mem_testlib.c @@ -0,0 +1,169 @@ +/* + * 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. + */ + +#include "testlib/s2n_mem_testlib.h" + +#include "stuffer/s2n_stuffer.h" + +/* Required to override memory callbacks at runtime */ +#include "utils/s2n_mem.c" + +struct s2n_mem_test_cb_ctx { + struct s2n_stuffer mallocs; +} s2n_mem_test_ctx = { 0 }; + +static s2n_mem_malloc_callback s2n_mem_malloc_cb_backup = NULL; +static s2n_mem_free_callback s2n_mem_free_cb_backup = NULL; + +static int s2n_mem_test_malloc_cb(void **ptr, uint32_t requested, uint32_t *allocated); +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; + } + 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; + } + return S2N_RESULT_OK; +} + +static S2N_RESULT s2n_mem_test_unset_callbacks() +{ + if (s2n_mem_malloc_cb_backup != NULL) { + s2n_mem_malloc_cb = s2n_mem_malloc_cb_backup; + } + if (s2n_mem_free_cb_backup != NULL) { + s2n_mem_free_cb = s2n_mem_free_cb_backup; + } + return S2N_RESULT_OK; +} + +static int s2n_mem_test_malloc_cb(void **ptr, uint32_t requested, uint32_t *allocated) +{ + int result = s2n_mem_malloc_cb_backup(ptr, requested, allocated); + POSIX_GUARD(result); + + struct s2n_mem_test_malloc new_info = { + .requested = requested, + .ptr = *ptr, + .allocated = *allocated, + .freed = false, + }; + + POSIX_GUARD_RESULT(s2n_mem_test_unset_callbacks()); + POSIX_GUARD(s2n_stuffer_write_bytes(&s2n_mem_test_ctx.mallocs, + (uint8_t *) &new_info, sizeof(new_info))); + POSIX_GUARD_RESULT(s2n_mem_test_set_callbacks()); + + return result; +} + +static int s2n_mem_test_free_cb(void *ptr, uint32_t size) +{ + int result = s2n_mem_free_cb_backup(ptr, size); + POSIX_GUARD(result); + + struct s2n_stuffer read_copy = s2n_mem_test_ctx.mallocs; + while (s2n_stuffer_data_available(&read_copy)) { + uint8_t *mem = s2n_stuffer_raw_read(&read_copy, sizeof(struct s2n_mem_test_malloc)); + struct s2n_mem_test_malloc *info = (struct s2n_mem_test_malloc *) (void *) mem; + if (info->ptr == ptr) { + info->freed = true; + } + } + return result; +} + +S2N_RESULT s2n_mem_test_init_callbacks(void *ctx) +{ + RESULT_GUARD(s2n_mem_test_free_callbacks(ctx)); + RESULT_GUARD_POSIX(s2n_stuffer_growable_alloc(&s2n_mem_test_ctx.mallocs, 0)); + RESULT_GUARD(s2n_mem_test_set_callbacks()); + return S2N_RESULT_OK; +} + +S2N_CLEANUP_RESULT s2n_mem_test_free_callbacks(void *ctx) +{ + (void) ctx; + RESULT_GUARD(s2n_mem_test_unset_callbacks()); + RESULT_GUARD_POSIX(s2n_stuffer_free(&s2n_mem_test_ctx.mallocs)); + return S2N_RESULT_OK; +} + +S2N_RESULT s2n_mem_test_wipe_callbacks() +{ + RESULT_GUARD_POSIX(s2n_stuffer_wipe(&s2n_mem_test_ctx.mallocs)); + return S2N_RESULT_OK; +} + +static S2N_RESULT s2n_mem_test_get_malloc_count(uint32_t *count) +{ + RESULT_ENSURE_REF(count); + size_t size = s2n_stuffer_data_available(&s2n_mem_test_ctx.mallocs); + const size_t sizeof_malloc = sizeof(struct s2n_mem_test_malloc); + *count = size / sizeof_malloc; + RESULT_ENSURE_EQ(size % sizeof_malloc, 0); + return S2N_RESULT_OK; +} + +static S2N_RESULT s2n_mem_test_get_malloc(uint32_t idx, struct s2n_mem_test_malloc *info) +{ + RESULT_ENSURE_REF(info); + struct s2n_stuffer read_copy = s2n_mem_test_ctx.mallocs; + const size_t sizeof_malloc = sizeof(struct s2n_mem_test_malloc); + RESULT_GUARD_POSIX(s2n_stuffer_skip_read(&read_copy, idx * sizeof_malloc)); + RESULT_GUARD_POSIX(s2n_stuffer_read_bytes(&read_copy, (uint8_t *) info, sizeof_malloc)); + return S2N_RESULT_OK; +} + +S2N_RESULT s2n_mem_test_assert_malloc_count(uint32_t count) +{ + uint32_t actual_count = 0; + RESULT_GUARD(s2n_mem_test_get_malloc_count(&actual_count)); + RESULT_ENSURE(count == actual_count, S2N_ERR_TEST_ASSERTION); + return S2N_RESULT_OK; +} + +S2N_RESULT s2n_mem_test_assert_malloc(uint32_t requested) +{ + uint32_t count = 0; + RESULT_GUARD(s2n_mem_test_get_malloc_count(&count)); + + struct s2n_mem_test_malloc mem_info = { 0 }; + for (size_t i = 0; i < count; i++) { + RESULT_GUARD(s2n_mem_test_get_malloc(i, &mem_info)); + if (mem_info.requested == requested) { + return S2N_RESULT_OK; + } + } + RESULT_BAIL(S2N_ERR_TEST_ASSERTION); +} + +S2N_RESULT s2n_mem_test_assert_all_freed() +{ + uint32_t count = 0; + RESULT_GUARD(s2n_mem_test_get_malloc_count(&count)); + + struct s2n_mem_test_malloc mem_info = { 0 }; + for (size_t i = 0; i < count; i++) { + RESULT_GUARD(s2n_mem_test_get_malloc(i, &mem_info)); + RESULT_ENSURE(mem_info.freed, S2N_ERR_TEST_ASSERTION); + } + return S2N_RESULT_OK; +} diff --git a/tests/testlib/s2n_mem_testlib.h b/tests/testlib/s2n_mem_testlib.h new file mode 100644 index 00000000000..c02d3759321 --- /dev/null +++ b/tests/testlib/s2n_mem_testlib.h @@ -0,0 +1,37 @@ +/* + * 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. + */ + +#pragma once + +#include "utils/s2n_result.h" + +struct s2n_mem_test_malloc { + void *ptr; + uint32_t requested; + uint32_t allocated; + bool freed; +}; + +struct s2n_mem_test_cb_scope { + uint8_t _reserve; +}; + +S2N_RESULT s2n_mem_test_init_callbacks(void *ctx); +S2N_CLEANUP_RESULT s2n_mem_test_free_callbacks(void *ctx); +S2N_RESULT s2n_mem_test_wipe_callbacks(); + +S2N_RESULT s2n_mem_test_assert_malloc_count(uint32_t count); +S2N_RESULT s2n_mem_test_assert_malloc(uint32_t requested); +S2N_RESULT s2n_mem_test_assert_all_freed(); diff --git a/tests/unit/s2n_mem_testlib_test.c b/tests/unit/s2n_mem_testlib_test.c new file mode 100644 index 00000000000..a44e822e379 --- /dev/null +++ b/tests/unit/s2n_mem_testlib_test.c @@ -0,0 +1,198 @@ +/* +* 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. +*/ + +#include "testlib/s2n_mem_testlib.h" + +#include "s2n_test.h" + +int main() +{ + BEGIN_TEST(); + + /* Test: No mallocs */ + for (size_t init = 0; init < 2; init++) { + /* The getters and asserts should behave the same whether the callbacks + * were never invoked or never initialized at all + */ + if (init) { + DEFER_CLEANUP(struct s2n_mem_test_cb_scope scope = { 0 }, + s2n_mem_test_free_callbacks); + EXPECT_OK(s2n_mem_test_init_callbacks(&scope)); + } + + /* Test: s2n_mem_test_assert_malloc_count */ + { + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc_count(1), + S2N_ERR_TEST_ASSERTION); + }; + + /* Test: s2n_mem_test_assert_malloc */ + { + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc(0), S2N_ERR_TEST_ASSERTION); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc(1), S2N_ERR_TEST_ASSERTION); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc(UINT32_MAX), + S2N_ERR_TEST_ASSERTION); + }; + + /* Test: s2n_mem_test_assert_all_freed */ + EXPECT_OK(s2n_mem_test_assert_all_freed()); + }; + + /* Test: Single malloc */ + { + DEFER_CLEANUP(struct s2n_mem_test_cb_scope scope = { 0 }, + s2n_mem_test_free_callbacks); + EXPECT_OK(s2n_mem_test_init_callbacks(&scope)); + + const size_t requested = 17; + struct s2n_blob mem = { 0 }; + EXPECT_SUCCESS(s2n_alloc(&mem, requested)); + EXPECT_NOT_NULL(mem.data); + EXPECT_EQUAL(mem.size, requested); + + /* Test: s2n_mem_test_assert_malloc_count */ + { + EXPECT_OK(s2n_mem_test_assert_malloc_count(1)); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc_count(0), + S2N_ERR_TEST_ASSERTION); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc_count(2), + S2N_ERR_TEST_ASSERTION); + }; + + /* Test: s2n_mem_test_assert_malloc */ + { + EXPECT_OK(s2n_mem_test_assert_malloc(requested)); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc(0), S2N_ERR_TEST_ASSERTION); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc(1), S2N_ERR_TEST_ASSERTION); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc(UINT32_MAX), + S2N_ERR_TEST_ASSERTION); + }; + + /* Test: s2n_mem_test_assert_all_freed */ + { + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_all_freed(), S2N_ERR_TEST_ASSERTION); + EXPECT_SUCCESS(s2n_free(&mem)); + EXPECT_OK(s2n_mem_test_assert_all_freed()); + }; + }; + + /* Test: Multiple mallocs */ + { + DEFER_CLEANUP(struct s2n_mem_test_cb_scope scope = { 0 }, + s2n_mem_test_free_callbacks); + EXPECT_OK(s2n_mem_test_init_callbacks(&scope)); + + const size_t requested[] = { 10, 1, 17, 17, 111 }; + const size_t count = s2n_array_len(requested); + struct s2n_blob mem = { 0 }; + for (size_t i = 0; i < count; i++) { + EXPECT_SUCCESS(s2n_alloc(&mem, requested[i])); + EXPECT_NOT_NULL(mem.data); + EXPECT_EQUAL(mem.size, requested[i]); + EXPECT_SUCCESS(s2n_free(&mem)); + } + + /* Test: s2n_mem_test_assert_malloc_count */ + { + EXPECT_OK(s2n_mem_test_assert_malloc_count(count)); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc_count(0), + S2N_ERR_TEST_ASSERTION); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc_count(count - 1), + S2N_ERR_TEST_ASSERTION); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc_count(count + 1), + S2N_ERR_TEST_ASSERTION); + }; + + /* Test: s2n_mem_test_assert_malloc */ + { + for (size_t i = 0; i < count; i++) { + EXPECT_OK(s2n_mem_test_assert_malloc(requested[i])); + } + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc(0), S2N_ERR_TEST_ASSERTION); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_malloc(UINT32_MAX), + S2N_ERR_TEST_ASSERTION); + }; + + /* Test: s2n_mem_test_assert_all_freed */ + { + EXPECT_OK(s2n_mem_test_assert_all_freed()); + EXPECT_SUCCESS(s2n_alloc(&mem, 1)); + EXPECT_ERROR_WITH_ERRNO(s2n_mem_test_assert_all_freed(), S2N_ERR_TEST_ASSERTION); + EXPECT_SUCCESS(s2n_free(&mem)); + EXPECT_OK(s2n_mem_test_assert_all_freed()); + }; + }; + + /* Test: s2n_mem_test_wipe_callbacks */ + { + DEFER_CLEANUP(struct s2n_mem_test_cb_scope scope = { 0 }, + s2n_mem_test_free_callbacks); + EXPECT_OK(s2n_mem_test_init_callbacks(&scope)); + + DEFER_CLEANUP(struct s2n_blob mem = { 0 }, s2n_free); + EXPECT_SUCCESS(s2n_alloc(&mem, 1)); + EXPECT_NOT_NULL(mem.data); + EXPECT_OK(s2n_mem_test_assert_malloc_count(1)); + + EXPECT_OK(s2n_mem_test_wipe_callbacks()); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); + }; + + /* Test: s2n_mem_test_init_callbacks */ + { + EXPECT_OK(s2n_mem_test_init_callbacks(NULL)); + for (size_t i = 0; i < 5; i++) { + DEFER_CLEANUP(struct s2n_blob mem = { 0 }, s2n_free); + EXPECT_SUCCESS(s2n_alloc(&mem, 1)); + EXPECT_NOT_NULL(mem.data); + EXPECT_OK(s2n_mem_test_assert_malloc_count(1)); + + /* If already initialized, we just wipe the state */ + EXPECT_OK(s2n_mem_test_init_callbacks(NULL)); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); + } + EXPECT_OK(s2n_mem_test_free_callbacks(NULL)); + }; + + /* Test: s2n_mem_test_free_callbacks */ + { + for (size_t i = 1; i < 5; i++) { + EXPECT_OK(s2n_mem_test_free_callbacks(NULL)); + + DEFER_CLEANUP(struct s2n_blob mem = { 0 }, s2n_free); + EXPECT_SUCCESS(s2n_alloc(&mem, 1)); + EXPECT_NOT_NULL(mem.data); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); + } + + EXPECT_OK(s2n_mem_test_init_callbacks(NULL)); + DEFER_CLEANUP(struct s2n_blob mem = { 0 }, s2n_free); + EXPECT_SUCCESS(s2n_alloc(&mem, 1)); + EXPECT_OK(s2n_mem_test_assert_malloc_count(1)); + + for (size_t i = 1; i < 5; i++) { + EXPECT_OK(s2n_mem_test_free_callbacks(NULL)); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); + + DEFER_CLEANUP(struct s2n_blob mem2 = { 0 }, s2n_free); + EXPECT_SUCCESS(s2n_alloc(&mem2, 1)); + EXPECT_NOT_NULL(mem2.data); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); + } + }; + + END_TEST(); +} diff --git a/tests/unit/s2n_post_handshake_recv_test.c b/tests/unit/s2n_post_handshake_recv_test.c index 00961e0e0e6..9789c785bff 100644 --- a/tests/unit/s2n_post_handshake_recv_test.c +++ b/tests/unit/s2n_post_handshake_recv_test.c @@ -19,6 +19,7 @@ #include "api/unstable/renegotiate.h" #include "error/s2n_errno.h" #include "s2n_test.h" +#include "testlib/s2n_mem_testlib.h" #include "testlib/s2n_testlib.h" #include "tls/s2n_connection.h" #include "tls/s2n_key_update.h" @@ -27,21 +28,10 @@ #include "tls/s2n_tls13_handshake.h" #include "utils/s2n_safety.h" -/* Required to override memory callbacks at runtime */ -#include "utils/s2n_mem.c" - #define S2N_TEST_MESSAGE_COUNT 5 int s2n_key_update_write(struct s2n_blob *out); -static s2n_mem_malloc_callback original_malloc_callback = NULL; -size_t mallocs_count = 0; -static int s2n_count_mallocs_cb(void **ptr, uint32_t requested, uint32_t *allocated) -{ - mallocs_count++; - return original_malloc_callback(ptr, requested, allocated); -} - size_t tickets_count = 0; static int s2n_ticket_count_cb(struct s2n_connection *conn, void *ctx, struct s2n_session_ticket *ticket) { @@ -91,7 +81,7 @@ static S2N_RESULT s2n_test_basic_recv(struct s2n_connection *sender, struct s2n_ RESULT_ENSURE_EQ(send_ret, sizeof(app_data)); /* Reset all counters */ - mallocs_count = 0; + RESULT_GUARD(s2n_mem_test_wipe_callbacks()); tickets_count = 0; hello_request_count = 0; @@ -118,7 +108,7 @@ static S2N_RESULT s2n_test_blocking_recv(struct s2n_connection *sender, struct s RESULT_ENSURE_EQ(send_ret, sizeof(app_data)); /* Reset all counters */ - mallocs_count = 0; + RESULT_GUARD(s2n_mem_test_wipe_callbacks()); tickets_count = 0; hello_request_count = 0; @@ -178,12 +168,6 @@ int main(int argc, char **argv) { BEGIN_TEST(); - /* Some tests need to make assertions about whether or not memory was allocated. - * Override the current memory callback with one that counts allocations. - */ - original_malloc_callback = s2n_mem_malloc_cb; - s2n_mem_malloc_cb = s2n_count_mallocs_cb; - const uint8_t unknown_message_type = UINT8_MAX; const uint32_t test_large_message_size = 3001; @@ -251,8 +235,10 @@ int main(int argc, char **argv) * if we successfully decrypt all records. If they were not processed, * then we would try to use the wrong key to decrypt the next record. */ + DEFER_CLEANUP(struct s2n_mem_test_cb_scope mem_ctx = { 0 }, s2n_mem_test_free_callbacks); + EXPECT_OK(s2n_mem_test_init_callbacks(&mem_ctx)); EXPECT_OK(s2n_test_basic_recv(sender, receiver)); - EXPECT_EQUAL(mallocs_count, 0); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); } } @@ -352,9 +338,11 @@ int main(int argc, char **argv) EXPECT_OK(s2n_tls13_server_nst_write(client, &messages)); EXPECT_OK(s2n_test_send_records(client, messages, fragment_size)); + DEFER_CLEANUP(struct s2n_mem_test_cb_scope mem_ctx = { 0 }, s2n_mem_test_free_callbacks); + EXPECT_OK(s2n_mem_test_init_callbacks(&mem_ctx)); EXPECT_ERROR_WITH_ERRNO(s2n_test_basic_recv(client, server), S2N_ERR_BAD_MESSAGE); EXPECT_EQUAL(tickets_count, 0); - EXPECT_EQUAL(mallocs_count, 0); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); } /* Test: server rejects fragmented post-handshake message (KeyUpdate) with an invalid size @@ -383,10 +371,12 @@ int main(int argc, char **argv) EXPECT_OK(s2n_test_send_records(client, message, fragment_size)); EXPECT_SUCCESS(s2n_update_application_traffic_keys(client, client->mode, SENDING)); + DEFER_CLEANUP(struct s2n_mem_test_cb_scope mem_ctx = { 0 }, s2n_mem_test_free_callbacks); + EXPECT_OK(s2n_mem_test_init_callbacks(&mem_ctx)); EXPECT_ERROR_WITH_ERRNO(s2n_test_basic_recv(client, server), S2N_ERR_BAD_MESSAGE); /* No post-handshake message should trigger the server to allocate memory */ - EXPECT_EQUAL(mallocs_count, 0); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); }; /* Test: client receives empty post-handshake messages (HelloRequests) @@ -413,15 +403,18 @@ int main(int argc, char **argv) EXPECT_SUCCESS(s2n_stuffer_write_uint24(&messages, 0)); } + DEFER_CLEANUP(struct s2n_mem_test_cb_scope mem_ctx = { 0 }, s2n_mem_test_free_callbacks); + EXPECT_OK(s2n_mem_test_init_callbacks(&mem_ctx)); + EXPECT_OK(s2n_test_send_records(server, messages, fragment_size)); EXPECT_OK(s2n_test_basic_recv(server, client)); EXPECT_EQUAL(hello_request_count, S2N_TEST_MESSAGE_COUNT); - EXPECT_EQUAL(mallocs_count, 0); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); EXPECT_OK(s2n_test_send_records(server, messages, fragment_size)); EXPECT_OK(s2n_test_blocking_recv(server, client, &io_pair)); EXPECT_EQUAL(hello_request_count, S2N_TEST_MESSAGE_COUNT); - EXPECT_EQUAL(mallocs_count, 0); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); } /* Test: client and server reject known, invalid messages (ClientHellos) */ @@ -444,11 +437,13 @@ int main(int argc, char **argv) EXPECT_SUCCESS(s2n_stuffer_skip_write(&message, test_large_message_size)); EXPECT_OK(s2n_test_send_records(sender, message, fragment_size)); + DEFER_CLEANUP(struct s2n_mem_test_cb_scope mem_ctx = { 0 }, s2n_mem_test_free_callbacks); + EXPECT_OK(s2n_mem_test_init_callbacks(&mem_ctx)); EXPECT_ERROR_WITH_ERRNO(s2n_test_basic_recv(sender, receiver), S2N_ERR_BAD_MESSAGE); /* No post-handshake message should trigger the server to allocate memory */ if (mode == S2N_SERVER) { - EXPECT_EQUAL(mallocs_count, 0); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); } } } @@ -473,11 +468,13 @@ int main(int argc, char **argv) EXPECT_SUCCESS(s2n_stuffer_skip_write(&message, test_large_message_size)); EXPECT_OK(s2n_test_send_records(sender, message, fragment_size)); + DEFER_CLEANUP(struct s2n_mem_test_cb_scope mem_ctx = { 0 }, s2n_mem_test_free_callbacks); + EXPECT_OK(s2n_mem_test_init_callbacks(&mem_ctx)); EXPECT_ERROR_WITH_ERRNO(s2n_test_basic_recv(sender, receiver), S2N_ERR_BAD_MESSAGE); /* No post-handshake message should trigger the server to allocate memory */ if (mode == S2N_SERVER) { - EXPECT_EQUAL(mallocs_count, 0); + EXPECT_OK(s2n_mem_test_assert_malloc_count(0)); } } }