Merge pull request #2 from zhanghaitao3/feat/gaussdb-sha256-auth

feat(protocol): Refactoring sha256 authentication code to comply with…

Author: liubao68

asyncpg/protocol/coreproto.pyx

@@ -9,8 +9,7 @@ import hashlib
 import hmac
 from hashlib import pbkdf2_hmac
 include "scram.pyx"
-
-
+include "sha256.pyx"
 
 
 cdef dict AUTH_METHOD_NAME = {
@@ -23,248 +22,6 @@ cdef dict AUTH_METHOD_NAME = {
 }
 
 
-cdef hex_string_to_bytes(str hex_string):
-    """
-    将hex字符串转换为bytes，对应Go的hexStringToBytes函数
-    """
-    if not hex_string:
-        return b''
-    
-    cdef str upper_string = hex_string.upper()
-    cdef int bytes_len = len(upper_string) // 2
-    cdef bytearray result = bytearray(bytes_len)
-    cdef int i, pos
-    cdef str high_char, low_char
-    cdef int high_val, low_val
-    
-    for i in range(bytes_len):
-        pos = i * 2
-        high_char = upper_string[pos]
-        low_char = upper_string[pos + 1]
-        
-        # 将字符转换为数值
-        high_val = "0123456789ABCDEF".index(high_char)
-        low_val = "0123456789ABCDEF".index(low_char)
-        
-        result[i] = (high_val << 4) | low_val
-    
-    return bytes(result)
-
-cdef generate_k_from_pbkdf2(str password, str random64code, int server_iteration):
-    """
-    对应Go的generateKFromPBKDF2函数
-    注意：Go代码使用的是SHA1，不是SHA256
-    """
-    cdef bytes random32code = hex_string_to_bytes(random64code)
-    # Go代码使用sha1.New，所以这里使用'sha1'
-    cdef bytes pwd_encoded = pbkdf2_hmac('sha1', password.encode('utf-8'), random32code, server_iteration, 32)
-    return pwd_encoded
-
-cdef bytes_to_hex_string(bytes src):
-    """
-    对应Go的bytesToHexString函数
-    """
-    cdef str s = ""
-    cdef int byte_val, v
-    cdef str hv
-    
-    for byte_val in src:
-        v = byte_val & 0xFF
-        hv = format(v, 'x')
-        if len(hv) < 2:
-            s += "0" + hv
-        else:
-            s += hv
-    return s
-
-cdef get_key_from_hmac(bytes key, bytes data):
-    """
-    对应Go的getKeyFromHmac函数，使用SHA256
-    """
-    cdef object h = hmac.new(key, data, hashlib.sha256)
-    return h.digest()
-
-cdef get_sha256(bytes message):
-    """
-    对应Go的getSha256函数
-    """
-    cdef object hash_obj = hashlib.sha256()
-    hash_obj.update(message)
-    return hash_obj.digest()
-
-cdef get_sm3(bytes message):
-    """
-    对应Go的getSm3函数 (这里用SHA256代替，因为Python标准库没有SM3)
-    实际项目中需要安装gmssl库来支持SM3
-    """
-    # 注意：这里用SHA256代替SM3，实际使用时需要proper的SM3实现
-    cdef object hash_obj = hashlib.sha256()  # 临时用SHA256代替
-    hash_obj.update(message)
-    return hash_obj.digest()
-
-cdef xor_between_password(bytes password1, bytes password2, int length):
-    """
-    对应Go的XorBetweenPassword函数
-    """
-    cdef bytearray result = bytearray(length)
-    cdef int i
-    
-    for i in range(length):
-        result[i] = password1[i] ^ password2[i]
-    return bytes(result)
-
-cdef bytes_to_hex(bytes source_bytes, bytearray result_array=None, int start_pos=0, int length=-1):
-    """
-    对应Go的bytesToHex函数，支持Java风格的4参数调用
-    但Go代码只传1个参数，所以做兼容处理
-    """
-    cdef bytes lookup = b'0123456789abcdef'
-    cdef int pos, i, c, j
-    cdef int byte_val
-    cdef bytearray result
-    
-    if result_array is not None:
-        # Java风格：4个参数 bytesToHex(hValue, result, 0, hValue.length)
-        if length == -1:
-            length = len(source_bytes)
-        
-        pos = start_pos
-        
-        for i in range(length):
-            if i >= len(source_bytes):
-                break
-            byte_val = source_bytes[i]
-            c = int(byte_val & 0xFF)
-            j = c >> 4
-            result_array[pos] = lookup[j]
-            pos += 1
-            j = c & 0xF
-            result_array[pos] = lookup[j]
-            pos += 1
-        return result_array
-    else:
-        # Go风格：1个参数，返回新的bytes
-        result = bytearray(len(source_bytes) * 2)
-        pos = 0
-        
-        for byte_val in source_bytes:
-            c = int(byte_val & 0xFF)
-            j = c >> 4
-            result[pos] = lookup[j]
-            pos += 1
-            j = c & 0xF
-            result[pos] = lookup[j]
-            pos += 1
-        
-        return bytes(result)
-
-cpdef rfc5802_algorithm(str password, str random64code, str token, str server_signature="", int server_iteration=4096, str method="sha256"):
-    """
-    RFC5802算法实现，完全对应Go代码逻辑
-    """
-    cdef bytes k, server_key, client_key, stored_key, token_byte
-    cdef bytes client_signature, hmac_result, h_value
-    cdef bytearray result
-    cdef int h_value_len
-    
-    try:
-        # Step 1: 生成K (SaltedPassword)
-        k = generate_k_from_pbkdf2(password, random64code, server_iteration)
-        
-        # Step 2: 生成ServerKey和ClientKey
-        server_key = get_key_from_hmac(k, b"Sever Key")  # 保持"Sever Key"拼写
-        client_key = get_key_from_hmac(k, b"Client Key")
-        
-        # Step 3: 生成StoredKey
-        if method.lower() == "sha256":
-            stored_key = get_sha256(client_key)
-        elif method.lower() == "sm3":
-            stored_key = get_sm3(client_key)
-        else:
-            stored_key = get_sha256(client_key)  # 默认使用SHA256
-        
-        # Step 4: 转换token为bytes
-        token_byte = hex_string_to_bytes(token)
-        
-        # Step 5: 计算clientSignature (实际上是ServerSignature，用于验证)
-        client_signature = get_key_from_hmac(server_key, token_byte)
-        
-        # Step 6: 验证serverSignature (如果提供)
-        if server_signature and server_signature != bytes_to_hex_string(client_signature):
-            return b""
-        
-        # Step 7: 计算真正的ClientSignature
-        hmac_result = get_key_from_hmac(stored_key, token_byte)
-        
-        # Step 8: XOR操作得到ClientProof
-        h_value = xor_between_password(hmac_result, client_key, len(client_key))
-        
-        # Step 9: 转换为hex bytes格式 (对应Java的 bytesToHex(hValue, result, 0, hValue.length))
-        h_value_len = len(h_value)
-        result = bytearray(h_value_len * 2)
-        bytes_to_hex(h_value, result, 0, h_value_len)
-        
-        return bytes(result)
-        
-    except Exception as e:
-        raise ValueError(f"RFC5802Algorithm failed: {e}")
-
-
-
-
-
-
-
-
-
-
-def bytes_to_hex(src_bytes, dst_bytes, offset, length):
-    """
-    Java: bytesToHex(byte[] src, byte[] dst, int offset, int length)
-    - src: 源字节数组
-    - dst: 目标字节数组
-    - offset: dst写入起始位置
-    - length: 需要转换的src字节数量
-    写入的输出是十六进制ASCII字节（不是16进制数值），每个字节转换成2个字母。
-    """
-    HEX_DIGITS = b'0123456789abcdef'
-    for i in range(length):
-        v = src_bytes[i]
-        if isinstance(v, str):
-            v = ord(v)
-        if v < 0:
-            v += 256
-        dst_bytes[offset + (i * 2)] = HEX_DIGITS[v >> 4]
-        dst_bytes[offset + (i * 2) + 1] = HEX_DIGITS[v & 0x0F]
-
-def SHA256_MD5encode(user: bytes, password: bytes, salt: bytes) -> bytes:
-    try:
-        md = hashlib.md5()
-        md.update(password)
-        md.update(user)
-        temp_digest = md.digest()  # 16 bytes
-
-        # hex_digest 70字节（实际前6和后64有效）
-        hex_digest = bytearray(70)
-
-        # 前32个字节为temp_digest的hex(16字节*2)
-        bytes_to_hex(temp_digest, hex_digest, 0, 16)
-
-        # 取前32字节(hex后缀): hex_digest[0:32], 作为SHA256输入
-        sha = hashlib.sha256()
-        sha.update(hex_digest[0:32])
-        sha.update(salt)
-        pass_digest = sha.digest()  # 32 bytes
-
-        # pass_digest的hex写到hex_digest[6:]
-        bytes_to_hex(pass_digest, hex_digest, 6, 32)
-
-        # 填入ASCII签名'sha256'
-        hex_digest[0:6] = b'sha256'
-
-    except Exception as e:
-        raise ValueError('SHA256_MD5encode failed: %s' % str(e))
-    return bytes(hex_digest)
 
 cdef class CoreProtocol:
 
@@ -282,6 +39,8 @@ cdef class CoreProtocol:
         self.encoding = 'utf-8'
         # type of `scram` is `SCRAMAuthentcation`
         self.scram = None
+        # type of `sha256` is `RFC5802Authentication`
+        self.sha256 = None
         # type of `gss_ctx` is `gssapi.SecurityContext` or
         # `sspilib.SecurityContext`
         self.gss_ctx = None
@@ -814,7 +573,6 @@ cdef class CoreProtocol:
             bytes token
             int32_t server_iteration
         status = self.buffer.read_int32()
-
         if status == AUTH_SUCCESSFUL:
             # AuthenticationOk
             self.result_type = RESULT_OK
@@ -842,15 +600,15 @@ cdef class CoreProtocol:
                     'The server requested password-based authentication, '
                     'but no password was provided.')
             if password_stored_method in (SHA256_PASSWORD,PLAIN_PASSWORD):
