[yaml2obj] - Add a way to override sh_flags section field.

Currently we have the `Flags` property that allows to
set flags for a section. The problem is that it does not
allow us to set an arbitrary value, because of bit fields
validation under the hood. An arbitrary values can be used
to test specific broken cases.

We probably do not want to relax the validation, so this
patch adds a `ShSize` property that allows to
override the `sh_size`. It is inline with others `Sh*` properties
we have already.

Differential revision: https://reviews.llvm.org/D71411

4 files changed, 117 insertions, 22 deletions

diff --git a/llvm/include/llvm/ObjectYAML/ELFYAML.h b/llvm/include/llvm/ObjectYAML/ELFYAML.h
index 7e0244ed08e..c9f14589906 100644
--- a/llvm/include/llvm/ObjectYAML/ELFYAML.h
+++ b/llvm/include/llvm/ObjectYAML/ELFYAML.h
@@ -192,6 +192,9 @@ struct Section : public Chunk {
   // This can be used to override the sh_size field. It does not affect the
   // content written.
   Optional<llvm::yaml::Hex64> ShSize;
+
+  // This can be used to override the sh_flags field.
+  Optional<llvm::yaml::Hex64> ShFlags;
 };
 
 // Fill is a block of data which is placed outside of sections. It is
diff --git a/llvm/lib/ObjectYAML/ELFEmitter.cpp b/llvm/lib/ObjectYAML/ELFEmitter.cpp
index d2e6fdfea00..69f5510fc7f 100644
--- a/llvm/lib/ObjectYAML/ELFEmitter.cpp
+++ b/llvm/lib/ObjectYAML/ELFEmitter.cpp
@@ -350,6 +350,20 @@ unsigned ELFState<ELFT>::toSymbolIndex(StringRef S, StringRef LocSec,
 }
 
 template <class ELFT>
+static void overrideFields(ELFYAML::Section *From, typename ELFT::Shdr &To) {
+  if (!From)
+    return;
+  if (From->ShFlags)
+    To.sh_flags = *From->ShFlags;
+  if (From->ShName)
+    To.sh_name = *From->ShName;
+  if (From->ShOffset)
+    To.sh_offset = *From->ShOffset;
+  if (From->ShSize)
+    To.sh_size = *From->ShSize;
+}
+
+template <class ELFT>
 bool ELFState<ELFT>::initImplicitHeader(ContiguousBlobAccumulator &CBA,
                                         Elf_Shdr &Header, StringRef SecName,
                                         ELFYAML::Section *YAMLSec) {
@@ -370,16 +384,8 @@ bool ELFState<ELFT>::initImplicitHeader(ContiguousBlobAccumulator &CBA,
   else
     return false;
 
-  // Override the fields if requested.
-  if (YAMLSec) {
-    if (YAMLSec->ShName)
-      Header.sh_name = *YAMLSec->ShName;
-    if (YAMLSec->ShOffset)
-      Header.sh_offset = *YAMLSec->ShOffset;
-    if (YAMLSec->ShSize)
-      Header.sh_size = *YAMLSec->ShSize;
-  }
-
+  // Override section fields if requested.
+  overrideFields<ELFT>(YAMLSec, Header);
   return true;
 }
 
@@ -484,15 +490,8 @@ void ELFState<ELFT>::initSectionHeaders(std::vector<Elf_Shdr> &SHeaders,
       llvm_unreachable("Unknown section type");
     }
 
-    // Override the fields if requested.
-    if (Sec) {
-      if (Sec->ShName)
-        SHeader.sh_name = *Sec->ShName;
-      if (Sec->ShOffset)
-        SHeader.sh_offset = *Sec->ShOffset;
-      if (Sec->ShSize)
-        SHeader.sh_size = *Sec->ShSize;
-    }
+    // Override section fields if requested.
+    overrideFields<ELFT>(Sec, SHeader);
   }
 }
 
diff --git a/llvm/lib/ObjectYAML/ELFYAML.cpp b/llvm/lib/ObjectYAML/ELFYAML.cpp
index c8de7a662fc..aa871af79b1 100644
--- a/llvm/lib/ObjectYAML/ELFYAML.cpp
+++ b/llvm/lib/ObjectYAML/ELFYAML.cpp
@@ -1017,10 +1017,12 @@ static void commonSectionMapping(IO &IO, ELFYAML::Section &Section) {
   // are producing YAML, because yaml2obj sets appropriate values for them
   // automatically when they are not explicitly defined.
   assert(!IO.outputting() ||
-         (!Section.ShOffset.hasValue() && !Section.ShSize.hasValue()));
+         (!Section.ShOffset.hasValue() && !Section.ShSize.hasValue() &&
+          !Section.ShName.hasValue() && !Section.ShFlags.hasValue()));
   IO.mapOptional("ShName", Section.ShName);
   IO.mapOptional("ShOffset", Section.ShOffset);
   IO.mapOptional("ShSize", Section.ShSize);
