[dfsan] Add origin ABI wrappers

[stephan.yichao.zhao]

Summary

supported: bcmp, fstat, memcmp, stat, strcasecmp, strchr, strcmp,
strncasecmp, strncp, strpbrk

This is a part of https://reviews.llvm.org/D95835.

[morehouse]

compiler-rt/lib/dfsan/dfsan_custom.cpp

119

126–129

174

This function is always called after dfsan_strchr. Should we just combine them?

195

207

214–215

215

217–220

219

225

Is the casting necessary?

301

310

These lines are unreachable.

[stephan.yichao.zhao]

updated

compiler-rt/lib/dfsan/dfsan_custom.cpp

195

Thank you.

215

This is similar to dfsan_get_label_of_string_prefixs. It is used by multiple wrappers.
Fixed typo.

219

This function is called by other wrappers too: bcmp, strcmp, strncmp, ... so we give it a general name.
I fixed the typo prefixs to prefixes.

301

done. replaced all pos -> bytes_read

310

removed. and fixed other similar cases.

[morehouse]

compiler-rt/lib/dfsan/dfsan_custom.cpp

219

Yes, but ultimately it is getting the label from comparing two byte sequences (i.e. memcmp). I was just suggesting a shorter name, but don't feel too strongly about it.

302–303

348–349

369

For simplicity, can we combine this with dfsan_strcmp? For example we could make n a ssize_t and do a normal strcmp when it's -1.

452

Can we also combine this with dfsan_strcasecmp?

compiler-rt/test/dfsan/custom.cpp

1646–1647

Can we move test_fork up to keep the lexicographic ordering?

[stephan.yichao.zhao]

compiler-rt/lib/dfsan/dfsan_custom.cpp

219

I misunderstood that. renamed to dfsan_get_memcmp_label and dfsan_get_memcmp_origin.

369

Using ssize_t reduces the size of n this function can accept.

__dfsw_strncmp always ensures n > 0 before calling dfsan_strncmp. so we could use n==0 as a flag to distinguish the two cases.

[stephan.yichao.zhao]

refactored code

[morehouse]

LGTM

compiler-rt/lib/dfsan/dfsan_custom.cpp

-//===-- dfsan.cpp ---------------------------------------------------------===//
+//===-- dfsan_custom.cpp --------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file is a part of DataFlowSanitizer.
__dfsw_stat(const char *path, struct stat *buf, dfsan_label path_label,
             dfsan_label buf_label, dfsan_label *ret_label) {
   int ret = stat(path, buf);
   if (ret == 0)
     dfsan_set_label(0, buf, sizeof(struct stat));
   *ret_label = 0;
   return ret;
 }
 
+SANITIZER_INTERFACE_ATTRIBUTE int __dfso_stat(
+    const char *path, struct stat *buf, dfsan_label path_label,
+    dfsan_label buf_label, dfsan_label *ret_label, dfsan_origin path_origin,
+    dfsan_origin buf_origin, dfsan_origin *ret_origin) {
+  int ret = __dfsw_stat(path, buf, path_label, buf_label, ret_label);
+  return ret;
+}
+
 SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_fstat(int fd, struct stat *buf,
                                                dfsan_label fd_label,
                                                dfsan_label buf_label,
                                                dfsan_label *ret_label) {
   int ret = fstat(fd, buf);
   if (ret == 0)
     dfsan_set_label(0, buf, sizeof(struct stat));
   *ret_label = 0;
   return ret;
 }
 
