[lld] [LLD][COFF] Check load config size before setting its DependentLoadFlags (PR #118535)

[Jacek Caban via llvm-commits]

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

Merge prepareLoadConfig and checkLoadConfigGuardData to share helper macros.

>From f840fe78ab75d7b438bc4fb435ae02a27eb35914 Mon Sep 17 00:00:00 2001
From: Jacek Caban <jacek at codeweavers.com>
Date: Tue, 3 Dec 2024 20:22:37 +0100
Subject: [PATCH] [LLD][COFF] Check load config size before setting its
 DependentLoadFlags

Merge prepareLoadConfig and checkLoadConfigGuardData to share helper macros.
---
 lld/COFF/Writer.cpp                   | 14 +++++---------
 lld/test/COFF/deploadflag-cfg-short.s | 12 ++++++++++++
 2 files changed, 17 insertions(+), 9 deletions(-)
 create mode 100644 lld/test/COFF/deploadflag-cfg-short.s

diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
index d3e326378ed2d4..fe78b8cf4ecff7 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -277,7 +277,6 @@ class Writer {
 
   void prepareLoadConfig();
   template <typename T> void prepareLoadConfig(T *loadConfig);
-  template <typename T> void checkLoadConfigGuardData(const T *loadConfig);
 
   std::unique_ptr<FileOutputBuffer> &buffer;
   std::map<PartialSectionKey, PartialSection *> partialSections;
@@ -2631,14 +2630,6 @@ void Writer::prepareLoadConfig() {
 }
 
 template <typename T> void Writer::prepareLoadConfig(T *loadConfig) {
-  if (ctx.config.dependentLoadFlags)
-    loadConfig->DependentLoadFlags = ctx.config.dependentLoadFlags;
-
-  checkLoadConfigGuardData(loadConfig);
-}
-
-template <typename T>
-void Writer::checkLoadConfigGuardData(const T *loadConfig) {
   size_t loadConfigSize = loadConfig->Size;
 
 #define RETURN_IF_NOT_CONTAINS(field)                                          \
@@ -2660,6 +2651,11 @@ void Writer::checkLoadConfigGuardData(const T *loadConfig) {
     if (loadConfig->field != s->getVA())                                       \
       warn(#field " not set correctly in '_load_config_used'");
 
+  if (ctx.config.dependentLoadFlags) {
+    RETURN_IF_NOT_CONTAINS(DependentLoadFlags)
+    loadConfig->DependentLoadFlags = ctx.config.dependentLoadFlags;
+  }
+
   if (ctx.config.guardCF == GuardCFLevel::Off)
     return;
   RETURN_IF_NOT_CONTAINS(GuardFlags)
diff --git a/lld/test/COFF/deploadflag-cfg-short.s b/lld/test/COFF/deploadflag-cfg-short.s
new file mode 100644
index 00000000000000..9cc5248044d8a3
--- /dev/null
+++ b/lld/test/COFF/deploadflag-cfg-short.s
@@ -0,0 +1,12 @@
+# REQUIRES: x86
+
+# RUN: llvm-mc -triple x86_64-windows-msvc -filetype=obj %s -o %t.obj
+# RUN: lld-link %t.obj -out:%t.dll -dll -noentry -nodefaultlib -dependentloadflag:0x800 2>&1 | FileCheck %s
+# CHECK: lld-link: warning: '_load_config_used' structure too small to include DependentLoadFlags
+
+        .section .rdata,"dr"
+        .balign 8
+.globl _load_config_used
+_load_config_used:
+        .long 0x4c
+        .fill 0x48, 1, 0