lld-link: Store comdat selection in SectionChunk, reject more invalid associated comdats

I need the comdat selection for PR40094. To keep the patch for that smaller,
I'm adding it here, and as a first application I'm using it to reject
associative comdats referring to earlier associative comdats. Depends on
D56929; together with that all associative comdats referring to other
associative comdats are now rejected.

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

llvm-svn: 352254

3 files changed, 32 insertions, 12 deletions

diff --git a/lld/COFF/Chunks.h b/lld/COFF/Chunks.h
index 0185da17b2b..60dbf1ff7a4 100644
--- a/lld/COFF/Chunks.h
+++ b/lld/COFF/Chunks.h
@@ -223,6 +223,9 @@ public:
   // The COMDAT leader symbol if this is a COMDAT chunk.
   DefinedRegular *Sym = nullptr;
 
+  // The COMDAT selection if this is a COMDAT chunk.
+  llvm::COFF::COMDATType Selection;
+
   ArrayRef<coff_relocation> Relocs;
 
   // Used by the garbage collector.
diff --git a/lld/COFF/InputFiles.cpp b/lld/COFF/InputFiles.cpp
index f291a9dbf83..4b87e02a6dd 100644
--- a/lld/COFF/InputFiles.cpp
+++ b/lld/COFF/InputFiles.cpp
@@ -226,10 +226,7 @@ void ObjFile::readAssociativeDefinition(COFFSymbolRef Sym,
                                         uint32_t ParentIndex) {
   SectionChunk *Parent = SparseChunks[ParentIndex];
 
-  if (Parent == PendingComdat) {
-    // This can happen if an associative comdat refers to another associative
-    // comdat that appears after it (invalid per COFF spec) or to a section
-    // without any symbols.
+  auto Diag = [&]() {
     StringRef Name, ParentName;
     COFFObj->getSymbolName(Sym, Name);
 
@@ -238,6 +235,13 @@ void ObjFile::readAssociativeDefinition(COFFSymbolRef Sym,
     COFFObj->getSectionName(ParentSec, ParentName);
     error(toString(this) + ": associative comdat " + Name +
           " has invalid reference to section " + ParentName);
+  };
+
+  if (Parent == PendingComdat) {
+    // This can happen if an associative comdat refers to another associative
+    // comdat that appears after it (invalid per COFF spec) or to a section
+    // without any symbols.
+    Diag();
     return;
   }
 
@@ -245,7 +249,14 @@ void ObjFile::readAssociativeDefinition(COFFSymbolRef Sym,
   // the section; otherwise mark it as discarded.
   int32_t SectionNumber = Sym.getSectionNumber();
   if (Parent) {
-    SparseChunks[SectionNumber] = readSection(SectionNumber, Def, "");
+    if (Parent->Selection == IMAGE_COMDAT_SELECT_ASSOCIATIVE) {
+      Diag();
+      return;
+    }
+
+    SectionChunk *C = readSection(SectionNumber, Def, "");
+    C->Selection = IMAGE_COMDAT_SELECT_ASSOCIATIVE;
+    SparseChunks[SectionNumber] = C;
     if (SparseChunks[SectionNumber])
       Parent->addAssociative(SparseChunks[SectionNumber]);
   } else {
diff --git a/lld/test/COFF/associative-comdat-order.test b/lld/test/COFF/associative-comdat-order.test
index 31105d51268..09b1b460215 100644
--- a/lld/test/COFF/associative-comdat-order.test
+++ b/lld/test/COFF/associative-comdat-order.test
@@ -1,10 +1,16 @@
-# RUN: yaml2obj < %s > %t.obj
-# RUN: not lld-link /include:symbol /dll /noentry /nodefaultlib %t.obj /out:%t.exe 2>&1 | FileCheck %s
-
 # Tests that an associative comdat being associated with another
 # associated comdat later in the file produces an error.
-# CHECK: lld-link: error: {{.*}}: associative comdat .text$ac1 has invalid reference to section .text$ac2
-# CHECK-NOT: lld-link: error:
+# RUN: sed -e s/ASSOC1/2/ -e s/ASSOC2/3/ %s | yaml2obj > %t.obj
+# RUN: not lld-link /include:symbol /dll /noentry /nodefaultlib %t.obj /out:%t.exe 2>&1 | FileCheck --check-prefix=FORWARD %s
+# FORWARD: lld-link: error: {{.*}}: associative comdat .text$ac1 has invalid reference to section .text$ac2
+# FORWARD-NOT: lld-link: error:
+
+# Tests that an associative comdat being associated with another
+# associated comdat earlier in the file produces an error.
+# RUN: sed -e s/ASSOC1/3/ -e s/ASSOC2/1/ %s | yaml2obj > %t.obj
+# RUN: not lld-link /include:symbol /dll /noentry /nodefaultlib %t.obj /out:%t.exe 2>&1 | FileCheck --check-prefix=BACKWARD %s
+# BACKWARD: lld-link: error: {{.*}}: associative comdat .text$ac2 has invalid reference to section .text$ac1
+# BACKWARD-NOT: lld-link: error:
 
 --- !COFF
 header:          
@@ -35,7 +41,7 @@ symbols:
       NumberOfRelocations: 0
       NumberOfLinenumbers: 0
       CheckSum:        3099354981
-      Number:          2
+      Number:          ASSOC1
       Selection:       IMAGE_COMDAT_SELECT_ASSOCIATIVE
   - Name:            '.text$ac2'
     Value:           0
@@ -48,7 +54,7 @@ symbols:
       NumberOfRelocations: 0
       NumberOfLinenumbers: 0
       CheckSum:        3099354981
-      Number:          3
+      Number:          ASSOC2
       Selection:       IMAGE_COMDAT_SELECT_ASSOCIATIVE
   - Name:            '.text$nm'
     Value:           0