-SANITIZER_INTERFACE_ATTRIBUTE char *__dfsw_strchr(const char *s, int c,
-                                                  dfsan_label s_label,
-                                                  dfsan_label c_label,
-                                                  dfsan_label *ret_label) {
+SANITIZER_INTERFACE_ATTRIBUTE int __dfso_fstat(
+    int fd, struct stat *buf, dfsan_label fd_label, dfsan_label buf_label,
+    dfsan_label *ret_label, dfsan_origin fd_origin, dfsan_origin buf_origin,
+    dfsan_origin *ret_origin) {
+  int ret = __dfsw_fstat(fd, buf, fd_label, buf_label, ret_label);
+  return ret;
+}
+
+static char *dfsan_strchr_with_label(const char *s, int c, size_t *bytes_read,
+                                     dfsan_label s_label, dfsan_label c_label,
+  char *match_pos = nullptr;
   for (size_t i = 0;; ++i) {
     if (s[i] == c || s[i] == 0) {
-      if (flags().strict_data_dependencies) {
-        *ret_label = s_label;
-      } else {
-        *ret_label = dfsan_union(dfsan_read_label(s, i + 1),
-                                 dfsan_union(s_label, c_label));
-      }
+      // If s[i] is the \0 at the end of the string, and \0 is not the
+      // character we are searching for, then return null.
+      *bytes_read = i + 1;
+      match_pos = s[i] == 0 && c != 0 ? nullptr : const_cast<char *>(s + i);
+      break;
+    }
+  }
+  if (flags().strict_data_dependencies)
+  else
+    *ret_label = dfsan_union(dfsan_read_label(s, *bytes_read),
+  return match_pos;
 
-      // If s[i] is the \0 at the end of the string, and \0 is not the
-      // character we are searching for, then return null.
-      if (s[i] == 0 && c != 0) {
-        return nullptr;
-      }
-      return const_cast<char *>(s + i);
-    }
+SANITIZER_INTERFACE_ATTRIBUTE char *__dfsw_strchr(const char *s, int c,
+                                                  dfsan_label s_label,
+                                                  dfsan_label c_label,
+                                                  dfsan_label *ret_label) {
+  size_t bytes_read;
+  return dfsan_strchr_with_label(s, c, &bytes_read, s_label, c_label,
+                                 ret_label);
+
+SANITIZER_INTERFACE_ATTRIBUTE char *__dfso_strchr(
+    const char *s, int c, dfsan_label s_label, dfsan_label c_label,
+    dfsan_label *ret_label, dfsan_origin s_origin, dfsan_origin c_origin,
+    dfsan_origin *ret_origin) {
+  size_t bytes_read;
+  char *r =
+      dfsan_strchr_with_label(s, c, &bytes_read, s_label, c_label, ret_label);
+  if (flags().strict_data_dependencies) {
+    *ret_origin = s_origin;
+  } else if (*ret_label) {
+    dfsan_origin o = dfsan_read_origin_of_first_taint(s, bytes_read);
+    *ret_origin = o ? o : (s_label ? s_origin : c_origin);
   }
+  return r;
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE char *__dfsw_strpbrk(const char *s,
                                                    const char *accept,
                                                    dfsan_label s_label,
                                                    dfsan_label accept_label,
                                                    dfsan_label *ret_label) {
   const char *ret = strpbrk(s, accept);
   if (flags().strict_data_dependencies) {
     *ret_label = ret ? s_label : 0;
   } else {
     size_t s_bytes_read = (ret ? ret - s : strlen(s)) + 1;
     *ret_label =
         dfsan_union(dfsan_read_label(s, s_bytes_read),
                     dfsan_union(dfsan_read_label(accept, strlen(accept) + 1),
                                 dfsan_union(s_label, accept_label)));
   }
   return const_cast<char *>(ret);
 }
 
+SANITIZER_INTERFACE_ATTRIBUTE char *__dfso_strpbrk(
+    const char *s, const char *accept, dfsan_label s_label,
+    dfsan_label accept_label, dfsan_label *ret_label, dfsan_origin s_origin,
+    dfsan_origin accept_origin, dfsan_origin *ret_origin) {
+  const char *ret = __dfsw_strpbrk(s, accept, s_label, accept_label, ret_label);
+  if (flags().strict_data_dependencies) {
+    if (ret)
+      *ret_origin = s_origin;
+  } else {
+    if (*ret_label) {
+      size_t s_bytes_read = (ret ? ret - s : strlen(s)) + 1;
+      dfsan_origin o = dfsan_read_origin_of_first_taint(s, s_bytes_read);
+      if (o) {
+        *ret_origin = o;
+      } else {
+        o = dfsan_read_origin_of_first_taint(accept, strlen(accept) + 1);
+        *ret_origin = o ? o : (s_label ? s_origin : accept_origin);
+      }
+    }
+  }
+  return const_cast<char *>(ret);
+}
+
 static int dfsan_memcmp_bcmp(const void *s1, const void *s2, size_t n,
-                             dfsan_label s1_label, dfsan_label s2_label,
-                             dfsan_label n_label, dfsan_label *ret_label) {
+                             size_t *bytes_read) {
   const char *cs1 = (const char *) s1, *cs2 = (const char *) s2;
   for (size_t i = 0; i != n; ++i) {
     if (cs1[i] != cs2[i]) {
-      if (flags().strict_data_dependencies) {
-        *ret_label = 0;
-      } else {
-        *ret_label = dfsan_union(dfsan_read_label(cs1, i + 1),
-                                 dfsan_read_label(cs2, i + 1));
-      }
+      *bytes_read = i + 1;
       return cs1[i] - cs2[i];
     }
   }
-
-  if (flags().strict_data_dependencies) {
-    *ret_label = 0;
-  } else {
-    *ret_label = dfsan_union(dfsan_read_label(cs1, n),
-                             dfsan_read_label(cs2, n));
-  }
+  *bytes_read = n;
+  return 0;
+
+static dfsan_label dfsan_get_memcmp_label(const void *s1, const void *s2,
+                                          size_t pos) {
+  if (flags().strict_data_dependencies)
     return 0;
+  return dfsan_union(dfsan_read_label(s1, pos), dfsan_read_label(s2, pos));
+}
+
+static void dfsan_get_memcmp_origin(const void *s1, const void *s2, size_t pos,
+                                    dfsan_label *ret_label,
+                                    dfsan_origin *ret_origin) {
+  *ret_label = dfsan_get_memcmp_label(s1, s2, pos);
+  if (*ret_label == 0)
+    return;
+  dfsan_origin o = dfsan_read_origin_of_first_taint(s1, pos);
+  *ret_origin = o ? o : dfsan_read_origin_of_first_taint(s2, pos);
+}
+
+static int dfsan_memcmp_bcmp_label(const void *s1, const void *s2, size_t n,
+                                   dfsan_label *ret_label) {
+  size_t bytes_read;
+  int r = dfsan_memcmp_bcmp(s1, s2, n, &bytes_read);
+  *ret_label = dfsan_get_memcmp_label(s1, s2, bytes_read);
+  return r;
+}
+
+static int dfsan_memcmp_bcmp_origin(const void *s1, const void *s2, size_t n,
+                                    dfsan_label *ret_label,
+                                    dfsan_origin *ret_origin) {
+  size_t bytes_read;
+  int r = dfsan_memcmp_bcmp(s1, s2, n, &bytes_read);
+  dfsan_get_memcmp_origin(s1, s2, bytes_read, ret_label, ret_origin);
+  return r;
 }
 
 DECLARE_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_memcmp, uptr caller_pc,
                               const void *s1, const void *s2, size_t n,
                               dfsan_label s1_label, dfsan_label s2_label,
                               dfsan_label n_label)
 
+DECLARE_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_origin_memcmp, uptr caller_pc,
+                              const void *s1, const void *s2, size_t n,
+                              dfsan_label s1_label, dfsan_label s2_label,
+                              dfsan_label n_label, dfsan_origin s1_origin,
+                              dfsan_origin s2_origin, dfsan_origin n_origin)
+
 SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_memcmp(const void *s1, const void *s2,
                                                 size_t n, dfsan_label s1_label,
                                                 dfsan_label s2_label,
                                                 dfsan_label n_label,
                                                 dfsan_label *ret_label) {
   CALL_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_memcmp, GET_CALLER_PC(), s1, s2, n,
                              s1_label, s2_label, n_label);
-  return dfsan_memcmp_bcmp(s1, s2, n, s1_label, s2_label, n_label, ret_label);
+  return dfsan_memcmp_bcmp_label(s1, s2, n, ret_label);
+}
+
+SANITIZER_INTERFACE_ATTRIBUTE int __dfso_memcmp(
+    const void *s1, const void *s2, size_t n, dfsan_label s1_label,
+    dfsan_label s2_label, dfsan_label n_label, dfsan_label *ret_label,
+    dfsan_origin s1_origin, dfsan_origin s2_origin, dfsan_origin n_origin,
+    dfsan_origin *ret_origin) {
+  CALL_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_origin_memcmp, GET_CALLER_PC(), s1,
+                             s2, n, s1_label, s2_label, n_label, s1_origin,
+                             s2_origin, n_origin);
+  return dfsan_memcmp_bcmp_origin(s1, s2, n, ret_label, ret_origin);
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_bcmp(const void *s1, const void *s2,
                                               size_t n, dfsan_label s1_label,
                                               dfsan_label s2_label,
                                               dfsan_label n_label,
                                               dfsan_label *ret_label) {
-  return dfsan_memcmp_bcmp(s1, s2, n, s1_label, s2_label, n_label, ret_label);
+  return dfsan_memcmp_bcmp_label(s1, s2, n, ret_label);
+}
+
+SANITIZER_INTERFACE_ATTRIBUTE int __dfso_bcmp(
+    const void *s1, const void *s2, size_t n, dfsan_label s1_label,
+    dfsan_label s2_label, dfsan_label n_label, dfsan_label *ret_label,
+    dfsan_origin s1_origin, dfsan_origin s2_origin, dfsan_origin n_origin,
+    dfsan_origin *ret_origin) {
+  return dfsan_memcmp_bcmp_origin(s1, s2, n, ret_label, ret_origin);
+}
+
+// When n == 0, compare strings without byte limit.
+// When n > 0, compare the first (at most) n bytes of s1 and s2.
+static int dfsan_strncmp(const char *s1, const char *s2, size_t n,
+                         size_t *bytes_read) {
+  for (size_t i = 0;; ++i) {
+    if (s1[i] != s2[i] || s1[i] == 0 || s2[i] == 0 || (n > 0 && i == n - 1)) {
+      *bytes_read = i + 1;
+      return s1[i] - s2[i];
+    }
+  }
 }
 
 DECLARE_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_strcmp, uptr caller_pc,
                               const char *s1, const char *s2,
                               dfsan_label s1_label, dfsan_label s2_label)
 
+DECLARE_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_origin_strcmp, uptr caller_pc,
+                              const char *s1, const char *s2,
+                              dfsan_label s1_label, dfsan_label s2_label,
+                              dfsan_origin s1_origin, dfsan_origin s2_origin)
+
 SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_strcmp(const char *s1, const char *s2,
                                                 dfsan_label s1_label,
                                                 dfsan_label s2_label,
                                                 dfsan_label *ret_label) {
   CALL_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_strcmp, GET_CALLER_PC(), s1, s2,
                              s1_label, s2_label);
