Fix/branch hint parse issue (#4878)

* test,security: added example code for invalid branch hints
* security: fixes parsing issue for invalid branch_hint sections
* fix: improve comments and logging in branch hint handling
* feat: add support for branch hints in build configuration and documentation
* test: add regression tests for invalid branch hints handling

```bash
$ pwd
/workspaces/wasm-micro-runtime/tests/regression/ba-issues

$ ./run.py -i 980002,980003
```

---------

Signed-off-by: Stephen Berard <stephen.berard@outlook.com>
Co-authored-by: liang.he@intel.com <liang.he@intel.com>

build-scripts/config_common.cmake

@@ -775,6 +775,11 @@ endif ()
 if (WAMR_BUILD_LIME1 EQUAL 1)
   message ("     Lime1 enabled")
 endif ()
+if (WAMR_BUILD_BRANCH_HINTS EQUAL 1)
+  message ("     Branch hints enabled")
+  add_definitions(-DWASM_ENABLE_BRANCH_HINTS=1)
+endif ()
+
 ########################################
 # Show Phase4 Wasm proposals status.
 ########################################
@@ -787,8 +792,8 @@ message (
 "       \"Non-trapping float-to-int Conversions\"\n"
 "       \"Sign-extension Operators\"\n"
 "       \"WebAssembly C and C++ API\"\n"
-"       \"Branch Hinting\"\n"
 "     Configurable. 0 is OFF. 1 is ON:\n"
+"       \"Branch Hinting\" via WAMR_BUILD_BRANCH_HINTS: ${WAMR_BUILD_BRANCH_HINTS}\n"
 "       \"Bulk Memory Operation\" via WAMR_BUILD_BULK_MEMORY: ${WAMR_BUILD_BULK_MEMORY}\n"
 "       \"Bulk-memory-opt\" via WAMR_BUILD_BULK_MEMORY_OPT: ${WAMR_BUILD_BULK_MEMORY_OPT}\n"
 "       \"Call-indirect-overlong\" via WAMR_BUILD_CALL_INDIRECT_OVERLONG: ${WAMR_BUILD_CALL_INDIRECT_OVERLONG}\n"

core/iwasm/interpreter/wasm_loader.c

@@ -5577,6 +5577,27 @@ fail:
 #endif
 
 #if WASM_ENABLE_BRANCH_HINTS != 0
+/**
+ * Count the number of branch instructions for the specified function.
+ */
+static uint32
+calculate_num_branch_instructions(const WASMFunction *func)
+{
+    const uint8 *code = func->code;
+    const uint8 *code_end = code + func->code_size;
+    uint32 max_hints = 0;
+
+    while (code < code_end) {
+        uint8 opcode = *code++;
+
+        if (opcode == WASM_OP_IF || opcode == WASM_OP_BR_IF) {
+            max_hints++;
+        }
+    }
+
+    return max_hints;
+}
+
 static bool
 handle_branch_hint_section(const uint8 *buf, const uint8 *buf_end,
                            WASMModule *module, char *error_buf,
@@ -5611,22 +5632,52 @@ handle_branch_hint_section(const uint8 *buf, const uint8 *buf_end,
 
         uint32 num_hints;
         read_leb_uint32(buf, buf_end, num_hints);
+
+        /* Ensure that num_hints doesn't exceed the actual number of branch
+         * instructions */
+        WASMFunction *func =
+            module->functions[func_idx - module->import_function_count];
+        uint32 max_branch_instructions =
+            calculate_num_branch_instructions(func);
+        if (num_hints > max_branch_instructions) {
+            set_error_buf_v(
+                error_buf, error_buf_size,
+                "invalid number of branch hints: expected at most %u, got %u",
+                max_branch_instructions, num_hints);
+            goto fail;
+        }
+
         struct WASMCompilationHintBranchHint *new_hints = loader_malloc(
             sizeof(struct WASMCompilationHintBranchHint) * num_hints, error_buf,
             error_buf_size);
+        if (!new_hints) {
+            goto fail;
+        }
         for (uint32 j = 0; j < num_hints; ++j) {
             struct WASMCompilationHintBranchHint *new_hint = &new_hints[j];
             new_hint->next = NULL;
             new_hint->type = WASM_COMPILATION_BRANCH_HINT;
             read_leb_uint32(buf, buf_end, new_hint->offset);
 
+            /* Validate offset is within the function's code bounds */
+            if (new_hint->offset >= func->code_size) {
+                set_error_buf_v(
+                    error_buf, error_buf_size,
+                    "invalid branch hint offset: %u exceeds function "
+                    "code size %u",
+                    new_hint->offset, func->code_size);
+                goto fail;
+            }
+
             uint32 size;
             read_leb_uint32(buf, buf_end, size);
             if (size != 1) {
                 set_error_buf_v(error_buf, error_buf_size,
                                 "invalid branch hint size, expected 1, got %d.",
                                 size);
-                wasm_runtime_free(new_hint);
+                /* Do not free new_hints here - any hints already linked into
+                 * the module structure will be freed during module cleanup.
+                 * Freeing here would cause a double-free. */
                 goto fail;
             }
 
@@ -5639,7 +5690,9 @@ handle_branch_hint_section(const uint8 *buf, const uint8 *buf_end,
                 set_error_buf_v(error_buf, error_buf_size,
                                 "invalid branch hint, expected 0 or 1, got %d",
                                 data);
-                wasm_runtime_free(new_hint);
+                /* Do not free new_hints here - any hints already linked into
+                 * the module structure will be freed during module cleanup.
+                 * Freeing here would cause a double-free. */
                 goto fail;
             }
 
@@ -5720,7 +5773,7 @@ load_user_section(const uint8 *buf, const uint8 *buf_end, WASMModule *module,
 #else
     if (name_len == 25
         && memcmp((const char *)p, "metadata.code.branch_hint", 25) == 0) {
-        LOG_VERBOSE("Found branch hint section, but branch hints are disabled "
+        LOG_WARNING("Found branch hint section, but branch hints are disabled "
                     "in this build, skipping.");
     }
 #endif

doc/build_wamr.md

@@ -616,6 +616,15 @@ SIMDE (SIMD Everywhere) implements SIMD operations in fast interpreter mode.
 > [!WARNING]
 > This is only supported in classic interpreter mode.
 
+## **Branch hints**
+
+- **WAMR_BUILD_BRANCH_HINTS**=1/0, default to disable if not set
+
+> [!NOTE]
+> Enabling this feature allows the runtime to utilize branch hints for better performance during aot/jit execution.
+
+## **Combination of configurations:**
+
 ### **Invoke general FFI**
 
 - **WAMR_BUILD_INVOKE_NATIVE_GENERAL**=1/0, default to off.

doc/tiered_support.md

@@ -165,6 +165,7 @@ This tier indicates experimental features with foundational support levels. Thes
 | GC (Garbage Collection)        | [WAMR_BUILD_GC](./build_wamr.md#garbage-collection)                                                   | Wasm Proposal      |
 | Legacy Exception Handling      | [WAMR_BUILD_EXCE_HANDLING](./build_wamr.md#exception-handling)                                        | Wasm Proposal      |
 | Multi-memory                   | [WAMR_BUILD_MULTI_MEMORY](./build_wamr.md#multi-memory)                                               | Wasm Proposal      |
+| Branch Hints                   | [WAMR_BUILD_BRANCH_HINTS](./build_wamr.md#branch-hints-feature)                                       | Wasm Proposal      |
 | Fast JIT                       | [WAMR_BUILD_FAST_JIT](./build_wamr.md#configure-fast-jit)                                             | Running mode       |
 | Multi-tier JIT                 | [Combination of flags](./build_wamr.md#configure-multi-tier-jit)                                      | Running mode       |
 | AoT Validator                  | [WAMR_BUILD_AOT_VALIDATOR](./build_wamr.md#aot-validator)                                             | Runtime Extensions |

tests/regression/ba-issues/build_wamr.sh

@@ -60,4 +60,7 @@ build_iwasm "-DWAMR_BUILD_REF_TYPES=1 -DWAMR_BUILD_JIT=1 -DWAMR_BUILD_LIBC_WASI=
 # build fast-jit iwasm for testing fast-jit with libc-wasi disabled
 build_iwasm "-DWAMR_BUILD_REF_TYPES=1 -DWAMR_BUILD_FAST_JIT=1 -DWAMR_BUILD_SIMD=0 -DWAMR_BUILD_LIBC_WASI=0" "fast-jit-wasi-disabled"
 
+# build default iwasm for testing wasm loader with branch hints enabled
+build_iwasm "-DWAMR_BUILD_BRANCH_HINTS=1" "default-branch-hints-enabled"
+
 # TODO: add more version of iwasm, for example, sgx version

tests/regression/ba-issues/issues/issue-980002/create_samples.py

@@ -0,0 +1,48 @@
+from pathlib import Path
+
+def u32leb(n):
+    out = bytearray()
+    while True:
+        b = n & 0x7f
+        n >>= 7
+        if n:
+            b |= 0x80
+        out.append(b)
+        if not n:
+            break
+    return bytes(out)
+name = b"metadata.code.branch_hint"
+assert len(name) == 25
+def build_module(payload_tail, out_path):
+    payload = b"".join([
+        u32leb(len(name)),
+        name,
+        payload_tail
+    ])
+    custom_section = b"\x00" + u32leb(len(payload)) + payload
+    payload_type = u32leb(1) + b"\x60" + u32leb(0) + u32leb(0)
+    sec_type = b"\x01" + u32leb(len(payload_type)) + payload_type
+    payload_func = u32leb(1) + u32leb(0)
+    sec_func = b"\x03" + u32leb(len(payload_func)) + payload_func
+    body = u32leb(0) + b"\x0b"
+    payload_code = u32leb(1) + u32leb(len(body)) + body
+    sec_code = b"\x0a" + u32leb(len(payload_code)) + payload_code
+    module = b"\x00asm" + b"\x01\x00\x00\x00" + sec_type + sec_func + sec_code + custom_section
+    Path(out_path).write_bytes(module)
+payload_invalid_free = b"".join([
+    b"\x01",            # numFunctionHints
+    b"\x00",            # func_idx
+    b"\x02",            # num_hints
+    b"\x00",            # hint0 offset
+    b"\x01",            # hint0 size
+    b"\x00",            # hint0 data
+    b"\x00",            # hint1 offset
+    b"\x02",            # hint1 size (invalid)
+])
+build_module(payload_invalid_free, "branch_hint_invalid_free.wasm")
+payload_dos = b"".join([
+    b"\x01",
+    b"\x00",
+    b"\xff\xff\xff\xff\x0f",
+])
+build_module(payload_dos, "branch_hint_null_deref.wasm")

tests/regression/ba-issues/running_config.json

@@ -1818,6 +1818,40 @@
                 "stdout content": "Exception: unsupported opcode",
                 "description": "classic-interp will exit gracefully when meeting simd opcodes"
             }
+        },
+        {
+            "deprecated": false,
+            "ids": [
+                980002
+            ],
+            "runtime": "iwasm-default-branch-hints-enabled",
+            "file": "branch_hint_invalid_free.wasm",
+            "mode": "fast-interp",
+            "options": "",
+            "argument": "",
+            "expected return": {
+                "ret code": 255,
+                "stdout content": "WASM module load failed: invalid number of branch hints: expected at most 0, got 2",
+                "description": ""
+            }
+
+        },
+        {
+            "deprecated": false,
+            "ids": [
+                980003
+            ],
+            "runtime": "iwasm-default-branch-hints-enabled",
+            "file": "branch_hint_null_deref.wasm",
+            "mode": "fast-interp",
+            "options": "",
+            "argument": "",
+            "expected return": {
+                "ret code": 255,
+                "stdout content": "WASM module load failed: invalid number of branch hints: expected at most 0, got 42949672",
+                "description": ""
+            }
+
         }
     ]
 }