coreproto.pyx

# Copyright (C) 2016-present the asyncpg authors and contributors
# <see AUTHORS file>
#
# This module is part of asyncpg and is released under
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0


import hashlib
import socket


include "scram.pyx"


cdef dict AUTH_METHOD_NAME = {
    AUTH_REQUIRED_KERBEROS: 'kerberosv5',
    AUTH_REQUIRED_PASSWORD: 'password',
    AUTH_REQUIRED_PASSWORDMD5: 'md5',
    AUTH_REQUIRED_GSS: 'gss',
    AUTH_REQUIRED_SASL: 'scram-sha-256',
    AUTH_REQUIRED_SSPI: 'sspi',
}


cdef class CoreProtocol:

    def __init__(self, addr, con_params):
        self.address = addr
        # type of `con_params` is `_ConnectionParameters`
        self.buffer = ReadBuffer()
        self.user = con_params.user
        self.password = con_params.password
        self.auth_msg = None
        self.con_params = con_params
        self.con_status = CONNECTION_BAD
        self.state = PROTOCOL_IDLE
        self.xact_status = PQTRANS_IDLE
        self.encoding = 'utf-8'
        # type of `scram` is `SCRAMAuthentcation`
        self.scram = None
        # type of `gss_ctx` is `gssapi.SecurityContext`
        self.gss_ctx = None

        self._reset_result()

    cpdef is_in_transaction(self):
        # PQTRANS_INTRANS = idle, within transaction block
        # PQTRANS_INERROR = idle, within failed transaction
        return self.xact_status in (PQTRANS_INTRANS, PQTRANS_INERROR)

    cdef _read_server_messages(self):
        cdef:
            char mtype
            ProtocolState state
            pgproto.take_message_method take_message = \
                <pgproto.take_message_method>self.buffer.take_message
            pgproto.get_message_type_method get_message_type= \
                <pgproto.get_message_type_method>self.buffer.get_message_type

        while take_message(self.buffer) == 1:
            mtype = get_message_type(self.buffer)
            state = self.state

            try:
                if mtype == b'S':
                    # ParameterStatus
                    self._parse_msg_parameter_status()

                elif mtype == b'A':
                    # NotificationResponse
                    self._parse_msg_notification()

                elif mtype == b'N':
                    # 'N' - NoticeResponse
                    self._on_notice(self._parse_msg_error_response(False))

                elif state == PROTOCOL_AUTH:
                    self._process__auth(mtype)

                elif state == PROTOCOL_PREPARE:
                    self._process__prepare(mtype)

                elif state == PROTOCOL_BIND_EXECUTE:
                    self._process__bind_execute(mtype)

                elif state == PROTOCOL_BIND_EXECUTE_MANY:
                    self._process__bind_execute_many(mtype)

                elif state == PROTOCOL_EXECUTE:
                    self._process__bind_execute(mtype)

                elif state == PROTOCOL_BIND:
                    self._process__bind(mtype)

                elif state == PROTOCOL_CLOSE_STMT_PORTAL:
                    self._process__close_stmt_portal(mtype)

                elif state == PROTOCOL_SIMPLE_QUERY:
                    self._process__simple_query(mtype)

                elif state == PROTOCOL_COPY_OUT:
                    self._process__copy_out(mtype)

                elif (state == PROTOCOL_COPY_OUT_DATA or
                        state == PROTOCOL_COPY_OUT_DONE):
                    self._process__copy_out_data(mtype)

                elif state == PROTOCOL_COPY_IN:
                    self._process__copy_in(mtype)

                elif state == PROTOCOL_COPY_IN_DATA:
                    self._process__copy_in_data(mtype)

                elif state == PROTOCOL_CANCELLED:
                    # discard all messages until the sync message
                    if mtype == b'E':
                        self._parse_msg_error_response(True)
                    elif mtype == b'Z':
                        self._parse_msg_ready_for_query()
                        self._push_result()
                    else:
                        self.buffer.discard_message()

                elif state == PROTOCOL_ERROR_CONSUME:
                    # Error in protocol (on asyncpg side);
                    # discard all messages until sync message

                    if mtype == b'Z':
                        # Sync point, self to push the result
                        if self.result_type != RESULT_FAILED:
                            self.result_type = RESULT_FAILED
                            self.result = apg_exc.InternalClientError(
                                'unknown error in protocol implementation')

                        self._parse_msg_ready_for_query()
                        self._push_result()

                    else:
                        self.buffer.discard_message()

                elif state == PROTOCOL_TERMINATING:
                    # The connection is being terminated.
                    # discard all messages until connection
                    # termination.
                    self.buffer.discard_message()

                else:
                    raise apg_exc.InternalClientError(
                        f'cannot process message {chr(mtype)!r}: '
                        f'protocol is in an unexpected state {state!r}.')

            except Exception as ex:
                self.result_type = RESULT_FAILED
                self.result = ex

                if mtype == b'Z':
                    self._push_result()
                else:
                    self.state = PROTOCOL_ERROR_CONSUME

            finally:
                self.buffer.finish_message()

    cdef _process__auth(self, char mtype):
        if mtype == b'R':
            # Authentication...
            try:
                self._parse_msg_authentication()
            except Exception as ex:
                # Exception in authentication parsing code
                # is usually either malformed authentication data
                # or missing support for cryptographic primitives
                # in the hashlib module.
                self.result_type = RESULT_FAILED
                self.result = apg_exc.InternalClientError(
                    f"unexpected error while performing authentication: {ex}")
                self.result.__cause__ = ex
                self.con_status = CONNECTION_BAD
                self._push_result()
            else:
                if self.result_type != RESULT_OK:
                    self.con_status = CONNECTION_BAD
                    self._push_result()

                elif self.auth_msg is not None:
                    # Server wants us to send auth data, so do that.
                    self._write(self.auth_msg)
                    self.auth_msg = None

        elif mtype == b'K':
            # BackendKeyData
            self._parse_msg_backend_key_data()

        elif mtype == b'E':
            # ErrorResponse
            self.con_status = CONNECTION_BAD
            self._parse_msg_error_response(True)
            self._push_result()

        elif mtype == b'Z':
            # ReadyForQuery
            self._parse_msg_ready_for_query()
            self.con_status = CONNECTION_OK
            self._push_result()

    cdef _process__prepare(self, char mtype):
        if mtype == b't':
            # Parameters description
            self.result_param_desc = self.buffer.consume_message()

        elif mtype == b'1':
            # ParseComplete
            self.buffer.discard_message()

        elif mtype == b'T':
            # Row description
            self.result_row_desc = self.buffer.consume_message()
            self._push_result()

        elif mtype == b'E':
            # ErrorResponse
            self._parse_msg_error_response(True)
            # we don't send a sync during the parse/describe sequence
            # but send a FLUSH instead. If an error happens we need to
            # send a SYNC explicitly in order to mark the end of the transaction.
            # this effectively clears the error and we then wait until we get a
            # ready for new query message
            self._write(SYNC_MESSAGE)
            self.state = PROTOCOL_ERROR_CONSUME

        elif mtype == b'n':
            # NoData
            self.buffer.discard_message()
            self._push_result()

    cdef _process__bind_execute(self, char mtype):
        if mtype == b'D':
            # DataRow
            self._parse_data_msgs()

        elif mtype == b's':
            # PortalSuspended
            self.buffer.discard_message()

        elif mtype == b'C':
            # CommandComplete
            self.result_execute_completed = True
            self._parse_msg_command_complete()

        elif mtype == b'E':
            # ErrorResponse
            self._parse_msg_error_response(True)

        elif mtype == b'1':
            # ParseComplete, in case `_bind_execute()` is reparsing
            self.buffer.discard_message()

        elif mtype == b'2':
            # BindComplete
            self.buffer.discard_message()

        elif mtype == b'Z':
            # ReadyForQuery
            self._parse_msg_ready_for_query()
            self._push_result()

        elif mtype == b'I':
            # EmptyQueryResponse
            self.buffer.discard_message()

    cdef _process__bind_execute_many(self, char mtype):
        cdef WriteBuffer buf

        if mtype == b'D':
            # DataRow
            self._parse_data_msgs()

        elif mtype == b's':
            # PortalSuspended
            self.buffer.discard_message()

        elif mtype == b'C':
            # CommandComplete
            self._parse_msg_command_complete()

        elif mtype == b'E':
            # ErrorResponse
            self._parse_msg_error_response(True)

        elif mtype == b'1':
            # ParseComplete, in case `_bind_execute_many()` is reparsing
            self.buffer.discard_message()

        elif mtype == b'2':
            # BindComplete
            self.buffer.discard_message()

        elif mtype == b'Z':
            # ReadyForQuery
            self._parse_msg_ready_for_query()
            self._push_result()

        elif mtype == b'I':
            # EmptyQueryResponse
            self.buffer.discard_message()

        elif mtype == b'1':
            # ParseComplete
            self.buffer.discard_message()

    cdef _process__bind(self, char mtype):
        if mtype == b'E':
            # ErrorResponse
            self._parse_msg_error_response(True)

        elif mtype == b'2':
            # BindComplete
            self.buffer.discard_message()

        elif mtype == b'Z':
            # ReadyForQuery
            self._parse_msg_ready_for_query()
            self._push_result()

    cdef _process__close_stmt_portal(self, char mtype):
        if mtype == b'E':
            # ErrorResponse
            self._parse_msg_error_response(True)

        elif mtype == b'3':
            # CloseComplete
            self.buffer.discard_message()

        elif mtype == b'Z':
            # ReadyForQuery
            self._parse_msg_ready_for_query()
            self._push_result()

    cdef _process__simple_query(self, char mtype):
        if mtype in {b'D', b'I', b'T'}:
            # 'D' - DataRow
            # 'I' - EmptyQueryResponse
            # 'T' - RowDescription
            self.buffer.discard_message()

        elif mtype == b'E':
            # ErrorResponse
            self._parse_msg_error_response(True)

        elif mtype == b'Z':
            # ReadyForQuery
            self._parse_msg_ready_for_query()
            self._push_result()

        elif mtype == b'C':
            # CommandComplete
            self._parse_msg_command_complete()

        else:
            # We don't really care about COPY IN etc
            self.buffer.discard_message()

    cdef _process__copy_out(self, char mtype):
        if mtype == b'E':
            self._parse_msg_error_response(True)

        elif mtype == b'H':
            # CopyOutResponse
            self._set_state(PROTOCOL_COPY_OUT_DATA)
            self.buffer.discard_message()

        elif mtype == b'Z':
            # ReadyForQuery
            self._parse_msg_ready_for_query()
            self._push_result()

    cdef _process__copy_out_data(self, char mtype):
        if mtype == b'E':
            self._parse_msg_error_response(True)

        elif mtype == b'd':
            # CopyData
            self._parse_copy_data_msgs()

        elif mtype == b'c':
            # CopyDone
            self.buffer.discard_message()
            self._set_state(PROTOCOL_COPY_OUT_DONE)

        elif mtype == b'C':
            # CommandComplete
            self._parse_msg_command_complete()

        elif mtype == b'Z':
            # ReadyForQuery
            self._parse_msg_ready_for_query()
            self._push_result()

    cdef _process__copy_in(self, char mtype):
        if mtype == b'E':
            self._parse_msg_error_response(True)

        elif mtype == b'G':
            # CopyInResponse
            self._set_state(PROTOCOL_COPY_IN_DATA)
            self.buffer.discard_message()

        elif mtype == b'Z':
            # ReadyForQuery
            self._parse_msg_ready_for_query()
            self._push_result()

    cdef _process__copy_in_data(self, char mtype):
        if mtype == b'E':
            self._parse_msg_error_response(True)

        elif mtype == b'C':
            # CommandComplete
            self._parse_msg_command_complete()

        elif mtype == b'Z':
            # ReadyForQuery
            self._parse_msg_ready_for_query()
            self._push_result()

    cdef _parse_msg_command_complete(self):
        cdef:
            const char* cbuf
            ssize_t cbuf_len

        cbuf = self.buffer.try_consume_message(&cbuf_len)
        if cbuf != NULL and cbuf_len > 0:
            msg = cpython.PyBytes_FromStringAndSize(cbuf, cbuf_len - 1)
        else:
            msg = self.buffer.read_null_str()
        self.result_status_msg = msg

    cdef _parse_copy_data_msgs(self):
        cdef:
            ReadBuffer buf = self.buffer

        self.result = buf.consume_messages(b'd')

        # By this point we have consumed all CopyData messages
        # in the inbound buffer.  If there are no messages left
        # in the buffer, we need to push the accumulated data
        # out to the caller in anticipation of the new CopyData
        # batch.  If there _are_ non-CopyData messages left,
        # we must not push the result here and let the
        # _process__copy_out_data subprotocol do the job.
        if not buf.take_message():
            self._on_result()
            self.result = None
        else:
            # If there is a message in the buffer, put it back to
            # be processed by the next protocol iteration.
            buf.put_message()

    cdef _write_copy_data_msg(self, object data):
        cdef:
            WriteBuffer buf
            object mview
            Py_buffer *pybuf

        mview = cpythonx.PyMemoryView_GetContiguous(
            data, cpython.PyBUF_READ, b'C')

        try:
            pybuf = cpythonx.PyMemoryView_GET_BUFFER(mview)

            buf = WriteBuffer.new_message(b'd')
            buf.write_cstr(<const char *>pybuf.buf, pybuf.len)
            buf.end_message()
        finally:
            mview.release()

        self._write(buf)

    cdef _write_copy_done_msg(self):
        cdef:
            WriteBuffer buf

        buf = WriteBuffer.new_message(b'c')
        buf.end_message()
        self._write(buf)

    cdef _write_copy_fail_msg(self, str cause):
        cdef:
            WriteBuffer buf

        buf = WriteBuffer.new_message(b'f')
        buf.write_str(cause or '', self.encoding)
        buf.end_message()
        self._write(buf)

    cdef _parse_data_msgs(self):
        cdef:
            ReadBuffer buf = self.buffer
            list rows

            decode_row_method decoder = <decode_row_method>self._decode_row
            pgproto.try_consume_message_method try_consume_message = \
                <pgproto.try_consume_message_method>buf.try_consume_message
            pgproto.take_message_type_method take_message_type = \
                <pgproto.take_message_type_method>buf.take_message_type

            const char* cbuf
            ssize_t cbuf_len
            object row
            bytes mem

        if PG_DEBUG:
            if buf.get_message_type() != b'D':
                raise apg_exc.InternalClientError(
                    '_parse_data_msgs: first message is not "D"')

        if self._discard_data:
            while take_message_type(buf, b'D'):
                buf.discard_message()
            return

        if PG_DEBUG:
            if type(self.result) is not list:
                raise apg_exc.InternalClientError(
                    '_parse_data_msgs: result is not a list, but {!r}'.
                    format(self.result))

        rows = self.result
        while take_message_type(buf, b'D'):
            cbuf = try_consume_message(buf, &cbuf_len)
            if cbuf != NULL:
                row = decoder(self, cbuf, cbuf_len)
            else:
                mem = buf.consume_message()
                row = decoder(
                    self,
                    cpython.PyBytes_AS_STRING(mem),
                    cpython.PyBytes_GET_SIZE(mem))

            cpython.PyList_Append(rows, row)

    cdef _parse_msg_backend_key_data(self):
        self.backend_pid = self.buffer.read_int32()
        self.backend_secret = self.buffer.read_int32()

    cdef _parse_msg_parameter_status(self):
        name = self.buffer.read_null_str()
        name = name.decode(self.encoding)

        val = self.buffer.read_null_str()
        val = val.decode(self.encoding)

        self._set_server_parameter(name, val)

    cdef _parse_msg_notification(self):
        pid = self.buffer.read_int32()
        channel = self.buffer.read_null_str().decode(self.encoding)
        payload = self.buffer.read_null_str().decode(self.encoding)
        self._on_notification(pid, channel, payload)

    cdef _parse_msg_authentication(self):
        cdef:
            int32_t status
            bytes md5_salt
            list sasl_auth_methods
            list unsupported_sasl_auth_methods

        status = self.buffer.read_int32()

        if status == AUTH_SUCCESSFUL:
            # AuthenticationOk
            self.result_type = RESULT_OK

        elif status == AUTH_REQUIRED_PASSWORD:
            # AuthenticationCleartextPassword
            self.result_type = RESULT_OK
            self.auth_msg = self._auth_password_message_cleartext()

        elif status == AUTH_REQUIRED_PASSWORDMD5:
            # AuthenticationMD5Password
            # Note: MD5 salt is passed as a four-byte sequence
            md5_salt = self.buffer.read_bytes(4)
            self.auth_msg = self._auth_password_message_md5(md5_salt)

        elif status == AUTH_REQUIRED_SASL:
            # AuthenticationSASL
            # This requires making additional requests to the server in order
            # to follow the SCRAM protocol defined in RFC 5802.
            # get the SASL authentication methods that the server is providing
            sasl_auth_methods = []
            unsupported_sasl_auth_methods = []
            # determine if the advertised authentication methods are supported,
            # and if so, add them to the list
            auth_method = self.buffer.read_null_str()
            while auth_method:
                if auth_method in SCRAMAuthentication.AUTHENTICATION_METHODS:
                    sasl_auth_methods.append(auth_method)
                else:
                    unsupported_sasl_auth_methods.append(auth_method)
                auth_method = self.buffer.read_null_str()

            # if none of the advertised authentication methods are supported,
            # raise an error
            # otherwise, initialize the SASL authentication exchange
            if not sasl_auth_methods:
                unsupported_sasl_auth_methods = [m.decode("ascii")
                    for m in unsupported_sasl_auth_methods]
                self.result_type = RESULT_FAILED
                self.result = apg_exc.InterfaceError(
                    'unsupported SASL Authentication methods requested by the '
                    'server: {!r}'.format(
                        ", ".join(unsupported_sasl_auth_methods)))
            else:
                self.auth_msg = self._auth_password_message_sasl_initial(
                    sasl_auth_methods)

        elif status == AUTH_SASL_CONTINUE:
            # AUTH_SASL_CONTINUE
            # this requeires sending the second part of the SASL exchange, where
            # the client parses information back from the server and determines
            # if this is valid.
            # The client builds a challenge response to the server
            server_response = self.buffer.consume_message()
            self.auth_msg = self._auth_password_message_sasl_continue(
                server_response)

        elif status == AUTH_SASL_FINAL:
            # AUTH_SASL_FINAL
            server_response = self.buffer.consume_message()
            if not self.scram.verify_server_final_message(server_response):
                self.result_type = RESULT_FAILED
                self.result = apg_exc.InterfaceError(
                    'could not verify server signature for '
                    'SCRAM authentciation: scram-sha-256',
                )
            self.scram = None

        elif status == AUTH_REQUIRED_GSS:
            self._auth_gss_init()
            self.auth_msg = self._auth_gss_step(None)

        elif status == AUTH_REQUIRED_GSS_CONTINUE:
            server_response = self.buffer.consume_message()
            self.auth_msg = self._auth_gss_step(server_response)

        elif status in (AUTH_REQUIRED_KERBEROS, AUTH_REQUIRED_SCMCRED,
                        AUTH_REQUIRED_SSPI):
            self.result_type = RESULT_FAILED
            self.result = apg_exc.InterfaceError(
                'unsupported authentication method requested by the '
                'server: {!r}'.format(AUTH_METHOD_NAME[status]))

        else:
            self.result_type = RESULT_FAILED
            self.result = apg_exc.InterfaceError(
                'unsupported authentication method requested by the '
                'server: {}'.format(status))

        if status not in [AUTH_SASL_CONTINUE, AUTH_SASL_FINAL,
                          AUTH_REQUIRED_GSS_CONTINUE]:
            self.buffer.discard_message()

    cdef _auth_password_message_cleartext(self):
        cdef:
            WriteBuffer msg

        msg = WriteBuffer.new_message(b'p')
        msg.write_bytestring(self.password.encode(self.encoding))
        msg.end_message()

        return msg

    cdef _auth_password_message_md5(self, bytes salt):
        cdef:
            WriteBuffer msg

        msg = WriteBuffer.new_message(b'p')

        # 'md5' + md5(md5(password + username) + salt))
        userpass = (self.password or '') + (self.user or '')
        md5_1 = hashlib.md5(userpass.encode(self.encoding)).hexdigest()
        md5_2 = hashlib.md5(md5_1.encode('ascii') + salt).hexdigest()

        msg.write_bytestring(b'md5' + md5_2.encode('ascii'))
        msg.end_message()

        return msg

    cdef _auth_password_message_sasl_initial(self, list sasl_auth_methods):
        cdef:
            WriteBuffer msg

        # use the first supported advertized mechanism
        self.scram = SCRAMAuthentication(sasl_auth_methods[0])
        # this involves a call and response with the server
        msg = WriteBuffer.new_message(b'p')
        msg.write_bytes(self.scram.create_client_first_message(self.user or ''))
        msg.end_message()

        return msg

    cdef _auth_password_message_sasl_continue(self, bytes server_response):
        cdef:
            WriteBuffer msg

        # determine if there is a valid server response
        self.scram.parse_server_first_message(server_response)
        # this involves a call and response with the server
        msg = WriteBuffer.new_message(b'p')
        client_final_message = self.scram.create_client_final_message(
            self.password or '')
        msg.write_bytes(client_final_message)
        msg.end_message()

        return msg

    cdef _auth_gss_init(self):
        try:
            import gssapi
        except ModuleNotFoundError:
            raise RuntimeError(
                'gssapi module not found; please install asyncpg[gssapi] to '
                'use asyncpg with Kerberos or GSSAPI authentication'
            ) from None

        service_name = self.con_params.krbsrvname or 'postgres'
        # find the canonical name of the server host
        if isinstance(self.address, str):
            raise RuntimeError('GSSAPI authentication is only supported for '
                               'TCP/IP connections')

        host = self.address[0]
        host_cname = socket.gethostbyname_ex(host)[0]
        gss_name = gssapi.Name(f'{service_name}/{host_cname}')
        self.gss_ctx = gssapi.SecurityContext(name=gss_name, usage='initiate')

    cdef _auth_gss_step(self, bytes server_response):
        cdef:
            WriteBuffer msg

        token = self.gss_ctx.step(server_response)
        if not token:
            self.gss_ctx = None
            return None
        msg = WriteBuffer.new_message(b'p')
        msg.write_bytes(token)
        msg.end_message()

        return msg

    cdef _parse_msg_ready_for_query(self):
        cdef char status = self.buffer.read_byte()

        if status == b'I':
            self.xact_status = PQTRANS_IDLE
        elif status == b'T':
            self.xact_status = PQTRANS_INTRANS
        elif status == b'E':
            self.xact_status = PQTRANS_INERROR
        else:
            self.xact_status = PQTRANS_UNKNOWN

    cdef _parse_msg_error_response(self, is_error):
        cdef:
            char code
            bytes message
            dict parsed = {}

        while True:
            code = self.buffer.read_byte()
            if code == 0:
                break

            message = self.buffer.read_null_str()

            parsed[chr(code)] = message.decode()

        if is_error:
            self.result_type = RESULT_FAILED
            self.result = parsed
        else:
            return parsed

    cdef _push_result(self):
        try:
            self._on_result()
        finally:
            self._set_state(PROTOCOL_IDLE)
            self._reset_result()

    cdef _reset_result(self):
        self.result_type = RESULT_OK
        self.result = None
        self.result_param_desc = None
        self.result_row_desc = None
        self.result_status_msg = None
        self.result_execute_completed = False
        self._discard_data = False

        # executemany support data
        self._execute_iter = None
        self._execute_portal_name = None
        self._execute_stmt_name = None

    cdef _set_state(self, ProtocolState new_state):
        if new_state == PROTOCOL_IDLE:
            if self.state == PROTOCOL_FAILED:
                raise apg_exc.InternalClientError(
                    'cannot switch to "idle" state; '
                    'protocol is in the "failed" state')
            elif self.state == PROTOCOL_IDLE:
                pass
            else:
                self.state = new_state

        elif new_state == PROTOCOL_FAILED:
            self.state = PROTOCOL_FAILED

        elif new_state == PROTOCOL_CANCELLED:
            self.state = PROTOCOL_CANCELLED

        elif new_state == PROTOCOL_TERMINATING:
            self.state = PROTOCOL_TERMINATING

        else:
            if self.state == PROTOCOL_IDLE:
                self.state = new_state

            elif (self.state == PROTOCOL_COPY_OUT and
                    new_state == PROTOCOL_COPY_OUT_DATA):
                self.state = new_state

            elif (self.state == PROTOCOL_COPY_OUT_DATA and
                    new_state == PROTOCOL_COPY_OUT_DONE):
                self.state = new_state

            elif (self.state == PROTOCOL_COPY_IN and
                    new_state == PROTOCOL_COPY_IN_DATA):
                self.state = new_state

            elif self.state == PROTOCOL_FAILED:
                raise apg_exc.InternalClientError(
                    'cannot switch to state {}; '
                    'protocol is in the "failed" state'.format(new_state))
            else:
                raise apg_exc.InternalClientError(
                    'cannot switch to state {}; '
                    'another operation ({}) is in progress'.format(
                        new_state, self.state))

    cdef _ensure_connected(self):
        if self.con_status != CONNECTION_OK:
            raise apg_exc.InternalClientError('not connected')

    cdef WriteBuffer _build_parse_message(self, str stmt_name, str query):
        cdef WriteBuffer buf

        buf = WriteBuffer.new_message(b'P')
        buf.write_str(stmt_name, self.encoding)
        buf.write_str(query, self.encoding)
        buf.write_int16(0)

        buf.end_message()
        return buf

    cdef WriteBuffer _build_bind_message(self, str portal_name,
                                         str stmt_name,
                                         WriteBuffer bind_data):
        cdef WriteBuffer buf

        buf = WriteBuffer.new_message(b'B')
        buf.write_str(portal_name, self.encoding)
        buf.write_str(stmt_name, self.encoding)

        # Arguments
        buf.write_buffer(bind_data)

        buf.end_message()
        return buf

    cdef WriteBuffer _build_empty_bind_data(self):
        cdef WriteBuffer buf
        buf = WriteBuffer.new()
        buf.write_int16(0)  # The number of parameter format codes
        buf.write_int16(0)  # The number of parameter values
        buf.write_int16(0)  # The number of result-column format codes
        return buf

    cdef WriteBuffer _build_execute_message(self, str portal_name,
                                            int32_t limit):
        cdef WriteBuffer buf

        buf = WriteBuffer.new_message(b'E')
        buf.write_str(portal_name, self.encoding)  # name of the portal
        buf.write_int32(limit)  # number of rows to return; 0 - all

        buf.end_message()
        return buf

    # API for subclasses

    cdef _connect(self):
        cdef:
            WriteBuffer buf
            WriteBuffer outbuf

        if self.con_status != CONNECTION_BAD:
            raise apg_exc.InternalClientError('already connected')

        self._set_state(PROTOCOL_AUTH)
        self.con_status = CONNECTION_STARTED

        # Assemble a startup message
        buf = WriteBuffer()

        # protocol version
        buf.write_int16(3)
        buf.write_int16(0)

        buf.write_bytestring(b'client_encoding')
        buf.write_bytestring("'{}'".format(self.encoding).encode('ascii'))

        buf.write_str('user', self.encoding)
        buf.write_str(self.con_params.user, self.encoding)

        buf.write_str('database', self.encoding)
        buf.write_str(self.con_params.database, self.encoding)

        if self.con_params.server_settings is not None:
            for k, v in self.con_params.server_settings.items():
                buf.write_str(k, self.encoding)
                buf.write_str(v, self.encoding)

        buf.write_bytestring(b'')

        # Send the buffer
        outbuf = WriteBuffer()
        outbuf.write_int32(buf.len() + 4)
        outbuf.write_buffer(buf)
        self._write(outbuf)

    cdef _send_parse_message(self, str stmt_name, str query):
        cdef:
            WriteBuffer msg

        self._ensure_connected()
        msg = self._build_parse_message(stmt_name, query)
        self._write(msg)

    cdef _prepare_and_describe(self, str stmt_name, str query):
        cdef:
            WriteBuffer packet
            WriteBuffer buf

        self._ensure_connected()
        self._set_state(PROTOCOL_PREPARE)

        packet = self._build_parse_message(stmt_name, query)

        buf = WriteBuffer.new_message(b'D')
        buf.write_byte(b'S')
        buf.write_str(stmt_name, self.encoding)
        buf.end_message()
        packet.write_buffer(buf)

        packet.write_bytes(FLUSH_MESSAGE)

        self._write(packet)

    cdef _send_bind_message(self, str portal_name, str stmt_name,
                            WriteBuffer bind_data, int32_t limit):

        cdef:
            WriteBuffer packet
            WriteBuffer buf

        buf = self._build_bind_message(portal_name, stmt_name, bind_data)
        packet = buf

        buf = self._build_execute_message(portal_name, limit)
        packet.write_buffer(buf)

        packet.write_bytes(SYNC_MESSAGE)

        self._write(packet)

    cdef _bind_execute(self, str portal_name, str stmt_name,
                       WriteBuffer bind_data, int32_t limit):

        cdef WriteBuffer buf

        self._ensure_connected()
        self._set_state(PROTOCOL_BIND_EXECUTE)

        self.result = []

        self._send_bind_message(portal_name, stmt_name, bind_data, limit)

    cdef bint _bind_execute_many(self, str portal_name, str stmt_name,
                                 object bind_data):
        self._ensure_connected()
        self._set_state(PROTOCOL_BIND_EXECUTE_MANY)

        self.result = None
        self._discard_data = True
        self._execute_iter = bind_data
        self._execute_portal_name = portal_name
        self._execute_stmt_name = stmt_name
        return self._bind_execute_many_more(True)

    cdef bint _bind_execute_many_more(self, bint first=False):
        cdef:
            WriteBuffer packet
            WriteBuffer buf
            list buffers = []

        # as we keep sending, the server may return an error early
        if self.result_type == RESULT_FAILED:
            self._write(SYNC_MESSAGE)
            return False

        # collect up to four 32KB buffers to send
        # https://github.com/MagicStack/asyncpg/pull/289#issuecomment-391215051
        while len(buffers) < _EXECUTE_MANY_BUF_NUM:
            packet = WriteBuffer.new()

            # fill one 32KB buffer
            while packet.len() < _EXECUTE_MANY_BUF_SIZE:
                try:
                    # grab one item from the input
                    buf = <WriteBuffer>next(self._execute_iter)

                # reached the end of the input
                except StopIteration:
                    if first:
                        # if we never send anything, simply set the result
                        self._push_result()
                    else:
                        # otherwise, append SYNC and send the buffers
                        packet.write_bytes(SYNC_MESSAGE)
                        buffers.append(memoryview(packet))
                        self._writelines(buffers)
                    return False

                # error in input, give up the buffers and cleanup
                except Exception as ex:
                    self._bind_execute_many_fail(ex, first)
                    return False

                # all good, write to the buffer
                first = False
                packet.write_buffer(
                    self._build_bind_message(
                        self._execute_portal_name,
                        self._execute_stmt_name,
                        buf,
                    )
                )
                packet.write_buffer(
                    self._build_execute_message(self._execute_portal_name, 0,
                    )
                )

            # collected one buffer
            buffers.append(memoryview(packet))

        # write to the wire, and signal the caller for more to send
        self._writelines(buffers)
        return True

    cdef _bind_execute_many_fail(self, object error, bint first=False):
        cdef WriteBuffer buf

        self.result_type = RESULT_FAILED
        self.result = error
        if first:
            self._push_result()
        elif self.is_in_transaction():
            # we're in an explicit transaction, just SYNC
            self._write(SYNC_MESSAGE)
        else:
            # In an implicit transaction, if `ignore_till_sync` is set,
            # `ROLLBACK` will be ignored and `Sync` will restore the state;
            # or the transaction will be rolled back with a warning saying
            # that there was no transaction, but rollback is done anyway,
            # so we could safely ignore this warning.
            # GOTCHA: cannot use simple query message here, because it is
            # ignored if `ignore_till_sync` is set.
            buf = self._build_parse_message('', 'ROLLBACK')
            buf.write_buffer(self._build_bind_message(
                '', '', self._build_empty_bind_data()))
            buf.write_buffer(self._build_execute_message('', 0))
            buf.write_bytes(SYNC_MESSAGE)
            self._write(buf)

    cdef _execute(self, str portal_name, int32_t limit):
        cdef WriteBuffer buf

        self._ensure_connected()
        self._set_state(PROTOCOL_EXECUTE)

        self.result = []

        buf = self._build_execute_message(portal_name, limit)

        buf.write_bytes(SYNC_MESSAGE)

        self._write(buf)

    cdef _bind(self, str portal_name, str stmt_name,
               WriteBuffer bind_data):

        cdef WriteBuffer buf

        self._ensure_connected()
        self._set_state(PROTOCOL_BIND)

        buf = self._build_bind_message(portal_name, stmt_name, bind_data)

        buf.write_bytes(SYNC_MESSAGE)

        self._write(buf)

    cdef _close(self, str name, bint is_portal):
        cdef WriteBuffer buf

        self._ensure_connected()
        self._set_state(PROTOCOL_CLOSE_STMT_PORTAL)

        buf = WriteBuffer.new_message(b'C')

        if is_portal:
            buf.write_byte(b'P')
        else:
            buf.write_byte(b'S')

        buf.write_str(name, self.encoding)
        buf.end_message()

        buf.write_bytes(SYNC_MESSAGE)

        self._write(buf)

    cdef _simple_query(self, str query):
        cdef WriteBuffer buf
        self._ensure_connected()
        self._set_state(PROTOCOL_SIMPLE_QUERY)
        buf = WriteBuffer.new_message(b'Q')
        buf.write_str(query, self.encoding)
        buf.end_message()
        self._write(buf)

    cdef _copy_out(self, str copy_stmt):
        cdef WriteBuffer buf

        self._ensure_connected()
        self._set_state(PROTOCOL_COPY_OUT)

        # Send the COPY .. TO STDOUT using the SimpleQuery protocol.
        buf = WriteBuffer.new_message(b'Q')
        buf.write_str(copy_stmt, self.encoding)
        buf.end_message()
        self._write(buf)

    cdef _copy_in(self, str copy_stmt):
        cdef WriteBuffer buf

        self._ensure_connected()
        self._set_state(PROTOCOL_COPY_IN)

        buf = WriteBuffer.new_message(b'Q')
        buf.write_str(copy_stmt, self.encoding)
        buf.end_message()
        self._write(buf)

    cdef _terminate(self):
        cdef WriteBuffer buf
        self._ensure_connected()
        self._set_state(PROTOCOL_TERMINATING)
        buf = WriteBuffer.new_message(b'X')
        buf.end_message()
        self._write(buf)

    cdef _write(self, buf):
        raise NotImplementedError

    cdef _writelines(self, list buffers):
        raise NotImplementedError

    cdef _decode_row(self, const char* buf, ssize_t buf_len):
        pass

    cdef _set_server_parameter(self, name, val):
        pass

    cdef _on_result(self):
        pass

    cdef _on_notice(self, parsed):
        pass

    cdef _on_notification(self, pid, channel, payload):
        pass

    cdef _on_connection_lost(self, exc):
        pass


cdef bytes SYNC_MESSAGE = bytes(WriteBuffer.new_message(b'S').end_message())
cdef bytes FLUSH_MESSAGE = bytes(WriteBuffer.new_message(b'H').end_message())