diff --git a/api/controllers/files/upload.py b/api/controllers/files/upload.py index 6820cc7b75..6c3e7f4af2 100644 --- a/api/controllers/files/upload.py +++ b/api/controllers/files/upload.py @@ -23,8 +23,12 @@ class PluginUploadFileApi(Resource): timestamp = request.args.get("timestamp") nonce = request.args.get("nonce") sign = request.args.get("sign") + tenant_id = request.args.get("tenant_id") + if not tenant_id: + raise Forbidden("Invalid request.") + user_id = request.args.get("user_id") - user = get_user(user_id) + user = get_user(tenant_id, user_id) filename = file.filename mimetype = file.mimetype @@ -38,6 +42,7 @@ class PluginUploadFileApi(Resource): if not verify_plugin_file_signature( filename=filename, mimetype=mimetype, + tenant_id=tenant_id, user_id=user_id, timestamp=timestamp, nonce=nonce, diff --git a/api/controllers/inner_api/plugin/plugin.py b/api/controllers/inner_api/plugin/plugin.py index 5ea2af8e84..fd317519bb 100644 --- a/api/controllers/inner_api/plugin/plugin.py +++ b/api/controllers/inner_api/plugin/plugin.py @@ -259,7 +259,7 @@ class PluginUploadFileRequestApi(Resource): @plugin_data(payload_type=RequestRequestUploadFile) def post(self, user_model: Account | EndUser, tenant_model: Tenant, payload: RequestRequestUploadFile): # generate signed url - url = get_signed_file_url_for_plugin(payload.filename, payload.mimetype, user_model.id) + url = get_signed_file_url_for_plugin(payload.filename, payload.mimetype, tenant_model.id, user_model.id) return BaseBackwardsInvocationResponse(data={"url": url}).model_dump() diff --git a/api/core/file/helpers.py b/api/core/file/helpers.py index 34fcdef5f6..d6c426438b 100644 --- a/api/core/file/helpers.py +++ b/api/core/file/helpers.py @@ -20,7 +20,7 @@ def get_signed_file_url(upload_file_id: str) -> str: return f"{url}?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}" -def get_signed_file_url_for_plugin(filename: str, mimetype: str, user_id: str) -> str: +def get_signed_file_url_for_plugin(filename: str, mimetype: str, tenant_id: str, user_id: str) -> str: url = f"{dify_config.FILES_URL}/files/upload/for-plugin" if user_id is None: @@ -29,7 +29,7 @@ def get_signed_file_url_for_plugin(filename: str, mimetype: str, user_id: str) - timestamp = str(int(time.time())) nonce = os.urandom(16).hex() key = dify_config.SECRET_KEY.encode() - msg = f"upload|{filename}|{mimetype}|{user_id}|{timestamp}|{nonce}" + msg = f"upload|{filename}|{mimetype}|{tenant_id}|{user_id}|{timestamp}|{nonce}" sign = hmac.new(key, msg.encode(), hashlib.sha256).digest() encoded_sign = base64.urlsafe_b64encode(sign).decode() @@ -37,12 +37,12 @@ def get_signed_file_url_for_plugin(filename: str, mimetype: str, user_id: str) - def verify_plugin_file_signature( - *, filename: str, mimetype: str, user_id: str | None, timestamp: str, nonce: str, sign: str + *, filename: str, mimetype: str, tenant_id: str, user_id: str | None, timestamp: str, nonce: str, sign: str ) -> bool: if user_id is None: user_id = "DEFAULT-USER" - data_to_sign = f"upload|{filename}|{mimetype}|{user_id}|{timestamp}|{nonce}" + data_to_sign = f"upload|{filename}|{mimetype}|{tenant_id}|{user_id}|{timestamp}|{nonce}" secret_key = dify_config.SECRET_KEY.encode() recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode()