From cf1b681b10250c8c0931b20d804409e51447a13d Mon Sep 17 00:00:00 2001
From: Dmitriy Musatkin <63878209+DmitriyMusatkin@users.noreply.github.com>
Date: Wed, 10 Apr 2024 13:37:58 -0700
Subject: [PATCH] Remove setenv (#634)

---
 source/s2n/s2n_tls_channel_handler.c | 31 ++++++++++++++++++++++++++--
 tests/socket_test.c                  |  4 +++-
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/source/s2n/s2n_tls_channel_handler.c b/source/s2n/s2n_tls_channel_handler.c
index 80770b62b..14839d19f 100644
--- a/source/s2n/s2n_tls_channel_handler.c
+++ b/source/s2n/s2n_tls_channel_handler.c
@@ -174,11 +174,34 @@ AWS_IO_API const char *aws_determine_default_pki_ca_file(void) {
     return NULL;
 }
 
+static struct aws_allocator *s_library_allocator = NULL;
+
+static int s_s2n_mem_init(void) {
+    return S2N_SUCCESS;
+}
+
+static int s_s2n_mem_cleanup(void) {
+    return S2N_SUCCESS;
+}
+
+static int s_s2n_mem_malloc(void **ptr, uint32_t requested, uint32_t *allocated) {
+    *ptr = aws_mem_acquire(s_library_allocator, requested);
+    *allocated = requested;
+
+    return S2N_SUCCESS;
+}
+
+static int s_s2n_mem_free(void *ptr, uint32_t size) {
+    (void)size;
+    aws_mem_release(s_library_allocator, ptr);
+    return S2N_SUCCESS;
+}
+
 /* If s2n is already initialized, then we don't call s2n_init() or s2n_cleanup() ourselves */
 static bool s_s2n_initialized_externally = false;
 
 void aws_tls_init_static_state(struct aws_allocator *alloc) {
-    (void)alloc;
+    AWS_FATAL_ASSERT(alloc);
     AWS_LOGF_INFO(AWS_LS_IO_TLS, "static: Initializing TLS using s2n.");
 
     /* Disable atexit behavior, so that s2n_cleanup() fully cleans things up.
@@ -196,7 +219,11 @@ void aws_tls_init_static_state(struct aws_allocator *alloc) {
     }
 
     if (!s_s2n_initialized_externally) {
-        setenv("S2N_DONT_MLOCK", "1", 1);
+        s_library_allocator = alloc;
+        if (S2N_SUCCESS != s2n_mem_set_callbacks(s_s2n_mem_init, s_s2n_mem_cleanup, s_s2n_mem_malloc, s_s2n_mem_free)) {
+            fprintf(stderr, "s2n_mem_set_callbacks() failed: %d (%s)\n", s2n_errno, s2n_strerror(s2n_errno, "EN"));
+            AWS_FATAL_ASSERT(0 && "s2n_mem_set_callbacks() failed");
+        }
 
         if (s2n_init() != S2N_SUCCESS) {
             fprintf(stderr, "s2n_init() failed: %d (%s)\n", s2n_errno, s2n_strerror(s2n_errno, "EN"));
diff --git a/tests/socket_test.c b/tests/socket_test.c
index b23ed084b..07740fc21 100644
--- a/tests/socket_test.c
+++ b/tests/socket_test.c
@@ -790,7 +790,9 @@ static int s_test_outgoing_tcp_sock_error(struct aws_allocator *allocator, void
 
     struct aws_socket_endpoint endpoint = {
         .address = "127.0.0.1",
-        .port = 8567,
+        /* note: the port is completely random from testing perspective, but
+         * freebsd seems to firewall higher numbered ports so keeping it low */
+        .port = 1567,
     };
 
     struct error_test_args args = {