Addressed copilot's comments

Author: danielinux

src/delta.c

@@ -397,6 +397,8 @@ int wb_diff(WB_DIFF_CTX *ctx, uint8_t *patch, uint32_t len)
 
         if (!found) {
             if (*(ctx->src_b + ctx->off_b) == ESC) {
+                if ((p_off + 1) >= (len - BLOCK_HDR_SIZE))
+                    break;
                 *(patch + p_off++) = ESC;
                 *(patch + p_off++) = ESC;
             } else {

src/store_sbrk.c

@@ -1,4 +1,26 @@
+/* store_sbrk.c
+ *
+ * Copyright (C) 2025 wolfSSL Inc.
+ *
+ * This file is part of wolfBoot.
+ *
+ * wolfBoot is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * wolfBoot is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1335, USA
+ */
+
 #include <stddef.h>
+#include <limits.h>
 
 #include "store_sbrk.h"
 
@@ -8,8 +30,11 @@ void *wolfboot_store_sbrk(unsigned int incr, uint8_t **heap,
     uint8_t *heap_limit = heap_base + heap_size;
     void *old_heap = *heap;
 
-    if (((incr >> 2) << 2) != incr)
-        incr = ((incr >> 2) + 1) << 2;
+    if ((incr & 3U) != 0U) {
+        if (incr > (UINT_MAX - 3U))
+            return (void *)-1;
+        incr = (incr + 3U) & ~3U;
+    }
 
     if (*heap == NULL) {
         *heap = heap_base;

src/store_sbrk.h

@@ -1,3 +1,24 @@
+/* store_sbrk.h
+ *
+ * Copyright (C) 2025 wolfSSL Inc.
+ *
+ * This file is part of wolfBoot.
+ *
+ * wolfBoot is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * wolfBoot is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1335, USA
+ */
+
 #ifndef WOLFBOOT_STORE_SBRK_H
 #define WOLFBOOT_STORE_SBRK_H
 

tools/unit-tests/unit-delta.c

@@ -217,6 +217,27 @@ START_TEST(test_wb_diff_preserves_trailing_header_margin_for_escape)
 }
 END_TEST
 
+START_TEST(test_wb_diff_preserves_main_loop_header_margin_for_escape)
+{
+    WB_DIFF_CTX diff_ctx;
+    uint8_t src_a[64] = {0};
+    uint8_t src_b[64] = {0};
+    uint8_t patch[BLOCK_HDR_SIZE + 2] = {0};
+    int ret;
+
+    memset(src_b, 0x5a, BLOCK_HDR_SIZE + 1);
+    src_b[0] = ESC;
+
+    ret = wb_diff_init(&diff_ctx, src_a, sizeof(src_a), src_b, BLOCK_HDR_SIZE + 1);
+    ck_assert_int_eq(ret, 0);
+
+    ret = wb_diff(&diff_ctx, patch, BLOCK_HDR_SIZE + 1);
+
+    ck_assert_int_eq(ret, 0);
+    ck_assert_uint_eq(patch[0], 0);
+}
+END_TEST
+
 static void initialize_buffers(uint8_t *src_a, uint8_t *src_b, size_t size)
 {
     uint32_t pseudo_rand = 0;
@@ -326,6 +347,7 @@ Suite *patch_diff_suite(void)
     tcase_add_test(tc_wolfboot_delta, test_wb_diff_match_extends_to_src_b_end);
     tcase_add_test(tc_wolfboot_delta, test_wb_diff_self_match_extends_to_src_b_end);
     tcase_add_test(tc_wolfboot_delta, test_wb_diff_preserves_trailing_header_margin_for_escape);
+    tcase_add_test(tc_wolfboot_delta, test_wb_diff_preserves_main_loop_header_margin_for_escape);
     tcase_add_test(tc_wolfboot_delta, test_wb_patch_and_diff);
     suite_add_tcase(s, tc_wolfboot_delta);
 

tools/unit-tests/unit-store-sbrk.c

@@ -1,9 +1,28 @@
 /* unit-store-sbrk.c
  *
  * Unit tests for store allocator helper.
+ *
+ * Copyright (C) 2025 wolfSSL Inc.
+ *
+ * This file is part of wolfBoot.
+ *
+ * wolfBoot is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * wolfBoot is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1335, USA
  */
 
 #include <check.h>
+#include <limits.h>
 #include <stdint.h>
 
 #include "../../src/store_sbrk.h"
@@ -36,13 +55,27 @@ START_TEST(test_sbrk_rejects_overflow)
 }
 END_TEST
 
+START_TEST(test_sbrk_rejects_alignment_overflow)
+{
+    uint8_t heap_buf[16];
+    uint8_t *heap = NULL;
+    void *ret;
+
+    ret = wolfboot_store_sbrk(UINT_MAX - 1U, &heap, heap_buf, sizeof(heap_buf));
+
+    ck_assert_ptr_eq(ret, (void *)-1);
+    ck_assert_ptr_eq(heap, NULL);
+}
+END_TEST
+
 Suite *wolfboot_suite(void)
 {
     Suite *s = suite_create("store-sbrk");
     TCase *tcase = tcase_create("store_sbrk");
 
     tcase_add_test(tcase, test_sbrk_first_call_advances_heap);
     tcase_add_test(tcase, test_sbrk_rejects_overflow);
+    tcase_add_test(tcase, test_sbrk_rejects_alignment_overflow);
     suite_add_tcase(s, tcase);
     return s;
 }