From ad5bbdbe942a05f02c502247eec3c2d1204234ad Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333@gmail.com>
Date: Fri, 6 Jun 2025 12:40:44 +0800
Subject: [PATCH] Add basic support

---
 clang/include/clang/Basic/CodeGenOptions.def | 1 +
 clang/include/clang/Driver/Options.td        | 5 +++++
 clang/lib/CodeGen/CGCall.cpp                 | 3 +++
 llvm/include/llvm/IR/Attributes.td           | 3 +++
 llvm/lib/CodeGen/SelectOptimize.cpp          | 7 +++++--
 llvm/lib/Target/X86/X86CmovConversion.cpp    | 4 ++++
 6 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/clang/include/clang/Basic/CodeGenOptions.def b/clang/include/clang/Basic/CodeGenOptions.def
index aad4e107cbeb3..b70b7b03bfc77 100644
--- a/clang/include/clang/Basic/CodeGenOptions.def
+++ b/clang/include/clang/Basic/CodeGenOptions.def
@@ -316,6 +316,7 @@ CODEGENOPT(SanitizeStats     , 1, 0) ///< Collect statistics for sanitizers.
 CODEGENOPT(SimplifyLibCalls  , 1, 1) ///< Set when -fbuiltin is enabled.
 CODEGENOPT(SoftFloat         , 1, 0) ///< -soft-float.
 CODEGENOPT(SpeculativeLoadHardening, 1, 0) ///< Enable speculative load hardening.
+CODEGENOPT(TimingAttackHardening, 1, 0) ///< Enable timing attack hardening.
 CODEGENOPT(FineGrainedBitfieldAccesses, 1, 0) ///< Enable fine-grained bitfield accesses.
 CODEGENOPT(StrictEnums       , 1, 0) ///< Optimize based on strict enum definition.
 CODEGENOPT(StrictVTablePointers, 1, 0) ///< Optimize based on the strict vtable pointers
diff --git a/clang/include/clang/Driver/Options.td b/clang/include/clang/Driver/Options.td
index 5ca31c253ed8f..bc96c2917e1e9 100644
--- a/clang/include/clang/Driver/Options.td
+++ b/clang/include/clang/Driver/Options.td
@@ -5012,6 +5012,11 @@ defm speculative_load_hardening : BoolMOption<"speculative-load-hardening",
   CodeGenOpts<"SpeculativeLoadHardening">, DefaultFalse,
   PosFlag<SetTrue, [], [ClangOption, CC1Option]>,
   NegFlag<SetFalse>, BothFlags<[], [ClangOption, CLOption]>>;
+defm timing_attack_hardening
+    : BoolMOption<"timing-attack-hardening",
+                  CodeGenOpts<"TimingAttackHardening">, DefaultTrue,
+                  PosFlag<SetTrue, [], [ClangOption, CC1Option]>,
+                  NegFlag<SetFalse>, BothFlags<[], [ClangOption, CLOption]>>;
 def mlvi_hardening : Flag<["-"], "mlvi-hardening">, Group<m_Group>,
   Visibility<[ClangOption, CLOption]>,
   HelpText<"Enable all mitigations for Load Value Injection (LVI)">;
diff --git a/clang/lib/CodeGen/CGCall.cpp b/clang/lib/CodeGen/CGCall.cpp
index 056caf64525a5..8c497cea27319 100644
--- a/clang/lib/CodeGen/CGCall.cpp
+++ b/clang/lib/CodeGen/CGCall.cpp
@@ -2049,6 +2049,9 @@ static void getTrivialDefaultFunctionAttributes(
     if (CodeGenOpts.SpeculativeLoadHardening)
       FuncAttrs.addAttribute(llvm::Attribute::SpeculativeLoadHardening);
 
+    if (CodeGenOpts.TimingAttackHardening)
+      FuncAttrs.addAttribute(llvm::Attribute::TimingAttackHardening);
+
     // Add zero-call-used-regs attribute.
     switch (CodeGenOpts.getZeroCallUsedRegs()) {
     case llvm::ZeroCallUsedRegs::ZeroCallUsedRegsKind::Skip:
diff --git a/llvm/include/llvm/IR/Attributes.td b/llvm/include/llvm/IR/Attributes.td
index d488c5f419b82..6bac2e8450185 100644
--- a/llvm/include/llvm/IR/Attributes.td
+++ b/llvm/include/llvm/IR/Attributes.td
@@ -349,6 +349,9 @@ def SpeculativeLoadHardening : EnumAttr<"speculative_load_hardening",
                                         IntersectPreserve,
                                         [FnAttr]>;
 
+def TimingAttackHardening
+    : EnumAttr<"timing_attack_hardening", IntersectPreserve, [FnAttr]>;
+
 /// Argument is swift error.
 def SwiftError : EnumAttr<"swifterror", IntersectPreserve, [ParamAttr]>;
 
diff --git a/llvm/lib/CodeGen/SelectOptimize.cpp b/llvm/lib/CodeGen/SelectOptimize.cpp
index 13ed8f28d5507..65e1096752551 100644
--- a/llvm/lib/CodeGen/SelectOptimize.cpp
+++ b/llvm/lib/CodeGen/SelectOptimize.cpp
@@ -25,6 +25,7 @@
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSchedule.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
@@ -363,7 +364,8 @@ PreservedAnalyses SelectOptimizeImpl::run(Function &F,
     return PreservedAnalyses::all();
 
   TTI = &FAM.getResult<TargetIRAnalysis>(F);
-  if (!TTI->enableSelectOptimize())
+  if (!TTI->enableSelectOptimize() ||
+      F.hasFnAttribute(Attribute::TimingAttackHardening))
     return PreservedAnalyses::all();
 
   PSI = FAM.getResult<ModuleAnalysisManagerFunctionProxy>(F)
@@ -398,7 +400,8 @@ bool SelectOptimizeImpl::runOnFunction(Function &F, Pass &P) {
 
   TTI = &P.getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 
-  if (!TTI->enableSelectOptimize())
+  if (!TTI->enableSelectOptimize() ||
+      F.hasFnAttribute(Attribute::TimingAttackHardening))
     return false;
 
   LI = &P.getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
diff --git a/llvm/lib/Target/X86/X86CmovConversion.cpp b/llvm/lib/Target/X86/X86CmovConversion.cpp
index e25d7551e3890..59e5951b87d62 100644
--- a/llvm/lib/Target/X86/X86CmovConversion.cpp
+++ b/llvm/lib/Target/X86/X86CmovConversion.cpp
@@ -60,6 +60,7 @@
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSchedule.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/IR/Attributes.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/MC/MCSchedule.h"
@@ -168,6 +169,9 @@ bool X86CmovConverterPass::runOnMachineFunction(MachineFunction &MF) {
   if (!EnableCmovConverter)
     return false;
 
+  if (MF.getFunction().hasFnAttribute(Attribute::TimingAttackHardening))
+    return false;
+
   // If the SelectOptimize pass is enabled, cmovs have already been optimized.
   if (!getCGPassBuilderOption().DisableSelectOptimize)
     return false;