[DFSAN] Add support for strnlen

This patch adds a support for the libc strnlen() function in DFSAN

Reviewed by: browneee

Differential Revision: https://reviews.llvm.org/D149459

3 files changed, 60 insertions, 0 deletions

diff --git a/compiler-rt/lib/dfsan/dfsan_custom.cpp b/compiler-rt/lib/dfsan/dfsan_custom.cpp
index 8bb5d39ee8f2..e63e68084fa3 100644
--- a/compiler-rt/lib/dfsan/dfsan_custom.cpp
+++ b/compiler-rt/lib/dfsan/dfsan_custom.cpp
@@ -535,6 +535,36 @@ SANITIZER_INTERFACE_ATTRIBUTE size_t __dfso_strlen(const char *s,
   return ret;
 }
 
+SANITIZER_INTERFACE_ATTRIBUTE size_t __dfsw_strnlen(const char *s,
+                                                    size_t maxlen,
+                                                    dfsan_label s_label,
+                                                    dfsan_label maxlen_label,
+                                                    dfsan_label *ret_label) {
+  size_t ret = strnlen(s, maxlen);
+  if (flags().strict_data_dependencies) {
+    *ret_label = 0;
+  } else {
+    size_t full_len = strlen(s);
+    size_t covered_len = maxlen > (full_len + 1) ? (full_len + 1) : maxlen;
+    *ret_label = dfsan_union(maxlen_label, dfsan_read_label(s, covered_len));
+  }
+  return ret;
+}
+
+SANITIZER_INTERFACE_ATTRIBUTE size_t __dfso_strnlen(
+    const char *s, size_t maxlen, dfsan_label s_label, dfsan_label maxlen_label,
+    dfsan_label *ret_label, dfsan_origin s_origin, dfsan_origin maxlen_origin,
+    dfsan_origin *ret_origin) {
+  size_t ret = __dfsw_strnlen(s, maxlen, s_label, maxlen_label, ret_label);
+  if (!flags().strict_data_dependencies) {
+    size_t full_len = strlen(s);
+    size_t covered_len = maxlen > (full_len + 1) ? (full_len + 1) : maxlen;
+    dfsan_origin o = dfsan_read_origin_of_first_taint(s, covered_len);
+    *ret_origin = o ? o : maxlen_origin;
+  }
+  return ret;
+}
+
 static void *dfsan_memmove(void *dest, const void *src, size_t n) {
   dfsan_label *sdest = shadow_for(dest);
   const dfsan_label *ssrc = shadow_for(src);
diff --git a/compiler-rt/lib/dfsan/done_abilist.txt b/compiler-rt/lib/dfsan/done_abilist.txt
index 88ec5cf504ba..cc3578b95cd4 100644
--- a/compiler-rt/lib/dfsan/done_abilist.txt
+++ b/compiler-rt/lib/dfsan/done_abilist.txt
@@ -278,6 +278,7 @@ fun:strcasecmp=custom
 fun:strchr=custom
 fun:strcmp=custom
 fun:strlen=custom
+fun:strnlen=custom
 fun:strncasecmp=custom
 fun:strncmp=custom
 fun:strpbrk=custom
diff --git a/compiler-rt/test/dfsan/custom.cpp b/compiler-rt/test/dfsan/custom.cpp
index 6dbf0d71c966..62483c1ed96f 100644
--- a/compiler-rt/test/dfsan/custom.cpp
+++ b/compiler-rt/test/dfsan/custom.cpp
@@ -381,6 +381,34 @@ void test_strlen() {
 #endif
 }
 
+void test_strnlen() {
+  char str1[] = "str1";
+  dfsan_set_label(i_label, &str1[3], 1);
+
+  int maxlen = 4;
+  dfsan_set_label(j_label, &maxlen, sizeof(maxlen));
+
+  int rv = strnlen(str1, maxlen);
+  assert(rv == 4);
+#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
+
+  maxlen = 2;
+  dfsan_set_label(j_label, &maxlen, sizeof(maxlen));
+  rv = strnlen(str1, maxlen);
+  assert(rv == 2);
+#ifdef STRICT_DATA_DEPENDENCIES
+  ASSERT_ZERO_LABEL(rv);
+#else
+  ASSERT_LABEL(rv, j_label);
+  ASSERT_EQ_ORIGIN(rv, maxlen);
+#endif
+}
+
 void test_strdup() {
   char str1[] = "str1";
   dfsan_set_label(i_label, &str1[3], 1);
@@ -2085,6 +2113,7 @@ int main(void) {
   test_strcpy();
   test_strdup();
   test_strlen();
+  test_strnlen();
   test_strncasecmp();
   test_strncmp();
   test_strncpy();