Add support for password authentication (password and md5)

Other authentication protocols that require frontend intervention are
not supported yet.

Author: elprans

asyncpg/_testbase.py

@@ -96,7 +96,7 @@ def assertRunUnder(self, delta):
 _default_cluster = None
 
 
-def _start_cluster():
+def _start_cluster(server_settings={}):
     global _default_cluster
 
     if _default_cluster is None:
@@ -107,7 +107,8 @@ def _start_cluster():
         else:
             _default_cluster = pg_cluster.TempCluster()
             _default_cluster.init()
-            _default_cluster.start(port=12345)
+            _default_cluster.trust_local_connections()
+            _default_cluster.start(port=12345, server_settings=server_settings)
             atexit.register(_shutdown_cluster, _default_cluster)
 
     return _default_cluster
@@ -121,7 +122,9 @@ def _shutdown_cluster(cluster):
 class ClusterTestCase(TestCase):
     def setUp(self):
         super().setUp()
-        self.cluster = _start_cluster()
+        self.cluster = _start_cluster({
+            'log_connections': 'on'
+        })
 
 
 class ConnectedTestCase(ClusterTestCase):

asyncpg/cluster.py

@@ -131,6 +131,23 @@ def start(self, wait=60, *, server_settings={}, **opts):
 
         self._test_connection(timeout=wait)
 
