From 4db0d955ab24e6fee03d294a36dc8553c8ac89be Mon Sep 17 00:00:00 2001
From: Stephen Crane <sjc@immunant.com>
Date: Mon, 14 Oct 2024 17:18:13 -0700
Subject: [PATCH] Add static addr taken function test and fix bug.

Fixes a bug when the static function and its use are in different files (i.e. function definition is in a header).
---
 tests/CMakeLists.txt                         |  1 +
 tests/static_addr_taken/CMakeLists.txt       | 14 +++++++
 tests/static_addr_taken/include/static_fns.h | 10 +++++
 tests/static_addr_taken/lib.c                | 38 +++++++++++++++++++
 tests/static_addr_taken/main.c               | 39 ++++++++++++++++++++
 tools/rewriter/SourceRewriter.cpp            |  5 ++-
 6 files changed, 105 insertions(+), 2 deletions(-)
 create mode 100644 tests/static_addr_taken/CMakeLists.txt
 create mode 100644 tests/static_addr_taken/include/static_fns.h
 create mode 100644 tests/static_addr_taken/lib.c
 create mode 100644 tests/static_addr_taken/main.c

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 9fd534215..fb3997b70 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -59,6 +59,7 @@ add_subdirectory(rewrite_fn_ptr_eq)
 add_subdirectory(rewrite_macros)
 add_subdirectory(sighandler)
 add_subdirectory(simple1)
+add_subdirectory(static_addr_taken)
 add_subdirectory(structs)
 
 # The following tests are not supported on ARM64 yet
diff --git a/tests/static_addr_taken/CMakeLists.txt b/tests/static_addr_taken/CMakeLists.txt
new file mode 100644
index 000000000..41723b1b6
--- /dev/null
+++ b/tests/static_addr_taken/CMakeLists.txt
@@ -0,0 +1,14 @@
+define_shared_lib(
+    SRCS lib.c
+    NEEDS_LD_WRAP
+    PKEY 2
+)
+
+define_test(
+    SRCS main.c
+    NEEDS_LD_WRAP
+    PKEY 1
+    CRITERION_TEST
+)
+
+define_ia2_wrapper()
diff --git a/tests/static_addr_taken/include/static_fns.h b/tests/static_addr_taken/include/static_fns.h
new file mode 100644
index 000000000..8c9263562
--- /dev/null
+++ b/tests/static_addr_taken/include/static_fns.h
@@ -0,0 +1,10 @@
+#pragma once
+
+typedef void (*fn_ptr_ty)(void);
+
+static void inline_noop(void) {
+    printf("called %s defined in header\n", __func__);
+}
+
+fn_ptr_ty *get_ptrs_in_main(void);
+fn_ptr_ty *get_ptrs_in_lib(void);
diff --git a/tests/static_addr_taken/lib.c b/tests/static_addr_taken/lib.c
new file mode 100644
index 000000000..f8f8de550
--- /dev/null
+++ b/tests/static_addr_taken/lib.c
@@ -0,0 +1,38 @@
+#include <criterion/criterion.h>
+#include <criterion/logging.h>
+#include <ia2.h>
+
+#define IA2_COMPARTMENT 2
+#include <ia2_compartment_init.inc>
+
+#include "static_fns.h"
+
+static void duplicate_noop(void) {
+    printf("called %s in library\n", __func__);
+}
+
+static void identical_name(void) {
+    static int x = 4;
+    printf("%s in library read x = %d\n", __func__, x);
+}
+
+static fn_ptr_ty ptrs[3] IA2_SHARED_DATA = {
+    inline_noop, duplicate_noop, identical_name
+};
+
+fn_ptr_ty *get_ptrs_in_lib(void) {
+    return ptrs;
+}
+
+Test(static_addr_taken, call_ptrs_in_lib) {
+    for (int i = 0; i < 3; i++) {
+        ptrs[i]();
+    }
+}
+
+Test(static_addr_taken, call_ptr_from_main) {
+    fn_ptr_ty *main_ptrs = get_ptrs_in_main();
+    for (int i = 0; i < 3; i++) {
+        main_ptrs[i]();
+    }
+}
\ No newline at end of file
diff --git a/tests/static_addr_taken/main.c b/tests/static_addr_taken/main.c
new file mode 100644
index 000000000..adc88b239
--- /dev/null
+++ b/tests/static_addr_taken/main.c
@@ -0,0 +1,39 @@
+#include <criterion/criterion.h>
+#include <criterion/logging.h>
+#include <ia2.h>
+
+INIT_RUNTIME(2);
+#define IA2_COMPARTMENT 1
+#include <ia2_compartment_init.inc>
+
+#include "static_fns.h"
+
+static void duplicate_noop(void) {
+    printf("called %s in main binary\n", __func__);
+}
+
+static void identical_name(void) {
+    static int x = 3;
+    printf("%s in main binary read x = %d\n", __func__, x);
+}
+
+static fn_ptr_ty ptrs[3] IA2_SHARED_DATA = {
+    inline_noop, duplicate_noop, identical_name
+};
+
+fn_ptr_ty *get_ptrs_in_main(void) {
+    return ptrs;
+}
+
+Test(static_addr_taken, call_ptrs_in_main) {
+    for (int i = 0; i < 3; i++) {
+        ptrs[i]();
+    }
+}
+
+Test(static_addr_taken, call_ptr_from_lib) {
+    fn_ptr_ty *lib_ptrs = get_ptrs_in_lib();
+    for (int i = 0; i < 3; i++) {
+        lib_ptrs[i]();
+    }
+}
\ No newline at end of file
diff --git a/tools/rewriter/SourceRewriter.cpp b/tools/rewriter/SourceRewriter.cpp
index 07cf45da5..87f6d60db 100644
--- a/tools/rewriter/SourceRewriter.cpp
+++ b/tools/rewriter/SourceRewriter.cpp
@@ -675,10 +675,11 @@ class FnPtrExpr : public RefactoringCallback {
         if (!decl_start.isFileID()) {
           llvm::errs() << "Error: non-file loc for function " << fn_name << '\n';
         } else {
+          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(filename, old_used_attr);
-          auto err = file_replacements[filename].add(used_attr);
+          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';
           }