Fix taint analysis with sret

lib/PhasarLLVM/DataFlow/IfdsIde/Problems/IFDSTaintAnalysis.cpp

@@ -266,6 +266,11 @@ transferAndKillTwoFlows(d_t To, d_t From1, d_t From2) {
 auto IFDSTaintAnalysis::getNormalFlowFunction(n_t Curr,
                                               [[maybe_unused]] n_t Succ)
     -> FlowFunctionPtrType {
+  PHASAR_LOG_LEVEL_CAT(DEBUG, "IFDSTaintAnalysis",
+                       "getNormalFlowFunction: " << llvmIRToString(Curr)
+                                                 << " --> "
+                                                 << llvmIRToString(Succ));
+
   // If a tainted value is stored, the store location must be tainted too
   if (const auto *Store = llvm::dyn_cast<llvm::StoreInst>(Curr)) {
     container_type Gen;
@@ -396,14 +401,6 @@ auto IFDSTaintAnalysis::getSummaryFlowFunction([[maybe_unused]] n_t CallSite,
   // populateWithMayAliases(Leak, CallSite);
   populateWithMustAliases(Kill, CallSite);
 
-  if (CS->hasStructRetAttr()) {
-    const auto *SRet = CS->getArgOperand(0);
-    if (!Gen.count(SRet)) {
-      // SRet is guaranteed to be written to by the call. If it does not
-      // generate it, we can freely kill it
-      Kill.insert(SRet);
-    }
-  }
   if (Gen.empty() && Leak.empty() && Kill.empty()) {
     if (Llvmfdff.contains(DestFun)) {
       // Note: The LLVMfdff is constant during the lifetime of the analysis, so
@@ -437,6 +434,15 @@ auto IFDSTaintAnalysis::getSummaryFlowFunction([[maybe_unused]] n_t CallSite,
     // not found
     return nullptr;
   }
+
+  if (CS->hasStructRetAttr()) {
+    const auto *SRet = CS->getArgOperand(0);
+    if (!Gen.count(SRet)) {
+      // SRet is guaranteed to be written to by the call. If it does not
+      // generate it, we can freely kill it
+      Kill.insert(SRet);
+    }
+  }
   if (Gen.empty()) {
     if (!Leak.empty() || !Kill.empty()) {
       return lambdaFlow([Leak{std::move(Leak)}, Kill{std::move(Kill)}, this,

test/llvm_test_code/taint_analysis/dummy_source_sink/CMakeLists.txt

@@ -15,6 +15,7 @@ set(taint_tests
   taint_exception_09.cpp
   taint_exception_10.cpp
   taint_lib_sum_01.cpp
+  sret.c
 )
 
 set(taint_tests_mem2reg

test/llvm_test_code/taint_analysis/dummy_source_sink/sret.c

@@ -0,0 +1,23 @@
+#include <stdio.h>
+
+extern char *source(void); // dummy source
+extern void sink(char c);  // dummy sink
+
+typedef struct {
+  char data[64];
+  int x;
+} BigStruct;
+
+BigStruct get_data(char *input) {
+  BigStruct s;
+  s.data[0] = input[0];
+  s.x = 42;
+  return s;
+}
+
+int main() {
+  char *tainted = source();
+  BigStruct bs = get_data(tainted); // sret call
+  sink(bs.data[0]);                 // Should be flagged as leak!
+  return 0;
+}

unittests/PhasarLLVM/DataFlow/IfdsIde/Problems/IFDSTaintAnalysisTest.cpp

@@ -34,12 +34,18 @@ class IFDSTaintAnalysisTest : public ::testing::Test {
   std::optional<IFDSTaintAnalysis> TaintProblem;
   std::optional<LLVMTaintConfig> TSF;
 
+  static bool isDummySrcFun(llvm::StringRef Name) {
+    return Name == "_Z6sourcev" || Name == "source";
+  }
+  static bool isDummySinkFun(llvm::StringRef Name) {
+    return Name == "_Z4sinki" || Name == "sink";
+  }
   static LLVMTaintConfig getDefaultConfig() {
     auto SourceCB = [](const llvm::Instruction *Inst) {
       std::set<const llvm::Value *> Ret;
       if (const auto *Call = llvm::dyn_cast<llvm::CallBase>(Inst);
           Call && Call->getCalledFunction() &&
-          Call->getCalledFunction()->getName() == "_Z6sourcev") {
+          isDummySrcFun(Call->getCalledFunction()->getName())) {
         Ret.insert(Call);
       }
       return Ret;
@@ -48,7 +54,7 @@ class IFDSTaintAnalysisTest : public ::testing::Test {
       std::set<const llvm::Value *> Ret;
       if (const auto *Call = llvm::dyn_cast<llvm::CallBase>(Inst);
           Call && Call->getCalledFunction() &&
-          Call->getCalledFunction()->getName() == "_Z4sinki") {
+          isDummySinkFun(Call->getCalledFunction()->getName())) {
         assert(Call->arg_size() > 0);
         Ret.insert(Call->getArgOperand(0));
       }
@@ -193,6 +199,18 @@ TEST_F(IFDSTaintAnalysisTest, TaintTest_06) {
   compare(TaintProblem->Leaks, GroundTruth);
 }
 
+TEST_F(IFDSTaintAnalysisTest, SRetTest_01) {
+  initialize({PathToLlFiles + "dummy_source_sink/sret_c_dbg.ll"});
+  IFDSSolver TaintSolver(*TaintProblem, &HA->getICFG());
+  TaintSolver.solve();
+
+  auto SinkCall = LineColFun{21, 3, "main"};
+  auto BsdataAt0 = LineColFunOp{21, 8, "main", llvm::Instruction::Load};
+  GroundTruthTy GroundTruth{{SinkCall, {BsdataAt0}}};
+
+  compare(TaintProblem->Leaks, GroundTruth);
+}
+
 TEST_F(IFDSTaintAnalysisTest, TaintTest_ExceptionHandling_01) {
   initialize(
       {PathToLlFiles + "dummy_source_sink/taint_exception_01_cpp_dbg.ll"});