Use ssize_t for buffer sizes consistently

Author: elprans

asyncpg/protocol/buffer.pxd

@@ -9,12 +9,12 @@ cdef class Memory:
     cdef:
         char* buf
         object owner
-        int length
+        ssize_t length
 
     cdef as_bytes(self)
 
     @staticmethod
-    cdef inline Memory new(char* buf, object owner, int length)
+    cdef inline Memory new(char* buf, object owner, ssize_t length)
 
 
 cdef class WriteBuffer:
@@ -26,10 +26,10 @@ cdef class WriteBuffer:
         char *_buf
 
         # Allocated size
-        size_t _size
+        ssize_t _size
 
         # Length of data in the buffer
-        size_t _length
+        ssize_t _length
 
         # Number of memoryviews attached to the buffer
         int _view_count
@@ -39,8 +39,8 @@ cdef class WriteBuffer:
 
     cdef inline _check_readonly(self)
     cdef inline len(self)
-    cdef inline _ensure_alloced(self, size_t extra_length)
-    cdef _reallocate(self, new_size)
+    cdef inline _ensure_alloced(self, ssize_t extra_length)
+    cdef _reallocate(self, ssize_t new_size)
     cdef inline start_message(self, char type)
     cdef inline end_message(self)
     cdef write_buffer(self, WriteBuffer buf)
@@ -81,31 +81,31 @@ cdef class ReadBuffer:
         int32_t _bufs_len
 
         # A read position in the first buffer in `_bufs`
-        int32_t _pos0
+        ssize_t _pos0
 
         # Length of the first buffer in `_bufs`
-        int32_t _len0
+        ssize_t _len0
 
         # A total number of buffered bytes in ReadBuffer
-        int32_t _length
+        ssize_t _length
 
         char _current_message_type
         int _current_message_len
-        int32_t _current_message_len_unread
+        ssize_t _current_message_len_unread
         bint _current_message_ready
 
     cdef feed_data(self, data)
     cdef inline _ensure_first_buf(self)
     cdef _switch_to_next_buf(self)
     cdef inline read_byte(self)
-    cdef inline char* _try_read_bytes(self, int nbytes)
-    cdef inline read(self, int nbytes)
+    cdef inline char* _try_read_bytes(self, ssize_t nbytes)
+    cdef inline read(self, ssize_t nbytes)
     cdef inline read_bytes(self, ssize_t n)
     cdef inline read_int32(self)
     cdef inline read_int16(self)
     cdef inline read_cstr(self)
     cdef int32_t has_message(self) except -1
-    cdef inline char* try_consume_message(self, int32_t* len)
+    cdef inline char* try_consume_message(self, ssize_t* len)
     cdef Memory consume_message(self)
     cdef discard_message(self)
     cdef inline _discard_message(self)
@@ -119,13 +119,13 @@ cdef class ReadBuffer:
 cdef class FastReadBuffer:
     cdef:
         const char* buf
-        size_t len
+        ssize_t len
 
-    cdef inline const char* read(self, size_t n) except NULL
+    cdef inline const char* read(self, ssize_t n) except NULL
     cdef inline const char* read_all(self)
     cdef inline FastReadBuffer slice_from(self, FastReadBuffer source,
-                                          size_t len)
-    cdef _raise_ins_err(self, size_t n, size_t len)
+                                          ssize_t len)
+    cdef _raise_ins_err(self, ssize_t n, ssize_t len)
 
     @staticmethod
     cdef FastReadBuffer new()

asyncpg/protocol/buffer.pyx

@@ -21,7 +21,7 @@ cdef class Memory:
         return cpython.PyBytes_FromStringAndSize(self.buf, self.length)
 
     @staticmethod
-    cdef inline Memory new(char* buf, object owner, int length):
+    cdef inline Memory new(char* buf, object owner, ssize_t length):
         cdef Memory mem
         mem = Memory.__new__(Memory)
         mem.buf = buf
@@ -70,13 +70,13 @@ cdef class WriteBuffer:
     cdef inline len(self):
         return self._length
 
-    cdef inline _ensure_alloced(self, size_t extra_length):
-        cdef size_t new_size = extra_length + self._length
+    cdef inline _ensure_alloced(self, ssize_t extra_length):
+        cdef ssize_t new_size = extra_length + self._length
 
         if new_size > self._size:
             self._reallocate(new_size)
 
-    cdef _reallocate(self, new_size):
+    cdef _reallocate(self, ssize_t new_size):
         cdef char *new_buf
 
         if new_size < _BUFFER_MAX_GROW:
@@ -86,18 +86,18 @@ cdef class WriteBuffer:
             new_size += _BUFFER_INITIAL_SIZE
 
         if self._smallbuf_inuse:
-            new_buf = <char*>PyMem_Malloc(sizeof(char) * new_size)
+            new_buf = <char*>PyMem_Malloc(sizeof(char) * <size_t>new_size)
             if new_buf is NULL:
                 self._buf = NULL
                 self._size = 0
                 self._length = 0
                 raise MemoryError
