From 372836114743aa9f3feb55620a1b35fed586ed3a Mon Sep 17 00:00:00 2001
From: Boquan Fang <boquanfang3@gmail.com>
Date: Fri, 23 Aug 2024 11:50:48 -0700
Subject: [PATCH] fix: resolve UBSAN violations in the codebase (#4722)

---
 stuffer/s2n_stuffer.c                  | 12 +++++++++++-
 tests/s2n_test.h                       |  8 +++++++-
 tests/testlib/s2n_test_server_client.c |  7 ++++++-
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/stuffer/s2n_stuffer.c b/stuffer/s2n_stuffer.c
index a76a2f41546..134ec0ba9c8 100644
--- a/stuffer/s2n_stuffer.c
+++ b/stuffer/s2n_stuffer.c
@@ -323,6 +323,10 @@ int s2n_stuffer_write(struct s2n_stuffer *stuffer, const struct s2n_blob *in)
 
 int s2n_stuffer_write_bytes(struct s2n_stuffer *stuffer, const uint8_t *data, const uint32_t size)
 {
+    if (size == 0) {
+        return S2N_SUCCESS;
+    }
+
     POSIX_ENSURE(S2N_MEM_IS_READABLE(data, size), S2N_ERR_SAFETY);
     POSIX_PRECONDITION(s2n_stuffer_validate(stuffer));
     POSIX_GUARD(s2n_stuffer_skip_write(stuffer, size));
@@ -437,7 +441,13 @@ int s2n_stuffer_shift(struct s2n_stuffer *stuffer)
     POSIX_ENSURE_REF(stuffer);
     struct s2n_stuffer copy = *stuffer;
     POSIX_GUARD(s2n_stuffer_rewrite(&copy));
-    uint8_t *data = stuffer->blob.data + stuffer->read_cursor;
+
+    uint8_t *data = stuffer->blob.data;
+    /* Adding 0 to a NULL value is undefined behavior */
+    if (stuffer->read_cursor != 0) {
+        data += stuffer->read_cursor;
+    }
+
     uint32_t data_size = s2n_stuffer_data_available(stuffer);
     POSIX_GUARD(s2n_stuffer_write_bytes(&copy, data, data_size));
     *stuffer = copy;
diff --git a/tests/s2n_test.h b/tests/s2n_test.h
index 7b79981355d..696ba68ed0d 100644
--- a/tests/s2n_test.h
+++ b/tests/s2n_test.h
@@ -200,7 +200,13 @@ int test_count;
 /* for use with S2N_RESULT */
 #define EXPECT_OK( function_call )  EXPECT_TRUE( s2n_result_is_ok(function_call) )
 
-#define EXPECT_BYTEARRAY_EQUAL( p1, p2, l ) EXPECT_EQUAL( memcmp( (p1), (p2), (l) ), 0 )
+#define EXPECT_BYTEARRAY_EQUAL( p1, p2, l ) \
+    do { \
+        if (l != 0) { \
+            EXPECT_EQUAL( memcmp( (p1), (p2), (l) ), 0 ); \
+        } \
+    } while (0)
+
 #define EXPECT_BYTEARRAY_NOT_EQUAL( p1, p2, l ) EXPECT_NOT_EQUAL( memcmp( (p1), (p2), (l) ), 0 )
 
 #define EXPECT_STRING_EQUAL( p1, p2 ) EXPECT_EQUAL( strcmp( (p1), (p2) ), 0 )
diff --git a/tests/testlib/s2n_test_server_client.c b/tests/testlib/s2n_test_server_client.c
index ecaec2a752f..8ed01012c11 100644
--- a/tests/testlib/s2n_test_server_client.c
+++ b/tests/testlib/s2n_test_server_client.c
@@ -77,8 +77,13 @@ S2N_RESULT s2n_negotiate_test_server_and_client_with_early_data(struct s2n_conne
     bool server_early_done = false, client_early_done = false;
     do {
         if (!client_early_done) {
+            /* Adding 0 to a NULL value is undefined behavior */
+            uint8_t *next_early_data_to_send = early_data_to_send->data;
+            if (total_data_sent != 0) {
+                next_early_data_to_send += total_data_sent;
+            }
             bool success = s2n_send_early_data(client_conn,
-                                   early_data_to_send->data + total_data_sent,
+                                   next_early_data_to_send,
                                    early_data_to_send->size - total_data_sent,
                                    &data_sent, &blocked)
                     == S2N_SUCCESS;