2 files changed, 60 insertions, 0 deletions

diff --git a/compiler-rt/lib/dfsan/dfsan_interceptors.cpp b/compiler-rt/lib/dfsan/dfsan_interceptors.cpp
index 12b74df..5ab7c2b 100644
--- a/compiler-rt/lib/dfsan/dfsan_interceptors.cpp
+++ b/compiler-rt/lib/dfsan/dfsan_interceptors.cpp
@@ -46,12 +46,25 @@ INTERCEPTOR(void *, mmap64, void *addr, SIZE_T length, int prot, int flags,
   return res;
 }
 
+INTERCEPTOR(int, munmap, void *addr, SIZE_T length) {
+  int res = REAL(munmap)(addr, length);
+  if (res != -1) {
+    uptr beg_shadow_addr = (uptr)__dfsan::shadow_for(addr);
+    void *end_addr =
+        (void *)((uptr)addr + RoundUpTo(length, GetPageSizeCached()));
+    uptr end_shadow_addr = (uptr)__dfsan::shadow_for(end_addr);
+    ReleaseMemoryPagesToOS(beg_shadow_addr, end_shadow_addr);
+  }
+  return res;
+}
+
 namespace __dfsan {
 void InitializeInterceptors() {
   CHECK(!interceptors_initialized);
 
   INTERCEPT_FUNCTION(mmap);
   INTERCEPT_FUNCTION(mmap64);
+  INTERCEPT_FUNCTION(munmap);
 
   interceptors_initialized = true;
 }
diff --git a/compiler-rt/test/dfsan/munmap_release_shadow.c b/compiler-rt/test/dfsan/munmap_release_shadow.c
new file mode 100644
index 0000000..085844d
--- /dev/null
+++ b/compiler-rt/test/dfsan/munmap_release_shadow.c
@@ -0,0 +1,47 @@
+// RUN: %clang_dfsan %s -o %t && %run %t
+
+#include <assert.h>
+#include <sanitizer/dfsan_interface.h>
+#include <stdbool.h>
+#include <stdio.h>
+#include <string.h>
+#include <sys/mman.h>
+#include <unistd.h>
+
+size_t get_rss_kb() {
+  long rss = 0L;
+  FILE *f = NULL;
+  assert((f = fopen("/proc/self/statm", "r")));
+  assert(fscanf(f, "%*s%ld", &rss) == 1);
+  fclose(f);
+  return ((size_t)rss * (size_t)sysconf(_SC_PAGESIZE)) >> 10;
+}
+
+int main(int argc, char **argv) {
+  const size_t map_size = 100 << 20;
+  size_t before = get_rss_kb();
+
+  char *p = mmap(NULL, map_size, PROT_READ | PROT_WRITE,
+                 MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  const dfsan_label label = dfsan_create_label("l", 0);
+  char val = 0xff;
+  dfsan_set_label(label, &val, sizeof(val));
+  memset(p, val, map_size);
+  size_t after_mmap = get_rss_kb();
+
+  munmap(p, map_size);
+  size_t after_munmap = get_rss_kb();
+
+  fprintf(stderr, "RSS at start: %td, after mmap: %td, after mumap: %td\n",
+          before, after_mmap, after_munmap);
+
+  // The memory after mmap increases 3 times of map_size because the overhead of
+  // shadow memory is 2x.
+  const size_t mmap_cost_kb = 3 * (map_size >> 10);
+  assert(after_mmap >= before + mmap_cost_kb);
+  // OS does not release memory to the same level as the start of the program.
+  // The assert checks the memory after munmap up to a delta.
+  const size_t delta = 5000;
+  assert(after_munmap + mmap_cost_kb <= after_mmap + delta);
+  return 0;
+}