test aws_http_request OOM errors (#93)

Author: graebm

tests/CMakeLists.txt

@@ -11,6 +11,7 @@ add_test_case(request_method)
 add_test_case(request_path)
 add_test_case(request_add_headers)
 add_test_case(request_erase_headers)
+add_test_case(request_handles_oom)
 
 add_test_case(h1_test_get_request)
 add_test_case(h1_test_request_bad_version)

tests/test_request.c

@@ -16,7 +16,7 @@
 #include <aws/http/request_response.h>
 
 #include <aws/common/string.h>
-#include <aws/testing/aws_test_harness.h>
+#include <aws/testing/aws_test_allocators.h>
 
 #define TEST_CASE(NAME)                                                                                                \
     AWS_TEST_CASE(NAME, s_test_##NAME);                                                                                \
@@ -132,6 +132,11 @@ TEST_CASE(request_add_headers) {
     ASSERT_SUCCESS(aws_http_request_get_header(request, &get, 0));
     ASSERT_SUCCESS(s_check_header_eq(get, "Host", "example.com"));
 
+    /* Overwrite header and check values */
+    ASSERT_SUCCESS(aws_http_request_set_header(request, s_make_header("Connection", "Upgrade"), 0));
+    ASSERT_SUCCESS(aws_http_request_get_header(request, &get, 0));
+    ASSERT_SUCCESS(s_check_header_eq(get, "Connection", "Upgrade"));
+
     aws_http_request_destroy(request);
     return AWS_OP_SUCCESS;
 }
@@ -194,3 +199,78 @@ TEST_CASE(request_erase_headers) {
     aws_http_request_destroy(request);
     return AWS_OP_SUCCESS;
 }
+
+/* Do every operation that involves allocating some memory */
+int s_request_handles_oom_attempt(struct aws_http_request *request) {
+    ASSERT_NOT_NULL(request);
+
+    /* Set, and then overwrite, method and path */
+    ASSERT_SUCCESS(aws_http_request_set_method(request, aws_byte_cursor_from_c_str("POST")));
+    ASSERT_SUCCESS(aws_http_request_set_path(request, aws_byte_cursor_from_c_str("/")));
+    ASSERT_SUCCESS(aws_http_request_set_method(request, aws_byte_cursor_from_c_str("GET")));
+    ASSERT_SUCCESS(aws_http_request_set_path(request, aws_byte_cursor_from_c_str("/chat")));
+
+    /* Add a lot of headers, enough to force the underlying array-list to expand.
+     * (just loop through the list above again and again) */
+    char name_buf[16];
+    char value_buf[16];
+    for (size_t i = 0; i < 128; ++i) {
+        snprintf(name_buf, sizeof(name_buf), "Name-%zu", i);
+        snprintf(name_buf, sizeof(name_buf), "Value-%zu", i);
+        struct aws_http_header header = {.name = aws_byte_cursor_from_c_str(name_buf),
+                                         .value = aws_byte_cursor_from_c_str(value_buf)};
+        ASSERT_SUCCESS(aws_http_request_add_header(request, header));
+    }
+
+    /* Overwrite all the headers */
+    for (size_t i = 0; i < 128; ++i) {
+        snprintf(name_buf, sizeof(name_buf), "New-Name-%zu", i);
+        snprintf(name_buf, sizeof(name_buf), "New-Value-%zu", i);
+        struct aws_http_header header = {.name = aws_byte_cursor_from_c_str(name_buf),
+                                         .value = aws_byte_cursor_from_c_str(value_buf)};
+        ASSERT_SUCCESS(aws_http_request_set_header(request, header, i));
+    }
+
+    return AWS_OP_SUCCESS;
+}
+
+TEST_CASE(request_handles_oom) {
+    (void)ctx;
+    struct aws_allocator timebomb_alloc;
+    ASSERT_SUCCESS(aws_timebomb_allocator_init(&timebomb_alloc, allocator, SIZE_MAX));
+
+    bool test_succeeded = false;
+    size_t allocations_until_failure;
+    for (allocations_until_failure = 0; allocations_until_failure < 10000; ++allocations_until_failure) {
+        /* Allow one more allocation each time we loop. */
+        aws_timebomb_allocator_reset_countdown(&timebomb_alloc, allocations_until_failure);
+
+        /* Create a request, then do a bunch of stuff with it. */
+        struct aws_http_request *request = aws_http_request_new(&timebomb_alloc);
+        int err = 0;
+        if (request) {
+            err = s_request_handles_oom_attempt(request);
+            if (err) {
+                /* Ensure failure was due to OOM */
+                ASSERT_INT_EQUALS(AWS_ERROR_OOM, aws_last_error());
+            } else {
+                test_succeeded = true;
+            }
+
+            aws_http_request_destroy(request);
+        } else {
+            /* Ensure failure was due to OOM */
+            ASSERT_INT_EQUALS(AWS_ERROR_OOM, aws_last_error());
+        }
+
+        if (test_succeeded) {
+            break;
+        }
+    }
+
+    ASSERT_TRUE(test_succeeded);
+    ASSERT_TRUE(allocations_until_failure > 2); /* Assert that this did fail a few times */
+
+    aws_timebomb_allocator_clean_up(&timebomb_alloc);
+    return AWS_OP_SUCCESS;
+}