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/_testbase/__init__.py

@@ -205,16 +205,21 @@ def _format_loop_exception(self, context, n):
 _default_cluster = None
 
 
-def _start_cluster(ClusterCls, cluster_kwargs, server_settings,
-                   initdb_options=None):
+def _init_cluster(ClusterCls, cluster_kwargs, initdb_options=None):
     cluster = ClusterCls(**cluster_kwargs)
     cluster.init(**(initdb_options or {}))
     cluster.trust_local_connections()
-    cluster.start(port='dynamic', server_settings=server_settings)
     atexit.register(_shutdown_cluster, cluster)
     return cluster
 
 
+def _start_cluster(ClusterCls, cluster_kwargs, server_settings,
+                   initdb_options=None):
+    cluster = _init_cluster(ClusterCls, cluster_kwargs, initdb_options)
+    cluster.start(port='dynamic', server_settings=server_settings)
+    return cluster
+
+
 def _get_initdb_options(initdb_options=None):
     if not initdb_options:
         initdb_options = {}
@@ -228,7 +233,7 @@ def _get_initdb_options(initdb_options=None):
     return initdb_options
 
 
-def _start_default_cluster(server_settings={}, initdb_options=None):
+def _init_default_cluster(initdb_options=None):
     global _default_cluster
 
     if _default_cluster is None:
@@ -237,9 +242,8 @@ def _start_default_cluster(server_settings={}, initdb_options=None):
             # Using existing cluster, assuming it is initialized and running
             _default_cluster = pg_cluster.RunningCluster()
         else:
-            _default_cluster = _start_cluster(
+            _default_cluster = _init_cluster(
                 pg_cluster.TempCluster, cluster_kwargs={},
-                server_settings=server_settings,
                 initdb_options=_get_initdb_options(initdb_options))
 
     return _default_cluster
@@ -248,7 +252,8 @@ def _start_default_cluster(server_settings={}, initdb_options=None):
 def _shutdown_cluster(cluster):
     if cluster.get_status() == 'running':
         cluster.stop()
-    cluster.destroy()
+    if cluster.get_status() != 'not-initialized':
+        cluster.destroy()
 
 
 def create_pool(dsn=None, *,
@@ -278,15 +283,40 @@ def get_server_settings(cls):
             'log_connections': 'on'
         }
 
+    @classmethod
+    def new_cluster(cls, ClusterCls, *, cluster_kwargs={}, initdb_options={}):
+        cluster = _init_cluster(ClusterCls, cluster_kwargs,
+                                _get_initdb_options(initdb_options))
+        cls._clusters.append(cluster)
+        return cluster
+
+    @classmethod
+    def start_cluster(cls, cluster, *, server_settings={}):
+        cluster.start(port='dynamic', server_settings=server_settings)
+
     @classmethod
     def setup_cluster(cls):
-        cls.cluster = _start_default_cluster(cls.get_server_settings())
+        cls.cluster = _init_default_cluster()
+
+        if cls.cluster.get_status() != 'running':
+            cls.cluster.start(
+                port='dynamic', server_settings=cls.get_server_settings())
 
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
+        cls._clusters = []
         cls.setup_cluster()
 
+    @classmethod
+    def tearDownClass(cls):
+        super().tearDownClass()
+        for cluster in cls._clusters:
+            if cluster is not _default_cluster:
+                cluster.stop()
+                cluster.destroy()
+        cls._clusters = []
+
     @classmethod
     def get_connection_spec(cls, kwargs={}):
         conn_spec = cls.cluster.get_connection_spec()
@@ -309,14 +339,6 @@ def connect(cls, **kwargs):
         conn_spec = cls.get_connection_spec(kwargs)
         return pg_connection.connect(**conn_spec, loop=cls.loop)
 
-    @classmethod
-    def start_cluster(cls, ClusterCls, *,
-                      cluster_kwargs={}, server_settings={},
-                      initdb_options={}):
-        return _start_cluster(
-            ClusterCls, cluster_kwargs,
-            server_settings, _get_initdb_options(initdb_options))
-
 
 class ProxiedClusterTestCase(ClusterTestCase):
     @classmethod

asyncpg/cluster.py

@@ -106,7 +106,8 @@ def get_status(self):
             stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         stdout, stderr = process.stdout, process.stderr
 
-        if process.returncode == 4 or not os.listdir(self._data_dir):
+        if (process.returncode == 4 or not os.path.exists(self._data_dir) or
+                not os.listdir(self._data_dir)):
             return 'not-initialized'
         elif process.returncode == 3:
             return 'stopped'
@@ -299,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/connection.py

@@ -912,7 +912,7 @@ async def set_type_codec(self, typename, *,
             raise ValueError('unknown type: {}.{}'.format(schema, typename))
 
         oid = typeinfo['oid']
-        if typeinfo['kind'] != b'b' or typeinfo['elemtype']:
+        if typeinfo['kind'] not in (b'b', b'd') or typeinfo['elemtype']:
             raise ValueError(
                 'cannot use custom codec on non-scalar type {}.{}'.format(
                     schema, typename))

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? 0xFFFFFFFF:
+    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 > 4294967295L):
+        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

@@ -62,6 +62,11 @@ cdef int4_decode(ConnectionSettings settings, FastReadBuffer buf):
     return cpython.PyLong_FromLong(hton.unpack_int32(buf.read(4)))
 
 
+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

asyncpg/protocol/codecs/misc.pyx

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

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()