[ELF] Support INSERT [AFTER|BEFORE] for orphan sections

D43468+D44380 added INSERT [AFTER|BEFORE] for non-orphan sections. This patch
makes INSERT work for orphan sections as well.

`SECTIONS {...} INSERT [AFTER|BEFORE] .foo` does not set `hasSectionCommands`, so the result
will be similar to a regular link without a linker script. The differences when `hasSectionCommands` is set include:

* image base is different
* -z noseparate-code/-z noseparate-loadable-segments are unavailable
* some special symbols such as `_end _etext _edata` are not defined

The behavior is similar to GNU ld:
INSERT is not considered an external linker script.

This feature makes the section layout more flexible. It can be used to:

* Place .nv_fatbin before other readonly SHT_PROGBITS sections to mitigate relocation overflows.
* Disturb the layout to expose address sensitive application bugs.

Reviewed By: grimar

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

9 files changed, 160 insertions, 87 deletions

diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
index 60f3e88..b31da37 100644
--- a/lld/ELF/Driver.cpp
+++ b/lld/ELF/Driver.cpp
@@ -1865,10 +1865,6 @@ template <class ELFT> void LinkerDriver::link(opt::InputArgList &args) {
   if (errorCount())
     return;
 
-  // Now when we read all script files, we want to finalize order of linker
-  // script commands, which can be not yet final because of INSERT commands.
-  script->processInsertCommands();
-
   // We want to declare linker script's symbols early,
   // so that we can version them.
   // They also might be exported if referenced by DSOs.
diff --git a/lld/ELF/LinkerScript.cpp b/lld/ELF/LinkerScript.cpp
index edb5cac..9643da2 100644
--- a/lld/ELF/LinkerScript.cpp
+++ b/lld/ELF/LinkerScript.cpp
@@ -246,32 +246,30 @@ getChangedSymbolAssignment(const SymbolAssignmentMap &oldValues) {
   return changed;
 }
 
-// This method is used to handle INSERT AFTER statement. Here we rebuild
-// the list of script commands to mix sections inserted into.
+// Process INSERT [AFTER|BEFORE] commands. For each command, we move the
+// specified output section to the designated place.
 void LinkerScript::processInsertCommands() {
-  std::vector<BaseCommand *> v;
-  auto insert = [&](std::vector<BaseCommand *> &from) {
-    v.insert(v.end(), from.begin(), from.end());
-    from.clear();
-  };
-
-  for (BaseCommand *base : sectionCommands) {
-    if (auto *os = dyn_cast<OutputSection>(base)) {
-      insert(insertBeforeCommands[os->name]);
-      v.push_back(base);
-      insert(insertAfterCommands[os->name]);
+  for (const InsertCommand &cmd : insertCommands) {
+    // If cmd.os is empty, it may have been discarded by
+    // adjustSectionsBeforeSorting(). We do not handle such output sections.
+    auto from = llvm::find(sectionCommands, cmd.os);
+    if (from == sectionCommands.end())
       continue;
+    sectionCommands.erase(from);
+
+    auto insertPos = llvm::find_if(sectionCommands, [&cmd](BaseCommand *base) {
+      auto *to = dyn_cast<OutputSection>(base);
+      return to != nullptr && to->name == cmd.where;
+    });
+    if (insertPos == sectionCommands.end()) {
+      error("unable to insert " + cmd.os->name +
+            (cmd.isAfter ? " after " : " before ") + cmd.where);
+    } else {
+      if (cmd.isAfter)
+        ++insertPos;
+      sectionCommands.insert(insertPos, cmd.os);
     }
-    v.push_back(base);
   }
-
-  for (auto &cmds : {insertBeforeCommands, insertAfterCommands})
-    for (const std::pair<StringRef, std::vector<BaseCommand *>> &p : cmds)
-      if (!p.second.empty())
-        error("unable to INSERT AFTER/BEFORE " + p.first +
-              ": section not defined");
-
-  sectionCommands = std::move(v);
 }
 
 // Symbols defined in script should not be inlined by LTO. At the same time
diff --git a/lld/ELF/LinkerScript.h b/lld/ELF/LinkerScript.h
index f020552..848e3e4 100644
--- a/lld/ELF/LinkerScript.h
+++ b/lld/ELF/LinkerScript.h
@@ -208,6 +208,12 @@ struct ByteCommand : BaseCommand {
   unsigned size;
 };
 
+struct InsertCommand {
+  OutputSection *os;
+  bool isAfter;
+  StringRef where;
+};
+
 struct PhdrsCommand {
   StringRef name;
   unsigned type = llvm::ELF::PT_NULL;
@@ -311,10 +317,9 @@ public:
   // A list of symbols referenced by the script.
   std::vector<llvm::StringRef> referencedSymbols;
 
-  // Used to implement INSERT [AFTER|BEFORE]. Contains commands that need
-  // to be inserted into SECTIONS commands list.
-  llvm::DenseMap<StringRef, std::vector<BaseCommand *>> insertAfterCommands;
-  llvm::DenseMap<StringRef, std::vector<BaseCommand *>> insertBeforeCommands;
+  // Used to implement INSERT [AFTER|BEFORE]. Contains output sections that need
+  // to be reordered.
+  std::vector<InsertCommand> insertCommands;
 };
 
 extern LinkerScript *script;
diff --git a/lld/ELF/ScriptParser.cpp b/lld/ELF/ScriptParser.cpp
index 635d281..8186420 100644
--- a/lld/ELF/ScriptParser.cpp
+++ b/lld/ELF/ScriptParser.cpp
@@ -523,13 +523,6 @@ std::vector<BaseCommand *> ScriptParser::readOverlay() {
 }
 
 void ScriptParser::readSections() {
-  script->hasSectionsCommand = true;
-
-  // -no-rosegment is used to avoid placing read only non-executable sections in
-  // their own segment. We do the same if SECTIONS command is present in linker
-  // script. See comment for computeFlags().
-  config->singleRoRx = true;
-
   expect("{");
   std::vector<BaseCommand *> v;
   while (!errorCount() && !consume("}")) {
@@ -548,22 +541,29 @@ void ScriptParser::readSections() {
     else
       v.push_back(readOutputSectionDescription(tok));
   }
+  script->sectionCommands.insert(script->sectionCommands.end(), v.begin(),
+                                 v.end());
 
-  if (!atEOF() && consume("INSERT")) {
-    std::vector<BaseCommand *> *dest = nullptr;
-    if (consume("AFTER"))
-      dest = &script->insertAfterCommands[next()];
-    else if (consume("BEFORE"))
-      dest = &script->insertBeforeCommands[next()];
-    else
-      setError("expected AFTER/BEFORE, but got '" + next() + "'");
-    if (dest)
-      dest->insert(dest->end(), v.begin(), v.end());
+  if (atEOF() || !consume("INSERT")) {
+    // --no-rosegment is used to avoid placing read only non-executable sections
+    // in their own segment. We do the same if SECTIONS command is present in
+    // linker script. See comment for computeFlags().
+    // TODO This rule will be dropped in the future.
+    config->singleRoRx = true;
+
+    script->hasSectionsCommand = true;
     return;
   }
 
-  script->sectionCommands.insert(script->sectionCommands.end(), v.begin(),
-                                 v.end());
+  bool isAfter = false;
+  if (consume("AFTER"))
+    isAfter = true;
+  else if (!consume("BEFORE"))
+    setError("expected AFTER/BEFORE, but got '" + next() + "'");
+  StringRef where = next();
+  for (BaseCommand *cmd : v)
+    if (auto *os = dyn_cast<OutputSection>(cmd))
+      script->insertCommands.push_back({os, isAfter, where});
 }
 
 void ScriptParser::readTarget() {
diff --git a/lld/ELF/Writer.cpp b/lld/ELF/Writer.cpp
index 7c2a3fe..e3e39ac 100644
--- a/lld/ELF/Writer.cpp
+++ b/lld/ELF/Writer.cpp
@@ -1420,9 +1420,15 @@ template <class ELFT> void Writer<ELFT>::sortSections() {
         llvm::find_if(script->sectionCommands, isSection),
         llvm::find_if(llvm::reverse(script->sectionCommands), isSection).base(),
         compareSections);
+
+    // Process INSERT commands. From this point onwards the order of
+    // script->sectionCommands is fixed.
+    script->processInsertCommands();
     return;
   }
 
+  script->processInsertCommands();
+
   // Orphan sections are sections present in the input files which are
   // not explicitly placed into the output file by the linker script.
   //
diff --git a/lld/test/ELF/linkerscript/insert-after.test b/lld/test/ELF/linkerscript/insert-after.test
index 005c150..65751c7 100644
--- a/lld/test/ELF/linkerscript/insert-after.test
+++ b/lld/test/ELF/linkerscript/insert-after.test
@@ -5,25 +5,39 @@
 ## we check that can use INSERT AFTER to insert sections .foo.data
 ## and .foo.text at the right places.
 
-SECTIONS {
-  .foo.data : { *(.foo.data) }
-} INSERT AFTER .data;
+# RUN: ld.lld %t1.o -o %t1 --script %p/Inputs/insert-after.script --script %s
+# RUN: llvm-readelf -S -l %t1 | FileCheck %s
+# CHECK:      Name      Type     Address          Off
+# CHECK-NEXT:           NULL     0000000000000000 000000
+# CHECK-NEXT: .text     PROGBITS 0000000000000000 001000
+# CHECK-NEXT: .foo.text PROGBITS 0000000000000008 001008
+# CHECK-NEXT: .data     PROGBITS 0000000000000010 001010
+# CHECK-NEXT: .foo.data PROGBITS 0000000000000018 001018
+# CHECK:      Type
+# CHECK-NEXT: LOAD {{.*}} R E
+# CHECK-NEXT: LOAD {{.*}} RW
+# CHECK-NEXT: GNU_STACK {{.*}} RW
 
-SECTIONS {
-  .foo.text : { *(.foo.text) }
-} INSERT AFTER .text;
+## There is no main linker script. INSERT AFTER just reorders output sections,
+## without making more layout changes. Address/offset assignments are different
+## with a main linker script.
 
-# RUN: ld.lld %t1.o -o %t1 --script %p/Inputs/insert-after.script --script %s
-# RUN: llvm-objdump -section-headers %t1 | FileCheck %s
-# CHECK:      Sections:
-# CHECK-NEXT: Idx Name          Size     VMA              Type
-# CHECK-NEXT:   0               00000000 0000000000000000
-# CHECK-NEXT:   1 .text         00000008 0000000000000000 TEXT
-# CHECK-NEXT:   2 .foo.text     00000008 0000000000000008 TEXT
-# CHECK-NEXT:   3 .data         00000008 0000000000000010 DATA
-# CHECK-NEXT:   4 .foo.data     00000008 0000000000000018 DATA
+# RUN: ld.lld --script %s %t1.o -o %t2
+# RUN: llvm-readelf -S -l %t2 | FileCheck --check-prefix=CHECK2 %s
+# CHECK2:      Name      Type     Address          Off
+# CHECK2-NEXT:           NULL     0000000000000000 000000
+# CHECK2-NEXT: .text     PROGBITS 0000000000201158 000158
+# CHECK2-NEXT: .foo.text PROGBITS 0000000000201160 000160
+# CHECK2-NEXT: .data     PROGBITS 0000000000202168 000168
+# CHECK2-NEXT: .foo.data PROGBITS 0000000000202170 000170
+# CHECK2:      Type
+# CHECK2-NEXT: PHDR {{.*}} R
+# CHECK2-NEXT: LOAD {{.*}} R
+# CHECK2-NEXT: LOAD {{.*}} R E
+# CHECK2-NEXT: LOAD {{.*}} RW
+# CHECK2-NEXT: GNU_STACK {{.*}} RW
+
+SECTIONS { .foo.data : { *(.foo.data) } } INSERT AFTER .data;
 
-# RUN: not ld.lld %t1.o -o %t1 --script %s 2>&1 \
-# RUN:   | FileCheck %s --check-prefix=ERR
-# ERR-DAG: error: unable to INSERT AFTER/BEFORE .text: section not defined
-# ERR-DAG: error: unable to INSERT AFTER/BEFORE .data: section not defined
+## The input section .foo.text is an orphan. It will be placed in .foo.text
+SECTIONS { .foo.text : {} } INSERT AFTER .text;
diff --git a/lld/test/ELF/linkerscript/insert-before.test b/lld/test/ELF/linkerscript/insert-before.test
index 24fdd69..8fe912f 100644
--- a/lld/test/ELF/linkerscript/insert-before.test
+++ b/lld/test/ELF/linkerscript/insert-before.test
@@ -5,25 +5,38 @@
 ## we check that can use INSERT BEFORE to insert sections .foo.data
 ## and .foo.text at the right places.
 
-SECTIONS {
-  .foo.data : { *(.foo.data) }
-} INSERT BEFORE .data;
+# RUN: ld.lld %t1.o -o %t1 --script %p/Inputs/insert-after.script --script %s
+# RUN: llvm-readelf -S -l %t1 | FileCheck %s
+# CHECK:      Name      Type     Address          Off
+# CHECK-NEXT:           NULL     0000000000000000 000000
+# CHECK-NEXT: .foo.text PROGBITS 0000000000000000 001000
+# CHECK-NEXT: .text     PROGBITS 0000000000000008 001008
+# CHECK-NEXT: .foo.data PROGBITS 0000000000000010 001010
+# CHECK-NEXT: .data     PROGBITS 0000000000000018 001018
+# CHECK:      Type
+# CHECK-NEXT: LOAD {{.*}} R E
+# CHECK-NEXT: LOAD {{.*}} RW
+# CHECK-NEXT: GNU_STACK {{.*}} RW
 
-SECTIONS {
-  .foo.text : { *(.foo.text) }
-} INSERT BEFORE .text;
+## There is no main linker script. INSERT BEFORE just reorders output sections,
+## without making more layout changes. Address/offset assignments are different
+## with a main linker script.
 
-# RUN: ld.lld %t1.o -o %t1 --script %p/Inputs/insert-after.script --script %s
-# RUN: llvm-objdump -section-headers %t1 | FileCheck %s
-# CHECK:      Sections:
-# CHECK-NEXT: Idx Name          Size     VMA         Type
-# CHECK-NEXT:   0               00000000 0000000000000000
-# CHECK-NEXT:   1 .foo.text     00000008 0000000000000000 TEXT
-# CHECK-NEXT:   2 .text         00000008 0000000000000008 TEXT
-# CHECK-NEXT:   3 .foo.data     00000008 0000000000000010 DATA
-# CHECK-NEXT:   4 .data         00000008 0000000000000018 DATA
+# RUN: ld.lld --script %s %t1.o -o %t2
+# RUN: llvm-readelf -S -l %t2 | FileCheck --check-prefix=CHECK2 %s
+# CHECK2:      Name      Type     Address          Off
+# CHECK2-NEXT:           NULL     0000000000000000 000000
+# CHECK2-NEXT: .foo.text PROGBITS 0000000000201158 000158
+# CHECK2-NEXT: .text     PROGBITS 0000000000201160 000160
+# CHECK2-NEXT: .foo.data PROGBITS 0000000000202168 000168
+# CHECK2-NEXT: .data     PROGBITS 0000000000202170 000170
+# CHECK2:      Type
+# CHECK2-NEXT: PHDR {{.*}} R
+# CHECK2-NEXT: LOAD {{.*}} R
+# CHECK2-NEXT: LOAD {{.*}} R E
+# CHECK2-NEXT: LOAD {{.*}} RW
+
+SECTIONS { .foo.data : { *(.foo.data) } } INSERT BEFORE .data;
 
-# RUN: not ld.lld %t1.o -o %t1 --script %s 2>&1 \
-# RUN:   | FileCheck %s --check-prefix=ERR
-# ERR-DAG: error: unable to INSERT AFTER/BEFORE .text: section not defined
-# ERR-DAG: error: unable to INSERT AFTER/BEFORE .data: section not defined
+## The input section .foo.text is an orphan. It will be placed in .foo.text
+SECTIONS { .foo.text : {} } INSERT BEFORE .text;
diff --git a/lld/test/ELF/linkerscript/insert-duplicate.test b/lld/test/ELF/linkerscript/insert-duplicate.test
new file mode 100644
index 0000000..8fa7e2c
--- /dev/null
+++ b/lld/test/ELF/linkerscript/insert-duplicate.test
@@ -0,0 +1,32 @@
+# REQUIRES: x86
+## Test that we can handle cases where an output section is specified by multiple
+## INSERT commands. Each output section description creates a new instance.
+## A redundant description matches no input sections and thus is a no-op.
+
+# RUN: llvm-mc -filetype=obj -triple=x86_64 %p/Inputs/insert-after.s -o %t.o
+# RUN: ld.lld -T %s %t.o -o %t
+# RUN: llvm-readelf -S -l %t | FileCheck %s
+
+# CHECK:      Name      Type     Address          Off
+# CHECK-NEXT:           NULL     0000000000000000 000000
+# CHECK-NEXT: .text     PROGBITS 00000000002011c8 0001c8
+# CHECK-NEXT: .foo.data PROGBITS 00000000002021d0 0001d0
+# CHECK-NEXT: .foo.text PROGBITS 00000000002031d8 0001d8
+# CHECK:      Type
+# CHECK-NEXT: PHDR {{.*}} R
+# CHECK-NEXT: LOAD {{.*}} R
+# CHECK-NEXT: LOAD {{.*}} R E
+# CHECK-NEXT: LOAD {{.*}} RW
+# CHECK-NEXT: LOAD {{.*}} R E
+# CHECK-NEXT: LOAD {{.*}} RW
+# CHECK-NEXT: GNU_STACK {{.*}} RW
+
+## First, move .foo.data after .foo.text
+SECTIONS { .foo.data : { *(.foo.data) } } INSERT AFTER .foo.text;
+
+## Next, move .foo.text after .foo.data
+SECTIONS { .foo.text : { *(.foo.text) } } INSERT AFTER .foo.data;
+
+## No-op. The .foo.data output section is a different instance and matches no
+## input sections.
+SECTIONS { .foo.data : { *(.foo.data) } } INSERT AFTER .foo.text;
diff --git a/lld/test/ELF/linkerscript/insert-not-exist.test b/lld/test/ELF/linkerscript/insert-not-exist.test
new file mode 100644
index 0000000..1f95070
--- /dev/null
+++ b/lld/test/ELF/linkerscript/insert-not-exist.test
@@ -0,0 +1,9 @@
+# REQUIRES: x86
+# RUN: llvm-mc -filetype=obj -triple=x86_64 %p/Inputs/insert-after.s -o %t.o
+# RUN: not ld.lld -T %s %t.o 2>&1 | FileCheck %s
+
+# CHECK: error: unable to insert .foo.data after .not_exist
+# CHECK: error: unable to insert .foo.text before .not_exist
+
+SECTIONS { .foo.data : {} } INSERT AFTER .not_exist;
+SECTIONS { .foo.text : {} } INSERT BEFORE .not_exist;