Support mb sequences across calls to mb*to*wcs* functions

Bug: 13077905
Change-Id: I5abdc7cc3c27c109b7900c94b112f18a95c35763
diff --git a/libc/bionic/wchar.cpp b/libc/bionic/wchar.cpp
index b46ad49..5da882f 100644
--- a/libc/bionic/wchar.cpp
+++ b/libc/bionic/wchar.cpp
@@ -32,23 +32,69 @@
 #include <wchar.h>
 
 //
-// This file is basically OpenBSD's citrus_utf8.c but rewritten to not require a 12-byte mbstate_t
-// so we're backwards-compatible with our LP32 ABI where mbstate_t was only 4 bytes. An additional
-// advantage of this is that callers who don't supply their own mbstate_t won't be accessing shared
-// state.
+// This file is basically OpenBSD's citrus_utf8.c but rewritten to not require a
+// 12-byte mbstate_t so we're backwards-compatible with our LP32 ABI where
+// mbstate_t was only 4 bytes.
 //
-// We also implement the POSIX interface directly rather than being accessed via function pointers.
+// The state is the UTF-8 sequence. We only support <= 4-bytes sequences so LP32
+// mbstate_t already has enough space (out of the 4 available bytes we only
+// need 3 since we should never need to store the entire sequence in the
+// intermediary state).
+//
+// The C standard leaves the conversion state undefined after a bad conversion.
+// To avoid unexpected failures due to the possible use of the internal private
+// state we always reset the conversion state when encountering illegal
+// sequences.
+//
+// We also implement the POSIX interface directly rather than being accessed via
+// function pointers.
 //
 
 #define ERR_ILLEGAL_SEQUENCE static_cast<size_t>(-1)
 #define ERR_INCOMPLETE_SEQUENCE static_cast<size_t>(-2)
 
