MSAN mark memory initialized only when it's actually copied

This reduces performance for MSAN builds in exchange for better sanitizer coverage.

PiperOrigin-RevId: 790741800
Change-Id: I05c653a648a86ce92149e9f3085609595f836335

Author: happyCoder92

sandboxed_api/call.h

@@ -40,6 +40,7 @@ constexpr uint32_t kMsgRecvFd = 0x107;
 constexpr uint32_t kMsgClose = 0x108;
 constexpr uint32_t kMsgReallocate = 0x109;
 constexpr uint32_t kMsgStrlen = 0x10A;
+constexpr uint32_t kMsgMarkMemoryInit = 0x10B;
 // Return:
 constexpr uint32_t kMsgReturn = 0x201;
 

sandboxed_api/client.cc

@@ -241,10 +241,6 @@ void HandleAllocMsg(const size_t size, FuncRet* ret) {
   VLOG(1) << "HandleAllocMsg: size=" << size;
 
   const void* allocated = malloc(size);
-  // Memory is copied to the pointer using an API that the memory sanitizer
-  // is blind to (process_vm_writev). Mark the memory as initialized here, so
-  // that the sandboxed code can still be tested using MSAN.
-  ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(allocated, size);
 
   ret->ret_type = v::Type::kPointer;
   ret->int_val = reinterpret_cast<uintptr_t>(allocated);
@@ -257,16 +253,20 @@ void HandleReallocMsg(uintptr_t ptr, size_t size, FuncRet* ret) {
           << ")";
 
   const void* reallocated = realloc(reinterpret_cast<void*>(ptr), size);
-  // Memory is copied to the pointer using an API that the memory sanitizer
-  // is blind to (process_vm_writev). Mark the memory as initialized here, so
-  // that the sandboxed code can still be tested using MSAN.
-  ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(reallocated, size);
 
   ret->ret_type = v::Type::kPointer;
   ret->int_val = reinterpret_cast<uintptr_t>(reallocated);
   ret->success = true;
 }
 
+void HandleMarkMemoryInit(uintptr_t ptr, size_t size, FuncRet* ret) {
+  ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(reinterpret_cast<void*>(ptr), size);
+
+  ret->ret_type = v::Type::kVoid;
+  ret->success = true;
+  ret->int_val = 0ULL;
+}
+
 // Handles requests to free memory previously allocated by HandleAllocMsg() and
 // HandleReallocMsg().
 void HandleFreeMsg(uintptr_t ptr, FuncRet* ret) {
@@ -399,6 +399,12 @@ void ServeRequest(sandbox2::Comms* comms) {
       VLOG(1) << "Received Client::kMsgStrlen message";
       HandleStrlen(comms, BytesAs<const char*>(bytes), &ret);
       break;
+    case comms::kMsgMarkMemoryInit:
+      VLOG(1) << "Received Client::kMsgMarkMemoryInit message";
+      {
+        auto req = BytesAs<comms::ReallocRequest>(bytes);
+        HandleMarkMemoryInit(req.old_addr, req.size, &ret);
+      }
       break;
     default:
       LOG(FATAL) << "Received unknown tag: " << tag;

sandboxed_api/rpcchannel.cc

@@ -196,4 +196,23 @@ absl::StatusOr<size_t> RPCChannel::Strlen(void* str) {
   return fret.int_val;
 }
 
+absl::Status RPCChannel::MarkMemoryInit(void* addr, size_t size) {
+#ifdef MEMORY_SANITIZER
+  absl::MutexLock lock(&mutex_);
+  comms::ReallocRequest req = {
+      .old_addr = reinterpret_cast<uintptr_t>(addr),
+      .size = size,
+  };
+  if (!comms_->SendTLV(comms::kMsgMarkMemoryInit, sizeof(comms::ReallocRequest),
+                       &req)) {
+    return absl::UnavailableError("Sending TLV value failed");
+  }
+  SAPI_ASSIGN_OR_RETURN(auto fret, Return(v::Type::kVoid));
+  if (!fret.success) {
+    return absl::UnavailableError("MarkMemoryInit() failed on the remote side");
+  }
+#endif
+  return absl::OkStatus();
+}
+
 }  // namespace sapi

sandboxed_api/rpcchannel.h

@@ -43,6 +43,9 @@ class RPCChannel {
   // Reallocates memory.
   absl::Status Reallocate(void* old_addr, size_t size, void** new_addr);
 
+  // Marks the memory as initialized (used with MSAN).
+  absl::Status MarkMemoryInit(void* addr, size_t size);
+
   // Frees memory.
   absl::Status Free(void* addr);
 

sandboxed_api/var_abstract.cc

@@ -131,6 +131,8 @@ absl::Status Var::TransferToSandboxee(RPCChannel* rpc_channel, pid_t pid) {
     return absl::UnavailableError("process_vm_writev: partial success");
   }
 
+  SAPI_RETURN_IF_ERROR(rpc_channel->MarkMemoryInit(GetRemote(), GetSize()));
+
   return absl::OkStatus();
 }