From fecbf0e03a2240d8de2fcda38632f2eeaa3f069a Mon Sep 17 00:00:00 2001
From: Ali Mohammad Pur <ali.mpfard@gmail.com>
Date: Sat, 4 Dec 2021 17:57:48 +0330
Subject: [PATCH] LibWasm: Make blocks that take arguments actually work

Previously we were ignoring the actual parameters and setting the arity
to an incorrect value, which could cause crashes (or unexpected traps).
---
 .../AbstractMachine/BytecodeInterpreter.cpp   | 49 ++++++++++++++++---
 .../LibWasm/AbstractMachine/Validator.cpp     | 11 ++++-
 2 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/Userland/Libraries/LibWasm/AbstractMachine/BytecodeInterpreter.cpp b/Userland/Libraries/LibWasm/AbstractMachine/BytecodeInterpreter.cpp
index 1b64d8704e..5017916d88 100644
--- a/Userland/Libraries/LibWasm/AbstractMachine/BytecodeInterpreter.cpp
+++ b/Userland/Libraries/LibWasm/AbstractMachine/BytecodeInterpreter.cpp
@@ -376,25 +376,60 @@ void BytecodeInterpreter::interpret(Configuration& configuration, InstructionPoi
         return;
     case Instructions::block.value(): {
         size_t arity = 0;
+        size_t parameter_count = 0;
         auto& args = instruction.arguments().get<Instruction::StructuredInstructionArgs>();
-        if (args.block_type.kind() != BlockType::Empty)
+        switch (args.block_type.kind()) {
+        case BlockType::Empty:
+            break;
+        case BlockType::Type:
             arity = 1;
-        configuration.stack().push(Label(arity, args.end_ip));
+            break;
+        case BlockType::Index: {
+            auto& type = configuration.frame().module().types()[args.block_type.type_index().value()];
+            arity = type.results().size();
+            parameter_count = type.parameters().size();
+        }
+        }
+
+        configuration.stack().entries().insert(configuration.stack().size() - parameter_count, Label(arity, args.end_ip));
         return;
     }
     case Instructions::loop.value(): {
         size_t arity = 0;
+        size_t parameter_count = 0;
         auto& args = instruction.arguments().get<Instruction::StructuredInstructionArgs>();
-        if (args.block_type.kind() != BlockType::Empty)
+        switch (args.block_type.kind()) {
+        case BlockType::Empty:
+            break;
+        case BlockType::Type:
             arity = 1;
-        configuration.stack().push(Label(arity, ip.value() + 1));
+            break;
+        case BlockType::Index: {
+            auto& type = configuration.frame().module().types()[args.block_type.type_index().value()];
+            arity = type.results().size();
+            parameter_count = type.parameters().size();
+        }
+        }
+
+        configuration.stack().entries().insert(configuration.stack().size() - parameter_count, Label(arity, ip.value() + 1));
         return;
     }
     case Instructions::if_.value(): {
         size_t arity = 0;
+        size_t parameter_count = 0;
         auto& args = instruction.arguments().get<Instruction::StructuredInstructionArgs>();
-        if (args.block_type.kind() != BlockType::Empty)
+        switch (args.block_type.kind()) {
+        case BlockType::Empty:
+            break;
+        case BlockType::Type:
             arity = 1;
+            break;
+        case BlockType::Index: {
+            auto& type = configuration.frame().module().types()[args.block_type.type_index().value()];
+            arity = type.results().size();
+            parameter_count = type.parameters().size();
+        }
+        }
 
         auto entry = configuration.stack().pop();
         auto value = entry.get<Value>().to<i32>();
@@ -402,12 +437,12 @@ void BytecodeInterpreter::interpret(Configuration& configuration, InstructionPoi
         if (value.value() == 0) {
             if (args.else_ip.has_value()) {
                 configuration.ip() = args.else_ip.value();
-                configuration.stack().push(move(end_label));
+                configuration.stack().entries().insert(configuration.stack().size() - parameter_count, end_label);
             } else {
                 configuration.ip() = args.end_ip.value() + 1;
             }
         } else {
-            configuration.stack().push(move(end_label));
+            configuration.stack().entries().insert(configuration.stack().size() - parameter_count, end_label);
         }
         return;
     }
diff --git a/Userland/Libraries/LibWasm/AbstractMachine/Validator.cpp b/Userland/Libraries/LibWasm/AbstractMachine/Validator.cpp
index e573e11dfd..6a5d6bd9b1 100644
--- a/Userland/Libraries/LibWasm/AbstractMachine/Validator.cpp
+++ b/Userland/Libraries/LibWasm/AbstractMachine/Validator.cpp
@@ -2424,11 +2424,14 @@ VALIDATE_INSTRUCTION(block)
     if (stack.size() < parameters.size())
         return Errors::invalid_stack_state();
 
-    for (size_t i = 0; i < parameters.size(); ++i) {
+    for (size_t i = 1; i <= parameters.size(); ++i) {
         if (stack.take_last() != parameters[parameters.size() - i])
             return Errors::invalid_stack_state();
     }
 
+    for (auto& parameter : parameters)
+        stack.append(parameter);
+
     m_entered_scopes.append(ChildScopeKind::Block);
     m_block_details.empend(stack.actual_size(), Empty {});
     m_parent_contexts.append(m_context);
@@ -2451,6 +2454,9 @@ VALIDATE_INSTRUCTION(loop)
             return Errors::invalid_stack_state();
     }
 
+    for (auto& parameter : parameters)
+        stack.append(parameter);
+
     m_entered_scopes.append(ChildScopeKind::Block);
     m_block_details.empend(stack.actual_size(), Empty {});
     m_parent_contexts.append(m_context);
@@ -2476,6 +2482,9 @@ VALIDATE_INSTRUCTION(if_)
             return Errors::invalid_stack_state();
     }
 
+    for (auto& parameter : parameters)
+        stack.append(parameter);
+
     m_entered_scopes.append(args.else_ip.has_value() ? ChildScopeKind::IfWithElse : ChildScopeKind::IfWithoutElse);
     m_block_details.empend(stack.actual_size(), BlockDetails::IfDetails { stack, {} });
     m_parent_contexts.append(m_context);