Address missing arch() serialization

Previous work assumed everything as native. That doesn't work when
debugging a 32-bit mode application on a 64-bit machine. That means some
`SupportedArch` fields needs to be serialized as well, so that if the
process was a 32-bit application it gets re-created correctly. This was
instantly made visible by the provided test (although it did not test it
explicitly).

Author: theIDinside

src/CheckpointInfo.cc

@@ -3,6 +3,7 @@
 #include "ReplayTimeline.h"
 #include "ScopedFd.h"
 #include "rr_pcp.capnp.h"
+#include "util.h"
 #include <algorithm>
 #include <capnp/blob.h>
 #include <capnp/message.h>
@@ -22,9 +23,10 @@ MarkData::MarkData(const ReplayTimeline::Mark& m)
       extra_regs(m.extra_regs()),
       return_addresses(m.get_internal()->proto.return_addresses),
       singlestep_to_next_mark_no_signal(
-          m.get_internal()->singlestep_to_next_mark_no_signal) {}
+          m.get_internal()->singlestep_to_next_mark_no_signal),
+      arch(m.get_internal()->extra_regs.arch()) {}
 
-MarkData::MarkData(rr::pcp::MarkData::Reader reader, SupportedArch arch,
+MarkData::MarkData(rr::pcp::MarkData::Reader reader,
                    const CPUIDRecords& cpuid_recs)
     : time(reader.getTime()),
       ticks(reader.getTicks()),
@@ -34,7 +36,8 @@ MarkData::MarkData(rr::pcp::MarkData::Reader reader, SupportedArch arch,
       extra_regs(),
       return_addresses(),
       singlestep_to_next_mark_no_signal(
-          reader.getSinglestepToNextMarkNoSignal()) {
+          reader.getSinglestepToNextMarkNoSignal()),
+      arch(from_trace_arch(reader.getArch())) {
   regs.set_arch(arch);
   regs.set_from_trace(arch, reader.getRegs().getRaw().begin(),
                       reader.getRegs().getRaw().size());
@@ -69,8 +72,7 @@ static std::vector<string> checkpoint_directories(const string& trace_dir) {
 }
 
 std::vector<CheckpointInfo> get_checkpoint_infos(
-    const std::string& trace_dir, SupportedArch arch,
-    const CPUIDRecords& cpuid_recs) {
+    const std::string& trace_dir, const CPUIDRecords& cpuid_recs) {
 
   std::vector<CheckpointInfo> checkpoints;
   for (auto checkpoint_dir : checkpoint_directories(trace_dir)) {
@@ -81,8 +83,8 @@ std::vector<CheckpointInfo> get_checkpoint_infos(
     }
     capnp::PackedFdMessageReader reader(fd);
     auto checkpointsInfoReader = reader.getRoot<pcp::CheckpointInfo>();
-    auto info = CheckpointInfo{ checkpoint_dir, checkpointsInfoReader, arch,
-                                cpuid_recs };
+    auto info =
+        CheckpointInfo{ checkpoint_dir, checkpointsInfoReader, cpuid_recs };
     checkpoints.push_back(info);
   }
 
@@ -138,6 +140,7 @@ bool CheckpointInfo::serialize(ReplaySession& session) {
     builder.setTicksAtEventStart(mark_data.ticks_at_event_start);
     builder.setSinglestepToNextMarkNoSignal(
         mark_data.singlestep_to_next_mark_no_signal);
+    builder.setArch(to_trace_arch(mark_data.arch));
   };
 
   if (is_explicit()) {
@@ -220,18 +223,17 @@ CheckpointInfo::CheckpointInfo(const Checkpoint& non_explicit_cp,
 
 CheckpointInfo::CheckpointInfo(std::string metadata_file,
                                rr::pcp::CheckpointInfo::Reader reader,
-                               SupportedArch arch,
                                const CPUIDRecords& cpuid_recs)
     : capnp_directory(std::move(metadata_file)),
       unique_id(reader.getId()),
       where(data_to_str(reader.getWhere())),
       next_serial(reader.getNextSerial()),
       clone_data(reader.isExplicit() ? reader.getExplicit()
                                      : reader.getNonExplicit().getCloneMark(),
-                 arch, cpuid_recs),
+                 cpuid_recs),
       non_explicit_mark_data(
           reader.isNonExplicit()
-              ? new MarkData{ reader.getNonExplicit().getCheckpointMark(), arch,
+              ? new MarkData{ reader.getNonExplicit().getCheckpointMark(),
                               cpuid_recs }
               : nullptr),
       stats() {

src/CheckpointInfo.h

@@ -5,6 +5,7 @@
 #include "GdbServerConnection.h"
 #include "ReplayTimeline.h"
 #include "ReturnAddressList.h"
+#include "kernel_abi.h"
 #include "rr_pcp.capnp.h"
 #include "util.h"
 #include <dirent.h>
@@ -31,8 +32,7 @@ struct MarkData {
   // Constructor when serializing
   MarkData(const ReplayTimeline::Mark& m);
   // Constructor when de-serializing
-  MarkData(rr::pcp::MarkData::Reader reader, SupportedArch arch,
-           const CPUIDRecords& cpuid_recs);
+  MarkData(rr::pcp::MarkData::Reader reader, const CPUIDRecords& cpuid_recs);
 
   FrameTime time;
   Ticks ticks;
@@ -42,6 +42,7 @@ struct MarkData {
   ExtraRegisters extra_regs;
   ReturnAddressList return_addresses;
   bool singlestep_to_next_mark_no_signal;
+  SupportedArch arch;
 };
 
 class CheckpointInfo {
@@ -66,7 +67,7 @@ class CheckpointInfo {
                  const ReplayTimeline::Mark& mark_with_checkpoint);
   // When deserializing from capnproto stream
   CheckpointInfo(std::string metadata_file,
-                 rr::pcp::CheckpointInfo::Reader reader, SupportedArch arch,
+                 rr::pcp::CheckpointInfo::Reader reader,
                  const CPUIDRecords& cpuid_recs);
 
   bool serialize(ReplaySession& session);
@@ -116,7 +117,6 @@ std::string checkpoints_index_file(const std::string& trace_dir);
  * order by event time.
  */
 std::vector<CheckpointInfo> get_checkpoint_infos(
-    const std::string& trace_dir, SupportedArch arch,
-    const CPUIDRecords& cpuid_recs);
+    const std::string& trace_dir, const CPUIDRecords& cpuid_recs);
 
 } // namespace rr

src/PersistentCheckpointing.cc

@@ -24,6 +24,7 @@
 #include "TraceFrame.h"
 #include "TraceStream.h"
 #include "VirtualPerfCounterMonitor.h"
+#include "kernel_abi.h"
 #include "log.h"
 #include "replay_syscall.h"
 #include "rr_pcp.capnp.h"
@@ -165,7 +166,7 @@ static void write_map(const WriteVmConfig& cfg,
                           << map.map.size();
 
     auto bytes_read = 0ull;
-    while(static_cast<size_t>(bytes_read) < map.map.size()) {
+    while (static_cast<size_t>(bytes_read) < map.map.size()) {
       const auto current_read = cfg.pread(bytes_read, map.map);
       if (current_read == -1)
         FILE_OP_FATAL(file) << " couldn't read contents of " << map.map.str();
@@ -309,20 +310,17 @@ void map_private_anonymous(AutoRemoteSyscalls& remote,
 }
 
 Task::CapturedState reconstitute_captured_state(
-    ReplaySession& s, pcp::CapturedState::Reader reader) {
+    SupportedArch arch, const std::vector<CPUIDRecord>& cpuid_records,
+    pcp::CapturedState::Reader reader) {
   Task::CapturedState res;
   res.ticks = reader.getTicks();
-  {
-    auto register_raw = reader.getRegs().getRaw();
-    res.regs = Registers{ s.arch() };
-    res.regs.set_from_trace(s.arch(), register_raw.begin(),
-                            register_raw.size());
-  }
-  {
-    auto raw = reader.getExtraRegs().getRaw();
-    set_extra_regs_from_raw(s.arch(), s.trace_reader().cpuid_records(), raw,
-                            res.extra_regs);
-  }
+  auto register_raw = reader.getRegs().getRaw();
+  res.regs = Registers{ arch };
+  res.regs.restore_from_persistent_checkpoint(arch, register_raw.begin(),
+                                              register_raw.size());
+
+  auto raw = reader.getExtraRegs().getRaw();
+  set_extra_regs_from_raw(arch, cpuid_records, raw, res.extra_regs);
 
   res.prname = data_to_str(reader.getPrname());
   res.fdtable_identity = reader.getFdtableIdentity();
@@ -380,6 +378,13 @@ void init_scratch_memory(ReplayTask* t, const KernelMapping& km) {
   }
 }
 
+kj::Array<capnp::byte> prepare_user_desc(const X86Arch::user_desc& desc) {
+  kj::Array<capnp::byte> data =
+      kj::heapArray<capnp::byte>(sizeof(X86Arch::user_desc));
+  memcpy(data.begin(), &desc, sizeof(X86Arch::user_desc)); // Copy raw bytes
+  return data;
+}
+
 void write_capture_state(pcp::CapturedState::Builder& sb,
                          const Task::CapturedState& state) {
   sb.setTicks(state.ticks);
@@ -421,7 +426,9 @@ void write_capture_state(pcp::CapturedState::Builder& sb,
   auto thread_areas = sb.initThreadAreas(state.thread_areas.size());
   auto i = 0;
   for (const auto& ta : state.thread_areas) {
-    thread_areas[i++] = capnp::Data::Builder{ (std::uint8_t*)&ta, sizeof(ta) };
+    thread_areas.set(
+        i++, kj::ArrayPtr<const capnp::byte>(
+                 reinterpret_cast<const capnp::byte*>(&ta), sizeof(ta)));
   }
 }
 

src/PersistentCheckpointing.h

@@ -2,6 +2,7 @@
 
 #include "AddressSpace.h"
 #include "CheckpointInfo.h"
+#include "kernel_abi.h"
 #include "log.h"
 #include "rr_pcp.capnp.h"
 #include <capnp/blob.h>
@@ -57,7 +58,8 @@ void write_monitor(rr::pcp::FileMonitor::Builder& builder, int fd,
  * Restores Task::CapturedState from capnproto data.
  */
 Task::CapturedState reconstitute_captured_state(
-    ReplaySession& s, pcp::CapturedState::Reader reader);
+    SupportedArch arch, const std::vector<CPUIDRecord>& cpuid_records,
+    pcp::CapturedState::Reader reader);
 
 void map_private_anonymous(AutoRemoteSyscalls& remote, const KernelMapping& km);
 

src/Registers.cc

@@ -10,6 +10,7 @@
 
 #include "ReplayTask.h"
 #include "core.h"
+#include "kernel_abi.h"
 #include "log.h"
 
 using namespace std;
@@ -625,6 +626,25 @@ void Registers::set_from_ptrace_for_arch(SupportedArch a, const void* data,
   memcpy(&u.x86regs, data, sizeof(u.x86regs));
 }
 
+void Registers::restore_from_persistent_checkpoint(SupportedArch arch,
+                                                   const void* data,
+                                                   size_t size) {
+  switch (arch) {
+    case x86:
+      DEBUG_ASSERT(sizeof(u.x86regs) == size);
+      memcpy(&u.x86regs, data, size);
+      break;
+    case x86_64:
+      DEBUG_ASSERT(sizeof(u.x64regs) == size);
+      memcpy(&u.x64regs, data, size);
+      break;
+    case aarch64:
+      DEBUG_ASSERT(sizeof(u.x64regs) == size);
+      memcpy(&u.arm64regs, data, size);
+      break;
+  }
+}
+
 void Registers::set_from_trace(SupportedArch a, const void* data,
                                size_t size) {
   if (is_x86ish(a)) {

src/Registers.h

@@ -124,6 +124,9 @@ class Registers {
   void set_from_trace(SupportedArch arch, const void* data,
                       size_t size);
 
+  void restore_from_persistent_checkpoint(SupportedArch arch, const void* data,
+                                          size_t size);
+
 #define ARCH_SWITCH_CASE(rettype, x86case, x64case, arm64case)                 \
 (([=](void) -> rettype {                                                       \
   switch (arch()) {                                                            \

src/ReplaySession.cc

@@ -2260,6 +2260,7 @@ void ReplaySession::serialize_checkpoint(
     auto addr_space_clone = addr_space_builders[i];
     addr_space_clone.setAuxv(kj::ArrayPtr<const capnp::byte>{
         leader->vm()->saved_auxv().data(), leader->vm()->saved_auxv().size() });
+    addr_space_clone.setArch(to_trace_arch(as.clone_leader->arch()));
     auto cls = addr_space_clone.initCloneLeaderState();
     write_capture_state(cls, as.clone_leader_state);
     auto pspace = addr_space_builders[i].initProcessSpace();
@@ -2334,8 +2335,9 @@ void ReplaySession::load_checkpoint(const CheckpointInfo& cp_info) {
     const auto proc_space = as.getProcessSpace();
     const auto cleader_captured_state = as.getCloneLeaderState();
 
+    SupportedArch address_space_arch = from_trace_arch(as.getArch());
     leader->is_stopped_ = true;
-    leader->os_exec_stub(arch());
+    leader->os_exec_stub(address_space_arch);
     std::string exe_name = data_to_str(proc_space.getExe());
     std::string original_exe_name = data_to_str(proc_space.getOriginalExe());
     leader->post_exec(original_exe_name);
@@ -2465,9 +2467,9 @@ void ReplaySession::load_checkpoint(const CheckpointInfo& cp_info) {
     leader->vm()->restore_auxv(leader, std::move(auxv));
     syscall(SYS_rrcall_reload_auxv, leader->tid);
     std::vector<Task::CapturedState> member_states;
-
     for (const auto& member_state : as.getMemberState()) {
-      member_states.push_back(reconstitute_captured_state(*this, member_state));
+      member_states.push_back(reconstitute_captured_state(
+          address_space_arch, trace_in.cpuid_records(), member_state));
     }
 
     CapturedMemory captured_memory;
@@ -2478,9 +2480,8 @@ void ReplaySession::load_checkpoint(const CheckpointInfo& cp_info) {
       captured_memory.push_back(
           std::make_pair(captured_mem.getStartAddress(), std::move(mem)));
     }
-
-    Task::CapturedState cloneLeaderCaptureState =
-        reconstitute_captured_state(*this, as.getCloneLeaderState());
+    Task::CapturedState cloneLeaderCaptureState = reconstitute_captured_state(
+        address_space_arch, trace_in.cpuid_records(), as.getCloneLeaderState());
     auto fd_table_key = cloneLeaderCaptureState.fdtable_identity;
     leader->preload_globals = cloneLeaderCaptureState.preload_globals;
     partial_init_addr_spaces.push_back(CloneCompletion::AddressSpaceClone{
@@ -2574,7 +2575,7 @@ void ReplaySession::load_checkpoint(const CheckpointInfo& cp_info) {
 
 std::vector<CheckpointInfo> ReplaySession::get_persistent_checkpoints() {
   return rr::get_checkpoint_infos(resolve_trace_name(trace_reader().dir()),
-                                  arch(), trace_reader().cpuid_records());
+                                  trace_reader().cpuid_records());
 }
 
 void ReplaySession::restore_session_info(const CheckpointInfo& cp) {

src/TraceStream.cc

@@ -162,20 +162,6 @@ static trace::CpuTriState to_tristate(bool value) {
   return value ? trace::CpuTriState::KNOWN_TRUE : trace::CpuTriState::KNOWN_FALSE;
 }
 
-static SupportedArch from_trace_arch(trace::Arch arch) {
-  switch (arch) {
-    case trace::Arch::X86:
-      return x86;
-    case trace::Arch::X8664:
-      return x86_64;
-    case trace::Arch::AARCH64:
-      return aarch64;
-    default:
-      FATAL() << "Unknown arch";
-      return x86;
-  }
-}
-
 static trace::SignalDisposition to_trace_disposition(
     SignalResolvedDisposition disposition) {
   switch (disposition) {

src/rr_pcp.capnp

@@ -141,6 +141,8 @@ struct AddressSpaceClone {
   memberState @2 :List(CapturedState);
   capturedMemory @3 :List(CapturedMemory);
   auxv @4 :Data;
+  # We need to know how to reconstitute the Register/ExtraRegister's in CapturedState
+  arch @5 :Arch;
 }
 
 struct CloneCompletionInfo {
@@ -161,6 +163,8 @@ struct MarkData {
   returnAddresses @5 :List(RemotePtr);
   extraRegs @6: ExtraRegisters;
   singlestepToNextMarkNoSignal @7 :Bool;
+  # The arch required to configure regs and extraRegs with the peristed data
+  arch @8 :Arch;
 }
 
 

src/util.cc

@@ -2734,6 +2734,20 @@ std::string default_rr_trace_dir() {
   return cached_dir;
 }
 
+SupportedArch from_trace_arch(trace::Arch arch) {
+  switch (arch) {
+    case trace::Arch::X86:
+      return x86;
+    case trace::Arch::X8664:
+      return x86_64;
+    case trace::Arch::AARCH64:
+      return aarch64;
+    default:
+      FATAL() << "Unknown arch";
+      return x86;
+  }
+}
+
 trace::Arch to_trace_arch(SupportedArch arch) {
   switch (arch) {
     case x86: