[libc-commits] [libc] 0784e62 - [libc] Fix strtok_r crash when src and *saveptr are both nullptr

[Siva Chandra Reddy via libc-commits]

Author: Alfonso Gregory
Date: 2021-07-27T21:49:23Z
New Revision: 0784e62c3c4a4aaabdf29f6fa0b5f8f7598a90d4

URL: https://github.com/llvm/llvm-project/commit/0784e62c3c4a4aaabdf29f6fa0b5f8f7598a90d4
DIFF: https://github.com/llvm/llvm-project/commit/0784e62c3c4a4aaabdf29f6fa0b5f8f7598a90d4.diff

LOG: [libc] Fix strtok_r crash when src and *saveptr are both nullptr

While working and testing my refactoring of multiple string functions in libc, I came across a bug that needs to be addressed in a patch on its own: src is checked for nullptr and assigned to *saveptr if it is nullptr. However, saveptr is initially nullptr when it comes to reentry. This could cause a problem if both saveptr and src are null; we need to do the check first and return nullptr if both are nullptr.

Reviewed By: sivachandra

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

Added: 
    

Modified: 
    libc/src/string/string_utils.h
    libc/test/src/string/strtok_r_test.cpp

Removed: 
    


################################################################################
diff  --git a/libc/src/string/string_utils.h b/libc/src/string/string_utils.h
index dfb2c8af45279..4f179bb628d6c 100644
--- a/libc/src/string/string_utils.h
+++ b/libc/src/string/string_utils.h
@@ -9,6 +9,7 @@
 #ifndef LIBC_SRC_STRING_STRING_UTILS_H
 #define LIBC_SRC_STRING_STRING_UTILS_H
 
+#include "src/__support/common.h"
 #include "utils/CPP/Bitset.h"
 #include <stddef.h> // size_t
 
@@ -58,23 +59,27 @@ static inline size_t complementary_span(const char *src, const char *segment) {
 static inline char *string_token(char *__restrict src,
                                  const char *__restrict delimiter_string,
                                  char **__restrict saveptr) {
+  // Return nullptr immediately if both src AND saveptr are nullptr
+  if (unlikely(src == nullptr && ((src = *saveptr) == nullptr)))
+    return nullptr;
+
   cpp::Bitset<256> delimiter_set;
-  for (; *delimiter_string; ++delimiter_string)
+  for (; *delimiter_string != '\0'; ++delimiter_string)
     delimiter_set.set(*delimiter_string);
 
-  src = src ? src : *saveptr;
-  for (; *src && delimiter_set.test(*src); ++src)
+  for (; *src != '\0' && delimiter_set.test(*src); ++src)
     ;
-  if (!*src) {
+  if (*src == '\0') {
     *saveptr = src;
     return nullptr;
   }
   char *token = src;
-  for (; *src && !delimiter_set.test(*src); ++src)
-    ;
-  if (*src) {
-    *src = '\0';
-    ++src;
+  for (; *src != '\0'; ++src) {
+    if (delimiter_set.test(*src)) {
+      *src = '\0';
+      ++src;
+      break;
+    }
   }
   *saveptr = src;
   return token;

diff  --git a/libc/test/src/string/strtok_r_test.cpp b/libc/test/src/string/strtok_r_test.cpp
index 35d4b91810a74..e767b5ce0c838 100644
--- a/libc/test/src/string/strtok_r_test.cpp
+++ b/libc/test/src/string/strtok_r_test.cpp
@@ -80,6 +80,18 @@ TEST(LlvmLibcStrTokReentrantTest, ShouldNotGoPastNullTerminator) {
   ASSERT_STREQ(__llvm_libc::strtok_r(nullptr, ",", &reserve), nullptr);
 }
 
+TEST(LlvmLibcStrTokReentrantTest,
+     ShouldReturnNullptrWhenBothSrcAndSaveptrAreNull) {
+  char *src = nullptr;
+  char *reserve = nullptr;
+  // Ensure that instead of crashing if src and reserve are null, nullptr is
+  // returned
+  ASSERT_STREQ(__llvm_libc::strtok_r(src, ",", &reserve), nullptr);
+  // And that neither src nor reserve are changed when that happens
+  ASSERT_STREQ(src, nullptr);
+  ASSERT_STREQ(reserve, nullptr);
+}
+
 TEST(LlvmLibcStrTokReentrantTest,
      SubsequentCallsShouldFindFollowingDelimiters) {
   char src[] = "12,34.56";