-  for (size_t i = 0;; ++i) {
-    if (s1[i] != s2[i] || s1[i] == 0 || s2[i] == 0) {
-      if (flags().strict_data_dependencies) {
-        *ret_label = 0;
-      } else {
-        *ret_label = dfsan_union(dfsan_read_label(s1, i + 1),
-                                 dfsan_read_label(s2, i + 1));
-      }
-      return s1[i] - s2[i];
-    }
-  }
-  return 0;
+  size_t bytes_read;
+  int r = dfsan_strncmp(s1, s2, 0, &bytes_read);
+  *ret_label = dfsan_get_memcmp_label(s1, s2, bytes_read);
+  return r;
+
+SANITIZER_INTERFACE_ATTRIBUTE int __dfso_strcmp(
+    const char *s1, const char *s2, dfsan_label s1_label, dfsan_label s2_label,
+    dfsan_label *ret_label, dfsan_origin s1_origin, dfsan_origin s2_origin,
+    dfsan_origin *ret_origin) {
+  CALL_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_origin_strcmp, GET_CALLER_PC(), s1,
+                             s2, s1_label, s2_label, s1_origin, s2_origin);
+  size_t bytes_read;
+  int r = dfsan_strncmp(s1, s2, 0, &bytes_read);
+  dfsan_get_memcmp_origin(s1, s2, bytes_read, ret_label, ret_origin);
+  return r;
 }
 
