[Core] Fix heap overflow in LayoutPass.

Found this with asan. Code assumes that find doesn't return end, thus if
both atoms didn't have followon roots it would still compare their positions.

llvm-svn: 191865

5 files changed, 14 insertions, 13 deletions

diff --git a/lld/lib/Passes/LayoutPass.cpp b/lld/lib/Passes/LayoutPass.cpp
index b116496e036..c4c8fb9c33e 100644
--- a/lld/lib/Passes/LayoutPass.cpp
+++ b/lld/lib/Passes/LayoutPass.cpp
@@ -56,10 +56,12 @@ bool LayoutPass::CompareAtoms::operator()(const DefinedAtom *left,
 
   // Sort atoms by their ordinal overrides only if they fall in the same
   // chain.
-  const DefinedAtom *leftAtom = _layout._followOnRoots.find(left)->second;
-  const DefinedAtom *rightAtom = _layout._followOnRoots.find(right)->second;
+  auto leftAtom = _layout._followOnRoots.find(left);
+  auto rightAtom = _layout._followOnRoots.find(right);
 
-  if (leftAtom == rightAtom) {
+  if (leftAtom != _layout._followOnRoots.end() &&
+      rightAtom != _layout._followOnRoots.end() &&
+      leftAtom->second == rightAtom->second) {
     if ((lPos != end) && (rPos != end)) {
       return lPos->second < rPos->second;
     }
diff --git a/lld/test/elf/X86_64/largebss.test b/lld/test/elf/X86_64/largebss.test
index 985f2185cc0..8aea7878eee 100644
--- a/lld/test/elf/X86_64/largebss.test
+++ b/lld/test/elf/X86_64/largebss.test
@@ -5,17 +5,16 @@
 
 RUN: lld -flavor gnu -target x86_64 %p/Inputs/largebss.o --output-filetype=yaml --noinhibit-exec | FileCheck %s
 
-
-CHECK:  - name:            largecommon
+CHECK:  - name:            largebss
 CHECK:    scope:           global
 CHECK:    type:            zero-fill
 CHECK:    size:            4000
-CHECK:    merge:           as-tentative
 CHECK:    section-name:    .bss
-CHECK:  - name:            largebss
+CHECK:  - name:            largecommon
 CHECK:    scope:           global
 CHECK:    type:            zero-fill
 CHECK:    size:            4000
+CHECK:    merge:           as-tentative
 CHECK:    section-name:    .bss
 CHECK:  - name:            largetbss
 CHECK:    scope:           global
diff --git a/lld/test/elf/phdr.test b/lld/test/elf/phdr.test
index b3dfd10698d..e36d8ee521b 100644
--- a/lld/test/elf/phdr.test
+++ b/lld/test/elf/phdr.test
@@ -63,7 +63,7 @@ I386-NEXT:     Offset: 0x4000
 I386-NEXT:     VirtualAddress: 0x4000
 I386-NEXT:     PhysicalAddress: 0x4000
 I386-NEXT:     FileSize: 4
-I386-NEXT:     MemSize: 16392
+I386-NEXT:     MemSize: 16389
 I386-NEXT:     Flags [ (0x6)
 I386-NEXT:       PF_R (0x4)
 I386-NEXT:       PF_W (0x2)
diff --git a/lld/test/elf/quickdata.test b/lld/test/elf/quickdata.test
index d2506263350..9a3d966af0e 100644
--- a/lld/test/elf/quickdata.test
+++ b/lld/test/elf/quickdata.test
@@ -4,11 +4,11 @@ RUN: --noinhibit-exec | FileCheck %s -check-prefix hexagon
 hexagon:  - name:            init
 hexagon:    scope:           global
 hexagon:    type:            quick-data
+hexagon:  - name:            bss1
+hexagon:    scope:           global
+hexagon:    type:            zero-fill-quick
 hexagon:  - name:            ac1
 hexagon:    scope:           global
 hexagon:    type:            zero-fill-quick
 hexagon:    size:            1
 hexagon:    merge:           as-tentative
-hexagon:  - name:            bss1
-hexagon:    scope:           global
-hexagon:    type:            zero-fill-quick
diff --git a/lld/test/elf/sections.test b/lld/test/elf/sections.test
index 8839aa6d4e8..a1953d80b94 100644
--- a/lld/test/elf/sections.test
+++ b/lld/test/elf/sections.test
@@ -10,7 +10,7 @@ OBJDUMP:  1 .text         0000000a 0000000000000074 TEXT DATA
 OBJDUMP:  2 .data         00000004 0000000000001000 DATA
 OBJDUMP:  3 .special      00000004 0000000000001004 DATA
 OBJDUMP:  4 .anotherspecial 00000004 0000000000001008 DATA
-OBJDUMP:  5 .bss          00000004 000000000000100c BSS
+OBJDUMP:  5 .bss          00000001 000000000000100c BSS
 OBJDUMP:  6 .shstrtab     {{[0-9a-f]+}} 0000000000000000
 OBJDUMP:  7 .symtab       {{[0-9a-f]+}} 0000000000000000
 OBJDUMP:  8 .strtab       {{[0-9a-f]+}} 0000000000000000
@@ -90,7 +90,7 @@ READOBJ:       SHF_ALLOC
 READOBJ:       SHF_WRITE
 READOBJ:     ]
 READOBJ:     Address: 0x100C
-READOBJ:     Size: 4
+READOBJ:     Size: 1
 READOBJ:   }
 READOBJ:   Section {
 READOBJ:     Index: 6