Fix handling of OIDs >= 2**31

Currently asyncpg (incorrectly) assumes OIDs to be signed 32-bit
integers, whereas in reality they are unsigned.  As a result, things
would crash once the OID sequence reaches 2**31.

Fix this by decoding OID values as unsigned longs.

Fixes: #279

Author: elprans

asyncpg/cluster.py

@@ -300,6 +300,42 @@ def get_connection_spec(self):
     def override_connection_spec(self, **kwargs):
         self._connection_spec_override = kwargs
 
+    def reset_wal(self, *, oid=None, xid=None):
+        status = self.get_status()
+        if status == 'not-initialized':
+            raise ClusterError(
+                'cannot modify WAL status: cluster is not initialized')
+
+        if status == 'running':
+            raise ClusterError(
+                'cannot modify WAL status: cluster is running')
+
+        opts = []
+        if oid is not None:
+            opts.extend(['-o', str(oid)])
+        if xid is not None:
+            opts.extend(['-x', str(xid)])
+        if not opts:
+            return
+
+        opts.append(self._data_dir)
+
+        try:
+            reset_wal = self._find_pg_binary('pg_resetwal')
+        except ClusterError:
+            reset_wal = self._find_pg_binary('pg_resetxlog')
+
+        process = subprocess.run(
+            [reset_wal] + opts,
+            stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+        stderr = process.stderr
+
+        if process.returncode != 0:
+            raise ClusterError(
+                'pg_resetwal exited with status {:d}: {}'.format(
+                    process.returncode, stderr.decode()))
+
     def reset_hba(self):
         """Remove all records from pg_hba.conf."""
         status = self.get_status()

asyncpg/protocol/buffer.pxd

@@ -97,13 +97,13 @@ cdef class ReadBuffer:
     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 read_byte(self) except? -1
     cdef inline const char* _try_read_bytes(self, ssize_t nbytes)
     cdef inline _read(self, char *buf, ssize_t nbytes)
     cdef read(self, ssize_t nbytes)
     cdef inline const char* read_bytes(self, ssize_t n) except NULL
-    cdef inline read_int32(self)
-    cdef inline read_int16(self)
+    cdef inline int32_t read_int32(self) except? -1
+    cdef inline int16_t read_int16(self) except? -1
     cdef inline read_cstr(self)
     cdef int32_t has_message(self) except -1
     cdef inline int32_t has_message_type(self, char mtype) except -1

asyncpg/protocol/buffer.pyx

@@ -376,7 +376,7 @@ cdef class ReadBuffer:
 
         return Memory.new(buf, result, nbytes)
 
-    cdef inline read_byte(self):
+    cdef inline char read_byte(self) except? -1:
         cdef const char *first_byte
 
         if ASYNCPG_DEBUG:
@@ -404,7 +404,7 @@ cdef class ReadBuffer:
             mem = <Memory>(self.read(n))
             return mem.buf
 
-    cdef inline read_int32(self):
+    cdef inline int32_t read_int32(self) except? -1:
         cdef:
             Memory mem
             const char *cbuf
@@ -417,7 +417,7 @@ cdef class ReadBuffer:
             mem = <Memory>(self.read(4))
             return hton.unpack_int32(mem.buf)
 
-    cdef inline read_int16(self):
+    cdef inline int16_t read_int16(self) except? -1:
         cdef:
             Memory mem
             const char *cbuf

asyncpg/protocol/codecs/base.pyx

@@ -373,6 +373,22 @@ cdef codec_decode_func_ex(ConnectionSettings settings, FastReadBuffer buf,
     return (<Codec>arg).decode(settings, buf)
 
 
+cdef uint32_t pylong_as_oid(val) except? 0xFFFFFFFFl:
+    cdef:
+        int64_t oid = 0
+        bint overflow = False
+
+    try:
+        oid = cpython.PyLong_AsLongLong(val)
+    except OverflowError:
+        overflow = True
+
+    if overflow or (oid < 0 or oid > UINT32_MAX):
+        raise OverflowError('OID value too large: {!r}'.format(val))
+
+    return <uint32_t>val
+
+
 cdef class DataCodecConfig:
     def __init__(self, cache_key):
         try:
@@ -523,9 +539,10 @@ cdef class DataCodecConfig:
             Codec core_codec
             encode_func c_encoder = NULL
             decode_func c_decoder = NULL
+            uint32_t oid = pylong_as_oid(typeoid)
 
         if xformat == PG_XFORMAT_TUPLE:
-            core_codec = get_any_core_codec(typeoid, format, xformat)
+            core_codec = get_any_core_codec(oid, format, xformat)
             if core_codec is None:
                 raise ValueError(
                     "{} type does not support 'tuple' exchange format".format(
@@ -538,7 +555,7 @@ cdef class DataCodecConfig:
         self.remove_python_codec(typeoid, typename, typeschema)
 
         self._local_type_codecs[typeoid] = \
-            Codec.new_python_codec(typeoid, typename, typeschema, typekind,
+            Codec.new_python_codec(oid, typename, typeschema, typekind,
                                    encoder, decoder, c_encoder, c_decoder,
                                    format, xformat)
 
@@ -551,19 +568,18 @@ cdef class DataCodecConfig:
         cdef:
             Codec codec
             Codec target_codec
+            uint32_t oid = pylong_as_oid(typeoid)
+            uint32_t alias_pid
 
         if format == PG_FORMAT_ANY:
             formats = (PG_FORMAT_BINARY, PG_FORMAT_TEXT)
         else:
             formats = (format,)
 
         for format in formats:
-            if self.get_codec(typeoid, format) is not None:
-                raise ValueError('cannot override codec for type {}'.format(
-                    typeoid))
-
             if isinstance(alias_to, int):
-                target_codec = self.get_codec(alias_to, format)
+                alias_oid = pylong_as_oid(alias_to)
+                target_codec = self.get_codec(alias_oid, format)
             else:
                 target_codec = get_extra_codec(alias_to, format)
 

asyncpg/protocol/codecs/int.pyx

@@ -27,9 +27,9 @@ cdef int2_encode(ConnectionSettings settings, WriteBuffer buf, obj):
     except OverflowError:
         overflow = 1
 
-    if overflow or val < -32768 or val > 32767:
+    if overflow or val < INT16_MIN or val > INT16_MAX:
         raise OverflowError(
-            'int too big to be encoded as INT2: {!r}'.format(obj))
+            'int16 value out of range: {!r}'.format(obj))
 
     buf.write_int32(2)
     buf.write_int16(<int16_t>val)
@@ -49,10 +49,9 @@ cdef int4_encode(ConnectionSettings settings, WriteBuffer buf, obj):
         overflow = 1
 
     # "long" and "long long" have the same size for x86_64, need an extra check
-    if overflow or (sizeof(val) > 4 and (val < -2147483648 or
-                                         val > 2147483647)):
+    if overflow or (sizeof(val) > 4 and (val < INT32_MIN or val > INT32_MAX)):
         raise OverflowError(
-            'int too big to be encoded as INT4: {!r}'.format(obj))
+            'int32 value out of range: {!r}'.format(obj))
 
     buf.write_int32(4)
     buf.write_int32(<int32_t>val)
@@ -62,6 +61,29 @@ cdef int4_decode(ConnectionSettings settings, FastReadBuffer buf):
     return cpython.PyLong_FromLong(hton.unpack_int32(buf.read(4)))
 
 
+cdef uint4_encode(ConnectionSettings settings, WriteBuffer buf, obj):
+    cdef int overflow = 0
+    cdef unsigned long val = 0
+
+    try:
+        val = cpython.PyLong_AsUnsignedLong(obj)
+    except OverflowError:
+        overflow = 1
+
+    # "long" and "long long" have the same size for x86_64, need an extra check
+    if overflow or (sizeof(val) > 4 and val > UINT32_MAX):
+        raise OverflowError(
+            'uint32 value out of range: {!r}'.format(obj))
+
+    buf.write_int32(4)
+    buf.write_int32(<int32_t>val)
+
+
+cdef uint4_decode(ConnectionSettings settings, FastReadBuffer buf):
+    return cpython.PyLong_FromUnsignedLong(
+        <uint32_t>hton.unpack_int32(buf.read(4)))
+
+
 cdef int8_encode(ConnectionSettings settings, WriteBuffer buf, obj):
     cdef int overflow = 0
     cdef long long val
@@ -72,10 +94,9 @@ cdef int8_encode(ConnectionSettings settings, WriteBuffer buf, obj):
         overflow = 1
 
     # Just in case for systems with "long long" bigger than 8 bytes
-    if overflow or (sizeof(val) > 8 and (val < -9223372036854775808 or
-                                         val > 9223372036854775807)):
+    if overflow or (sizeof(val) > 8 and (val < INT64_MIN or val > INT64_MAX)):
         raise OverflowError(
-            'int too big to be encoded as INT8: {!r}'.format(obj))
+            'int64 value out of range: {!r}'.format(obj))
 
     buf.write_int32(8)
     buf.write_int64(<int64_t>val)

asyncpg/protocol/codecs/misc.pyx

@@ -36,8 +36,8 @@ cdef init_pseudo_codecs():
 
     for oid_type in oid_types:
         register_core_codec(oid_type,
-                            <encode_func>&int4_encode,
-                            <decode_func>&int4_decode,
+                            <encode_func>&uint4_encode,
+                            <decode_func>&uint4_decode,
                             PG_FORMAT_BINARY)
 
     # reg* types -- these are really system catalog OIDs, but

asyncpg/protocol/codecs/tid.pyx

@@ -23,9 +23,9 @@ cdef tid_encode(ConnectionSettings settings, WriteBuffer buf, obj):
         overflow = 1
 
     # "long" and "long long" have the same size for x86_64, need an extra check
-    if overflow or (sizeof(block) > 4 and block > 4294967295):
+    if overflow or (sizeof(block) > 4 and block > UINT32_MAX):
         raise OverflowError(
-            'block too big to be encoded as UINT4: {!r}'.format(obj[0]))
+            'tuple id block value out of range: {!r}'.format(obj[0]))
 
     try:
         offset = cpython.PyLong_AsUnsignedLong(obj[1])
@@ -35,7 +35,7 @@ cdef tid_encode(ConnectionSettings settings, WriteBuffer buf, obj):
 
     if overflow or offset > 65535:
         raise OverflowError(
-            'offset too big to be encoded as UINT2: {!r}'.format(obj[1]))
+            'tuple id offset value out of range: {!r}'.format(obj[1]))
 
     buf.write_int32(6)
     buf.write_int32(<int32_t>block)