-SANITIZER_INTERFACE_ATTRIBUTE int
-__dfsw_strcasecmp(const char *s1, const char *s2, dfsan_label s1_label,
-                  dfsan_label s2_label, dfsan_label *ret_label) {
+// When n == 0, compare strings without byte limit.
+// When n > 0, compare the first (at most) n bytes of s1 and s2.
+static int dfsan_strncasecmp(const char *s1, const char *s2, size_t n,
+                             size_t *bytes_read) {
   for (size_t i = 0;; ++i) {
     char s1_lower = tolower(s1[i]);
     char s2_lower = tolower(s2[i]);
 
-    if (s1_lower != s2_lower || s1[i] == 0 || s2[i] == 0) {
-      if (flags().strict_data_dependencies) {
-        *ret_label = 0;
-      } else {
-        *ret_label = dfsan_union(dfsan_read_label(s1, i + 1),
-                                 dfsan_read_label(s2, i + 1));
-      }
+    if (s1_lower != s2_lower || s1[i] == 0 || s2[i] == 0 ||
+        (n > 0 && i == n - 1)) {
+      *bytes_read = i + 1;
       return s1_lower - s2_lower;
     }
   }
-  return 0;
+}
+
+SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_strcasecmp(const char *s1,
+                                                    const char *s2,
+                                                    dfsan_label s1_label,
+                                                    dfsan_label s2_label,
+                                                    dfsan_label *ret_label) {
+  size_t bytes_read;
+  int r = dfsan_strncasecmp(s1, s2, 0, &bytes_read);
+  *ret_label = dfsan_get_memcmp_label(s1, s2, bytes_read);
+  return r;
+}
+
+SANITIZER_INTERFACE_ATTRIBUTE int __dfso_strcasecmp(
+    const char *s1, const char *s2, dfsan_label s1_label, dfsan_label s2_label,
+    dfsan_label *ret_label, dfsan_origin s1_origin, dfsan_origin s2_origin,
+    dfsan_origin *ret_origin) {
+  size_t bytes_read;
+  int r = dfsan_strncasecmp(s1, s2, 0, &bytes_read);
+  dfsan_get_memcmp_origin(s1, s2, bytes_read, ret_label, ret_origin);
+  return r;
 }
 
 DECLARE_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_strncmp, uptr caller_pc,
                               const char *s1, const char *s2, size_t n,
                               dfsan_label s1_label, dfsan_label s2_label,
                               dfsan_label n_label)
 
+DECLARE_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_origin_strncmp, uptr caller_pc,
+                              const char *s1, const char *s2, size_t n,
+                              dfsan_label s1_label, dfsan_label s2_label,
+                              dfsan_label n_label, dfsan_origin s1_origin,
+                              dfsan_origin s2_origin, dfsan_origin n_origin)
+
 SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_strncmp(const char *s1, const char *s2,
                                                  size_t n, dfsan_label s1_label,
                                                  dfsan_label s2_label,
                                                  dfsan_label n_label,
                                                  dfsan_label *ret_label) {
   if (n == 0) {
     *ret_label = 0;
     return 0;
   }
 
   CALL_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_strncmp, GET_CALLER_PC(), s1, s2,
                              n, s1_label, s2_label, n_label);
 
-  for (size_t i = 0;; ++i) {
-    if (s1[i] != s2[i] || s1[i] == 0 || s2[i] == 0 || i == n - 1) {
-      if (flags().strict_data_dependencies) {
-        *ret_label = 0;
-      } else {
-        *ret_label = dfsan_union(dfsan_read_label(s1, i + 1),
-                                 dfsan_read_label(s2, i + 1));
-      }
-      return s1[i] - s2[i];
-    }
-  }
-  return 0;
+  size_t bytes_read;
+  int r = dfsan_strncmp(s1, s2, n, &bytes_read);
+  *ret_label = dfsan_get_memcmp_label(s1, s2, bytes_read);
+  return r;
 }
 
