[WebAssembly] Check if the section order is correct

Summary:
This patch checks if the section order is correct when reading a wasm
object file in `WasmObjectFile` and converting YAML to wasm object in
yaml2wasm. (It is not possible to check when reading YAML because it is
handled exclusively by the YAML reader.)

This checks the ordering of all known sections (core sections + known
custom sections). This also adds section ID DataCount section that will
be scheduled to be added in near future.

Reviewers: sbc100

Subscribers: dschuff, mgorny, jgravelle-google, sunfish, llvm-commits

Differential Revision: https://reviews.llvm.org/D54924

7 files changed, 169 insertions, 16 deletions

diff --git a/llvm/include/llvm/BinaryFormat/Wasm.h b/llvm/include/llvm/BinaryFormat/Wasm.h
index 9b8ca28de66..ca752374e15 100644
--- a/llvm/include/llvm/BinaryFormat/Wasm.h
+++ b/llvm/include/llvm/BinaryFormat/Wasm.h
@@ -188,19 +188,20 @@ struct WasmLinkingData {
 };
 
 enum : unsigned {
-  WASM_SEC_CUSTOM = 0,   // Custom / User-defined section
-  WASM_SEC_TYPE = 1,     // Function signature declarations
-  WASM_SEC_IMPORT = 2,   // Import declarations
-  WASM_SEC_FUNCTION = 3, // Function declarations
-  WASM_SEC_TABLE = 4,    // Indirect function table and other tables
-  WASM_SEC_MEMORY = 5,   // Memory attributes
-  WASM_SEC_GLOBAL = 6,   // Global declarations
-  WASM_SEC_EXPORT = 7,   // Exports
-  WASM_SEC_START = 8,    // Start function declaration
-  WASM_SEC_ELEM = 9,     // Elements section
-  WASM_SEC_CODE = 10,    // Function bodies (code)
-  WASM_SEC_DATA = 11,    // Data segments
-  WASM_SEC_EVENT = 13    // Event declarations
+  WASM_SEC_CUSTOM = 0,     // Custom / User-defined section
+  WASM_SEC_TYPE = 1,       // Function signature declarations
+  WASM_SEC_IMPORT = 2,     // Import declarations
+  WASM_SEC_FUNCTION = 3,   // Function declarations
+  WASM_SEC_TABLE = 4,      // Indirect function table and other tables
+  WASM_SEC_MEMORY = 5,     // Memory attributes
+  WASM_SEC_GLOBAL = 6,     // Global declarations
+  WASM_SEC_EXPORT = 7,     // Exports
+  WASM_SEC_START = 8,      // Start function declaration
+  WASM_SEC_ELEM = 9,       // Elements section
+  WASM_SEC_CODE = 10,      // Function bodies (code)
+  WASM_SEC_DATA = 11,      // Data segments
+  WASM_SEC_DATACOUNT = 12, // Data segment count
+  WASM_SEC_EVENT = 13      // Event declarations
 };
 
 // Type immediate encodings used in various contexts.
diff --git a/llvm/include/llvm/Object/Wasm.h b/llvm/include/llvm/Object/Wasm.h
index 2edb93873da..ed857652a04 100644
--- a/llvm/include/llvm/Object/Wasm.h
+++ b/llvm/include/llvm/Object/Wasm.h
@@ -283,6 +283,49 @@ private:
   uint32_t EventSection = 0;
 };
 
