From 4a6d962ff22d6a34385dec0f3bd902d3a3b71675 Mon Sep 17 00:00:00 2001
From: dheaton-arm <Damian.Heaton@arm.com>
Date: Tue, 30 Jul 2024 13:52:20 +0100
Subject: [PATCH] shims/foreign_items: Add hook for updating the canonical
 address

Where a method modifies the address that should be considered canonical - such as via TBI or when an MTE tag has been set - Miri will need to be notified.

This adds a hook to inform Miri of the new address.
---
 src/alloc_addresses/mod.rs | 15 +++++++++++++++
 src/shims/foreign_items.rs |  7 +++++++
 tests/utils/miri_extern.rs |  5 +++++
 3 files changed, 27 insertions(+)

diff --git a/src/alloc_addresses/mod.rs b/src/alloc_addresses/mod.rs
index c19a962e65..3f76e984c7 100644
--- a/src/alloc_addresses/mod.rs
+++ b/src/alloc_addresses/mod.rs
@@ -460,6 +460,21 @@ impl<'tcx> MiriMachine<'tcx> {
             }
         })
     }
+
+    /// Updates the current canonical address for the allocation. Note that any access to this allocation *must* use this address from this point on.
+    pub fn set_alloc_address(&mut self, id: AllocId, new: u64) {
+        let global_state = self.alloc_addresses.get_mut();
+        if let Some(addr) = global_state.base_addr.insert(id, new) {
+            // Remove the old address' int->ptr mapping.
+            let pos =
+                global_state.int_to_ptr_map.binary_search_by_key(&addr, |(addr, _)| *addr).unwrap();
+            let removed = global_state.int_to_ptr_map.remove(pos);
+            assert_eq!(removed, (addr, id));
+        }
+        let new_pos =
+            global_state.int_to_ptr_map.binary_search_by_key(&new, |(addr, _)| *addr).unwrap_err();
+        global_state.int_to_ptr_map.insert(new_pos, (new, id));
+    }
 }
 
 #[cfg(test)]
diff --git a/src/shims/foreign_items.rs b/src/shims/foreign_items.rs
index 1b45bc2203..4b24671139 100644
--- a/src/shims/foreign_items.rs
+++ b/src/shims/foreign_items.rs
@@ -423,6 +423,13 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     }
                 }
             }
+            "miri_set_canonical_address" => {
+                let [old_ptr, new_ptr] = this.check_shim(abi, Abi::Rust, link_name, args)?;
+                let old_ptr = this.read_pointer(old_ptr)?;
+                let new_ptr = this.read_pointer(new_ptr)?;
+                let (alloc_id, _, _) = this.ptr_get_alloc_id(old_ptr, 0)?;
+                this.machine.set_alloc_address(alloc_id, new_ptr.addr().bytes());
+            }
 
             // Aborting the process.
             "exit" => {
diff --git a/tests/utils/miri_extern.rs b/tests/utils/miri_extern.rs
index d6c43b1882..8c8349d9e9 100644
--- a/tests/utils/miri_extern.rs
+++ b/tests/utils/miri_extern.rs
@@ -147,4 +147,9 @@ extern "Rust" {
     /// "symbolic" alignment checks. Will fail if the pointer is not actually aligned or `align` is
     /// not a power of two. Has no effect when alignment checks are concrete (which is the default).
     pub fn miri_promise_symbolic_alignment(ptr: *const (), align: usize);
+
+    /// Miri-provided extern function to specify that a new address is to be considered the
+    /// canonical address, where `new` is a valid alias to the `old` allocation,
+    /// usually due to them having different values in bits that are ignored by hardware.
+    pub fn miri_set_canonical_address(old: *const (), new: *const ());
 }