Model Runtime (#1858)

Co-authored-by: StyleZhang <jasonapring2015@outlook.com>
Co-authored-by: Garfield Dai <dai.hai@foxmail.com>
Co-authored-by: chenhe <guchenhe@gmail.com>
Co-authored-by: jyong <jyong@dify.ai>
Co-authored-by: Joel <iamjoel007@gmail.com>
Co-authored-by: Yeuoly <admin@srmxy.cn>

api/libs/rsa.py

@@ -5,6 +5,7 @@ from Crypto.Cipher import PKCS1_OAEP, AES
 from Crypto.PublicKey import RSA
 from Crypto.Random import get_random_bytes
 
+from core.helper.lru_cache import LRUCache
 from extensions.ext_redis import redis_client
 from extensions.ext_storage import storage
 
@@ -45,7 +46,15 @@ def encrypt(text, public_key):
     return prefix_hybrid + encrypted_data
 
 
-def decrypt(encrypted_text, tenant_id):
+tenant_rsa_keys = LRUCache(capacity=1000)
+
+
+def get_decrypt_decoding(tenant_id):
+    rsa_key = tenant_rsa_keys.get(tenant_id)
+    if rsa_key:
+        cipher_rsa = PKCS1_OAEP.new(rsa_key)
+        return rsa_key, cipher_rsa
+
     filepath = "privkeys/{tenant_id}".format(tenant_id=tenant_id) + "/private.pem"
 
     cache_key = 'tenant_privkey:{hash}'.format(hash=hashlib.sha3_256(filepath.encode()).hexdigest())
@@ -61,6 +70,12 @@ def decrypt(encrypted_text, tenant_id):
     rsa_key = RSA.import_key(private_key)
     cipher_rsa = PKCS1_OAEP.new(rsa_key)
 
+    tenant_rsa_keys.put(tenant_id, rsa_key)
+
+    return rsa_key, cipher_rsa
+
+
+def decrypt_token_with_decoding(encrypted_text, rsa_key, cipher_rsa):
     if encrypted_text.startswith(prefix_hybrid):
         encrypted_text = encrypted_text[len(prefix_hybrid):]
 
@@ -79,5 +94,11 @@ def decrypt(encrypted_text, tenant_id):
     return decrypted_text.decode()
 
 
+def decrypt(encrypted_text, tenant_id):
+    rsa_key, cipher_rsa = get_decrypt_decoding(tenant_id)
+
+    return decrypt_token_with_decoding(encrypted_text, rsa_key, cipher_rsa)
+
+
 class PrivkeyNotFoundError(Exception):
     pass