+class WasmSectionOrderChecker {
+public:
+  // We define orders for all core wasm sections and known custom sections.
+  enum : int {
+    // Core sections
+    // The order of standard sections is precisely given by the spec.
+    WASM_SEC_ORDER_TYPE = 1,
+    WASM_SEC_ORDER_IMPORT = 2,
+    WASM_SEC_ORDER_FUNCTION = 3,
+    WASM_SEC_ORDER_TABLE = 4,
+    WASM_SEC_ORDER_MEMORY = 5,
+    WASM_SEC_ORDER_GLOBAL = 6,
+    WASM_SEC_ORDER_EVENT = 7,
+    WASM_SEC_ORDER_EXPORT = 8,
+    WASM_SEC_ORDER_START = 9,
+    WASM_SEC_ORDER_ELEM = 10,
+    WASM_SEC_ORDER_DATACOUNT = 11,
+    WASM_SEC_ORDER_CODE = 12,
+    WASM_SEC_ORDER_DATA = 13,
+
+    // Custom sections
+    // "dylink" should be the very first section in the module
+    WASM_SEC_ORDER_DYLINK = 0,
+    // "linking" section requires DATA section in order to validate data symbols
+    WASM_SEC_ORDER_LINKING = 100,
+    // Must come after "linking" section in order to validate reloc indexes.
+    WASM_SEC_ORDER_RELOC = 101,
+    // "name" section must appear after DATA. Comes after "linking" to allow
+    // symbol table to set default function name.
+    WASM_SEC_ORDER_NAME = 102,
+    // "producers" section must appear after "name" section.
+    WASM_SEC_ORDER_PRODUCERS = 103
+  };
+
+  bool isValidSectionOrder(unsigned ID, StringRef CustomSectionName = "");
+
+private:
+  int LastOrder = -1; // Lastly seen known section's order
+
+  // Returns -1 for unknown sections.
+  int getSectionOrder(unsigned ID, StringRef CustomSectionName = "");
+};
+
 } // end namespace object
 
 inline raw_ostream &operator<<(raw_ostream &OS, const object::WasmSymbol &Sym) {
diff --git a/llvm/lib/Object/WasmObjectFile.cpp b/llvm/lib/Object/WasmObjectFile.cpp
index 1a687d94d7f..34e4d192e59 100644
--- a/llvm/lib/Object/WasmObjectFile.cpp
+++ b/llvm/lib/Object/WasmObjectFile.cpp
@@ -24,6 +24,7 @@
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/LEB128.h"
+#include "llvm/Support/ScopedPrinter.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
@@ -207,8 +208,8 @@ static wasm::WasmTable readTable(WasmObjectFile::ReadContext &Ctx) {
   return Table;
 }
 
-static Error readSection(WasmSection &Section,
-                         WasmObjectFile::ReadContext &Ctx) {
+static Error readSection(WasmSection &Section, WasmObjectFile::ReadContext &Ctx,
+                         WasmSectionOrderChecker &Checker) {
   Section.Offset = Ctx.Ptr - Ctx.Start;
   Section.Type = readUint8(Ctx);
   LLVM_DEBUG(dbgs() << "readSection type=" << Section.Type << "\n");
@@ -231,6 +232,13 @@ static Error readSection(WasmSection &Section,
     Ctx.Ptr += SectionNameSize;
     Size -= SectionNameSize;
   }
+
+  if (!Checker.isValidSectionOrder(Section.Type, Section.Name)) {
+    return make_error<StringError>("Out of order section type: " +
+                                       llvm::to_string(Section.Type),
+                                   object_error::parse_failed);
+  }
+
   Section.Content = ArrayRef<uint8_t>(Ctx.Ptr, Size);
   Ctx.Ptr += Size;
   return Error::success();
@@ -265,8 +273,9 @@ WasmObjectFile::WasmObjectFile(MemoryBufferRef Buffer, Error &Err)
   }
 
   WasmSection Sec;
+  WasmSectionOrderChecker Checker;
   while (Ctx.Ptr < Ctx.End) {
-    if ((Err = readSection(Sec, Ctx)))
+    if ((Err = readSection(Sec, Ctx, Checker)))
       return;
     if ((Err = parseSection(Sec)))
       return;
@@ -1433,3 +1442,58 @@ WasmObjectFile::getWasmRelocation(DataRefImpl Ref) const {
   assert(Ref.d.b < Sec.Relocations.size());
   return Sec.Relocations[Ref.d.b];
 }
+
+int WasmSectionOrderChecker::getSectionOrder(unsigned ID,
+                                             StringRef CustomSectionName) {
+  switch (ID) {
+  case wasm::WASM_SEC_CUSTOM:
+    return StringSwitch<unsigned>(CustomSectionName)
+        .Case("dylink", WASM_SEC_ORDER_DYLINK)
+        .Case("linking", WASM_SEC_ORDER_LINKING)
+        .StartsWith("reloc.", WASM_SEC_ORDER_RELOC)
+        .Case("name", WASM_SEC_ORDER_NAME)
+        .Case("producers", WASM_SEC_ORDER_PRODUCERS)
+        .Default(-1);
+  case wasm::WASM_SEC_TYPE:
+    return WASM_SEC_ORDER_TYPE;
+  case wasm::WASM_SEC_IMPORT:
+    return WASM_SEC_ORDER_IMPORT;
+  case wasm::WASM_SEC_FUNCTION:
+    return WASM_SEC_ORDER_FUNCTION;
+  case wasm::WASM_SEC_TABLE:
+    return WASM_SEC_ORDER_TABLE;
+  case wasm::WASM_SEC_MEMORY:
+    return WASM_SEC_ORDER_MEMORY;
+  case wasm::WASM_SEC_GLOBAL:
+    return WASM_SEC_ORDER_GLOBAL;
+  case wasm::WASM_SEC_EXPORT:
+    return WASM_SEC_ORDER_EXPORT;
+  case wasm::WASM_SEC_START:
+    return WASM_SEC_ORDER_START;
+  case wasm::WASM_SEC_ELEM:
+    return WASM_SEC_ORDER_ELEM;
+  case wasm::WASM_SEC_CODE:
+    return WASM_SEC_ORDER_CODE;
+  case wasm::WASM_SEC_DATA:
+    return WASM_SEC_ORDER_DATA;
+  case wasm::WASM_SEC_DATACOUNT:
+    return WASM_SEC_ORDER_DATACOUNT;
+  case wasm::WASM_SEC_EVENT:
+    return WASM_SEC_ORDER_EVENT;
+  default:
+    llvm_unreachable("invalid section");
+  }
+}
+
+bool WasmSectionOrderChecker::isValidSectionOrder(unsigned ID,
+                                                  StringRef CustomSectionName) {
+  int Order = getSectionOrder(ID, CustomSectionName);
+  if (Order == -1) // Skip unknown sections
+    return true;
+  // There can be multiple "reloc." sections. Otherwise there shouldn't be any
+  // duplicate section orders.
+  bool IsValid = (LastOrder == Order && Order == WASM_SEC_ORDER_RELOC) ||
+                 LastOrder < Order;
+  LastOrder = Order;
+  return IsValid;
+}
diff --git a/llvm/test/Object/Inputs/WASM/invalid-section-order.wasm b/llvm/test/Object/Inputs/WASM/invalid-section-order.wasm
new file mode 100644
index 00000000000..f2726ff176f
--- /dev/null
+++ b/llvm/test/Object/Inputs/WASM/invalid-section-order.wasm
diff --git a/llvm/test/Object/wasm-invalid-section-order.test b/llvm/test/Object/wasm-invalid-section-order.test
new file mode 100644
index 00000000000..bb008ffac49
--- /dev/null
+++ b/llvm/test/Object/wasm-invalid-section-order.test
@@ -0,0 +1,16 @@
+# RUN: not obj2yaml %p/Inputs/WASM/invalid-section-order.wasm 2>&1 | FileCheck %s
+# CHECK: {{.*}}: Out of order section type: 10
+
+# Inputs/WASM/invalid-section-order.wasm is generated from this ll file, by
+# modifying WasmObjectWriter to incorrectly write the data section before the
+# code section.
+#
+# target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+# target triple = "wasm32-unknown-unknown"
+#
+# @data = global i32 0, align 4
+#
+# define void @foo() {
+# entry:
+#   ret void
+# }
diff --git a/llvm/test/ObjectYAML/wasm/invalid_section_order.yaml b/llvm/test/ObjectYAML/wasm/invalid_section_order.yaml
new file mode 100644
index 00000000000..52ad1361962
--- /dev/null
+++ b/llvm/test/ObjectYAML/wasm/invalid_section_order.yaml
@@ -0,0 +1,20 @@
+# RUN: not yaml2obj %s -o /dev/null 2>&1 | FileCheck %s
+
+--- !WASM
+FileHeader:
+  Version:         0x00000001
+Sections:
+  - Type:            TYPE
+    Signatures:
+      - Index:           0
+        ReturnType:      NORESULT
+        ParamTypes:      []
+  - Type:            CODE
+    Functions:
+      - Index:           0
+        Locals:          []
+        Body:            0B
+  # CHECK: Out of order section type: 3
+  - Type:            FUNCTION
+    FunctionTypes:   [ 0 ]
+...
diff --git a/llvm/tools/yaml2obj/yaml2wasm.cpp b/llvm/tools/yaml2obj/yaml2wasm.cpp
index 7fd0bb06934..e8d998d9914 100644
--- a/llvm/tools/yaml2obj/yaml2wasm.cpp
+++ b/llvm/tools/yaml2obj/yaml2wasm.cpp
@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 //
 
+#include "llvm/Object/Wasm.h"
 #include "llvm/ObjectYAML/ObjectYAML.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/LEB128.h"
@@ -516,7 +517,15 @@ int WasmWriter::writeWasm(raw_ostream &OS) {
   writeUint32(OS, Obj.Header.Version);
 
   // Write each section
+  llvm::object::WasmSectionOrderChecker Checker;
   for (const std::unique_ptr<WasmYAML::Section> &Sec : Obj.Sections) {
+    StringRef SecName = "";
+    if (auto S = dyn_cast<WasmYAML::CustomSection>(Sec.get()))
+      SecName = S->Name;
+    if (!Checker.isValidSectionOrder(Sec->Type, SecName)) {
+      errs() << "Out of order section type: " << Sec->Type << "\n";
+      return 1;
+    }
     encodeULEB128(Sec->Type, OS);
     std::string OutString;
     raw_string_ostream StringStream(OutString);