4 files changed, 60 insertions, 10 deletions

diff --git a/include/wchar.h b/include/wchar.h
index 9e32720..2036245 100644
--- a/include/wchar.h
+++ b/include/wchar.h
@@ -75,16 +75,10 @@ _FORTIFY_FN(mbsnrtowcs) size_t mbsnrtowcs(wchar_t * _FORTIFY_POS0 __d,
         size_t __b = __bos(__d, 0);
         size_t __r;
 
-        if (__wn > __n / sizeof(wchar_t)) {
-                __b /= sizeof(wchar_t);
-                __r = __orig_mbsnrtowcs(__d, __s, __n, __wn > __b ? __b : __wn, __st);
-                if (__b < __wn && __d && *__s && __r != (size_t)-1)
-                        __builtin_trap();
-        } else {
-                __r = __orig_mbsnrtowcs(__d, __s, __n > __b ? __b : __n, __wn, __st);
-                if (__b < __n && __d && *__s && __r != (size_t)-1)
-                        __builtin_trap();
-        }
+        __b /= sizeof(wchar_t);
+        __r = __orig_mbsnrtowcs(__d, __s, __n, __wn > __b ? __b : __wn, __st);
+        if (__b < __wn && __d && *__s && __r != (size_t)-1)
+                __builtin_trap();
         return __r;
 }
 #endif
diff --git a/tests/Makefile b/tests/Makefile
index 9fc8287..6904b2d 100644
--- a/tests/Makefile
+++ b/tests/Makefile
@@ -50,6 +50,8 @@ RUNTIME_TARGETS= \
         test_mempcpy_static_write     \
         test_memset_dynamic_write     \
         test_memset_static_write      \
+        test_mbsnrtowcs_dynamic       \
+        test_mbsnrtowcs_static        \
         test_poll_dynamic             \
         test_poll_static              \
         test_ppoll_dynamic            \
diff --git a/tests/test_mbsnrtowcs_dynamic.c b/tests/test_mbsnrtowcs_dynamic.c
new file mode 100644
index 0000000..77b9082
--- /dev/null
+++ b/tests/test_mbsnrtowcs_dynamic.c
@@ -0,0 +1,28 @@
+#include "common.h"
+
+#include <wchar.h>
+#include <string.h>
+
+int main(int argc, char** argv) {
+  wchar_t buffer[4] = {0};
+  const char *src = "ABCDEFGHIJ";
+  const char *srcp = src;
+  mbstate_t st;
+  memset(&st, 0, sizeof(st));
+
+  /* Safe: convert up to 2 source bytes into at most 2 wide chars */
+  srcp = src;
+  mbsnrtowcs(buffer, &srcp, 2, 2, &st);
+
+  /* Unsafe: ask to write argc (10) wide chars into 4-element buffer.
+   * Before the fix, the else branch clamped source bytes instead of
+   * the output wide-char count, allowing destination overflow. */
+  CHK_FAIL_START
+  srcp = src;
+  memset(&st, 0, sizeof(st));
+  mbsnrtowcs(buffer, &srcp, 10, argc, &st);
+  CHK_FAIL_END
+
+  printf("%ls\n", buffer);
+  return ret;
+}
diff --git a/tests/test_mbsnrtowcs_static.c b/tests/test_mbsnrtowcs_static.c
new file mode 100644
index 0000000..755d453
--- /dev/null
+++ b/tests/test_mbsnrtowcs_static.c
@@ -0,0 +1,26 @@
+#include "common.h"
+
+#include <wchar.h>
+#include <string.h>
+
+int main(int argc, char** argv) {
+  wchar_t buffer[4] = {0};
+  const char *src = "ABCDEFGHIJKLMNOP";
+  const char *srcp = src;
+  mbstate_t st;
+  memset(&st, 0, sizeof(st));
+
+  /* Safe: convert up to 4 source bytes into at most 2 wide chars */
+  srcp = src;
+  mbsnrtowcs(buffer, &srcp, 4, 2, &st);
+
+  /* Unsafe: ask to write 16 wide chars into 4-element buffer */
+  CHK_FAIL_START
+  srcp = src;
+  memset(&st, 0, sizeof(st));
+  mbsnrtowcs(buffer, &srcp, 16, 16, &st);
+  CHK_FAIL_END
+
+  printf("%ls\n", buffer);
+  return ret;
+}