-            memcpy(new_buf, self._buf, self._size)
+            memcpy(new_buf, self._buf, <size_t>self._size)
             self._size = new_size
             self._buf = new_buf
             self._smallbuf_inuse = False
         else:
-            new_buf = <char*>PyMem_Realloc(<void*>self._buf, new_size)
+            new_buf = <char*>PyMem_Realloc(<void*>self._buf, <size_t>new_size)
             if new_buf is NULL:
                 PyMem_Free(self._buf)
                 self._buf = NULL
@@ -117,16 +117,18 @@ cdef class WriteBuffer:
 
     cdef inline end_message(self):
         # "length-1" to exclude the message type byte
-        cdef size_t mlen = self._length - 1
+        cdef ssize_t mlen = self._length - 1
 
         self._check_readonly()
         if not self._message_mode:
             raise BufferError(
                 'end_message can only be called with start_message')
         if self._length < 5:
             raise BufferError('end_message: buffer is too small')
+        if mlen > _MAXINT32:
+            raise BufferError('end_message: message is too large')
 
-        hton.pack_int32(&self._buf[1], mlen)
+        hton.pack_int32(&self._buf[1], <int32_t>mlen)
         return self
 
     cdef write_buffer(self, WriteBuffer buf):
@@ -138,7 +140,7 @@ cdef class WriteBuffer:
         self._ensure_alloced(buf._length)
         memcpy(self._buf + self._length,
                <void*>buf._buf,
-               buf._length)
+               <size_t>buf._length)
         self._length += buf._length
 
     cdef write_byte(self, char b):
@@ -171,7 +173,7 @@ cdef class WriteBuffer:
         self._check_readonly()
         self._ensure_alloced(len)
 
-        memcpy(self._buf + self._length, <void*>data, len)
+        memcpy(self._buf + self._length, <void*>data, <size_t>len)
         self._length += len
 
     cdef write_int16(self, int16_t i):
@@ -246,7 +248,7 @@ cdef class ReadBuffer:
 
     cdef feed_data(self, data):
         cdef:
-            int32_t dlen
+            ssize_t dlen
             bytes data_bytes
 
         if not cpython.PyBytes_CheckExact(data):
@@ -293,7 +295,7 @@ cdef class ReadBuffer:
                 raise RuntimeError(
                     'debug: second buffer of ReadBuffer is empty')
 
-    cdef inline char* _try_read_bytes(self, int nbytes):
+    cdef inline char* _try_read_bytes(self, ssize_t nbytes):
         # Important: caller must call _ensure_first_buf() prior
         # to calling try_read_bytes, and must not overread
 
@@ -319,10 +321,10 @@ cdef class ReadBuffer:
         else:
             return NULL
 
-    cdef inline read(self, int nbytes):
+    cdef inline read(self, ssize_t nbytes):
         cdef:
             object result
-            int nread
+            ssize_t nread
             char *buf
 
         self._ensure_first_buf()
@@ -419,8 +421,8 @@ cdef class ReadBuffer:
                 'to be in the buffer')
 
         cdef:
-            int pos
-            int nread
+            ssize_t pos
+            ssize_t nread
             bytes result
             char* buf
             char* buf_start
@@ -504,8 +506,8 @@ cdef class ReadBuffer:
         self._current_message_ready = 1
         return 1
 
-    cdef inline char* try_consume_message(self, int32_t* len):
-        cdef int32_t buf_len
+    cdef inline char* try_consume_message(self, ssize_t* len):
+        cdef ssize_t buf_len
 
         if not self._current_message_ready:
             return NULL
@@ -579,7 +581,7 @@ cdef class ReadBuffer:
 @cython.freelist(_BUFFER_FREELIST_SIZE)
 cdef class FastReadBuffer:
 
-    cdef inline const char* read(self, size_t n) except NULL:
+    cdef inline const char* read(self, ssize_t n) except NULL:
         cdef const char *result
 
         if n > self.len:
@@ -599,12 +601,12 @@ cdef class FastReadBuffer:
         return result
 
     cdef inline FastReadBuffer slice_from(self, FastReadBuffer source,
-                                          size_t len):
+                                          ssize_t len):
         self.buf = source.read(len)
         self.len = len
         return self
 
-    cdef _raise_ins_err(self, size_t n, size_t len):
+    cdef _raise_ins_err(self, ssize_t n, ssize_t len):
         raise BufferError(
             'insufficient data in buffer: requested {}, remaining {}'.
                 format(n, self.len))

asyncpg/protocol/codecs/array.pyx

@@ -38,21 +38,26 @@ cdef inline _is_sub_array(object obj):
 
 cdef _get_array_shape(object obj, int32_t *dims, int32_t *ndims):
     cdef:
-        int32_t mylen = len(obj)
-        int32_t elemlen = -2
+        ssize_t mylen = len(obj)
+        ssize_t elemlen = -2
         object it
 
+    if mylen > _MAXINT32:
+        raise ValueError('too many elements in array value')
+
     if ndims[0] > ARRAY_MAXDIM:
         raise ValueError(
             'number of array dimensions ({}) exceed the maximum expected ({})'.
                 format(ndims[0], ARRAY_MAXDIM))
 