+    def reload(self):
+        """Reload server configuration."""
+        status = self.get_status()
+        if status != 'running':
+            raise ClusterError('cannot reload: cluster is not running')
+
+        process = subprocess.run(
+            [self._pg_ctl, 'reload', '-D', self._data_dir],
+            stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+        stderr = process.stderr
+
+        if process.returncode != 0:
+            raise ClusterError(
+                'pg_ctl stop exited with status {:d}: {}'.format(
+                    process.returncode, stderr.decode()))
+
     def stop(self, wait=60):
         process = subprocess.run(
             [self._pg_ctl, 'stop', '-D', self._data_dir, '-t', str(wait),
@@ -165,6 +182,68 @@ def get_connection_addr(self):
 
         return self._connection_addr['host'], self._connection_addr['port']
 
+    def reset_hba(self):
+        """Remove all records from pg_hba.conf."""
+        status = self.get_status()
+        if status == 'not-initialized':
+            raise ClusterError(
+                'cannot modify HBA records: cluster is not initialized')
+
+        pg_hba = os.path.join(self._data_dir, 'pg_hba.conf')
+
+        try:
+            with open(pg_hba, 'w'):
+                pass
+        except IOError as e:
+            raise ClusterError(
+                'cannot modify HBA records: {}'.format(e)) from e
+
+    def add_hba_entry(self, *, type='host', database, user, address=None,
+                      auth_method, auth_options=None):
+        """Add a record to pg_hba.conf."""
+        status = self.get_status()
+        if status == 'not-initialized':
+            raise ClusterError(
+                'cannot modify HBA records: cluster is not initialized')
+
+        if type not in {'local', 'host', 'hostssl', 'hostnossl'}:
+            raise ValueError('invalid HBA record type: {!r}'.format(type))
+
+        pg_hba = os.path.join(self._data_dir, 'pg_hba.conf')
+
+        record = '{} {} {}'.format(type, database, user)
+
+        if type != 'local':
+            if address is None:
+                raise ValueError(
+                    '{!r} entry requires a valid address'.format(type))
+            else:
+                record += ' {}'.format(address)
+
+        record += ' {}'.format(auth_method)
+
+        if auth_options is not None:
+            record += ' ' + ' '.join(
+                '{}={}'.format(k, v) for k, v in auth_options)
+
+        try:
+            with open(pg_hba, 'a') as f:
+                print(record, file=f)
+        except IOError as e:
+            raise ClusterError(
+                'cannot modify HBA records: {}'.format(e)) from e
+
+    def trust_local_connections(self):
+        self.reset_hba()
+        self.add_hba_entry(type='local', database='all',
+                           user='all', auth_method='trust')
+        self.add_hba_entry(type='host', address='127.0.0.1/32',
+                           database='all', user='all',
+                           auth_method='trust')
+        status = self.get_status()
+        if status == 'running':
+            self.reload()
+
     def _init_env(self):
         self._pg_config = self._find_pg_config(self._pg_config_path)
         self._pg_config_data = self._run_pg_config(self._pg_config)

asyncpg/protocol/buffer.pxd

@@ -100,6 +100,7 @@ cdef class ReadBuffer:
     cdef inline read_byte(self)
     cdef inline char* _try_read_bytes(self, int nbytes)
     cdef inline read(self, int nbytes)
+    cdef inline read_bytes(self, ssize_t n)
     cdef inline read_int32(self)
     cdef inline read_int16(self)
     cdef inline read_cstr(self)

asyncpg/protocol/buffer.pyx

@@ -293,21 +293,6 @@ cdef class ReadBuffer:
                 raise RuntimeError(
                     'debug: second buffer of ReadBuffer is empty')
 
-    cdef inline read_byte(self):
-        cdef char* first_byte
-
-        IF DEBUG:
-            if not self._buf0:
-                raise RuntimeError(
-                    'debug: first buffer of ReadBuffer is empty')
-
-        self._ensure_first_buf()
-        first_byte = self._try_read_bytes(1)
-        if first_byte is NULL:
-            raise BufferError('not enough data to read one byte')
-
-        return first_byte[0]
-
     cdef inline char* _try_read_bytes(self, int nbytes):
         # Important: caller must call _ensure_first_buf() prior
         # to calling try_read_bytes, and must not overread
@@ -373,6 +358,34 @@ cdef class ReadBuffer:
                     result,
                     len(result))
 
+    cdef inline read_byte(self):
+        cdef char* first_byte
+
+        IF DEBUG:
+            if not self._buf0:
+                raise RuntimeError(
+                    'debug: first buffer of ReadBuffer is empty')
+
+        self._ensure_first_buf()
+        first_byte = self._try_read_bytes(1)
+        if first_byte is NULL:
+            raise BufferError('not enough data to read one byte')
+
+        return first_byte[0]
+
+    cdef inline read_bytes(self, ssize_t n):
+        cdef:
+            Memory mem
+            char *cbuf
+
+        self._ensure_first_buf()
+        cbuf = self._try_read_bytes(n)
+        if cbuf != NULL:
+            return cbuf
+        else:
+            mem = <Memory>(self.read(n))
+            return mem.buf
+
     cdef inline read_int32(self):
         cdef:
             Memory mem
@@ -508,7 +521,10 @@ cdef class ReadBuffer:
     cdef Memory consume_message(self):
         if not self._current_message_ready:
             raise BufferError('no message to consume')
-        mem = self.read(self._current_message_len_unread)
+        if self._current_message_len_unread > 0:
+            mem = self.read(self._current_message_len_unread)
+        else:
+            mem = None
         self._discard_message()
         return mem
 

asyncpg/protocol/coreproto.pxd

@@ -27,6 +27,26 @@ cdef enum ProtocolState:
     PROTOCOL_BIND = 16
 
 
+cdef enum AuthenticationMessage:
+    AUTH_SUCCESSFUL = 0
+    AUTH_REQUIRED_KERBEROS = 2
+    AUTH_REQUIRED_PASSWORD = 3
+    AUTH_REQUIRED_PASSWORDMD5 = 5
+    AUTH_REQUIRED_SCMCRED = 6
+    AUTH_REQUIRED_GSS = 7
+    AUTH_REQUIRED_GSS_CONTINUE = 8
+    AUTH_REQUIRED_SSPI = 9
+
+
+AUTH_METHOD_NAME = {
+    AUTH_REQUIRED_KERBEROS: 'kerberosv5',
+    AUTH_REQUIRED_PASSWORD: 'password',
+    AUTH_REQUIRED_PASSWORDMD5: 'md5',
+    AUTH_REQUIRED_GSS: 'gss',
+    AUTH_REQUIRED_SSPI: 'sspi',
+}
+
+
 cdef enum ResultType:
     RESULT_OK = 1
     RESULT_FAILED = 2
@@ -87,6 +107,9 @@ cdef class CoreProtocol:
     cdef _parse_msg_error_response(self, is_error)
     cdef _parse_msg_command_complete(self)
 
+    cdef _auth_password_message_cleartext(self)
+    cdef _auth_password_message_md5(self, bytes salt)
+
     cdef _write(self, buf)
     cdef inline _write_sync_message(self)
 

asyncpg/protocol/coreproto.pyx

@@ -5,10 +5,16 @@
 # the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
 
 
+from hashlib import md5 as hashlib_md5  # for MD5 authentication
+
+
 cdef class CoreProtocol:
 
     def __init__(self, con_args):
         self.buffer = ReadBuffer()
+        self.user = con_args.get('user')
+        self.password = con_args.pop('password', None)
+        self.auth_msg = None
         self.con_args = con_args
         self.transport = None
         self.con_status = CONNECTION_BAD
@@ -106,6 +112,11 @@ cdef class CoreProtocol:
                 self._push_result()
                 self.transport.close()
 
+            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()
@@ -294,18 +305,69 @@ cdef class CoreProtocol:
         self._set_server_parameter(name, val)
 
     cdef _parse_msg_authentication(self):
-        cdef int status
+        cdef:
+            int32_t status
+            bytes md5_salt
+
         status = self.buffer.read_int32()
-        if status != 0:
+
+        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 = cpython.PyBytes_FromStringAndSize(
+                self.buffer.read_bytes(4), 4)
+            self.auth_msg = self._auth_password_message_md5(md5_salt)
+
+        elif status in (AUTH_REQUIRED_KERBEROS, AUTH_REQUIRED_SCMCRED,
+                        AUTH_REQUIRED_GSS, AUTH_REQUIRED_GSS_CONTINUE,
+                        AUTH_REQUIRED_SSPI):
             self.result_type = RESULT_FAILED
             self.result = apg_exc.InterfaceError(
-                'unsupported status {} for Authentication (R) '
-                'message'.format(status))
+                'unsupported authentication method requested by the '
+                'server: {!r}'.format(AUTH_METHOD_NAME[status]))
+
         else:
-            # 0 == AuthenticationOk
-            self.result_type = RESULT_OK
+            self.result_type = RESULT_FAILED
+            self.result = apg_exc.InterfaceError(
+                'unsupported authentication method requested by the '
+                'server: {}'.format(status))
+
         self.buffer.consume_message()
 
+    cdef _auth_password_message_cleartext(self):
+        cdef:
+            WriteBuffer msg
+
+        msg = WriteBuffer.new_message(b'p')
+        msg.write_bytestring(self.password.encode('ascii'))
+        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 '')).encode('ascii')
+        hash = hashlib_md5(hashlib_md5(userpass).hexdigest().\
+                encode('ascii') + salt).hexdigest().encode('ascii')
+
+        msg.write_bytestring(b'md5' + hash)
+        msg.end_message()
+
+        return msg
 
     cdef _parse_msg_ready_for_query(self):
         cdef char status = self.buffer.read_byte()