asyncpg/protocol/prepared_stmt.pyx

@@ -163,7 +163,7 @@ cdef class PreparedStatementState:
             list cols_names
             object cols_mapping
             tuple row
-            int oid
+            uint32_t oid
             Codec codec
             list codecs
 
@@ -183,7 +183,7 @@ cdef class PreparedStatementState:
             cols_mapping[col_name] = i
             cols_names.append(col_name)
             oid = row[3]
-            codec = self.settings.get_data_codec(<uint32_t>oid)
+            codec = self.settings.get_data_codec(oid)
             if codec is None or not codec.has_decoder():
                 raise RuntimeError('no decoder for OID {}'.format(oid))
             if not codec.is_binary():
@@ -198,7 +198,7 @@ cdef class PreparedStatementState:
 
     cdef _ensure_args_encoder(self):
         cdef:
-            int p_oid
+            uint32_t p_oid
             Codec codec
             list codecs = []
 
@@ -207,7 +207,7 @@ cdef class PreparedStatementState:
 
         for i from 0 <= i < self.args_num:
             p_oid = self.parameters_desc[i]
-            codec = self.settings.get_data_codec(<uint32_t>p_oid)
+            codec = self.settings.get_data_codec(p_oid)
             if codec is None or not codec.has_encoder():
                 raise RuntimeError('no encoder for OID {}'.format(p_oid))
             if codec.type not in {}:
