From 23bab1eb43d39e7163eb55bcca6e412f68f930e3 Mon Sep 17 00:00:00 2001
From: Matt Morehouse <mascasa@google.com>
Date: Fri, 18 Sep 2020 08:47:31 -0700
Subject: [PATCH] [DFSan] Add strpbrk wrapper.

Reviewed By: vitalybuka

Differential Revision: https://reviews.llvm.org/D87849
---
 compiler-rt/lib/dfsan/dfsan_custom.cpp | 18 +++++++++++++++++
 compiler-rt/lib/dfsan/done_abilist.txt |  1 +
 compiler-rt/test/dfsan/custom.cpp      | 35 ++++++++++++++++++++++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/compiler-rt/lib/dfsan/dfsan_custom.cpp b/compiler-rt/lib/dfsan/dfsan_custom.cpp
index 77b93f8..80031d96 100644
--- a/compiler-rt/lib/dfsan/dfsan_custom.cpp
+++ b/compiler-rt/lib/dfsan/dfsan_custom.cpp
@@ -95,6 +95,24 @@ SANITIZER_INTERFACE_ATTRIBUTE char *__dfsw_strchr(const char *s, int c,
   }
 }
 
+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);
+}
+
 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) {
diff --git a/compiler-rt/lib/dfsan/done_abilist.txt b/compiler-rt/lib/dfsan/done_abilist.txt
index 85255f7..2146481 100644
--- a/compiler-rt/lib/dfsan/done_abilist.txt
+++ b/compiler-rt/lib/dfsan/done_abilist.txt
@@ -192,6 +192,7 @@ fun:strcmp=custom
 fun:strlen=custom
 fun:strncasecmp=custom
 fun:strncmp=custom
+fun:strpbrk=custom
 fun:strrchr=custom
 fun:strstr=custom
 
diff --git a/compiler-rt/test/dfsan/custom.cpp b/compiler-rt/test/dfsan/custom.cpp
index 6d5e06a..8f00b1e 100644
--- a/compiler-rt/test/dfsan/custom.cpp
+++ b/compiler-rt/test/dfsan/custom.cpp
@@ -758,6 +758,40 @@ void test_strstr() {
 #endif
 }
 
+void test_strpbrk() {
+  char s[] = "abcdefg";
+  char accept[] = "123fd";
+  dfsan_set_label(i_label, &s[5], 1);
+  dfsan_set_label(j_label, &accept[1], 1);
+
+  char *rv = strpbrk(s, accept);
+  assert(rv == &s[3]);
+#ifdef STRICT_DATA_DEPENDENCIES
+  ASSERT_ZERO_LABEL(rv);
+#else
+  ASSERT_LABEL(rv, j_label);
+#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);
+#endif
+
+  rv = strpbrk(ps, "123");
+  assert(rv == NULL);
+#ifdef STRICT_DATA_DEPENDENCIES
+  ASSERT_ZERO_LABEL(rv);
+#else
+  ASSERT_LABEL(rv, i_j_label);
+#endif
+}
+
 void test_memchr() {
   char str1[] = "str1";
   dfsan_set_label(i_label, &str1[3], 1);
@@ -1030,6 +1064,7 @@ int main(void) {
   test_strncasecmp();
   test_strncmp();
   test_strncpy();
+  test_strpbrk();
   test_strrchr();
   test_strstr();
   test_strtod();
-- 
2.7.4