Fix pem validation (#735)

Author: DmitriyMusatkin

source/pem.c

@@ -9,11 +9,7 @@
 
 #include <aws/io/logging.h>
 
-enum aws_pem_parse_state {
-    BEGIN,
-    ON_DATA,
-    END,
-};
+enum aws_pem_parse_state { BEGIN, ON_DATA };
 
 static const struct aws_byte_cursor begin_header = AWS_BYTE_CUR_INIT_FROM_STRING_LITERAL("-----BEGIN");
 static const struct aws_byte_cursor end_header = AWS_BYTE_CUR_INIT_FROM_STRING_LITERAL("-----END");
@@ -29,63 +25,48 @@ int aws_sanitize_pem(struct aws_byte_buf *pem, struct aws_allocator *allocator)
         return AWS_OP_ERR;
     }
     struct aws_byte_cursor pem_cursor = aws_byte_cursor_from_buf(pem);
-    enum aws_pem_parse_state state = BEGIN;
-
-    for (size_t i = 0; i < pem_cursor.len; i++) {
-        /* parse through the pem once */
-        char current = *(pem_cursor.ptr + i);
-        switch (state) {
-            case BEGIN:
-                if (current == '-') {
-                    struct aws_byte_cursor compare_cursor = pem_cursor;
-                    compare_cursor.len = begin_header.len;
-                    compare_cursor.ptr += i;
-                    if (aws_byte_cursor_eq(&compare_cursor, &begin_header)) {
-                        state = ON_DATA;
-                        i--;
-                    }
-                }
-                break;
-            case ON_DATA:
-                /* start copying everything */
-                if (current == '-') {
-                    struct aws_byte_cursor compare_cursor = pem_cursor;
-                    compare_cursor.len = end_header.len;
-                    compare_cursor.ptr += i;
-                    if (aws_byte_cursor_eq(&compare_cursor, &end_header)) {
-                        /* Copy the end header string and start to search for the end part of a pem */
-                        state = END;
-                        aws_byte_buf_append(&clean_pem_buf, &end_header);
-                        i += (end_header.len - 1);
-                        break;
-                    }
-                }
-                aws_byte_buf_append_byte_dynamic(&clean_pem_buf, (uint8_t)current);
-                break;
-            case END:
-                if (current == '-') {
-                    struct aws_byte_cursor compare_cursor = pem_cursor;
-                    compare_cursor.len = dashes.len;
-                    compare_cursor.ptr += i;
-                    if (aws_byte_cursor_eq(&compare_cursor, &dashes)) {
-                        /* End part of a pem, copy the last 5 dashes and a new line, then ignore everything before next
-                         * begin header */
-                        state = BEGIN;
-                        aws_byte_buf_append(&clean_pem_buf, &dashes);
-                        i += (dashes.len - 1);
+    bool is_on_data = false;
+
+    while (pem_cursor.len > 0) {
+        if (!is_on_data) {
+            if (aws_byte_cursor_starts_with(&pem_cursor, &begin_header)) {
+                /* mini optimization - just copy over begin to avoid all the starts with checks */
+                aws_byte_buf_append_dynamic(&clean_pem_buf, &begin_header);
+                aws_byte_cursor_advance(&pem_cursor, begin_header.len);
+                is_on_data = true;
+            } else {
+                aws_byte_cursor_advance(&pem_cursor, 1);
+            }
+        } else {
+            /* Note this does not validate that end label is same as begin label. */
+            if (aws_byte_cursor_starts_with(&pem_cursor, &end_header)) {
+                /* copy over end */
+                aws_byte_buf_append_dynamic(&clean_pem_buf, &end_header);
+                aws_byte_cursor_advance(&pem_cursor, end_header.len);
+
+                /* copy over label until the closing 5 dashes */
+                while (pem_cursor.len > 0) {
+                    aws_byte_buf_append_byte_dynamic(&clean_pem_buf, *pem_cursor.ptr);
+                    aws_byte_cursor_advance(&pem_cursor, 1);
+
+                    if (aws_byte_cursor_starts_with(&pem_cursor, &dashes)) {
+                        aws_byte_buf_append_dynamic(&clean_pem_buf, &dashes);
+                        aws_byte_cursor_advance(&pem_cursor, dashes.len);
                         aws_byte_buf_append_byte_dynamic(&clean_pem_buf, (uint8_t)'\n');
+                        is_on_data = false;
                         break;
                     }
                 }
-                aws_byte_buf_append_byte_dynamic(&clean_pem_buf, (uint8_t)current);
-                break;
-            default:
-                break;
+            } else {
+                aws_byte_buf_append_byte_dynamic(&clean_pem_buf, *pem_cursor.ptr);
+                aws_byte_cursor_advance(&pem_cursor, 1);
+            }
         }
     }
 
-    if (clean_pem_buf.len == 0) {
-        /* No valid data remains after sanitization. File might have been the wrong format */
+    if (clean_pem_buf.len == 0 || is_on_data) {
+        /* No valid data remains after sanitization or data block is left hanging.
+         File might have been the wrong format */
         aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
         goto error;
     }

tests/CMakeLists.txt

@@ -156,6 +156,7 @@ add_test_case(test_pem_invalid_in_chain_parse)
 add_test_case(pem_sanitize_comments_around_pem_object_removed)
 add_test_case(pem_sanitize_empty_file_rejected)
 add_test_case(pem_sanitize_wrong_format_rejected)
+add_test_case(pem_sanitize_bad_format_rejected)
 
 add_test_case(socket_handler_echo_and_backpressure)
 add_test_case(socket_handler_close)

tests/pem_test.c

@@ -59,6 +59,18 @@ static int s_test_pem_sanitize_empty_file_rejected(struct aws_allocator *allocat
 
 AWS_TEST_CASE(pem_sanitize_empty_file_rejected, s_test_pem_sanitize_empty_file_rejected)
 
+static int s_test_pem_sanitize_bad_format_rejected(struct aws_allocator *allocator, void *ctx) {
+    (void)ctx;
+    struct aws_byte_buf pem = aws_byte_buf_from_c_str("aaaaa-");
+
+    ASSERT_ERROR(AWS_ERROR_INVALID_ARGUMENT, aws_sanitize_pem(&pem, allocator));
+
+    aws_byte_buf_clean_up(&pem);
+    return AWS_OP_SUCCESS;
+}
+
+AWS_TEST_CASE(pem_sanitize_bad_format_rejected, s_test_pem_sanitize_bad_format_rejected)
+
 AWS_TEST_CASE(pem_sanitize_wrong_format_rejected, s_test_pem_sanitize_wrong_format_rejected)
 static int s_test_pem_sanitize_wrong_format_rejected(struct aws_allocator *allocator, void *ctx) {
     (void)ctx;