-                # 读取认证参数
+                # Read authentication parameters
                 random64code = self.buffer.read_bytes(64)
                 token = self.buffer.read_bytes(8)
                 server_iteration = self.buffer.read_int32()
-                # 调用_auth_password_message_sha256生成认证消息
+                # generate authentication message
                 self.auth_msg = self._auth_password_message_sha256(random64code, token, 
                                        server_iteration)
             elif password_stored_method == MD5_PASSWORD:
-                # MD5密码存储方式
+                # MD5 password storage method
                 salt = self.buffer.read_bytes(4)
                 self.auth_msg = self._auth_password_message_md5(salt)
 
@@ -938,19 +696,15 @@ cdef class CoreProtocol:
 
     cdef _auth_password_message_sha256(self, bytes random64code, bytes token, 
                                        int32_t server_iteration):
-        """
-        处理SHA256认证消息
-        """
         cdef:
             WriteBuffer msg
             bytes result
-
-        # 调用RFC5802算法计算认证结果
-        result = rfc5802_algorithm(
+        self.sha256 = RFC5802Authentication()
+        result = self.sha256.authenticate(
             self.password,
             random64code.decode('utf-8'),
             token.decode('utf-8'),
-            '',  # salt为空
+            '', 
             server_iteration,
             'sha256'
         )
@@ -959,8 +713,7 @@ cdef class CoreProtocol:
             self.result = apg_exc.InterfaceError(
                 'Invalid username/password, login denied.')
             return None
-        
-        # 构建认证响应消息
+
         msg = WriteBuffer.new_message(b'p')
         msg.write_bytes(result)
         msg.end_message()

asyncpg/protocol/sha256.pxd

@@ -0,0 +1,55 @@
+
+# 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
+
+
+@cython.final
+cdef class RFC5802Authentication:
+    """Cython header declarations for RFC5802 SCRAM authentication mechanism.
+
+    This header file defines the interface for the RFC5802Authentication class
+    which implements the RFC5802 SCRAM authentication mechanism for secure
+    password authentication.
+
+    The class provides cryptographic operations for:
+    - PBKDF2 key derivation
+    - HMAC-based key generation
+    - SHA256/SM3 hash computation
+    - Hexadecimal encoding/decoding
+    - XOR operations for password proofs
+
+    All methods follow the RFC5802 specification for SCRAM authentication.
+    """
+
+    # Class constants
+    cdef readonly int DEFAULT_ITERATIONS
+    cdef readonly int DEFAULT_KEY_LENGTH
+    cdef readonly list SUPPORTED_METHODS
+    cdef readonly str HEX_CHARS
+    cdef readonly bytes HEX_LOOKUP
+
+    # Hexadecimal conversion methods
+    cdef bytes hex_string_to_bytes(self, str hex_string)
+    cdef str _bytes_to_hex_string(self, bytes src)
+    cdef bytes _bytes_to_hex(self, bytes source_bytes, bytearray result_array=*, 
+                            int start_pos=*, int length=*)
+
+    # Cryptographic key generation methods
+    cdef bytes _generate_k_from_pbkdf2(self, str password, str random64code, 
+                                      int server_iteration)
+    cdef bytes _get_key_from_hmac(self, bytes key, bytes data)
+
+    # Hash computation methods
+    cdef bytes _get_sha256(self, bytes message)
+    cdef bytes _get_sm3(self, bytes message)
+
+    # Password operation methods
+    cdef bytes _xor_between_password(self, bytes password1, bytes password2, int length)
+
+    # Main authentication method
+    cpdef bytes authenticate(self, str password, str random64code, str token, 
+                           str server_signature=*, int server_iteration=*, 
+                           str method=*)