Test embedded NUL byte/character

Author: vstinner

Lib/test/test_codecs.py

@@ -4067,7 +4067,8 @@ class LocaleCodecTest(unittest.TestCase):
     STRINGS = ("ascii", "ulatin1:\xa7\xe9",
                "u255:\xff",
                "UCS:\xe9\u20ac\U0010ffff",
-               "surrogates:\uDC80\uDCFF")
+               "surrogates:\uDC80\uDCFF",
+               "embed\0char")
     BYTES_STRINGS = (b"blatin1:\xa7\xe9", b"b255:\xff")
     SURROGATES = "\uDC80\uDCFF"
 
@@ -4085,6 +4086,10 @@ def check_encode_strings(self, errors):
             with self.subTest(text=text):
                 try:
                     expected = text.encode(self.ENCODING, errors)
+                    if b"\0" in expected:
+                        # Py_EncodeLocale() and _Py_EncodeLocale()
+                        # truncate the input string at the first NUL character
+                        expected = expected.partition(b'\0')[0]
                 except UnicodeEncodeError:
                     for error_pos in range(len(text)):
                         try:
@@ -4169,8 +4174,12 @@ def check_decode_strings(self, errors):
             with self.subTest(encoded=encoded):
                 try:
                     expected = encoded.decode(self.ENCODING, errors)
+                    if "\0" in expected:
+                        # Py_DecodeLocale() and _Py_DecodeLocale() truncate
+                        # the input string at the first NUL byte
+                        expected = expected.partition('\0')[0]
                 except UnicodeDecodeError:
-                    for error_pos in range(len(text) - 1, -1, -1):
+                    for error_pos in range(len(encoded) - 1, -1, -1):
                         try:
                             encoded[:error_pos].decode(self.ENCODING, errors)
                         except UnicodeDecodeError:

Modules/_testinternalcapi.c

@@ -1059,7 +1059,9 @@ encode_locale(PyObject *self, PyObject *args)
         return NULL;
     }
 
-    wchar_t *wstr = PyUnicode_AsWideCharString(unicode, NULL);
+    // Accept embedded null characters
+    Py_ssize_t unused_wlen;
+    wchar_t *wstr = PyUnicode_AsWideCharString(unicode, &unused_wlen);
     if (wstr == NULL) {
         return NULL;
     }
@@ -1115,11 +1117,13 @@ static PyObject *
 decode_locale(PyObject *self, PyObject *args)
 {
     char *str;
+    Py_ssize_t unused_len;
     int current_locale = 0;
     PyObject *res = NULL;
     const char *errors = NULL;
-
-    if (!PyArg_ParseTuple(args, "y|is", &str, &current_locale, &errors)) {
+    // Accept embedded null bytes in str
+    if (!PyArg_ParseTuple(args, "y#|is",
+                          &str, &unused_len, &current_locale, &errors)) {
         return NULL;
     }
     _Py_error_handler error_handler = _Py_GetErrorHandler(errors);

Modules/_testlimitedcapi/codec.c

@@ -21,7 +21,9 @@ static PyObject *
 decode_locale(PyObject *Py_UNUSED(module), PyObject *arg)
 {
     const char *str;
-    if (PyArg_Parse(arg, "y", &str) < 0) {
+    Py_ssize_t unused_len;
+    // Accept embedded null bytes
+    if (PyArg_Parse(arg, "y#", &str, &unused_len) < 0) {
         return NULL;
     }
 
@@ -61,7 +63,9 @@ encode_locale(PyObject *Py_UNUSED(module), PyObject *arg)
         return NULL;
     }
 
-    wchar_t *wstr = PyUnicode_AsWideCharString(unicode, NULL);
+    // Accept embedded null characters
+    Py_ssize_t unused_wlen;
+    wchar_t *wstr = PyUnicode_AsWideCharString(unicode, &unused_wlen);
     if (wstr == NULL) {
         return NULL;
     }