@@ -290,14 +290,14 @@ cdef _decode_parameters_desc(object desc):
     cdef:
         ReadBuffer reader
         int16_t nparams
-        int32_t p_oid
+        uint32_t p_oid
         list result = []
 
     reader = ReadBuffer.new_message_parser(desc)
     nparams = reader.read_int16()
 
     for i from 0 <= i < nparams:
-        p_oid = reader.read_int32()
+        p_oid = <uint32_t>reader.read_int32()
         result.append(p_oid)
 
     return result
@@ -310,9 +310,9 @@ cdef _decode_row_desc(object desc):
         int16_t nfields
 
         bytes f_name
-        int32_t f_table_oid
+        uint32_t f_table_oid
         int16_t f_column_num
-        int32_t f_dt_oid
+        uint32_t f_dt_oid
         int16_t f_dt_size
         int32_t f_dt_mod
         int16_t f_format
@@ -325,9 +325,9 @@ cdef _decode_row_desc(object desc):
 
     for i from 0 <= i < nfields:
         f_name = reader.read_cstr()
-        f_table_oid = reader.read_int32()
+        f_table_oid = <uint32_t>reader.read_int32()
         f_column_num = reader.read_int16()
-        f_dt_oid = reader.read_int32()
+        f_dt_oid = <uint32_t>reader.read_int32()
         f_dt_size = reader.read_int16()
         f_dt_mod = reader.read_int32()
         f_format = reader.read_int16()

asyncpg/protocol/protocol.pyx

@@ -18,7 +18,9 @@ import socket
 import time
 
 from libc.stdint cimport int8_t, uint8_t, int16_t, uint16_t, \
-                         int32_t, uint32_t, int64_t, uint64_t
+                         int32_t, uint32_t, int64_t, uint64_t, \
+                         INT16_MIN, INT16_MAX, INT32_MIN, INT32_MAX, \
+                         UINT32_MAX, INT64_MIN, INT64_MAX
 
 from asyncpg.protocol cimport record