i#7585 AArch64 selfmod: Fix state corruption bug (#8128)

This commit fixes two related state corruption bugs.

When hw_cache_consistency is enabled DynamoRIO sets writable/executable
pages to read-only in order to detect any writes to code regions. If
these writes occurred whilst app state was spilled, the signal handler
failed to restore app state before returning to d_r_dispatch, which
caused app state corruption.

Similarly, if this fault happened on an instruction that uses the stolen
register, we would enter the signal handler with the stolen register
restored to its app state. The signal handler assumed it contained the
tls pointer and left things in a corrupt state.

Discovered while testing simple Java commands with hw_cache_consistency
enabled on AArch64:
    drrun -- javac HelloWorld.java

With this bug fixed I am not aware of any other issues with
hw_cache_consistency on AArch64 so this commit also turns
hw_cache_consistency on by deafult on AArch64.

Issue: #7585
Fixes #2055, #5771

core/drlibc/drlibc.c

@@ -200,12 +200,6 @@ clear_icache(void *beg, void *end)
     /* Data Synchronization Barrier */
     __asm__ __volatile__("dsb ish" : : : "memory");
 
-    /*TODO i#7585: The selfmod tests currently rely on seeing ic ivau instructions to
-     *            detect modified pages. Remove this line when we have fixed the
-     *            hw_cache_consistency bugs and enabled by default.
-     */
-    d_to_i_coherent = false;
-
     if (!d_to_i_coherent) {
         /* Invalidate instruction cache to point of unification, one line at a time. */
         addr = ALIGN_BACKWARD(beg_uint, icache_line_size);

core/optionsx.h

@@ -2381,8 +2381,7 @@ DYNAMIC_OPTION_DEFAULT(bool, allow_detach, true, "allow detaching from process")
 PC_OPTION_INTERNAL(bool, noasynch, "disable asynchronous event interceptions")
 #endif
 
-/* TODO i7585: Default to on for AArch64 when we have tested it more. */
-#if defined(X86) /* || defined(AARCH64) */
+#if defined(X86) || defined(AARCH64)
 #    define HW_CACHE_CONSISTENCY_DEFAULT true
 #else
 #    define HW_CACHE_CONSISTENCY_DEFAULT false

core/unix/signal.c

@@ -3152,12 +3152,15 @@ set_sigcontext_isa_mode(sig_full_cxt_t *sc_full, dr_isa_mode_t isa_mode)
 #    endif
 #endif
 
-/* Returns whether successful.  If avoid_failure, tries to translate
- * at least pc if not successful.  Pass f if known.
+/* Returns whether successful.  If avoid_failure, tries to translate at least pc if not
+ * successful.
+ * fragment_inout should point to the fragment to use for translation if known.
+ * If fragment_inout points to NULL, the fragment will be looked up from the PC in uc
+ * and fragment_inout will be set to point to the looked up fragment.
  */
 static bool
 translate_sigcontext(dcontext_t *dcontext, kernel_ucontext_t *uc, bool avoid_failure,
-                     fragment_t *f)
+                     DR_PARAM_INOUT fragment_t **fragment_inout, fragment_t *wrapper)
 {
     bool success = false;
     priv_mcontext_t mcontext;
@@ -3180,6 +3183,13 @@ translate_sigcontext(dcontext_t *dcontext, kernel_ucontext_t *uc, bool avoid_fai
     /* PR 214962: we assume we're going to relocate to this stored context,
      * so we restore memory now
      */
+
+    if (fragment_inout != NULL && *fragment_inout == NULL && wrapper != NULL) {
+        *fragment_inout = fragment_pclookup(dcontext, (cache_pc)sc->SC_XIP, wrapper);
+    }
+
+    fragment_t *f = fragment_inout == NULL ? NULL : *fragment_inout;
+
     if (translate_mcontext(dcontext->thread_record, &mcontext, true /*restore memory*/,
                            f)) {
         mcontext_to_ucontext(uc, &mcontext);
@@ -4112,35 +4122,28 @@ transfer_from_sig_handler_to_fcache_return(dcontext_t *dcontext, kernel_ucontext
      */
     sc->SC_XIP = (ptr_uint_t)fcache_return_routine(dcontext);
 #if defined(AARCHXX) || defined(RISCV64)
-    /* If !is_kernel_xfer then we have come here from main_signal_handler
-     * via check_for_modified_code and recreate_app_state_internal was
-     * called with just_pc so the stolen register still points at TLS.
+    /* We do not have to set dr_reg_stolen in dcontext's mcontext here
+     * because dcontext's mcontext is stale and we used the mcontext
+     * created from recreate_app_state_internal with the original sigcontext.
      */
-    if (is_kernel_xfer) {
-        /* We do not have to set dr_reg_stolen in dcontext's mcontext here
-         * because dcontext's mcontext is stale and we used the mcontext
-         * created from recreate_app_state_internal with the original sigcontext.
-         */
-        /* We restore dr_reg_stolen's app value in recreate_app_state_internal,
-         * so now we need set dr_reg_stolen to hold DR's TLS before sigreturn
-         * from DR's handler.
-         */
-        ASSERT(get_sigcxt_stolen_reg(sc) != (reg_t)*get_dr_tls_base_addr());
-        /* Preserve the translated value. */
-        dcontext->local_state->spill_space.reg_stolen = get_sigcxt_stolen_reg(sc);
-        /* Now put DR's base in the sigcontext. */
-        set_sigcxt_stolen_reg(sc, (reg_t)*get_dr_tls_base_addr());
+    /* We restore dr_reg_stolen's app value in recreate_app_state_internal,
+     * so now we need set dr_reg_stolen to hold DR's TLS before sigreturn
+     * from DR's handler.
+     */
+    /* Preserve the translated value. */
+    dcontext->local_state->spill_space.reg_stolen = get_sigcxt_stolen_reg(sc);
+    /* Now put DR's base in the sigcontext. */
+    set_sigcxt_stolen_reg(sc, (reg_t)*get_dr_tls_base_addr());
 #    ifdef RISCV64
-        os_set_app_tls_base(dcontext, TLS_REG_LIB, (void *)get_sigcxt_tp_reg(sc));
-        /* Now put host tp in the sigcontext. */
-        set_sigcxt_tp_reg(sc, (reg_t)read_thread_register(TLS_REG_LIB));
+    os_set_app_tls_base(dcontext, TLS_REG_LIB, (void *)get_sigcxt_tp_reg(sc));
+    /* Now put host tp in the sigcontext. */
+    set_sigcxt_tp_reg(sc, (reg_t)read_thread_register(TLS_REG_LIB));
 #    endif
 
 #    ifdef ARM
-        /* We're going to our fcache_return gencode which uses DEFAULT_ISA_MODE */
-        set_pc_mode_in_cpsr(sc, DEFAULT_ISA_MODE);
+    /* We're going to our fcache_return gencode which uses DEFAULT_ISA_MODE */
+    set_pc_mode_in_cpsr(sc, DEFAULT_ISA_MODE);
 #    endif
-    }
 #endif
 
 #if defined(X64) || defined(ARM)
@@ -5162,9 +5165,8 @@ record_pending_signal(dcontext_t *dcontext, int sig, kernel_ucontext_t *ucxt,
 
         sc_orig = *sc;
         ASSERT(!forged);
-        /* cache the fragment since pclookup is expensive for coarse (i#658) */
-        f = fragment_pclookup(dcontext, (cache_pc)sc->SC_XIP, &wrapper);
-        xl8_success = translate_sigcontext(dcontext, ucxt, !can_always_delay[sig], f);
+        xl8_success =
+            translate_sigcontext(dcontext, ucxt, !can_always_delay[sig], &f, &wrapper);
 
         if (can_always_delay[sig] && !xl8_success) {
             /* delay: we expect this for coarse fragments if alarm arrives
@@ -5208,10 +5210,8 @@ record_pending_signal(dcontext_t *dcontext, int sig, kernel_ucontext_t *ucxt,
          */
         if (blocked && !forged && !reroute && !can_always_delay[sig] &&
             safe_is_in_fcache(dcontext, pc, xsp)) {
-            /* cache the fragment since pclookup is expensive for coarse (i#658) */
-            f = fragment_pclookup(dcontext, (cache_pc)sc->SC_XIP, &wrapper);
             sc_orig = *sc;
-            translate_sigcontext(dcontext, ucxt, true /*shouldn't fail*/, f);
+            translate_sigcontext(dcontext, ucxt, true /*shouldn't fail*/, &f, &wrapper);
 
             if (!send_signal_to_client_and_handle_action(
                     dcontext, sig, frame, sc, &sc_orig, access_address, true /*blocked*/,
@@ -5235,7 +5235,7 @@ record_pending_signal(dcontext_t *dcontext, int sig, kernel_ucontext_t *ucxt,
             if (forged)
                 ASSERT(is_couldbelinking(dcontext));
             else
-                translate_sigcontext(dcontext, ucxt, true /*shouldn't fail*/, NULL);
+                translate_sigcontext(dcontext, ucxt, true /*shouldn't fail*/, NULL, NULL);
             /* the process should be terminated */
             execute_default_from_cache(dcontext, sig, frame, &sc_orig, forged);
             ASSERT_NOT_REACHED();
@@ -5520,31 +5520,49 @@ check_for_modified_code(dcontext_t *dcontext, cache_pc instr_cache_pc,
      * how can we tell the difference?
      */
     if (was_executable_area_writable(target)) {
-        /* translate instr_cache_pc to original app pc
-         * DO NOT use translate_sigcontext, don't want to change the
-         * signal frame or else we'll lose control when we try to
-         * return to signal pc!
+        /* There are three scenarios we need to consider here:
+         * 1. Fault was in native code, we return to the native PC,
+         * 2. Fault was in fcache code, we return to the fcache PC,
+         * 3. Fault was in fcache code, we return to d_r_dispatch.
+         *
+         * Cases 1 and 2 we need to preserve the sigcontext state because we are
+         * returning to that context but case 3 we need to restore app state.
+         *
+         * If native_state==true we are in case 1 and don't need to restore the app state
+         * but if native_state==false then theoretically we don't know whether we are in
+         * case 2 or 3 until handle_modified_code() returns. This is a problem because
+         * case 3 must restore state, and case 2 must not restore state. If we are
+         * restoring state it needs to be done before calling handle_modified_code()
+         * because it might flush the current fragment, but we don't know whether we need
+         * to until afterwards.
+         *
+         * In practice however case 2 never happens. Comments in handle_modified_code()
+         * describe the case 2 scenario as a possible future optimisation that hasn't
+         * been implemented. If we assume case 2 is impossible then we can distinguish
+         * case 1 and 3 using native_state.
          */
         app_pc next_pc, translated_pc = NULL;
         fragment_t *f = NULL;
         fragment_t wrapper;
         ASSERT((cache_pc)SIGCXT_FROM_UCXT(uc)->SC_XIP == instr_cache_pc);
+        bool context_is_translated = false;
         if (!native_state) {
-            /* For safe recreation we need to either be couldbelinking or hold
-             * the initexit lock (to keep someone from flushing current
-             * fragment), the initexit lock is easier
+            /* handle_modified_code() might flush f so we need to recreate app state
+             * before the call.
              */
-            d_r_mutex_lock(&thread_initexit_lock);
-            /* cache the fragment since pclookup is expensive for coarse units (i#658) */
-            f = fragment_pclookup(dcontext, instr_cache_pc, &wrapper);
-            translated_pc = recreate_app_pc(dcontext, instr_cache_pc, f);
+            context_is_translated =
+                translate_sigcontext(dcontext, uc, true, &f, &wrapper);
+            ASSERT(context_is_translated);
+            translated_pc = (app_pc)SIGCXT_FROM_UCXT(uc)->SC_XIP;
             ASSERT(translated_pc != NULL);
-            d_r_mutex_unlock(&thread_initexit_lock);
         }
 
         next_pc =
             handle_modified_code(dcontext, instr_cache_pc, translated_pc, target, f);
 
+        /* Confirm that we only see case 1 or case 3. */
+        ASSERT((native_state && next_pc == NULL) || (!native_state && next_pc != NULL));
+
         if (!native_state) {
             /* going to exit from middle of fragment (at the write) so will mess up
              * trace building
@@ -5559,7 +5577,6 @@ check_for_modified_code(dcontext_t *dcontext, cache_pc instr_cache_pc,
             /* re-execute the write -- just have main_signal_handler return */
             return true;
         } else {
-            ASSERT(!native_state);
             /* Do not resume execution in cache, go back to d_r_dispatch. */
             transfer_from_sig_handler_to_fcache_return(
                 dcontext, uc, NULL, SIGSEGV, next_pc,

core/vmareas.c

@@ -10974,8 +10974,12 @@ handle_modified_code(dcontext_t *dcontext, cache_pc instr_cache_pc, app_pc instr
      * flushing outside the requested region (entire vm_area_t). If we could tell
      * we could return NULL instead (which is a special flag that says redo the
      * write instead of going to d_r_dispatch) if f wasn't flushed.
-     * XXX - Redoing the write would be more efficient then going back to
-     * d_r_dispatch and should be the common case. */
+     * XXX i#7585 - Redoing the write would be more efficient than going back to
+     * d_r_dispatch and should be the common case.
+     * Note that check_for_modified_code() in unix/signal.c assumes that this function
+     * returns instr_app_pc even if f wasn't flushed. If we implement this optimisation
+     * check_for_modified_code() will need to be reworked.
+     */
     flush_fragments_in_region_finish(dcontext, false /*don't keep initexit_lock*/);
     if (DYNAMO_OPTION(opt_jit) && !TESTANY(MEMPROT_WRITE, prot) &&
         is_jit_managed_area(flush_start))

suite/tests/CMakeLists.txt

@@ -3178,6 +3178,12 @@ if (NOT RISCV64) # TODO i#3544: Port tests to RISC-V 64
   use_DynamoRIO_extension(client.drreg-cross.dll drreg)
   use_DynamoRIO_extension(client.drreg-cross.dll drutil)
 
+  if (AARCH64 AND UNIX AND NOT APPLE)
+    tobuild_ci(client.codemod-state client-interface/codemod-state.c "" "" "")
+    use_DynamoRIO_extension(client.codemod-state.dll drmgr)
+    use_DynamoRIO_extension(client.codemod-state.dll drreg)
+  endif ()
+
   tobuild_ci(client.drx-test client-interface/drx-test.c "" "" "")
   use_DynamoRIO_extension(client.drx-test.dll drx)
 
@@ -6787,9 +6793,8 @@ if (AARCH64)
 
   # Test that stolen register mangling does not interfere with selfmod mangling.
   # The -steal_reg <N> option must be kept in sync with ADDR_REG in selfmod.c.
-  # TODO i7585: Fix stolen reg corruption bug and enable this test.
-  #torunonly(security-common.selfmod-stolen_reg security-common.selfmod-coherent-icache
-  #  security-common/selfmod.c "-steal_reg 9 -hw_cache_consistency -sandbox_writes" "")
+  torunonly(security-common.selfmod-stolen_reg security-common.selfmod-coherent-icache
+    security-common/selfmod.c "-steal_reg 9 -hw_cache_consistency -sandbox_writes" "")
 endif ()
 
 if (X64 AND NOT AARCH64 AND NOT RISCV64) # XXX i#1569: get working on AArch64

suite/tests/client-interface/codemod-state.c

@@ -0,0 +1,130 @@
+/* ****************************************************
+ * Copyright (c) 2026 Arm Limited  All rights reserved.
+ * ***************************************************/
+
+/*
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * * Redistributions of source code must retain the above copyright notice,
+ *   this list of conditions and the following disclaimer.
+ *
+ * * Redistributions in binary form must reproduce the above copyright notice,
+ *   this list of conditions and the following disclaimer in the documentation
+ *   and/or other materials provided with the distribution.
+ *
+ * * Neither the name of Arm Limited nor the names of its contributors may be
+ *   used to endorse or promote products derived from this software without
+ *   specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED. IN NO EVENT SHALL ARM LIMITED OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+ * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+ * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+ * DAMAGE.
+ */
+
+#define CODEMOD_STATE_WRITER_INCREMENT 7
+
+#ifndef ASM_CODE_ONLY
+
+#    include "tools.h"
+
+/*
+ * codemod-state.c
+ * Test that the DynamoRIO signal handler restores spilled app state when an app writes to
+ * an rwx code page that DynamoRIO has made read-only for code modification detection.
+ */
+
+/*
+ * Write `value` to `*target` and return `input+CODEMOD_STATE_WRITER_INCREMENT`.
+ * We test that the value of `input` is not corrupted when its register is spilled
+ * by a client across the store to `target`.
+ */
+extern ptr_uint_t
+codemod_state_writer(uint *target, uint value, ptr_uint_t input);
+
+typedef void (*generated_func_t)(void);
+
+int
+main(int argc, char **argv)
+{
+#    ifdef X64
+#        define TEST_INPUT_VALUE 0x123456789abc0000ULL
+#    else
+#        error "Test does not support the target architecture."
+#    endif
+
+#    ifdef AARCH64
+    static const uint ret = 0xd65f03c0;
+#    else
+#        error "Test does not support the target architecture."
+#    endif
+
+    static const ptr_uint_t expected = TEST_INPUT_VALUE + CODEMOD_STATE_WRITER_INCREMENT;
+
+    /* Create a writable/executable page and copy our ret instruction there. */
+    uint *code = (uint *)allocate_mem(PAGE_SIZE, ALLOW_READ | ALLOW_WRITE | ALLOW_EXEC);
+    if (code == NULL) {
+        print("failed to allocate code\n");
+        return 1;
+    }
+
+    *code = ret;
+    tools_clear_icache(code, code + 1);
+
+    /* Call the generated code to make sure it is in the code cache. */
+    ((generated_func_t)code)();
+
+    ptr_uint_t result = codemod_state_writer(code, ret, TEST_INPUT_VALUE);
+    if (result != expected) {
+        print("register state corrupted: " PFX " != " PFX "\n", result, expected);
+        return 1;
+    }
+
+    print("state preserved\n");
+    free_mem((char *)code, PAGE_SIZE);
+    return 0;
+}
+
+#else /* ASM_CODE_ONLY */
+
+#    include "asm_defines.asm"
+/* clang-format off */
+START_FILE
+
+#define FUNCNAME codemod_state_writer
+    DECLARE_EXPORTED_FUNC(FUNCNAME)
+GLOBAL_LABEL(FUNCNAME:)
+#    ifdef AARCH64
+        /*
+        * ptr_uint_t
+        * codemod_state_writer(uint *target, uint value, ptr_uint_t input);
+        */
+#       define target x0
+#       define value w1
+#       define input x2
+        /* Client temporarily clobbers `input` here. This simulates a client spilling an
+         * app register to use as scratch.
+         * The address we are writing to is in the read/write/execute page which DR will
+         * have set to read-only to detect writes. The str will fault and DR returns from
+         * the signal handler back to the dispatcher (rather than directly back to fcache
+         * PC). We test that app state is correctly preserved in this situation.
+         */
+        str     value, [target]
+        /* Client restores `input` here, but the fault means the restore isn't reached.
+         */
+        add     x0, input, #(CODEMOD_STATE_WRITER_INCREMENT)
+        ret
+#    else
+#        error "Test does not support the target architecture."
+#    endif
+#undef FUNCNAME
+
+#endif

suite/tests/client-interface/codemod-state.dll.c

@@ -0,0 +1,120 @@
+/* ****************************************************
+ * Copyright (c) 2026 Arm Limited  All rights reserved.
+ * ***************************************************/
+
+/*
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * * Redistributions of source code must retain the above copyright notice,
+ *   this list of conditions and the following disclaimer.
+ *
+ * * Redistributions in binary form must reproduce the above copyright notice,
+ *   this list of conditions and the following disclaimer in the documentation
+ *   and/or other materials provided with the distribution.
+ *
+ * * Neither the name of Arm Limited nor the names of its contributors may be
+ *   used to endorse or promote products derived from this software without
+ *   specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED. IN NO EVENT SHALL ARM LIMITED OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+ * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+ * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+ * DAMAGE.
+ */
+
+#include "dr_api.h"
+#include "drmgr.h"
+#include "drreg.h"
+#include "client_tools.h"
+
+static app_pc test_store_pc;
+static bool found_test_store;
+
+static dr_emit_flags_t
+event_app_instruction(void *drcontext, void *tag, instrlist_t *bb, instr_t *instr,
+                      bool for_trace, bool translating, void *user_data)
+{
+    if (!instr_is_app(instr) || instr_get_app_pc(instr) != test_store_pc)
+        return DR_EMIT_DEFAULT;
+
+    CHECK(instr_get_opcode(instr) == OP_str, "test store has unexpected opcode");
+
+    /* Spill and clobber the add src register before the store instruction so that it
+     * does not have its app value when the store faults.
+     */
+
+    instr_t *add_instr = instr_get_next_app(instr);
+    CHECK(instr_get_opcode(add_instr) == OP_add,
+          "test store is not followed by the expected add.");
+    CHECK(instr_num_srcs(add_instr) >= 1, "add instr has wrong number of srcs.");
+    opnd_t add_src = instr_get_src(add_instr, 0);
+    CHECK(opnd_is_reg(add_src), "add src 0 is not a register.");
+    reg_id_t clobber_reg = opnd_get_reg(add_src);
+
+    drvector_t allowed;
+    reg_id_t reg = DR_REG_NULL;
+    drreg_init_and_fill_vector(&allowed, false);
+    drreg_set_vector_entry(&allowed, clobber_reg, true);
+    if (drreg_reserve_register(drcontext, bb, instr, &allowed, &reg) != DRREG_SUCCESS) {
+        CHECK(false, "failed to reserve clobber register");
+    }
+    ASSERT(reg == clobber_reg);
+    drvector_delete(&allowed);
+
+    /* Insert a meta instruction that clobbers the add src register. */
+    static const uint poison_value = 0x55;
+    instrlist_meta_preinsert(bb, instr,
+                             XINST_CREATE_load_int(drcontext,
+                                                   opnd_create_reg(clobber_reg),
+                                                   OPND_CREATE_INT16(poison_value)));
+
+    /* Now restore the app value after the store. */
+    if (drreg_unreserve_register(drcontext, bb, instr_get_next_app(instr), reg) !=
+        DRREG_SUCCESS) {
+        CHECK(false, "failed to unreserve register");
+    }
+
+    found_test_store = true;
+    return DR_EMIT_DEFAULT;
+}
+
+static void
+event_exit(void)
+{
+    CHECK(found_test_store, "failed to find test store");
+    if (!drmgr_unregister_bb_insertion_event(event_app_instruction) ||
+        drreg_exit() != DRREG_SUCCESS) {
+        CHECK(false, "exit failed");
+    }
+    drmgr_exit();
+}
+
+DR_EXPORT void
+dr_client_main(client_id_t id, int argc, const char *argv[])
+{
+    drreg_options_t ops = { sizeof(ops), 1 /*max slots needed*/, false };
+    module_data_t *module = dr_get_main_module();
+    CHECK(module != NULL, "failed to find main module");
+    test_store_pc = (app_pc)dr_get_proc_address(module->handle, "codemod_state_writer");
+    CHECK(test_store_pc != NULL, "failed to find codemod_state_writer");
+    dr_free_module_data(module);
+
+    if (!drmgr_init()) {
+        CHECK(false, "drmgr init failed");
+    }
+    if (drreg_init(&ops) != DRREG_SUCCESS) {
+        CHECK(false, "drreg init failed");
+    }
+    drmgr_register_exit_event(event_exit);
+    if (!drmgr_register_bb_instrumentation_event(NULL, event_app_instruction, NULL)) {
+        CHECK(false, "bb registration failed");
+    }
+}

suite/tests/client-interface/codemod-state.expect

@@ -0,0 +1 @@
+state preserved