Improve static address taken function rewriting (#447)

* 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).

* Handle static function declaration that starts with a macro expansion

* Disable identically named static function

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

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()

tests/static_addr_taken/include/static_fns.h

@@ -0,0 +1,12 @@
+#pragma once
+
+#define LOCAL static
+
+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);

tests/static_addr_taken/lib.c

@@ -0,0 +1,25 @@
+#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;
+}

tests/static_addr_taken/main.c

@@ -0,0 +1,43 @@
+#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__);
+}
+
+LOCAL void macro_attr_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[] IA2_SHARED_DATA = {
+    inline_noop, duplicate_noop, /* identical_name, */ macro_attr_noop,
+};
+
+fn_ptr_ty *get_ptrs_in_main(void) {
+    return ptrs;
+}
+
+Test(static_addr_taken, call_ptrs_in_main) {
+    for (int i = 0; i < sizeof(ptrs) / sizeof(ptrs[0]); 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]();
+    }
+}

tools/rewriter/SourceRewriter.cpp

@@ -671,14 +671,17 @@ class FnPtrExpr : public RefactoringCallback {
       // internal_addr_taken_fns map. To make the rewriter idempotent we should
       // check for an existing used attribute.
       if (new_fn) {
-        auto decl_start = fn_decl->getBeginLoc();
-        if (!decl_start.isFileID()) {
+        auto static_fn_range = fn_decl->getSourceRange();
+        auto expansion_range = sm.getExpansionRange(static_fn_range);
+        if (!expansion_range.getBegin().isFileID()) {
           llvm::errs() << "Error: non-file loc for function " << fn_name << '\n';
         } 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(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';
           }