-int mbsinit(const mbstate_t*) {
-  // We have no state, so we're always in the initial state.
-  return 1;
+static size_t mbstate_bytes_so_far(const mbstate_t* ps) {
+  return
+    (ps->__seq[2] != 0) ? 3 :
+    (ps->__seq[1] != 0) ? 2 :
+    (ps->__seq[0] != 0) ? 1 : 0;
 }
 
-size_t mbrtowc(wchar_t* pwc, const char* s, size_t n, mbstate_t*) {
+static void mbstate_set_byte(mbstate_t* ps, int i, char byte) {
+  ps->__seq[i] = static_cast<uint8_t>(byte);
+}
+
+static uint8_t mbstate_get_byte(const mbstate_t* ps, int n) {
+  return ps->__seq[n];
+}
+
+static size_t reset_and_return_illegal(int _errno, mbstate_t* ps) {
+  errno = _errno;
+  *(reinterpret_cast<uint32_t*>(ps->__seq)) = 0;
+  return ERR_ILLEGAL_SEQUENCE;
+}
+
+static size_t reset_and_return(int _return, mbstate_t* ps) {
+  *(reinterpret_cast<uint32_t*>(ps->__seq)) = 0;
+  return _return;
+}
+
+
+int mbsinit(const mbstate_t* ps) {
+  return (ps == NULL || (*(reinterpret_cast<const uint32_t*>(ps->__seq)) == 0));
+}
+
+size_t mbrtowc(wchar_t* pwc, const char* s, size_t n, mbstate_t* ps) {
+  static mbstate_t __private_state;
+  mbstate_t* state = (ps == NULL) ? &__private_state : ps;
+
+  // We should never get to a state which has all 4 bytes of the sequence set.
+  // Full state verification is done when decoding the sequence (after we have
+  // all the bytes).
+  if (mbstate_get_byte(state, 3) != 0) {
+    return reset_and_return_illegal(EINVAL, state);
+  }
+
   if (s == NULL) {
     s = "";
     n = 1;
@@ -59,8 +105,8 @@
     return 0;
   }
 
-  int ch;
-  if (((ch = static_cast<uint8_t>(*s)) & ~0x7f) == 0) {
+  uint8_t ch;
+  if (mbsinit(state) && (((ch = static_cast<uint8_t>(*s)) & ~0x7f) == 0)) {
     // Fast path for plain ASCII characters.
     if (pwc != NULL) {
       *pwc = ch;
@@ -82,7 +128,9 @@
   // between character codes and their multibyte representations.
   wchar_t lower_bound;
 
-  ch = static_cast<uint8_t>(*s);
+  // The first byte in the state (if any) tells the length.
+  size_t bytes_so_far = mbstate_bytes_so_far(state);
+  ch = bytes_so_far > 0 ? mbstate_get_byte(state, 0) : static_cast<uint8_t>(*s);
   if ((ch & 0x80) == 0) {
     mask = 0x7f;
     length = 1;
@@ -101,106 +149,144 @@
     lower_bound = 0x10000;
   } else {
     // Malformed input; input is not UTF-8. See RFC 3629.
-    errno = EILSEQ;
-    return ERR_ILLEGAL_SEQUENCE;
+    return reset_and_return_illegal(EILSEQ, state);
+  }
+
+  // Fill in the state.
+  size_t bytes_wanted = length - bytes_so_far;
+  size_t i;
+  for (i = 0; i < MIN(bytes_wanted, n); i++) {
+    if (!mbsinit(state) && ((*s & 0xc0) != 0x80)) {
+      // Malformed input; bad characters in the middle of a character.
+      return reset_and_return_illegal(EILSEQ, state);
+    }
+    mbstate_set_byte(state, bytes_so_far + i, *s++);
+  }
+  if (i < bytes_wanted) {
+    return ERR_INCOMPLETE_SEQUENCE;
   }
 
   // Decode the octet sequence representing the character in chunks
   // of 6 bits, most significant first.
-  wchar_t wch = static_cast<uint8_t>(*s++) & mask;
-  size_t i;
-  for (i = 1; i < MIN(length, n); i++) {
-    if ((*s & 0xc0) != 0x80) {
-      // Malformed input; bad characters in the middle of a character.
-      errno = EILSEQ;
-      return ERR_ILLEGAL_SEQUENCE;
-    }
+  wchar_t wch = mbstate_get_byte(state, 0) & mask;
+  for (i = 1; i < length; i++) {
     wch <<= 6;
-    wch |= *s++ & 0x3f;
+    wch |= mbstate_get_byte(state, i) & 0x3f;
   }
-  if (i < length) {
-    return ERR_INCOMPLETE_SEQUENCE;
-  }
+
   if (wch < lower_bound) {
     // Malformed input; redundant encoding.
-    errno = EILSEQ;
-    return ERR_ILLEGAL_SEQUENCE;
+    return reset_and_return_illegal(EILSEQ, state);
   }
   if ((wch >= 0xd800 && wch <= 0xdfff) || wch == 0xfffe || wch == 0xffff) {
     // Malformed input; invalid code points.
-    errno = EILSEQ;
-    return ERR_ILLEGAL_SEQUENCE;
+    return reset_and_return_illegal(EILSEQ, state);
   }
   if (pwc != NULL) {
     *pwc = wch;
   }
-  return (wch == L'\0' ? 0 : length);
+  return reset_and_return(wch == L'\0' ? 0 : bytes_wanted, state);
 }
 
 size_t mbsnrtowcs(wchar_t* dst, const char** src, size_t nmc, size_t len, mbstate_t* ps) {
+  static mbstate_t __private_state;
+  mbstate_t* state = (ps == NULL) ? &__private_state : ps;
   size_t i, o, r;
 
   if (dst == NULL) {
+    /*
+     * The fast path in the loop below is not safe if an ASCII
+     * character appears as anything but the first byte of a
+     * multibyte sequence. Check now to avoid doing it in the loop.
+     */
+    if ((nmc > 0) && (mbstate_bytes_so_far(state) > 0)
+        && (static_cast<uint8_t>((*src)[0]) < 0x80)) {
+      return reset_and_return_illegal(EILSEQ, state);
+    }
     for (i = o = 0; i < nmc; i += r, o++) {
       if (static_cast<uint8_t>((*src)[i]) < 0x80) {
         // Fast path for plain ASCII characters.
         if ((*src)[i] == '\0') {
-          return o;
+          return reset_and_return(o, state);
         }
         r = 1;
       } else {
-        r = mbrtowc(NULL, *src + i, nmc - i, ps);
+        r = mbrtowc(NULL, *src + i, nmc - i, state);
         if (r == ERR_ILLEGAL_SEQUENCE) {
-          return r;
+          return reset_and_return_illegal(EILSEQ, state);
         }
         if (r == ERR_INCOMPLETE_SEQUENCE) {
-          return o;
+          return reset_and_return_illegal(EILSEQ, state);
         }
         if (r == 0) {
-          return o;
+          return reset_and_return(o, state);
         }
       }
     }
-    return o;
+    return reset_and_return(o, state);
   }
 
+  /*
+   * The fast path in the loop below is not safe if an ASCII
+   * character appears as anything but the first byte of a
+   * multibyte sequence. Check now to avoid doing it in the loop.
+   */
+  if ((nmc > 0) && (mbstate_bytes_so_far(state) > 0)
+      && (static_cast<uint8_t>((*src)[0]) < 0x80)) {
+    return reset_and_return_illegal(EILSEQ, state);
+  }
   for (i = o = 0; i < nmc && o < len; i += r, o++) {
     if (static_cast<uint8_t>((*src)[i]) < 0x80) {
       // Fast path for plain ASCII characters.
       dst[o] = (*src)[i];
       if ((*src)[i] == '\0') {
         *src = NULL;
-        return o;
+        return reset_and_return_illegal(EILSEQ, state);
       }
       r = 1;
     } else {
-      r = mbrtowc(dst + o, *src + i, nmc - i, ps);
+      r = mbrtowc(dst + o, *src + i, nmc - i, state);
       if (r == ERR_ILLEGAL_SEQUENCE) {
         *src += i;
-        return r;
+        return reset_and_return_illegal(EILSEQ, state);
       }
       if (r == ERR_INCOMPLETE_SEQUENCE) {
         *src += nmc;
-        return o;
+        return reset_and_return(EILSEQ, state);
       }
       if (r == 0) {
         *src = NULL;
-        return o;
+        return reset_and_return(o, state);
       }
     }
   }
   *src += i;
-  return o;
+  return reset_and_return(o, state);
 }
 
 size_t mbsrtowcs(wchar_t* dst, const char** src, size_t len, mbstate_t* ps) {
   return mbsnrtowcs(dst, src, SIZE_MAX, len, ps);
 }
 
-size_t wcrtomb(char* s, wchar_t wc, mbstate_t*) {
+size_t wcrtomb(char* s, wchar_t wc, mbstate_t* ps) {
+  static mbstate_t __private_state;
+  mbstate_t* state = (ps == NULL) ? &__private_state : ps;
+
   if (s == NULL) {
-    // Reset to initial shift state (no-op).
-    return 1;
+    // Equivalent to wcrtomb(buf, L'\0', ps).
+    return reset_and_return(1, state);
+  }
+
+  // POSIX states that if wc is a null wide character, a null byte shall be
+  // stored, preceded by any shift sequence needed to restore the initial shift
+  // state. Since shift states are not supported, only the null byte is stored.
+  if (wc == L'\0') {
+    *s = '\0';
+    reset_and_return(1, state);
+  }
+
+  if (!mbsinit(state)) {
+    return reset_and_return_illegal(EILSEQ, state);
   }
 
   if ((wc & ~0x7f) == 0) {
@@ -246,6 +332,13 @@
 }
 
 size_t wcsnrtombs(char* dst, const wchar_t** src, size_t nwc, size_t len, mbstate_t* ps) {
+  static mbstate_t __private_state;
+  mbstate_t* state = (ps == NULL) ? &__private_state : ps;
+
+  if (!mbsinit(state)) {
+    return reset_and_return_illegal(EILSEQ, state);
+  }
+
   char buf[MB_LEN_MAX];
   size_t i, o, r;
   if (dst == NULL) {
@@ -258,7 +351,7 @@
         }
         r = 1;
       } else {
-        r = wcrtomb(buf, wc, ps);
+        r = wcrtomb(buf, wc, state);
         if (r == ERR_ILLEGAL_SEQUENCE) {
           return r;
         }
@@ -279,14 +372,14 @@
       r = 1;
     } else if (len - o >= sizeof(buf)) {
       // Enough space to translate in-place.
-      r = wcrtomb(dst + o, wc, ps);
+      r = wcrtomb(dst + o, wc, state);
       if (r == ERR_ILLEGAL_SEQUENCE) {
         *src += i;
         return r;
       }
     } else {
       // May not be enough space; use temp buffer.
-      r = wcrtomb(buf, wc, ps);
+      r = wcrtomb(buf, wc, state);
       if (r == ERR_ILLEGAL_SEQUENCE) {
         *src += i;
         return r;
diff --git a/libc/include/wchar.h b/libc/include/wchar.h
index fe2fe07..4ac468d 100644
--- a/libc/include/wchar.h
+++ b/libc/include/wchar.h
@@ -41,11 +41,9 @@
 
 typedef __WINT_TYPE__  wint_t;
 typedef struct {
-#ifdef __LP32__
-  int dummy;
-#else
-  // 8 bytes should be enough to support at least UTF-8
-  char __reserved[8];
+  uint8_t __seq[4];
+#ifdef __LP64__
+  char __reserved[4];
 #endif
 } mbstate_t;
 
diff --git a/tests/wchar_test.cpp b/tests/wchar_test.cpp
index 0d15f21..30d7bff 100644
--- a/tests/wchar_test.cpp
+++ b/tests/wchar_test.cpp
@@ -87,6 +87,29 @@
   EXPECT_EQ(EILSEQ, errno);
 }
 
+TEST(wchar, wcrtomb_start_state) {
+  char out[MB_LEN_MAX];
+  mbstate_t ps;
+
+  // Any non-initial state is invalid when calling wcrtomb.
+  memset(&ps, 0, sizeof(ps));
+  EXPECT_EQ(static_cast<size_t>(-2), mbrtowc(NULL, "\xc2", 1, &ps));
+  EXPECT_EQ(static_cast<size_t>(-1), wcrtomb(out, 0x00a2, &ps));
+  EXPECT_EQ(EILSEQ, errno);
+
+  // If the first argument to wcrtomb is NULL or the second is L'\0' the shift
+  // state should be reset.
+  memset(&ps, 0, sizeof(ps));
+  EXPECT_EQ(static_cast<size_t>(-2), mbrtowc(NULL, "\xc2", 1, &ps));
+  EXPECT_EQ(1U, wcrtomb(NULL, 0x00a2, &ps));
+  EXPECT_TRUE(mbsinit(&ps));
+
+  memset(&ps, 0, sizeof(ps));
+  EXPECT_EQ(static_cast<size_t>(-2), mbrtowc(NULL, "\xf0\xa4", 1, &ps));
+  EXPECT_EQ(1U, wcrtomb(out, L'\0', &ps));
+  EXPECT_TRUE(mbsinit(&ps));
+}
+
 TEST(wchar, wcstombs_wcrtombs) {
   const wchar_t chars[] = { L'h', L'e', L'l', L'l', L'o', 0 };
   const wchar_t bad_chars[] = { L'h', L'i', static_cast<wchar_t>(0xffffffff), 0 };
@@ -184,6 +207,14 @@
   EXPECT_EQ(EILSEQ, errno);
   bytes[3] = 0;
   EXPECT_STREQ("hix", bytes);
+
+  // Any non-initial state is invalid when calling wcsrtombs.
+  mbstate_t ps;
+  src = chars;
+  memset(&ps, 0, sizeof(ps));
+  ASSERT_EQ(static_cast<size_t>(-2), mbrtowc(NULL, "\xc2", 1, &ps));
+  EXPECT_EQ(static_cast<size_t>(-1), wcsrtombs(NULL, &src, 0, &ps));
+  EXPECT_EQ(EILSEQ, errno);
 }
 
 TEST(wchar, limits) {
@@ -267,6 +298,83 @@
   ASSERT_EQ(EILSEQ, errno);
 }
 
+void test_mbrtowc_incomplete(mbstate_t* ps) {
+  ASSERT_STREQ("C.UTF-8", setlocale(LC_CTYPE, "C.UTF-8"));
+  uselocale(LC_GLOBAL_LOCALE);
+
+  wchar_t out;
+  // 2-byte UTF-8.
+  ASSERT_EQ(static_cast<size_t>(-2), mbrtowc(&out, "\xc2", 1, ps));
+  ASSERT_EQ(1U, mbrtowc(&out, "\xa2" "cdef", 5, ps));
+  ASSERT_EQ(0x00a2, out);
+  ASSERT_TRUE(mbsinit(ps));
+  // 3-byte UTF-8.
+  ASSERT_EQ(static_cast<size_t>(-2), mbrtowc(&out, "\xe2", 1, ps));
+  ASSERT_EQ(static_cast<size_t>(-2), mbrtowc(&out, "\x82", 1, ps));
+  ASSERT_EQ(1U, mbrtowc(&out, "\xac" "def", 4, ps));
+  ASSERT_EQ(0x20ac, out);
+  ASSERT_TRUE(mbsinit(ps));
+  // 4-byte UTF-8.
+  ASSERT_EQ(static_cast<size_t>(-2), mbrtowc(&out, "\xf0", 1, ps));
+  ASSERT_EQ(static_cast<size_t>(-2), mbrtowc(&out, "\xa4\xad", 2, ps));
+  ASSERT_EQ(1U, mbrtowc(&out, "\xa2" "ef", 3, ps));
+  ASSERT_EQ(0x24b62, out);
+  ASSERT_TRUE(mbsinit(ps));
+
+  // Invalid 2-byte
+  ASSERT_EQ(static_cast<size_t>(-2), mbrtowc(&out, "\xc2", 1, ps));
+  ASSERT_EQ(static_cast<size_t>(-1), mbrtowc(&out, "\x20" "cdef", 5, ps));
+  ASSERT_EQ(EILSEQ, errno);
+}
+
+TEST(wchar, mbrtowc_incomplete) {
+  mbstate_t ps;
+  memset(&ps, 0, sizeof(ps));
+
+  test_mbrtowc_incomplete(&ps);
+  test_mbrtowc_incomplete(NULL);
+}
+
+void test_mbsrtowcs(mbstate_t* ps) {
+  wchar_t out[4];
+
+  const char* valid = "A" "\xc2\xa2" "\xe2\x82\xac" "\xf0\xa4\xad\xa2" "ef";
+  ASSERT_EQ(4U, mbsrtowcs(out, &valid, 4, ps));
+  ASSERT_EQ(L'A', out[0]);
+  ASSERT_EQ(0x00a2, out[1]);
+  ASSERT_EQ(0x20ac, out[2]);
+  ASSERT_EQ(0x24b62, out[3]);
+  ASSERT_EQ('e', *valid);
+
+  const char* invalid = "A" "\xc2\x20" "ef";
+  ASSERT_EQ(static_cast<size_t>(-1), mbsrtowcs(out, &invalid, 4, ps));
+  EXPECT_EQ(EILSEQ, errno);
+  ASSERT_EQ('\xc2', *invalid);
+
+  const char* incomplete = "A" "\xc2";
+  ASSERT_EQ(static_cast<size_t>(-1), mbsrtowcs(out, &incomplete, 2, ps));
+  EXPECT_EQ(EILSEQ, errno);
+  ASSERT_EQ('\xc2', *incomplete);
+}
+
+TEST(wchar, mbsrtowcs) {
+  ASSERT_STREQ("C.UTF-8", setlocale(LC_CTYPE, "C.UTF-8"));
+  uselocale(LC_GLOBAL_LOCALE);
+
+  mbstate_t ps;
+  memset(&ps, 0, sizeof(ps));
+  test_mbsrtowcs(&ps);
+  test_mbsrtowcs(NULL);
+
+  // Invalid multi byte continuation.
+  const char* invalid = "\x20";
+  wchar_t out;
+  ASSERT_EQ(static_cast<size_t>(-2), mbrtowc(&out, "\xc2", 1, &ps));
+  ASSERT_EQ(static_cast<size_t>(-1), mbsrtowcs(&out, &invalid, 1, &ps));
+  EXPECT_EQ(EILSEQ, errno);
+  ASSERT_EQ('\x20', *invalid);
+}
+
 TEST(wchar, wcstod) {
   ASSERT_DOUBLE_EQ(1.23, wcstod(L"1.23", NULL));
 }