+  IO.mapOptional("ShFlags", Section.ShFlags);
 }
 
 static void sectionMapping(IO &IO, ELFYAML::DynamicSection &Section) {
@@ -1283,11 +1285,12 @@ void MappingTraits<std::unique_ptr<ELFYAML::Chunk>>::mapping(
 
 StringRef MappingTraits<std::unique_ptr<ELFYAML::Chunk>>::validate(
     IO &io, std::unique_ptr<ELFYAML::Chunk> &C) {
-  if (const auto *RawSection =
-          dyn_cast<ELFYAML::RawContentSection>(C.get())) {
+  if (const auto *RawSection = dyn_cast<ELFYAML::RawContentSection>(C.get())) {
     if (RawSection->Size && RawSection->Content &&
         (uint64_t)(*RawSection->Size) < RawSection->Content->binary_size())
       return "Section size must be greater than or equal to the content size";
+    if (RawSection->Flags && RawSection->ShFlags)
+      return "ShFlags and Flags cannot be used together";
     return {};
   }
 
diff --git a/llvm/test/tools/yaml2obj/ELF/override-shflags.yaml b/llvm/test/tools/yaml2obj/ELF/override-shflags.yaml
new file mode 100644
index 00000000000..653deb5ef70
--- /dev/null
+++ b/llvm/test/tools/yaml2obj/ELF/override-shflags.yaml
@@ -0,0 +1,90 @@
+## Check we are able to set custom sh_flags field
+## for different sections.
+
+# RUN: yaml2obj --docnum=1 %s -o %t1
+# RUN: llvm-readobj --sections %t1 | FileCheck %s --check-prefix=UPDATE-FLAGS
+
+# UPDATE-FLAGS:      Name: .symtab
+# UPDATE-FLAGS-NEXT: Type: SHT_SYMTAB
+# UPDATE-FLAGS-NEXT: Flags [ (0x1)
+
+# UPDATE-FLAGS:      Name: .strtab
+# UPDATE-FLAGS-NEXT: Type: SHT_STRTAB
+# UPDATE-FLAGS-NEXT: Flags [ (0x2)
+
+# UPDATE-FLAGS:      Name: .shstrtab
+# UPDATE-FLAGS-NEXT: Type: SHT_STRTAB
+# UPDATE-FLAGS-NEXT: Flags [ (0x3)
+
+# UPDATE-FLAGS:      Name: .dynsym
+# UPDATE-FLAGS-NEXT: Type: SHT_DYNSYM
+# UPDATE-FLAGS-NEXT: Flags [ (0x4)
+
+# UPDATE-FLAGS:      Name: .dynstr
+# UPDATE-FLAGS-NEXT: Type: SHT_STRTAB
+# UPDATE-FLAGS-NEXT: Flags [ (0x5)
+
+# UPDATE-FLAGS:      Name: .regular.os
+# UPDATE-FLAGS-NEXT: Type: SHT_PROGBITS
+# UPDATE-FLAGS-NEXT: Flags [ (0xAB00000)
+
+# UPDATE-FLAGS:      Name: .regular.proc
+# UPDATE-FLAGS-NEXT: Type: SHT_PROGBITS
+# UPDATE-FLAGS-NEXT: Flags [ (0xC0000000)
+
+# UPDATE-FLAGS:      Name: .regular.all
+# UPDATE-FLAGS-NEXT: Type: SHT_PROGBITS
+# UPDATE-FLAGS-NEXT: Flags [ (0xFFFFFFFF)
+
+--- !ELF
+FileHeader:
+  Class:   ELFCLASS64
+  Data:    ELFDATA2LSB
+  Type:    ET_REL
+  Machine: EM_X86_64
+Sections:
+## Check we can override flags for implicit sections.
+  - Name: .symtab
+    Type: SHT_SYMTAB
+    ShFlags: 0x000000001
+  - Name: .strtab
+    Type: SHT_STRTAB
+    ShFlags: 0x000000002
+  - Name: .shstrtab
+    Type: SHT_STRTAB
+    ShFlags: 0x000000003
+  - Name: .dynsym
+    Type: SHT_DYNSYM
+    ShFlags: 0x000000004
+  - Name: .dynstr
+    Type: SHT_STRTAB
+    ShFlags: 0x000000005
+## Check we can override flags for regular sections.
+  - Name: .regular.os
+    Type: SHT_PROGBITS
+    ShFlags: 0x0AB00000
+  - Name: .regular.proc
+    Type: SHT_PROGBITS
+    ShFlags: 0xC0000000
+  - Name: .regular.all
+    Type: SHT_PROGBITS
+    ShFlags: 0xFFFFFFFF
+
+## ShFlags is used when we want to work around the flag values validation.
+## It does not make sense to allow both ShFlags and Flags in the document.
+
+# RUN: not yaml2obj --docnum=2 %s 2>&1 | FileCheck %s --check-prefix=ERR
+
+# ERR: error: ShFlags and Flags cannot be used together
+
+--- !ELF
+FileHeader:
+  Class:   ELFCLASS64
+  Data:    ELFDATA2LSB
+  Type:    ET_REL
+  Machine: EM_X86_64
+Sections:
+  - Name: .symtab
+    Type: SHT_SYMTAB
+    Flags: [ SHF_ALLOC ]
+    ShFlags: 0x1