From d7fbae286a666fa4cdb5edda74a0eaf49c0f542b Mon Sep 17 00:00:00 2001
From: Jyong <76649700+JohnJyong@users.noreply.github.com>
Date: Fri, 14 Jun 2024 15:19:59 +0800
Subject: [PATCH] add aws s3 iam check (#5174)

---
 api/.env.example                     |  1 +
 api/config.py                        |  2 ++
 api/extensions/storage/s3_storage.py | 20 ++++++++++++--------
 docker/docker-compose.yaml           |  2 ++
 4 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/api/.env.example b/api/.env.example
index 571f5b168ea5a7..42ed40e4c0e25a 100644
--- a/api/.env.example
+++ b/api/.env.example
@@ -42,6 +42,7 @@ DB_DATABASE=dify
 # storage type: local, s3, azure-blob
 STORAGE_TYPE=local
 STORAGE_LOCAL_PATH=storage
+S3_USE_AWS_MANAGED_IAM=false
 S3_ENDPOINT=https://your-bucket-name.storage.s3.clooudflare.com
 S3_BUCKET_NAME=your-bucket-name
 S3_ACCESS_KEY=your-access-key
diff --git a/api/config.py b/api/config.py
index 53cc6ce393b6ab..0c2c8d7cf679f2 100644
--- a/api/config.py
+++ b/api/config.py
@@ -24,6 +24,7 @@
     'APP_WEB_URL': 'https://udify.app',
     'FILES_URL': '',
     'FILES_ACCESS_TIMEOUT': 300,
+    'S3_USE_AWS_MANAGED_IAM': 'False',
     'S3_ADDRESS_STYLE': 'auto',
     'STORAGE_TYPE': 'local',
     'STORAGE_LOCAL_PATH': 'storage',
@@ -226,6 +227,7 @@ def __init__(self):
         self.STORAGE_LOCAL_PATH = get_env('STORAGE_LOCAL_PATH')
 
         # S3 Storage settings
+        self.S3_USE_AWS_MANAGED_IAM = get_bool_env('S3_USE_AWS_MANAGED_IAM')
         self.S3_ENDPOINT = get_env('S3_ENDPOINT')
         self.S3_BUCKET_NAME = get_env('S3_BUCKET_NAME')
         self.S3_ACCESS_KEY = get_env('S3_ACCESS_KEY')
diff --git a/api/extensions/storage/s3_storage.py b/api/extensions/storage/s3_storage.py
index 8aae68a740014a..787596fa791d4a 100644
--- a/api/extensions/storage/s3_storage.py
+++ b/api/extensions/storage/s3_storage.py
@@ -16,14 +16,18 @@ def __init__(self, app: Flask):
         super().__init__(app)
         app_config = self.app.config
         self.bucket_name = app_config.get('S3_BUCKET_NAME')
-        self.client = boto3.client(
-                    's3',
-                    aws_secret_access_key=app_config.get('S3_SECRET_KEY'),
-                    aws_access_key_id=app_config.get('S3_ACCESS_KEY'),
-                    endpoint_url=app_config.get('S3_ENDPOINT'),
-                    region_name=app_config.get('S3_REGION'),
-                    config=Config(s3={'addressing_style': app_config.get('S3_ADDRESS_STYLE')})
-                )
+        if app_config.get('S3_USE_AWS_MANAGED_IAM'):
+            session = boto3.Session()
+            self.client = session.client('s3')
+        else:
+            self.client = boto3.client(
+                        's3',
+                        aws_secret_access_key=app_config.get('S3_SECRET_KEY'),
+                        aws_access_key_id=app_config.get('S3_ACCESS_KEY'),
+                        endpoint_url=app_config.get('S3_ENDPOINT'),
+                        region_name=app_config.get('S3_REGION'),
+                        config=Config(s3={'addressing_style': app_config.get('S3_ADDRESS_STYLE')})
+                    )
 
     def save(self, filename, data):
         self.client.put_object(Bucket=self.bucket_name, Key=filename, Body=data)
diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml
index 5c67406bcb5eeb..f68e2aec3f153f 100644
--- a/docker/docker-compose.yaml
+++ b/docker/docker-compose.yaml
@@ -81,6 +81,7 @@ services:
       # only available when STORAGE_TYPE is `local`.
       STORAGE_LOCAL_PATH: storage
       # The S3 storage configurations, only available when STORAGE_TYPE is `s3`.
+      S3_USE_AWS_MANAGED_IAM: 'false'
       S3_ENDPOINT: 'https://xxx.r2.cloudflarestorage.com'
       S3_BUCKET_NAME: 'difyai'
       S3_ACCESS_KEY: 'ak-difyai'
@@ -236,6 +237,7 @@ services:
       STORAGE_TYPE: local
       STORAGE_LOCAL_PATH: storage
       # The S3 storage configurations, only available when STORAGE_TYPE is `s3`.
+      S3_USE_AWS_MANAGED_IAM: 'false'
       S3_ENDPOINT: 'https://xxx.r2.cloudflarestorage.com'
       S3_BUCKET_NAME: 'difyai'
       S3_ACCESS_KEY: 'ak-difyai'