aarch64: Save/Restore PAC keys

Should fix #3568.

src/Task.cc

@@ -1278,7 +1278,50 @@ bool Task::get_aarch64_debug_regs(int which, ARM64Arch::user_hwdebug_state *regs
   fallible_ptrace(PTRACE_GETREGSET, which, (void*)&iov);
   return errno == 0;
 }
+std::vector<uint8_t> Task::pac_keys(bool *ok)
+{
+  std::vector<uint8_t> pac_data(
+      sizeof(ARM64Arch::user_pac_address_keys) +
+      sizeof(ARM64Arch::user_pac_generic_keys));
+  struct iovec vec = { pac_data.data(), sizeof(ARM64Arch::user_pac_address_keys) };
+  if (fallible_ptrace(PTRACE_GETREGSET, NT_ARM_PACA_KEYS, &vec)) {
+    if (ok) {
+      *ok = false;
+    }
+    return std::vector<uint8_t>{};
+  }
+  vec = { pac_data.data() + sizeof(ARM64Arch::user_pac_address_keys),
+          sizeof(ARM64Arch::user_pac_generic_keys) };
+  if (fallible_ptrace(PTRACE_GETREGSET, NT_ARM_PACG_KEYS, &vec)) {
+    if (ok) {
+      *ok = false;
+    }
+    return std::vector<uint8_t>{};
+  }
+  return pac_data;
+}
+bool Task::set_pac_keys(const std::vector<uint8_t> &pac_data)
+{
+  if (pac_data.empty()) {
+    return true;
+  }
+  struct iovec vec = { (void*)pac_data.data(), sizeof(ARM64Arch::user_pac_address_keys) };
+  if (fallible_ptrace(PTRACE_SETREGSET, NT_ARM_PACA_KEYS, &vec)) {
+    return false;
+  }
+  vec = { (void*)(pac_data.data() + sizeof(ARM64Arch::user_pac_address_keys)),
+          sizeof(ARM64Arch::user_pac_generic_keys) };
+  return !fallible_ptrace(PTRACE_SETREGSET, NT_ARM_PACG_KEYS, &vec);
+}
 #else
+std::vector<uint8_t> Task::pac_keys(bool *ok)
+{
+  return std::vector<uint8_t>{};
+}
+bool Task::set_pac_keys(const std::vector<uint8_t> &pac_data)
+{
+  return true;
+}
 bool Task::set_aarch64_debug_regs(int, ARM64Arch::user_hwdebug_state *, size_t) {
   FATAL() << "Reached aarch64 code path on non-aarch64 system";
   return false;

src/Task.h

@@ -528,6 +528,16 @@ class Task {
    */
   void set_x86_debug_status(uintptr_t status);
 
+  /**
+   * Read the (architecture-specific) pointer authentication keys of the current task
+   */
+  std::vector<uint8_t> pac_keys(bool *ok = nullptr);
+
+  /**
+   * Set the (architecture-specific) pointer authentication keys for the current task
+   */
+  bool set_pac_keys(const std::vector<uint8_t> &data);
+
   /**
    * Determine why a SIGTRAP occurred. On x86, uses x86_debug_status() but doesn't
    * consume it.

src/TraceStream.cc

@@ -778,6 +778,9 @@ void TraceWriter::write_task_event(const TraceTaskEvent& event) {
       exec.setExeBase(event.exe_base().as_int());
       exec.setInterpBase(event.interp_base().as_int());
       exec.setInterpName(str_to_data(event.interp_name()));;
+      auto pac_data = exec.initPacData();
+      std::vector<uint8_t> pac_data_vec = event.pac_data();
+      pac_data.setRaw(Data::Reader(pac_data_vec.data(), pac_data_vec.size()));
       break;
     }
     case TraceTaskEvent::EXIT:
@@ -840,6 +843,8 @@ TraceTaskEvent TraceReader::read_task_event(FrameTime* time) {
       r.exe_base_ = exec.getExeBase();
       r.interp_base_ = exec.getInterpBase();
       r.interp_name_ = data_to_str(exec.getInterpName());
+      auto pac_data = exec.getPacData().getRaw();
+      r.pac_data_ = std::vector<uint8_t>(pac_data.begin(), pac_data.end());
       break;
     }
     case trace::TaskEvent::Which::EXIT:

src/TraceTaskEvent.h

@@ -105,6 +105,14 @@ class TraceTaskEvent {
     DEBUG_ASSERT(type() == EXEC);
     interp_name_ = name;
   }
+  std::vector<uint8_t> pac_data() const {
+    DEBUG_ASSERT(type() == EXEC);
+    return pac_data_;
+  }
+  void set_pac_data(std::vector<uint8_t> data) {
+    DEBUG_ASSERT(type() == EXEC);
+    pac_data_ = data;
+  }
   WaitStatus exit_status() const {
     DEBUG_ASSERT(type() == EXIT);
     return exit_status_;
@@ -124,6 +132,7 @@ class TraceTaskEvent {
   remote_ptr<void> exe_base_;         // EXEC only
   remote_ptr<void> interp_base_;      // EXEC only
   std::string interp_name_;           // EXEC only
+  std::vector<uint8_t> pac_data_;     // EXEC only
   WaitStatus exit_status_;            // EXIT only
 };
 

src/kernel_abi.h

@@ -2560,6 +2560,17 @@ struct ARM64Arch : public GenericArch<SupportedArch::aarch64, WordSize64Defs> {
     uint32_t rule_locs[0];
   };
   RR_VERIFY_TYPE_ARCH(SupportedArch::aarch64, struct ::ethtool_rxnfc, struct ethtool_rxnfc);
+
+  struct user_pac_address_keys {
+    __uint128_t apiakey;
+    __uint128_t apibkey;
+    __uint128_t	apdakey;
+    __uint128_t apdbkey;
+  };
+
+  struct user_pac_generic_keys {
+    __uint128_t apgakey;
+  };
 };
 
 #define RR_ARCH_FUNCTION(f, arch, args...)                                     \

src/record_syscall.cc

@@ -5772,6 +5772,7 @@ static void process_execve(RecordTask* t, TaskSyscallState& syscall_state) {
     }
   }
 
+  syscall_state.exec_saved_event->set_pac_data(t->pac_keys());
   t->session().trace_writer().write_task_event(*syscall_state.exec_saved_event);
 
   {

src/replay_syscall.cc

@@ -450,6 +450,9 @@ static void process_execve(ReplayTask* t, const TraceFrame& trace_frame,
   t->post_exec_syscall(exe_name, kms[exe_km].fsname());
   t->vm()->set_interp_base(tte.interp_base());
   t->vm()->set_interp_name(tte.interp_name());
+  if (!t->set_pac_keys(tte.pac_data())) {
+    LOG(warn) << "Failed to restore PAC keys. Replay may fail.";
+  }
 
   t->fd_table()->close_after_exec(
       t, t->current_trace_frame().event().Syscall().exec_fds_to_close);

src/rr_trace.capnp

@@ -218,6 +218,7 @@ struct TaskEvent {
       interpBase @10 :RemotePtr;
       # Not a Path since it is only meaningful during recording
       interpName @11 :CString;
+      pacData @12 :PACData;
     }
     # Most frame 'exit' events generate one of these, but these are not
     # generated if rr ends abnormally so the tasks did not in fact exit during
@@ -261,6 +262,11 @@ struct ExtraRegisters {
   raw @0 :Data;
 }
 
+struct PACData {
+  # Formay determined by Frame::arch
+  raw @0 :Data;
+}
+
 enum SyscallState {
   enteringPtrace @0;
   entering @1;