gh-111140: PyLong_From/AsNativeBytes: Take *flags* rather than just *endianness* (GH-116053)

1 files changed, 82 insertions, 24 deletions

diff --git a/Objects/longobject.c b/Objects/longobject.c
index cc2fe11..c4ab064 100644
--- a/Objects/longobject.c
+++ b/Objects/longobject.c
@@ -1083,18 +1083,17 @@ _fits_in_n_bits(Py_ssize_t v, Py_ssize_t n)
 static inline int
 _resolve_endianness(int *endianness)
 {
-    if (*endianness < 0) {
+    if (*endianness == -1 || (*endianness & 2)) {
         *endianness = PY_LITTLE_ENDIAN;
+    } else {
+        *endianness &= 1;
     }
-    if (*endianness != 0 && *endianness != 1) {
-        PyErr_SetString(PyExc_SystemError, "invalid 'endianness' value");
-        return -1;
-    }
+    assert(*endianness == 0 || *endianness == 1);
     return 0;
 }
 
 Py_ssize_t
-PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int endianness)
+PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int flags)
 {
     PyLongObject *v;
     union {
@@ -1109,7 +1108,7 @@ PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int endianness)
         return -1;
     }
 
-    int little_endian = endianness;
+    int little_endian = flags;
     if (_resolve_endianness(&little_endian) < 0) {
         return -1;
     }
@@ -1125,6 +1124,15 @@ PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int endianness)
         do_decref = 1;
     }
 
+    if ((flags != -1 && (flags & Py_ASNATIVEBYTES_REJECT_NEGATIVE))
+        && _PyLong_IsNegative(v)) {
+        PyErr_SetString(PyExc_ValueError, "Cannot convert negative int");
+        if (do_decref) {
+            Py_DECREF(v);
+        }
+        return -1;
+    }
+
     if (_PyLong_IsCompact(v)) {
         res = 0;
         cv.v = _PyLong_CompactValue(v);
@@ -1159,6 +1167,15 @@ PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int endianness)
             /* If we fit, return the requested number of bytes */
             if (_fits_in_n_bits(cv.v, n * 8)) {
                 res = n;
+            } else if (cv.v > 0 && _fits_in_n_bits(cv.v, n * 8 + 1)) {
+                /* Positive values with the MSB set do not require an
+                 * additional bit when the caller's intent is to treat them
+                 * as unsigned. */
+                if (flags == -1 || (flags & Py_ASNATIVEBYTES_UNSIGNED_BUFFER)) {
+                    res = n;
+                } else {
+                    res = n + 1;
+                }
             }
         }
         else {
@@ -1199,17 +1216,55 @@ PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int endianness)
             _PyLong_AsByteArray(v, buffer, (size_t)n, little_endian, 1, 0);
         }
 
-        // More efficient calculation for number of bytes required?
+        /* Calculates the number of bits required for the *absolute* value
+         * of v. This does not take sign into account, only magnitude. */
         size_t nb = _PyLong_NumBits((PyObject *)v);
-        /* Normally this would be((nb - 1) / 8) + 1 to avoid rounding up
-         * multiples of 8 to the next byte, but we add an implied bit for
-         * the sign and it cancels out. */
-        size_t n_needed = (nb / 8) + 1;
-        res = (Py_ssize_t)n_needed;
-        if ((size_t)res != n_needed) {
-            PyErr_SetString(PyExc_OverflowError,
-                "value too large to convert");
+        if (nb == (size_t)-1) {
             res = -1;
+        } else {
+            /* Normally this would be((nb - 1) / 8) + 1 to avoid rounding up
+             * multiples of 8 to the next byte, but we add an implied bit for
+             * the sign and it cancels out. */
+            res = (Py_ssize_t)(nb / 8) + 1;
+        }
+
+        /* Two edge cases exist that are best handled after extracting the
+         * bits. These may result in us reporting overflow when the value
+         * actually fits.
+         */
+        if (n > 0 && res == n + 1 && nb % 8 == 0) {
+            if (_PyLong_IsNegative(v)) {
+                /* Values of 0x80...00 from negative values that use every
+                 * available bit in the buffer do not require an additional
+                 * bit to store the sign. */
+                int is_edge_case = 1;
+                unsigned char *b = (unsigned char *)buffer;
+                for (Py_ssize_t i = 0; i < n && is_edge_case; ++i, ++b) {
+                    if (i == 0) {
+                        is_edge_case = (*b == (little_endian ? 0 : 0x80));
+                    } else if (i < n - 1) {
+                        is_edge_case = (*b == 0);
+                    } else {
+                        is_edge_case = (*b == (little_endian ? 0x80 : 0));
+                    }
+                }
+                if (is_edge_case) {
+                    res = n;
+                }
+            }
+            else {
+                /* Positive values with the MSB set do not require an
+                 * additional bit when the caller's intent is to treat them
+                 * as unsigned. */
+                unsigned char *b = (unsigned char *)buffer;
+                if (b[little_endian ? n - 1 : 0] & 0x80) {
+                    if (flags == -1 || (flags & Py_ASNATIVEBYTES_UNSIGNED_BUFFER)) {
+                        res = n;
+                    } else {
+                        res = n + 1;
+                    }
+                }
+            }
         }
     }
 
@@ -1222,38 +1277,41 @@ PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int endianness)
 
 
 PyObject *
-PyLong_FromNativeBytes(const void* buffer, size_t n, int endianness)
+PyLong_FromNativeBytes(const void* buffer, size_t n, int flags)
 {
     if (!buffer) {
         PyErr_BadInternalCall();
         return NULL;
     }
 
-    int little_endian = endianness;
+    int little_endian = flags;
     if (_resolve_endianness(&little_endian) < 0) {
         return NULL;
     }
 
-    return _PyLong_FromByteArray((const unsigned char *)buffer, n,
-                                 little_endian, 1);
+    return _PyLong_FromByteArray(
+        (const unsigned char *)buffer,
+        n,
+        little_endian,
+        (flags == -1 || !(flags & Py_ASNATIVEBYTES_UNSIGNED_BUFFER)) ? 1 : 0
+    );
 }
 
 
 PyObject *
-PyLong_FromUnsignedNativeBytes(const void* buffer, size_t n, int endianness)
+PyLong_FromUnsignedNativeBytes(const void* buffer, size_t n, int flags)
 {
     if (!buffer) {
         PyErr_BadInternalCall();
         return NULL;
     }
 
-    int little_endian = endianness;
+    int little_endian = flags;
     if (_resolve_endianness(&little_endian) < 0) {
         return NULL;
     }
 
-    return _PyLong_FromByteArray((const unsigned char *)buffer, n,
-                                 little_endian, 0);
+    return _PyLong_FromByteArray((const unsigned char *)buffer, n, little_endian, 0);
 }