feat(e2ee): fix b64 check

Author: liuhuiqi.7

volcenginesdkarkruntime/_utils/_key_agreement.py

@@ -100,17 +100,28 @@ def aes_gcm_decrypt_base64_list(key: bytes, nonce: bytes, ciphertext: str) -> st
         try:
             result.append(aes_gcm_decrypt_base64_string(key, nonce, b64))
         except Exception:
+            print("decrypt base64 string failed", b64)
             for i in range(20, len(b64), 4):
                 try:
-                    decrypted = aes_gcm_decrypt_base64_string(key, nonce, b64[:i+4])
+                    decrypted = aes_gcm_decrypt_base64_string(
+                        key, nonce, b64[:i+4])
                     result.append(decrypted)
-                    decrypted = aes_gcm_decrypt_base64_string(key, nonce, b64[i+4:])
+                    decrypted = aes_gcm_decrypt_base64_string(
+                        key, nonce, b64[i+4:])
                     result.append(decrypted)
                 except Exception:
-                    result.append('')
+                    print("decrypt base64 string failed",
+                          b64, b64[:i+4], b64[i+4:])
+    print('decrypt base64 list result', result)
     return ''.join(result)
 
 
+def decrypt_validate(ciphertext: str) -> bool:
+    cipher_bytes = ciphertext.encode()
+    cipher_b64_bytes = base64.decodebytes(cipher_bytes)
+    return len(cipher_bytes)/4 >= len(cipher_b64_bytes)/3 >= len(cipher_bytes)/4 - 1
+
+
 def marshal_cryptography_pub_key(key) -> bytes:
     # python version of crypto/elliptic/elliptic.go Marshal
     # without point on curve check

volcenginesdkarkruntime/resources/chat/completions.py

@@ -31,7 +31,7 @@
 
 from ..._types import Body, Query, Headers
 from ..._utils._utils import deepcopy_minimal, with_sts_token, async_with_sts_token
-from ..._utils._key_agreement import aes_gcm_decrypt_base64_string, aes_gcm_decrypt_base64_list
+from ..._utils._key_agreement import aes_gcm_decrypt_base64_string, aes_gcm_decrypt_base64_list, decrypt_validate
 from ..._base_client import make_request_options
 from ..._resource import SyncAPIResource, AsyncAPIResource
 from ..._compat import cached_property
@@ -149,7 +149,7 @@ def _decrypt(
                         content = aes_gcm_decrypt_base64_string(
                             key, nonce, choice.message.content
                         )
-                        if content == '':
+                        if not decrypt_validate(choice.message.content):
                             content = aes_gcm_decrypt_base64_list(
                                 key, nonce, choice.message.content
                             )
@@ -312,7 +312,7 @@ async def _decrypt(
                         content = aes_gcm_decrypt_base64_string(
                             key, nonce, choice.message.content
                         )
-                        if content == '':
+                        if not decrypt_validate(choice.message.content):
                             content = aes_gcm_decrypt_base64_list(
                                 key, nonce, choice.message.content
                             )