-SANITIZER_INTERFACE_ATTRIBUTE int
-__dfsw_strncasecmp(const char *s1, const char *s2, size_t n,
-                   dfsan_label s1_label, dfsan_label s2_label,
-                   dfsan_label n_label, dfsan_label *ret_label) {
+SANITIZER_INTERFACE_ATTRIBUTE int __dfso_strncmp(
+    const char *s1, const char *s2, size_t n, dfsan_label s1_label,
+    dfsan_label s2_label, dfsan_label n_label, dfsan_label *ret_label,
+    dfsan_origin s1_origin, dfsan_origin s2_origin, dfsan_origin n_origin,
+    dfsan_origin *ret_origin) {
   if (n == 0) {
     *ret_label = 0;
     return 0;
   }
 
-  for (size_t i = 0;; ++i) {
-    char s1_lower = tolower(s1[i]);
-    char s2_lower = tolower(s2[i]);
+  CALL_WEAK_INTERCEPTOR_HOOK(dfsan_weak_hook_origin_strncmp, GET_CALLER_PC(),
+                             s1, s2, n, s1_label, s2_label, n_label, s1_origin,
+                             s2_origin, n_origin);
+
+  size_t bytes_read;
+  int r = dfsan_strncmp(s1, s2, n, &bytes_read);
+  dfsan_get_memcmp_origin(s1, s2, bytes_read, ret_label, ret_origin);
+  return r;
+}
 
-    if (s1_lower != s2_lower || s1[i] == 0 || s2[i] == 0 || i == n - 1) {
-      if (flags().strict_data_dependencies) {
-        *ret_label = 0;
-      } else {
-        *ret_label = dfsan_union(dfsan_read_label(s1, i + 1),
-                                 dfsan_read_label(s2, i + 1));
-      }
-      return s1_lower - s2_lower;
-    }
-  }
+SANITIZER_INTERFACE_ATTRIBUTE int __dfsw_strncasecmp(
+    const char *s1, const char *s2, size_t n, dfsan_label s1_label,
+    dfsan_label s2_label, dfsan_label n_label, dfsan_label *ret_label) {
+  if (n == 0) {
+    return 0;
+
+  size_t bytes_read;
+  int r = dfsan_strncasecmp(s1, s2, n, &bytes_read);
+  *ret_label = dfsan_get_memcmp_label(s1, s2, bytes_read);
+  return r;
+
+SANITIZER_INTERFACE_ATTRIBUTE int __dfso_strncasecmp(
+    const char *s1, const char *s2, size_t n, dfsan_label s1_label,
+    dfsan_label s2_label, dfsan_label n_label, dfsan_label *ret_label,
+    dfsan_origin s1_origin, dfsan_origin s2_origin, dfsan_origin n_origin,
+    dfsan_origin *ret_origin) {
+  if (n == 0) {
+    *ret_label = 0;
     return 0;
   }
 
+  size_t bytes_read;
+  int r = dfsan_strncasecmp(s1, s2, n, &bytes_read);
+  dfsan_get_memcmp_origin(s1, s2, bytes_read, ret_label, ret_origin);
+  return r;
+}
+
 SANITIZER_INTERFACE_ATTRIBUTE void *__dfsw_calloc(size_t nmemb, size_t size,
                                                   dfsan_label nmemb_label,
                                                   dfsan_label size_label,
                                                   dfsan_label *ret_label) {
   void *p = calloc(nmemb, size);
   dfsan_set_label(0, p, nmemb * size);
   *ret_label = 0;
   return p;

compiler-rt/test/dfsan/custom.cpp

 #define ASSERT_ORIGIN(data, origin)
 #endif
 
 #ifdef ORIGIN_TRACKING
 #define ASSERT_ORIGINS(ptr, size, origin)                         \
   for (int i = 0; i < size; ++i) {                                \
     assert(origin == dfsan_get_origin((long)(((char *)ptr)[i]))); \
   }
-#define ASSERT_ORIGINS(ptr, size, origin)
 #else
+#define ASSERT_ORIGINS(ptr, size, origin)
 #endif
 
 #ifdef ORIGIN_TRACKING
 #define ASSERT_INIT_ORIGIN(ptr, origin) \
   assert(origin == dfsan_get_init_origin(ptr))
 #else
 #define ASSERT_INIT_ORIGIN(ptr, origin)
 #endif

 #ifdef ORIGIN_TRACKING
 #define ASSERT_SAVED_ORIGINS(val)       \
   for (int i = 0; i < sizeof(val); ++i) \
     ASSERT_ORIGIN(((char *)(&val))[i], val##_o[i]);
 #else
 #define ASSERT_SAVED_ORIGINS(val)
 #endif
 
-#ifdef ORIGIN_TRACKING
-#define ASSERT_SAVED_N_ORIGINS(val, n) \
-  for (int i = 0; i < n; ++i)          \
-    ASSERT_ORIGIN(val[i], val##_o[i]);
-#else
-#define ASSERT_SAVED_N_ORIGINS(val, n)
-#endif
-
-#if !defined(ORIGIN_TRACKING)
 void test_stat() {
   int i = 1;
   dfsan_set_label(i_label, &i, sizeof(i));
 
   struct stat s;
   s.st_dev = i;
-  assert(0 == stat("/", &s));
+  DEFINE_AND_SAVE_ORIGINS(s)
+  int ret = stat("/", &s);
+  assert(0 == ret);
+  ASSERT_ZERO_LABEL(ret);
   ASSERT_ZERO_LABEL(s.st_dev);
+  ASSERT_SAVED_ORIGINS(s)
 
   s.st_dev = i;
-  assert(-1 == stat("/nonexistent", &s));
+  SAVE_ORIGINS(s)
+  ret = stat("/nonexistent", &s);
+  assert(-1 == ret);
+  ASSERT_ZERO_LABEL(ret);
   ASSERT_LABEL(s.st_dev, i_label);
+  ASSERT_SAVED_ORIGINS(s)
 }
 
 void test_fstat() {
   int i = 1;
   dfsan_set_label(i_label, &i, sizeof(i));
 
   struct stat s;
   int fd = open("/dev/zero", O_RDONLY);
   s.st_dev = i;
+  DEFINE_AND_SAVE_ORIGINS(s)
   int rv = fstat(fd, &s);
   assert(0 == rv);
+  ASSERT_ZERO_LABEL(rv);
   ASSERT_ZERO_LABEL(s.st_dev);
+  ASSERT_SAVED_ORIGINS(s)
 }
 
 void test_memcmp() {
   char str1[] = "str1", str2[] = "str2";
   dfsan_set_label(i_label, &str1[3], 1);
   dfsan_set_label(j_label, &str2[3], 1);
 
   int rv = memcmp(str1, str2, sizeof(str1));
   assert(rv < 0);
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, i_j_label);
+  ASSERT_EQ_ORIGIN(rv, str1[3]);
 #endif
+
+  rv = memcmp(str1, str2, sizeof(str1) - 2);
+  assert(rv == 0);
+  ASSERT_ZERO_LABEL(rv);
 }
 
 void test_bcmp() {
   char str1[] = "str1", str2[] = "str2";
   dfsan_set_label(i_label, &str1[3], 1);
   dfsan_set_label(j_label, &str2[3], 1);
 
   int rv = bcmp(str1, str2, sizeof(str1));
   assert(rv != 0);
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, i_j_label);
+  ASSERT_EQ_ORIGIN(rv, str1[3]);
 #endif
 
   rv = bcmp(str1, str2, sizeof(str1) - 2);
   assert(rv == 0);
   ASSERT_ZERO_LABEL(rv);
 }
 
+#if !defined(ORIGIN_TRACKING)
 void test_memcpy() {
   char str1[] = "str1";
   char str2[sizeof(str1)];
   dfsan_set_label(i_label, &str1[3], 1);
 
   ASSERT_ZERO_LABEL(memcpy(str2, str1, sizeof(str1)));
   assert(0 == memcmp(str2, str1, sizeof(str1)));
   ASSERT_ZERO_LABEL(str2[0]);
void test_memset() {
   dfsan_set_label(j_label, &j, sizeof(j));
 
   ASSERT_ZERO_LABEL(memset(&buf, j, sizeof(buf)));
   for (int i = 0; i < 8; ++i) {
     ASSERT_LABEL(buf[i], j_label);
     assert(buf[i] == 'a');
   }
 }
+#endif // !defined(ORIGIN_TRACKING)
 
 void test_strcmp() {
   char str1[] = "str1", str2[] = "str2";
   dfsan_set_label(i_label, &str1[3], 1);
   dfsan_set_label(j_label, &str2[3], 1);
 
   int rv = strcmp(str1, str2);
   assert(rv < 0);
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, i_j_label);
+  ASSERT_EQ_ORIGIN(rv, str1[3]);
+#endif
+
+  rv = strcmp(str1, str1);
+  assert(rv == 0);
+#ifdef STRICT_DATA_DEPENDENCIES
+  ASSERT_ZERO_LABEL(rv);
+  ASSERT_ZERO_ORIGIN(rv);
+#else
+  ASSERT_LABEL(rv, i_label);
+  ASSERT_EQ_ORIGIN(rv, str1[3]);
 #endif
 }
 
+#if !defined(ORIGIN_TRACKING)
 void test_strcat() {
   char src[] = "world";
   char dst[] = "hello \0    ";
   char *p = dst;
   dfsan_set_label(k_label, &p, sizeof(p));
   dfsan_set_label(i_label, src, sizeof(src));
   dfsan_set_label(j_label, dst, sizeof(dst));
   char *ret = strcat(p, src);
void test_strncpy() {
   strd = strncpy(str2, str1, 3);
   assert(strd == str2);
   assert(strncmp(str1, str2, 3) == 0);
   ASSERT_ZERO_LABEL(strd);
   ASSERT_ZERO_LABEL(strd[0]);
   ASSERT_ZERO_LABEL(strd[1]);
   ASSERT_ZERO_LABEL(strd[2]);
 }
+#endif // !defined(ORIGIN_TRACKING)
 
 void test_strncmp() {
   char str1[] = "str1", str2[] = "str2";
   dfsan_set_label(i_label, &str1[3], 1);
   dfsan_set_label(j_label, &str2[3], 1);
 
   int rv = strncmp(str1, str2, sizeof(str1));
   assert(rv < 0);
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, dfsan_union(i_label, j_label));
+  ASSERT_EQ_ORIGIN(rv, str1[3]);
 #endif
 
+  rv = strncmp(str1, str2, 0);
+  assert(rv == 0);
+  ASSERT_ZERO_LABEL(rv);
+
   rv = strncmp(str1, str2, 3);
   assert(rv == 0);
   ASSERT_ZERO_LABEL(rv);
+
+  rv = strncmp(str1, str1, 4);
+  assert(rv == 0);
+#ifdef STRICT_DATA_DEPENDENCIES
+  ASSERT_ZERO_LABEL(rv);
+#else
+  ASSERT_LABEL(rv, i_label);
+  ASSERT_EQ_ORIGIN(rv, str1[3]);
+#endif
 }
 
 void test_strcasecmp() {
   char str1[] = "str1", str2[] = "str2", str3[] = "Str1";
   dfsan_set_label(i_label, &str1[3], 1);
   dfsan_set_label(j_label, &str2[3], 1);
   dfsan_set_label(j_label, &str3[2], 1);
 
   int rv = strcasecmp(str1, str2);
   assert(rv < 0);
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, dfsan_union(i_label, j_label));
+  ASSERT_EQ_ORIGIN(rv, str1[3]);
 #endif
 
   rv = strcasecmp(str1, str3);
   assert(rv == 0);
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, dfsan_union(i_label, j_label));
+  ASSERT_EQ_ORIGIN(rv, str1[3]);
 #endif
 
   char s1[] = "AbZ";
   char s2[] = "aBy";
   dfsan_set_label(i_label, &s1[2], 1);
   dfsan_set_label(j_label, &s2[2], 1);
 
   rv = strcasecmp(s1, s2);
   assert(rv > 0); // 'Z' > 'y'
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, dfsan_union(i_label, j_label));
+  ASSERT_EQ_ORIGIN(rv, s1[2]);
 #endif
 }
 
 void test_strncasecmp() {
   char str1[] = "Str1", str2[] = "str2";
   dfsan_set_label(i_label, &str1[3], 1);
   dfsan_set_label(j_label, &str2[3], 1);
 
   int rv = strncasecmp(str1, str2, sizeof(str1));
   assert(rv < 0);
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, dfsan_union(i_label, j_label));
+  ASSERT_EQ_ORIGIN(rv, str1[3]);
 #endif
 
   rv = strncasecmp(str1, str2, 3);
   assert(rv == 0);
   ASSERT_ZERO_LABEL(rv);
 
   char s1[] = "AbZ";
   char s2[] = "aBy";
#endif
   ASSERT_ZERO_LABEL(rv);
 
   rv = strncasecmp(s1, s2, 3);
   assert(rv > 0); // 'Z' > 'y'
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, dfsan_union(i_label, j_label));
+  ASSERT_EQ_ORIGIN(rv, s1[2]);
 #endif
 }
 
 void test_strchr() {
   char str1[] = "str1";
   dfsan_set_label(i_label, &str1[3], 1);
 
-  char *crv = strchr(str1, 'r');
+  char *p1 = str1;
+  char c = 'r';
+  dfsan_set_label(k_label, &c, sizeof(c));
+
+  char *crv = strchr(p1, c);
   assert(crv == &str1[2]);
+#ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(crv);
+#else
+  ASSERT_LABEL(crv, k_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&crv, c);
+#endif
+
+  dfsan_set_label(j_label, &p1, sizeof(p1));
+  crv = strchr(p1, 'r');
+  assert(crv == &str1[2]);
+  ASSERT_LABEL(crv, j_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&crv, p1);
 
-  crv = strchr(str1, '1');
+  crv = strchr(p1, '1');
   assert(crv == &str1[3]);
 #ifdef STRICT_DATA_DEPENDENCIES
-  ASSERT_ZERO_LABEL(crv);
+  ASSERT_LABEL(crv, j_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&crv, p1);
 #else
-  ASSERT_LABEL(crv, i_label);
+  ASSERT_LABEL(crv, i_j_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&crv, str1[3]);
 #endif
 
-  crv = strchr(str1, 'x');
+  crv = strchr(p1, 'x');
   assert(!crv);
 #ifdef STRICT_DATA_DEPENDENCIES
-  ASSERT_ZERO_LABEL(crv);
+  ASSERT_LABEL(crv, j_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&crv, p1);
 #else
-  ASSERT_LABEL(crv, i_label);
+  ASSERT_LABEL(crv, i_j_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&crv, str1[3]);
 #endif
 
   // `man strchr` says:
   // The terminating null byte is considered part of the string, so that if c
   // is specified as '\0', these functions return a pointer to the terminator.
-  crv = strchr(str1, '\0');
+  crv = strchr(p1, '\0');
   assert(crv == &str1[4]);
 #ifdef STRICT_DATA_DEPENDENCIES
-  ASSERT_ZERO_LABEL(crv);
+  ASSERT_LABEL(crv, j_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&crv, p1);
 #else
-  ASSERT_LABEL(crv, i_label);
+  ASSERT_LABEL(crv, i_j_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&crv, str1[3]);
 #endif
 }
 