-    dims[ndims[0] - 1] = mylen
+    dims[ndims[0] - 1] = <int32_t>mylen
 
     for elem in obj:
         if _is_sub_array(elem):
             if elemlen == -2:
                 elemlen = len(elem)
+                if elemlen > _MAXINT32:
+                    raise ValueError('too many elements in array value')
                 ndims[0] += 1
                 _get_array_shape(elem, dims, ndims)
             else:
@@ -123,7 +128,7 @@ cdef inline array_encode(ConnectionSettings settings, WriteBuffer buf,
     # flags
     buf.write_int32(0)
     # element type
-    buf.write_int32(elem_oid)
+    buf.write_int32(<int32_t>elem_oid)
     # upper / lower bounds
     for i in range(ndims):
         buf.write_int32(dims[i])
@@ -137,9 +142,9 @@ cdef inline array_decode(ConnectionSettings settings, FastReadBuffer buf,
     cdef:
         int32_t ndims = hton.unpack_int32(buf.read(4))
         int32_t flags = hton.unpack_int32(buf.read(4))
-        uint32_t elem_oid = hton.unpack_int32(buf.read(4))
+        uint32_t elem_oid = <uint32_t>hton.unpack_int32(buf.read(4))
         list result
-        uint32_t i
+        int i
         int32_t elem_len
         int64_t elem_count = 1
         FastReadBuffer elem_buf = FastReadBuffer.new()

asyncpg/protocol/codecs/base.pyx

@@ -95,20 +95,25 @@ cdef class Codec:
                           object obj):
         cdef:
             WriteBuffer elem_data
-            int32_t i
+            int i
             list elem_codecs = self.element_codecs
+            ssize_t count
+
+        count = len(obj)
+        if count > _MAXINT32:
+            raise ValueError('too many elements in composite type record')
 
         elem_data = WriteBuffer.new()
         i = 0
         for item in obj:
-            elem_data.write_int32(self.element_type_oids[i])
+            elem_data.write_int32(<int32_t>self.element_type_oids[i])
             if item is None:
                 elem_data.write_int32(-1)
             else:
                 (<Codec>elem_codecs[i]).encode(settings, elem_data, item)
             i += 1
 
-        record_encode_frame(settings, buf, elem_data, len(obj))
+        record_encode_frame(settings, buf, elem_data, <int32_t>count)
 
     cdef encode_in_python(self, ConnectionSettings settings, WriteBuffer buf,
                           object obj):
@@ -145,11 +150,11 @@ cdef class Codec:
             Codec elem_codec
             FastReadBuffer elem_buf = FastReadBuffer.new()
 
-        elem_count = hton.unpack_int32(buf.read(4))
+        elem_count = <uint32_t>hton.unpack_int32(buf.read(4))
         result = record.ApgRecord_New(self.element_names, elem_count)
         for i in range(elem_count):
             elem_typ = self.element_type_oids[i]
-            received_elem_typ = hton.unpack_int32(buf.read(4))
+            received_elem_typ = <uint32_t>hton.unpack_int32(buf.read(4))
 
             if received_elem_typ != elem_typ:
                 raise RuntimeError(

asyncpg/protocol/codecs/bits.pyx

@@ -14,7 +14,7 @@ cdef bits_encode(ConnectionSettings settings, WriteBuffer wbuf, obj):
         bint pybuf_used = False
         char *buf
         ssize_t len
-        int32_t bitlen
+        ssize_t bitlen
 
     if cpython.PyBytes_CheckExact(obj):
         buf = cpython.PyBytes_AS_STRING(obj)
@@ -31,7 +31,9 @@ cdef bits_encode(ConnectionSettings settings, WriteBuffer wbuf, obj):
         bitlen = len * 8
 
     try:
-        wbuf.write_int32(4 + len)
+        if bitlen > _MAXINT32:
+            raise ValueError('bit value too long')
+        wbuf.write_int32(4 + <int32_t>len)
         wbuf.write_int32(<int32_t>bitlen)
         wbuf.write_cstr(buf, len)
     finally:
@@ -42,7 +44,7 @@ cdef bits_encode(ConnectionSettings settings, WriteBuffer wbuf, obj):
 cdef bits_decode(ConnectionSettings settings, FastReadBuffer buf):
     cdef:
         int32_t bitlen = hton.unpack_int32(buf.read(4))
-        size_t buf_len = buf.len
+        ssize_t buf_len = buf.len
 
     bytes_ = cpython.PyBytes_FromStringAndSize(buf.read_all(), buf_len)
     return BitString.frombytes(bytes_, bitlen)

asyncpg/protocol/codecs/bytea.pyx

@@ -30,7 +30,7 @@ cdef bytea_encode(ConnectionSettings settings, WriteBuffer wbuf, obj):
 
 
 cdef bytea_decode(ConnectionSettings settings, FastReadBuffer buf):
-    cdef size_t buf_len = buf.len
+    cdef ssize_t buf_len = buf.len
     return cpython.PyBytes_FromStringAndSize(buf.read_all(), buf_len)