1 files changed, 154 insertions, 0 deletions
diff --git a/lib/dfsan/lit_tests/custom.c b/lib/dfsan/lit_tests/custom.c
new file mode 100644
index 0000000000000..c9fa9356154df
--- /dev/null
+++ b/lib/dfsan/lit_tests/custom.c
@@ -0,0 +1,154 @@
+// RUN: %clang_dfsan -m64 %s -o %t && %t
+// RUN: %clang_dfsan -mllvm -dfsan-args-abi -m64 %s -o %t && %t
+
+// Tests custom implementations of various libc functions.
+
+#define _GNU_SOURCE
+#include <sanitizer/dfsan_interface.h>
+#include <assert.h>
+#include <link.h>
+#include <pthread.h>
+#include <string.h>
+#include <stdlib.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <fcntl.h>
+#include <unistd.h>
+
+void *ptcb(void *p) {
+  assert(p == (void *)1);
+  assert(dfsan_get_label((uintptr_t)p) == 0);
+  return (void *)2;
+}
+
+int dlcb(struct dl_phdr_info *info, size_t size, void *data) {
+  assert(data == (void *)3);
+  assert(dfsan_get_label((uintptr_t)info) == 0);
+  assert(dfsan_get_label(size) == 0);
+  assert(dfsan_get_label((uintptr_t)data) == 0);
+  return 0;
+}
+
+int main(void) {
+  int i = 1;
+  dfsan_label i_label = dfsan_create_label("i", 0);
+  dfsan_set_label(i_label, &i, sizeof(i));
+
+  int j = 2;
+  dfsan_label j_label = dfsan_create_label("j", 0);
+  dfsan_set_label(j_label, &j, sizeof(j));
+
+  struct stat s;
+  s.st_dev = i;
+  int rv = stat("/", &s);
+  assert(rv == 0);
+  assert(dfsan_get_label(s.st_dev) == 0);
+
+  s.st_dev = i;
+  rv = stat("/nonexistent", &s);
+  assert(rv == -1);
+  assert(dfsan_get_label(s.st_dev) == i_label);
+
+  int fd = open("/dev/zero", O_RDONLY);
+  s.st_dev = i;
+  rv = fstat(fd, &s);
+  assert(rv == 0);
+  assert(dfsan_get_label(s.st_dev) == 0);
+
+  char str1[] = "str1", str2[] = "str2";
+  dfsan_set_label(i_label, &str1[3], 1);
+  dfsan_set_label(j_label, &str2[3], 1);
+
+  rv = memcmp(str1, str2, sizeof(str1));
+  assert(rv < 0);
+  assert(dfsan_get_label(rv) == dfsan_union(i_label, j_label));
+
+  char strc[sizeof(str1)];
+  memcpy(strc, str1, sizeof(str1));
+  assert(dfsan_get_label(strc[0]) == 0);
+  assert(dfsan_get_label(strc[3]) == i_label);
+
+  memset(strc, j, sizeof(strc));
+  assert(dfsan_get_label(strc[0]) == j_label);
+  assert(dfsan_get_label(strc[1]) == j_label);
+  assert(dfsan_get_label(strc[2]) == j_label);
+  assert(dfsan_get_label(strc[3]) == j_label);
+  assert(dfsan_get_label(strc[4]) == j_label);
+
+  rv = strcmp(str1, str2);
+  assert(rv < 0);
+  assert(dfsan_get_label(rv) == dfsan_union(i_label, j_label));
+
+  char *strd = strdup(str1);
+  assert(dfsan_get_label(strd[0]) == 0);
+  assert(dfsan_get_label(strd[3]) == i_label);
+  free(strd);
+
+  rv = strncmp(str1, str2, sizeof(str1));
+  assert(rv < 0);
+  assert(dfsan_get_label(rv) == dfsan_union(i_label, j_label));
+
+  rv = strncmp(str1, str2, 3);
+  assert(rv == 0);
+  assert(dfsan_get_label(rv) == 0);
+
+  str1[0] = 'S';
+
+  rv = strncasecmp(str1, str2, sizeof(str1));
+  assert(rv < 0);
+  assert(dfsan_get_label(rv) == dfsan_union(i_label, j_label));
+
+  rv = strncasecmp(str1, str2, 3);
+  assert(rv == 0);
+  assert(dfsan_get_label(rv) == 0);
+
+  char *crv = strchr(str1, 'r');
+  assert(crv == &str1[2]);
+  assert(dfsan_get_label((uintptr_t)crv) == 0);
+
+  crv = strchr(str1, '1');
+  assert(crv == &str1[3]);
+  assert(dfsan_get_label((uintptr_t)crv) == i_label);
+
+  crv = strchr(str1, 'x');
+  assert(crv == 0);
+  assert(dfsan_get_label((uintptr_t)crv) == i_label);
+
+  // 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.
+  crv = calloc(4096, 1);
+  assert(dfsan_get_label(crv[0]) == 0);
+  free(crv);
+
+  crv = calloc(4096, 1);
+  assert(dfsan_get_label(crv[0]) == 0);
+  free(crv);
+
+  char buf[16];
+  buf[0] = i;
+  buf[15] = j;
+  rv = read(fd, buf, sizeof(buf));
+  assert(rv == sizeof(buf));
+  assert(dfsan_get_label(buf[0]) == 0);
+  assert(dfsan_get_label(buf[15]) == 0);
+
+  close(fd);
+  fd = open("/bin/sh", O_RDONLY);
+  buf[0] = i;
+  buf[15] = j;
+  rv = pread(fd, buf, sizeof(buf), 0);
+  assert(rv == sizeof(buf));
+  assert(dfsan_get_label(buf[0]) == 0);
+  assert(dfsan_get_label(buf[15]) == 0);
+
+  pthread_t pt;
+  pthread_create(&pt, 0, ptcb, (void *)1);
+  void *cbrv;
+  pthread_join(pt, &cbrv);
+  assert(cbrv == (void *)2);
+
+  dl_iterate_phdr(dlcb, (void *)3);
+
+  return 0;
+}