+#if !defined(ORIGIN_TRACKING)
 void test_calloc() {
   // With any luck this sequence of calls will cause calloc to return the same
   // pointer both times.  This is probably the best we can do to test this
   // function.
   char *crv = (char *) calloc(4096, 1);
   ASSERT_ZERO_LABEL(crv[0]);
   dfsan_set_label(i_label, crv, 100);
   free(crv);
#endif
   rv = strstr(str1, "2s");
   assert(rv == NULL);
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, i_j_label);
 #endif
 }
+#endif // !defined(ORIGIN_TRACKING)
 
 void test_strpbrk() {
   char s[] = "abcdefg";
   char accept[] = "123fd";
+
+  char *p_s = s;
+  char *p_accept = accept;
+
+  dfsan_set_label(n_label, &p_accept, sizeof(p_accept));
+
+  char *rv = strpbrk(p_s, p_accept);
+  assert(rv == &s[3]);
+#ifdef STRICT_DATA_DEPENDENCIES
+  ASSERT_ZERO_LABEL(rv);
+#else
+  ASSERT_LABEL(rv, n_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&rv, p_accept);
+#endif
+
+  dfsan_set_label(m_label, &p_s, sizeof(p_s));
+
+  rv = strpbrk(p_s, p_accept);
+  assert(rv == &s[3]);
+#ifdef STRICT_DATA_DEPENDENCIES
+  ASSERT_LABEL(rv, m_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&rv, p_s);
+#else
+  ASSERT_LABEL(rv, dfsan_union(m_label, n_label));
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&rv, p_s);
+#endif
+
   dfsan_set_label(i_label, &s[5], 1);
   dfsan_set_label(j_label, &accept[1], 1);
 
