bpo-33578: Fix getstate/setstate for CJK decoder (GH-10290)

Previous version was casting to Py_ssize_t incorrectly
and exhibited unexpected behavior on big-endian systems.

2 files changed, 26 insertions, 5 deletions

diff --git a/Lib/test/test_multibytecodec.py b/Lib/test/test_multibytecodec.py
index 8e8362b..3cf5d7b 100644
--- a/Lib/test/test_multibytecodec.py
+++ b/Lib/test/test_multibytecodec.py
@@ -271,6 +271,10 @@ class Test_IncrementalDecoder(unittest.TestCase):
         pending4, _ = decoder.getstate()
         self.assertEqual(pending4, b'')
 
+        # Ensure state values are preserved correctly
+        decoder.setstate((b'abc', 123456789))
+        self.assertEqual(decoder.getstate(), (b'abc', 123456789))
+
     def test_setstate_validates_input(self):
         decoder = codecs.getincrementaldecoder('euc_jp')()
         self.assertRaises(TypeError, decoder.setstate, 123)
diff --git a/Modules/cjkcodecs/multibytecodec.c b/Modules/cjkcodecs/multibytecodec.c
index 4633499..9409456 100644
--- a/Modules/cjkcodecs/multibytecodec.c
+++ b/Modules/cjkcodecs/multibytecodec.c
@@ -1218,6 +1218,7 @@ _multibytecodec_MultibyteIncrementalDecoder_getstate_impl(MultibyteIncrementalDe
 /*[clinic end generated code: output=255009c4713b7f82 input=4006aa49bddbaa75]*/
 {
     PyObject *buffer;
+    PyObject *statelong;
 
     buffer = PyBytes_FromStringAndSize((const char *)self->pending,
                                        self->pendingsize);
@@ -1225,7 +1226,16 @@ _multibytecodec_MultibyteIncrementalDecoder_getstate_impl(MultibyteIncrementalDe
         return NULL;
     }
 
-    return make_tuple(buffer, (Py_ssize_t)*self->state.c);
+    statelong = (PyObject *)_PyLong_FromByteArray(self->state.c,
+                                                  sizeof(self->state.c),
+                                                  1 /* little-endian */ ,
+                                                  0 /* unsigned */ );
+    if (statelong == NULL) {
+        Py_DECREF(buffer);
+        return NULL;
+    }
+
+    return Py_BuildValue("NN", buffer, statelong);
 }
 
 /*[clinic input]
@@ -1240,16 +1250,23 @@ _multibytecodec_MultibyteIncrementalDecoder_setstate_impl(MultibyteIncrementalDe
 /*[clinic end generated code: output=106b2fbca3e2dcc2 input=e5d794e8baba1a47]*/
 {
     PyObject *buffer;
+    PyLongObject *statelong;
     Py_ssize_t buffersize;
     char *bufferstr;
-    unsigned long long flag;
+    unsigned char statebytes[8];
 
-    if (!PyArg_ParseTuple(state, "SK;setstate(): illegal state argument",
-                          &buffer, &flag))
+    if (!PyArg_ParseTuple(state, "SO!;setstate(): illegal state argument",
+                          &buffer, &PyLong_Type, &statelong))
     {
         return NULL;
     }
 
+    if (_PyLong_AsByteArray(statelong, statebytes, sizeof(statebytes),
+                            1 /* little-endian */ ,
+                            0 /* unsigned */ ) < 0) {
+        return NULL;
+    }
+
     buffersize = PyBytes_Size(buffer);
     if (buffersize == -1) {
         return NULL;
@@ -1266,7 +1283,7 @@ _multibytecodec_MultibyteIncrementalDecoder_setstate_impl(MultibyteIncrementalDe
     }
     self->pendingsize = buffersize;
     memcpy(self->pending, bufferstr, self->pendingsize);
-    memcpy(self->state.c, (unsigned char *)&flag, sizeof(flag));
+    memcpy(self->state.c, statebytes, sizeof(statebytes));
 
     Py_RETURN_NONE;
 }