From 30f4bfd822cbd72391327f503b9513aff162c596 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Mon, 28 Oct 2024 20:52:37 -0400
Subject: [PATCH] fix nginx

---
 tools/rewriter/SourceRewriter.cpp | 35 +++++++++++++++----------------
 1 file changed, 17 insertions(+), 18 deletions(-)

diff --git a/tools/rewriter/SourceRewriter.cpp b/tools/rewriter/SourceRewriter.cpp
index ab4df7b1c..1cd505296 100644
--- a/tools/rewriter/SourceRewriter.cpp
+++ b/tools/rewriter/SourceRewriter.cpp
@@ -660,9 +660,10 @@ class FnPtrExpr : public RefactoringCallback {
         mangle_type(ctxt, old_type->getCanonicalTypeInternal());
     OpaqueStruct new_type = kFnPtrTypePrefix + mangled_type;
 
+    auto caller_pkey = get_file_pkey(sm);
     auto linkage = fn_decl->getFormalLinkage();
     auto [it, new_fn] = addr_taken_fns[filename].insert(
-        std::make_tuple(fn_name, new_type, linkage));
+        std::make_tuple(fn_name, new_type, linkage, caller_pkey));
 
     // TODO: Note that this only checks if a function is added to the
     // addr_taken_fns map. To make the rewriter idempotent we should
@@ -675,12 +676,14 @@ class FnPtrExpr : public RefactoringCallback {
       } else {
         auto decl_start = expansion_range.getBegin();
         Filename decl_filename = get_filename(decl_start, sm);
-        Replacement old_used_attr(sm, decl_start, 0,
-                                  llvm::StringRef("__attribute__((used)) "));
-        Replacement used_attr = replace_new_file(decl_filename, old_used_attr);
-        auto err = file_replacements[decl_filename].add(used_attr);
-        if (err) {
-          llvm::errs() << "Error adding replacements: " << err << '\n';
+        if (!should_not_modify_file(decl_filename)) {
+          Replacement old_used_attr(sm, decl_start, 0,
+                                    llvm::StringRef("__attribute__((used)) "));
+          Replacement used_attr = replace_new_file(decl_filename, old_used_attr);
+          auto err = file_replacements[decl_filename].add(used_attr);
+          if (err) {
+            llvm::errs() << "Error adding replacements: " << err << '\n';
+          }
         }
       }
     }
@@ -710,7 +713,7 @@ class FnPtrExpr : public RefactoringCallback {
     return;
   }
 
-  std::map<Filename, std::set<std::tuple<Function, OpaqueStruct, clang::Linkage>>>
+  std::map<Filename, std::set<std::tuple<Function, OpaqueStruct, clang::Linkage, Pkey>>>
       addr_taken_fns;
 
 private:
@@ -1422,23 +1425,19 @@ int main(int argc, const char **argv) {
 
     // For each static function, define a macro to define the wrapper in the
     // output header and invoke it in the source file
-    for (const auto [fn_name, opaque, linkage] : addr_taken_fns_in_file) {
+    for (const auto [fn_name, opaque, linkage, caller_pkey] : addr_taken_fns_in_file) {
       std::string wrapper_name = "__ia2_"s + fn_name;
       std::string target_fn = fn_name;
 
-      if (clang::isExternallyVisible(linkage)) {// && filename pkey does not fn_name pkey) {
-        llvm::SmallString<256> input_file(filename);
-        llvm::sys::path::replace_path_prefix(input_file, OutputDirectory, RootDirectory);
-        Pkey file_pkey = file_pkeys.at(input_file.str().str());
-        Pkey fn_pkey = fn_decl_pass.fn_pkeys.at(fn_name);
-        if (file_pkey != fn_pkey) {
+      Pkey target_pkey = fn_decl_pass.fn_pkeys[fn_name];
+      if (clang::isExternallyVisible(linkage)) {
+        if (caller_pkey != target_pkey) {
           target_fn = "__real_"s + fn_name;
         }
       }
 
-      // TODO: These wrapper go from pkey 0 to the target pkey so if the target
-      // also has pkey 0 then it just needs call the original function
-      Pkey target_pkey = fn_decl_pass.fn_pkeys[fn_name];
+      // These wrapper go from pkey 0 to the target pkey so if the target
+      // also has pkey 0 then it just calls the original function
       header_out << "extern " << opaque << " " << wrapper_name << ";\n";
       if (target_pkey != 0) {
         AbiSignature c_abi_sig = fn_decl_pass.abi_signatures[fn_name];