[sanitizer] Support sandboxing in sanitizer coverage.

Summary:
Sandboxed code may now pass additional arguments to
__sanitizer_sandbox_on_notify() to force all coverage data to be dumped to a
single file (the default is one file per module). The user may supply a file or
socket to write to. The latter option can be used to broker out the file writing
functionality. If -1  is passed, we pre-open a file.

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@209121 91177308-0d34-0410-b5e6-96231b3b80d8

8 files changed, 159 insertions, 31 deletions

diff --git a/lib/sanitizer_common/sanitizer_common.cc b/lib/sanitizer_common/sanitizer_common.cc
index c0c7748a7..305056e24 100644
--- a/lib/sanitizer_common/sanitizer_common.cc
+++ b/lib/sanitizer_common/sanitizer_common.cc
@@ -300,9 +300,9 @@ void __sanitizer_set_report_path(const char *path) {
   }
 }
 
-void NOINLINE __sanitizer_sandbox_on_notify(void *reserved) {
-  (void)reserved;
-  PrepareForSandboxing();
+void NOINLINE
+__sanitizer_sandbox_on_notify(__sanitizer_sandbox_arguments *args) {
+  PrepareForSandboxing(args);
   if (sandboxing_callback)
     sandboxing_callback();
 }
diff --git a/lib/sanitizer_common/sanitizer_common.h b/lib/sanitizer_common/sanitizer_common.h
index 8fa41679f..ab3cc4522 100644
--- a/lib/sanitizer_common/sanitizer_common.h
+++ b/lib/sanitizer_common/sanitizer_common.h
@@ -183,7 +183,8 @@ void ReExec();
 bool StackSizeIsUnlimited();
 void SetStackSizeLimitInBytes(uptr limit);
 void AdjustStackSize(void *attr);
-void PrepareForSandboxing();
+void PrepareForSandboxing(__sanitizer_sandbox_arguments *args);
+void CovPrepareForSandboxing(__sanitizer_sandbox_arguments *args);
 void SetSandboxingCallback(void (*f)());
 
 void InitTlsSize();
diff --git a/lib/sanitizer_common/sanitizer_coverage.cc b/lib/sanitizer_common/sanitizer_coverage.cc
index a0fecfba6..cb0644f6d 100644
--- a/lib/sanitizer_common/sanitizer_coverage.cc
+++ b/lib/sanitizer_common/sanitizer_coverage.cc
@@ -43,7 +43,7 @@
 atomic_uint32_t dump_once_guard;  // Ensure that CovDump runs only once.
 
 // pc_array is the array containing the covered PCs.
-// To make the pc_array thread- and AS- safe it has to be large enough.
+// To make the pc_array thread- and async-signal-safe it has to be large enough.
 // 128M counters "ought to be enough for anybody" (4M on 32-bit).
 // pc_array is allocated with MmapNoReserveOrDie and so it uses only as
 // much RAM as it really needs.
@@ -51,6 +51,10 @@ static const uptr kPcArraySize = FIRST_32_SECOND_64(1 << 22, 1 << 27);
 static uptr *pc_array;
 static atomic_uintptr_t pc_array_index;
 
+static bool cov_sandboxed = false;
+static int cov_fd = kInvalidFd;
+static unsigned int cov_max_block_size = 0;
+
 namespace __sanitizer {
 
 // Simply add the pc into the vector under lock. If the function is called more
@@ -71,8 +75,56 @@ static inline bool CompareLess(const uptr &a, const uptr &b) {
   return a < b;
 }
 
+// Block layout for packed file format: header, followed by module name (no
+// trailing zero), followed by data blob.
+struct CovHeader {
+  int pid;
+  unsigned int module_name_length;
+  unsigned int data_length;
+};
+
+static void CovWritePacked(int pid, const char *module, const void *blob,
+                           unsigned int blob_size) {
+  CHECK_GE(cov_fd, 0);
+  unsigned module_name_length = internal_strlen(module);
+  CovHeader header = {pid, module_name_length, blob_size};
+
+  if (cov_max_block_size == 0) {
+    // Writing to a file. Just go ahead.
+    internal_write(cov_fd, &header, sizeof(header));
+    internal_write(cov_fd, module, module_name_length);
+    internal_write(cov_fd, blob, blob_size);
+  } else {
+    // Writing to a socket. We want to split the data into appropriately sized
+    // blocks.
+    InternalScopedBuffer<char> block(cov_max_block_size);
+    CHECK_EQ((uptr)block.data(), (uptr)(CovHeader *)block.data());
+    uptr header_size_with_module = sizeof(header) + module_name_length;
+    CHECK_LT(header_size_with_module, cov_max_block_size);
+    unsigned int max_payload_size =
+        cov_max_block_size - header_size_with_module;
+    char *block_pos = block.data();
+    internal_memcpy(block_pos, &header, sizeof(header));
+    block_pos += sizeof(header);
+    internal_memcpy(block_pos, module, module_name_length);
+    block_pos += module_name_length;
+    char *block_data_begin = block_pos;
+    char *blob_pos = (char *)blob;
+    while (blob_size > 0) {
+      unsigned int payload_size = Min(blob_size, max_payload_size);
+      blob_size -= payload_size;
+      internal_memcpy(block_data_begin, blob_pos, payload_size);
+      blob_pos += payload_size;
+      ((CovHeader *)block.data())->data_length = payload_size;
+      internal_write(cov_fd, block.data(),
+                     header_size_with_module + payload_size);
+    }
+  }
+}
+
 // Dump the coverage on disk.
-void CovDump() {
+static void CovDump() {
+  if (!common_flags()->coverage) return;
 #if !SANITIZER_WINDOWS
   if (atomic_fetch_add(&dump_once_guard, 1, memory_order_relaxed))
     return;
@@ -81,7 +133,7 @@ void CovDump() {
   InternalMmapVector<u32> offsets(size);
   const uptr *vb = pc_array;
   const uptr *ve = vb + size;
-  MemoryMappingLayout proc_maps(/*cache_enabled*/false);
+  MemoryMappingLayout proc_maps(/*cache_enabled*/true);
   uptr mb, me, off, prot;
   InternalScopedBuffer<char> module(4096);
   InternalScopedBuffer<char> path(4096 * 2);
@@ -102,22 +154,57 @@ void CovDump() {
         offsets.push_back(static_cast<u32>(diff));
       }
       char *module_name = StripModuleName(module.data());
-      internal_snprintf((char *)path.data(), path.size(), "%s.%zd.sancov",
-                        module_name, internal_getpid());
-      InternalFree(module_name);
-      uptr fd = OpenFile(path.data(), true);
-      if (internal_iserror(fd)) {
-        Report(" CovDump: failed to open %s for writing\n", path.data());
+      if (cov_sandboxed) {
+        CovWritePacked(internal_getpid(), module_name, offsets.data(),
+                       offsets.size() * sizeof(u32));
+        VReport(1, " CovDump: %zd PCs written to packed file\n", vb - old_vb);
       } else {
-        internal_write(fd, offsets.data(), offsets.size() * sizeof(u32));
-        internal_close(fd);
-        VReport(1, " CovDump: %s: %zd PCs written\n", path.data(), vb - old_vb);
+        // One file per module per process.
+        internal_snprintf((char *)path.data(), path.size(), "%s.%zd.sancov",
+                          module_name, internal_getpid());
+        uptr fd = OpenFile(path.data(), true);
+        if (internal_iserror(fd)) {
+          Report(" CovDump: failed to open %s for writing\n", path.data());
+        } else {
+          internal_write(fd, offsets.data(), offsets.size() * sizeof(u32));
+          internal_close(fd);
+          VReport(1, " CovDump: %s: %zd PCs written\n", path.data(),
+                  vb - old_vb);
+        }
       }
+      InternalFree(module_name);
     }
   }
+  if (cov_fd >= 0)
+    internal_close(cov_fd);
 #endif  // !SANITIZER_WINDOWS
 }
 
+static void OpenPackedFileForWriting() {
+  CHECK(cov_fd == kInvalidFd);
+  InternalScopedBuffer<char> path(1024);
+  internal_snprintf((char *)path.data(), path.size(), "%zd.sancov.packed",
+                    internal_getpid());
+  uptr fd = OpenFile(path.data(), true);
+  if (internal_iserror(fd)) {
+    Report(" Coverage: failed to open %s for writing\n", path.data());
+    Die();
+  }
+  cov_fd = fd;
+}
+
+void CovPrepareForSandboxing(__sanitizer_sandbox_arguments *args) {
+  if (!args) return;
+  if (!common_flags()->coverage) return;
+  cov_sandboxed = args->coverage_sandboxed;
+  if (!cov_sandboxed) return;
+  cov_fd = args->coverage_fd;
+  cov_max_block_size = args->coverage_max_block_size;
+  if (cov_fd < 0)
+    // Pre-open the file now. The sandbox won't allow us to do it later.
+    OpenPackedFileForWriting();
+}
+
 }  // namespace __sanitizer
 
 extern "C" {
diff --git a/lib/sanitizer_common/sanitizer_internal_defs.h b/lib/sanitizer_common/sanitizer_internal_defs.h
index d5d01f1d3..86e999df5 100644
--- a/lib/sanitizer_common/sanitizer_internal_defs.h
+++ b/lib/sanitizer_common/sanitizer_internal_defs.h
@@ -91,11 +91,15 @@ extern "C" {
   SANITIZER_INTERFACE_ATTRIBUTE
   void __sanitizer_set_report_path(const char *path);
 
-  // Notify the tools that the sandbox is going to be turned on. The reserved
-  // parameter will be used in the future to hold a structure with functions
-  // that the tools may call to bypass the sandbox.
-  SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE
-  void __sanitizer_sandbox_on_notify(void *reserved);
+  typedef struct {
+      int coverage_sandboxed;
+      __sanitizer::sptr coverage_fd;
+      unsigned int coverage_max_block_size;
+  } __sanitizer_sandbox_arguments;
+
+  // Notify the tools that the sandbox is going to be turned on.
+  SANITIZER_INTERFACE_ATTRIBUTE SANITIZER_WEAK_ATTRIBUTE void
+      __sanitizer_sandbox_on_notify(__sanitizer_sandbox_arguments *args);
 
   // This function is called by the tool when it has just finished reporting
   // an error. 'error_summary' is a one-line string that summarizes
diff --git a/lib/sanitizer_common/sanitizer_linux.cc b/lib/sanitizer_common/sanitizer_linux.cc
index 2edcf1fc4..4ca3b5b68 100644
--- a/lib/sanitizer_common/sanitizer_linux.cc
+++ b/lib/sanitizer_common/sanitizer_linux.cc
@@ -390,7 +390,7 @@ void GetThreadStackAndTls(bool main, uptr *stk_addr, uptr *stk_size,
 }
 #endif  // SANITIZER_GO
 
-void PrepareForSandboxing() {
+void PrepareForSandboxing(__sanitizer_sandbox_arguments *args) {
   // Some kinds of sandboxes may forbid filesystem access, so we won't be able
   // to read the file mappings from /proc/self/maps. Luckily, neither the
   // process will be able to load additional libraries, so it's fine to use the
@@ -401,6 +401,7 @@ void PrepareForSandboxing() {
   if (Symbolizer *sym = Symbolizer::GetOrNull())
     sym->PrepareForSandboxing();
 #endif
+  CovPrepareForSandboxing(args);
 }
 
 enum MutexState {
diff --git a/lib/sanitizer_common/sanitizer_mac.cc b/lib/sanitizer_common/sanitizer_mac.cc
index 7a4a0d586..683888a1a 100644
--- a/lib/sanitizer_common/sanitizer_mac.cc
+++ b/lib/sanitizer_common/sanitizer_mac.cc
@@ -196,7 +196,8 @@ void ReExec() {
   UNIMPLEMENTED();
 }
 
-void PrepareForSandboxing() {
+void PrepareForSandboxing(__sanitizer_sandbox_arguments *args) {
+  (void)args;
   // Nothing here for now.
 }
 
diff --git a/lib/sanitizer_common/sanitizer_win.cc b/lib/sanitizer_common/sanitizer_win.cc
index ec6ce994b..697c59f00 100644
--- a/lib/sanitizer_common/sanitizer_win.cc
+++ b/lib/sanitizer_common/sanitizer_win.cc
@@ -192,7 +192,8 @@ void ReExec() {
   UNIMPLEMENTED();
 }
 
-void PrepareForSandboxing() {
+void PrepareForSandboxing(__sanitizer_sandbox_arguments *args) {
+  (void)args;
   // Nothing here for now.
 }
 
diff --git a/lib/sanitizer_common/scripts/sancov.py b/lib/sanitizer_common/scripts/sancov.py
index aa791bc4e..4354359b0 100755
--- a/lib/sanitizer_common/scripts/sancov.py
+++ b/lib/sanitizer_common/scripts/sancov.py
@@ -4,6 +4,7 @@
 # We need to merge these integers into a set and then
 # either print them (as hex) or dump them into another file.
 import array
+import struct
 import sys
 
 prog_name = "";
@@ -11,16 +12,16 @@ prog_name = "";
 def Usage():
   print >> sys.stderr, "Usage: \n" + \
       " " + prog_name + " merge file1 [file2 ...]  > output\n" \
-      " " + prog_name + " print file1 [file2 ...]\n"
+      " " + prog_name + " print file1 [file2 ...]\n" \
+      " " + prog_name + " unpack file1 [file2 ...]\n"
   exit(1)
 
 def ReadOneFile(path):
-  f = open(path, mode="rb")
-  f.seek(0, 2)
-  size = f.tell()
-  f.seek(0, 0)
-  s = set(array.array('I', f.read(size)))
-  f.close()
+  with open(path, mode="rb") as f:
+    f.seek(0, 2)
+    size = f.tell()
+    f.seek(0, 0)
+    s = set(array.array('I', f.read(size)))
   print >>sys.stderr, "%s: read %d PCs from %s" % (prog_name, size / 4, path)
   return s
 
@@ -44,6 +45,36 @@ def MergeAndPrint(files):
   a = array.array('I', s)
   a.tofile(sys.stdout)
 
+
+def UnpackOneFile(path):
+  with open(path, mode="rb") as f:
+    print >> sys.stderr, "%s: unpacking %s" % (prog_name, path)
+    while True:
+      header = f.read(12)
+      if not header: return
+      if len(header) < 12:
+        break
+      pid, module_length, blob_size = struct.unpack('iII', header)
+      module = f.read(module_length)
+      blob = f.read(blob_size)
+      assert(len(module) == module_length)
+      assert(len(blob) == blob_size)
+      extracted_file = "%s.%d.sancov" % (module, pid)
+      print >> sys.stderr, "%s: extracting %s" % \
+        (prog_name, extracted_file)
+      # The packed file may contain multiple blobs for the same pid/module
+      # pair. Append to the end of the file instead of overwriting.
+      with open(extracted_file, 'ab') as f2:
+        f2.write(blob)
+    # fail
+    raise Exception('Error reading file %s' % path)
+
+
+def Unpack(files):
+  for f in files:
+    UnpackOneFile(f)
+
+
 if __name__ == '__main__':
   prog_name = sys.argv[0]
   if len(sys.argv) <= 2:
@@ -52,5 +83,7 @@ if __name__ == '__main__':
     PrintFiles(sys.argv[2:])
   elif sys.argv[1] == "merge":
     MergeAndPrint(sys.argv[2:])
+  elif sys.argv[1] == "unpack":
+    Unpack(sys.argv[2:])
   else:
     Usage()