-  char *rv = strpbrk(s, accept);
+  rv = strpbrk(s, accept);
   assert(rv == &s[3]);
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, j_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&rv, accept[1]);
 #endif
 
   char *ps = s;
   dfsan_set_label(j_label, &ps, sizeof(ps));
 
   rv = strpbrk(ps, "123gf");
   assert(rv == &s[5]);
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_LABEL(rv, j_label);
 #else
   ASSERT_LABEL(rv, i_j_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&rv, s[5]);
 #endif
 
   rv = strpbrk(ps, "123");
   assert(rv == NULL);
 #ifdef STRICT_DATA_DEPENDENCIES
   ASSERT_ZERO_LABEL(rv);
 #else
   ASSERT_LABEL(rv, i_j_label);
+  ASSERT_INIT_ORIGIN_EQ_ORIGIN(&rv, s[5]);
 #endif
 }
 
+#if !defined(ORIGIN_TRACKING)
 void test_memchr() {
   char str1[] = "str1";
   dfsan_set_label(i_label, &str1[3], 1);
   dfsan_set_label(j_label, &str1[4], 1);
 
   char *crv = (char *) memchr(str1, 'r', sizeof(str1));
   assert(crv == &str1[2]);
   ASSERT_ZERO_LABEL(crv);

 #endif
   i_j_label = dfsan_union(i_label, j_label);
   assert(i_j_label != i_label);
   assert(i_j_label != j_label);
   assert(i_j_label != k_label);
 
 #if !defined(ORIGIN_TRACKING)
   test__dl_get_tls_static_info();
+#endif // !defined(ORIGIN_TRACKING)
   test_bcmp();
+#if !defined(ORIGIN_TRACKING)
   test_calloc();
   test_clock_gettime();
   test_ctime_r();
   test_dfsan_set_write_callback();
   test_dl_iterate_phdr();
   test_dlopen();
   test_epoll_wait();
   test_fgets();
+#endif // !defined(ORIGIN_TRACKING)
+  test_fork();
   test_fstat();
+#if !defined(ORIGIN_TRACKING)
   test_get_current_dir_name();
   test_getcwd();
   test_gethostname();
   test_getpeername();
   test_getpwuid_r();
   test_getrlimit();
   test_getrusage();
   test_getsockname();
   test_getsockopt();
   test_gettimeofday();
   test_inet_pton();
   test_localtime_r();
   test_memchr();
+#endif // !defined(ORIGIN_TRACKING)
   test_memcmp();
+#if !defined(ORIGIN_TRACKING)
   test_memcpy();
   test_memmove();
   test_memset();
   test_nanosleep();
   test_poll();
   test_pread();
 #endif // !defined(ORIGIN_TRACKING)
   test_pthread_create();
   test_pthread_join();
 #if !defined(ORIGIN_TRACKING)
   test_read();
   test_recvmmsg();
   test_recvmsg();
   test_sched_getaffinity();
   test_select();
 #endif // !defined(ORIGIN_TRACKING)
   test_sigaction();
   test_signal();
   test_sigaltstack();
   test_sigemptyset();
 #if !defined(ORIGIN_TRACKING)
   test_snprintf();
   test_socketpair();
   test_sprintf();
+#endif // !defined(ORIGIN_TRACKING)
   test_stat();
   test_strcasecmp();
   test_strchr();
   test_strcmp();
+#if !defined(ORIGIN_TRACKING)
   test_strcat();
   test_strcpy();
   test_strdup();
 #endif // !defined(ORIGIN_TRACKING)
   test_strlen();
-#if !defined(ORIGIN_TRACKING)
   test_strncasecmp();
   test_strncmp();
+#if !defined(ORIGIN_TRACKING)
   test_strncpy();
+#endif // !defined(ORIGIN_TRACKING)
   test_strpbrk();
+#if !defined(ORIGIN_TRACKING)
   test_strrchr();
   test_strstr();
   test_strtod();
   test_strtol();
   test_strtoll();
   test_strtoul();
   test_strtoull();
   test_time();
 #endif // !defined(ORIGIN_TRACKING)
   test_write();
-  test_fork();
 }