[lld] [LLD][COFF] Detect weak reference cycles. (PR #104463)

[Jacek Caban via llvm-commits]

https://github.com/cjacek created https://github.com/llvm/llvm-project/pull/104463

Instead of entering an infinite loop.

>From b5976e4dcbda0dd481b6f40e060ed7ad45d9934b Mon Sep 17 00:00:00 2001
From: Jacek Caban <jacek at codeweavers.com>
Date: Sun, 10 Sep 2023 22:22:12 +0200
Subject: [PATCH] [LLD][COFF] Detect weak reference cycles.

---
 lld/COFF/Symbols.cpp              |  8 +++++-
 lld/test/COFF/weak-ref-cycle.test | 42 +++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 lld/test/COFF/weak-ref-cycle.test

diff --git a/lld/COFF/Symbols.cpp b/lld/COFF/Symbols.cpp
index f4efcf2266cd6b..2bc5476ff7f00a 100644
--- a/lld/COFF/Symbols.cpp
+++ b/lld/COFF/Symbols.cpp
@@ -13,6 +13,7 @@
 #include "lld/Common/Memory.h"
 #include "lld/Common/Strings.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/Demangle/Demangle.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@@ -126,9 +127,14 @@ DefinedImportThunk::DefinedImportThunk(COFFLinkerContext &ctx, StringRef name,
 
 Defined *Undefined::getWeakAlias() {
   // A weak alias may be a weak alias to another symbol, so check recursively.
-  for (Symbol *a = weakAlias; a; a = cast<Undefined>(a)->weakAlias)
+  SmallSet<Symbol *, 4> weakChain;
+  for (Symbol *a = weakAlias; a; a = cast<Undefined>(a)->weakAlias) {
+    if (weakChain.contains(a))
+      break; // We have a cycle.
     if (auto *d = dyn_cast<Defined>(a))
       return d;
+    weakChain.insert(a);
+  }
   return nullptr;
 }
 
diff --git a/lld/test/COFF/weak-ref-cycle.test b/lld/test/COFF/weak-ref-cycle.test
new file mode 100644
index 00000000000000..71d640fadaad07
--- /dev/null
+++ b/lld/test/COFF/weak-ref-cycle.test
@@ -0,0 +1,42 @@
+REQUIRES: x86
+RUN: split-file %s %t.dir && cd %t.dir
+
+RUN: llvm-mc -filetype=obj -triple=x86_64-windows test.s -o test.obj
+RUN: llvm-mc -filetype=obj -triple=x86_64-windows weak2.s -o weak2.obj
+RUN: llvm-mc -filetype=obj -triple=x86_64-windows weak-def.s -o weak-def.obj
+
+RUN: not lld-link -machine:amd64 -dll -noentry -out:test.dll test.obj weak2.obj 2>&1 | FileCheck -check-prefix=ERR %s
+
+ERR:      error: undefined symbol: testsym
+ERR-NEXT: >>> referenced by test.obj
+ERR-EMPTY:
+ERR-NEXT: error: undefined symbol: weak1
+ERR-NEXT: >>> referenced by test.obj
+ERR-NEXT: >>> referenced by weak2.obj
+ERR-EMPTY:
+ERR-NEXT: error: undefined symbol: weak2
+ERR-NEXT: >>> referenced by test.obj
+ERR-NEXT: >>> referenced by weak2.obj
+
+Depending on symbol processing order, we may have temporary weak reference cycles:
+
+RUN: lld-link -machine:amd64 -dll -noentry -out:test.dll test.obj weak2.obj weak-def.obj
+RUN: lld-link -machine:amd64 -dll -noentry -out:test.dll test.obj weak-def.obj weak2.obj
+RUN: lld-link -machine:amd64 -dll -noentry -out:test.dll weak-def.obj test.obj weak2.obj
+
+#--- test.s
+    .weak testsym
+.set testsym, weak1
+    .weak weak1
+.set weak1, weak2
+    .section .test, "r"
+
+#--- weak2.s
+    .weak weak2
+.set weak2, weak1
+
+#--- weak-def.s
+    .globl weak1
+    .data
+weak1:
+    .word 0