Procházet zdrojové kódy

feat: add redis ssl support (#65)

Yuhao před 1 rokem
rodič
revize
f8eefa31fe

+ 7 - 0
api/config.py

@@ -21,9 +21,11 @@ DEFAULTS = {
     'REDIS_HOST': 'localhost',
     'REDIS_PORT': '6379',
     'REDIS_DB': '0',
+    'REDIS_USE_SSL': 'False',
     'SESSION_REDIS_HOST': 'localhost',
     'SESSION_REDIS_PORT': '6379',
     'SESSION_REDIS_DB': '2',
+    'SESSION_REDIS_USE_SSL': 'False',
     'OAUTH_REDIRECT_PATH': '/console/api/oauth/authorize',
     'OAUTH_REDIRECT_INDEX_PATH': '/',
     'CONSOLE_URL': 'https://cloud.dify.ai',
@@ -105,14 +107,18 @@ class Config:
         # redis settings
         self.REDIS_HOST = get_env('REDIS_HOST')
         self.REDIS_PORT = get_env('REDIS_PORT')
+        self.REDIS_USERNAME = get_env('REDIS_USERNAME')
         self.REDIS_PASSWORD = get_env('REDIS_PASSWORD')
         self.REDIS_DB = get_env('REDIS_DB')
+        self.REDIS_USE_SSL = get_bool_env('REDIS_USE_SSL')
 
         # session redis settings
         self.SESSION_REDIS_HOST = get_env('SESSION_REDIS_HOST')
         self.SESSION_REDIS_PORT = get_env('SESSION_REDIS_PORT')
+        self.SESSION_REDIS_USERNAME = get_env('SESSION_REDIS_USERNAME')
         self.SESSION_REDIS_PASSWORD = get_env('SESSION_REDIS_PASSWORD')
         self.SESSION_REDIS_DB = get_env('SESSION_REDIS_DB')
+        self.SESSION_REDIS_USE_SSL = get_bool_env('SESSION_REDIS_USE_SSL')
 
         # storage settings
         self.STORAGE_TYPE = get_env('STORAGE_TYPE')
@@ -165,6 +171,7 @@ class Config:
         self.CELERY_BACKEND = get_env('CELERY_BACKEND')
         self.CELERY_RESULT_BACKEND = 'db+{}'.format(self.SQLALCHEMY_DATABASE_URI) \
             if self.CELERY_BACKEND == 'database' else self.CELERY_BROKER_URL
+        self.BROKER_USE_SSL = self.CELERY_BROKER_URL.startswith('rediss://')
 
         # hosted provider credentials
         self.OPENAI_API_KEY = get_env('OPENAI_API_KEY')

+ 15 - 0
api/extensions/ext_celery.py

@@ -15,9 +15,24 @@ def init_app(app: Flask) -> Celery:
         backend=app.config["CELERY_BACKEND"],
         task_ignore_result=True,
     )
+    
+    # Add SSL options to the Celery configuration
+    ssl_options = {
+        "ssl_cert_reqs": None,
+        "ssl_ca_certs": None,
+        "ssl_certfile": None,
+        "ssl_keyfile": None,
+    }
+
     celery_app.conf.update(
         result_backend=app.config["CELERY_RESULT_BACKEND"],
     )
+
+    if app.config["BROKER_USE_SSL"]:
+        celery_app.conf.update(
+            broker_use_ssl=ssl_options,  # Add the SSL options to the broker configuration
+        )
+        
     celery_app.set_default()
     app.extensions["celery"] = celery_app
     return celery_app

+ 7 - 2
api/extensions/ext_redis.py

@@ -1,18 +1,23 @@
 import redis
-
+from redis.connection import SSLConnection, Connection
 
 redis_client = redis.Redis()
 
 
 def init_app(app):
+    connection_class = Connection
+    if app.config.get('REDIS_USE_SSL', False):
+        connection_class = SSLConnection
+
     redis_client.connection_pool = redis.ConnectionPool(**{
         'host': app.config.get('REDIS_HOST', 'localhost'),
         'port': app.config.get('REDIS_PORT', 6379),
+        'username': app.config.get('REDIS_USERNAME', None),
         'password': app.config.get('REDIS_PASSWORD', None),
         'db': app.config.get('REDIS_DB', 0),
         'encoding': 'utf-8',
         'encoding_errors': 'strict',
         'decode_responses': False
-    })
+    }, connection_class=connection_class)
 
     app.extensions['redis'] = redis_client

+ 7 - 1
api/extensions/ext_session.py

@@ -1,4 +1,5 @@
 import redis
+from redis.connection import SSLConnection, Connection
 from flask import request
 from flask_session import Session, SqlAlchemySessionInterface, RedisSessionInterface
 from flask_session.sessions import total_seconds
@@ -23,16 +24,21 @@ def init_app(app):
     if session_type == 'sqlalchemy':
         app.session_interface = sqlalchemy_session_interface
     elif session_type == 'redis':
+        connection_class = Connection
+        if app.config.get('SESSION_REDIS_USE_SSL', False):
+            connection_class = SSLConnection
+
         sess_redis_client = redis.Redis()
         sess_redis_client.connection_pool = redis.ConnectionPool(**{
             'host': app.config.get('SESSION_REDIS_HOST', 'localhost'),
             'port': app.config.get('SESSION_REDIS_PORT', 6379),
+            'username': app.config.get('SESSION_REDIS_USERNAME', None),
             'password': app.config.get('SESSION_REDIS_PASSWORD', None),
             'db': app.config.get('SESSION_REDIS_DB', 2),
             'encoding': 'utf-8',
             'encoding_errors': 'strict',
             'decode_responses': False
-        })
+        }, connection_class=connection_class)
 
         app.extensions['session_redis'] = sess_redis_client
 

+ 6 - 0
docker/docker-compose.yaml

@@ -36,14 +36,18 @@ services:
       # It is consistent with the configuration in the 'redis' service below.
       REDIS_HOST: redis
       REDIS_PORT: 6379
+      REDIS_USERNAME: ''
       REDIS_PASSWORD: difyai123456
+      REDIS_USE_SSL: 'false'
       # use redis db 0 for redis cache
       REDIS_DB: 0
       # The configurations of session, Supported values are `sqlalchemy`. `redis`
       SESSION_TYPE: redis
       SESSION_REDIS_HOST: redis
       SESSION_REDIS_PORT: 6379
+      SESSION_REDIS_USERNAME: ''
       SESSION_REDIS_PASSWORD: difyai123456
+      SESSION_REDIS_USE_SSL: 'false'
       # use redis db 2 for session store
       SESSION_REDIS_DB: 2
       # The configurations of celery broker.
@@ -129,8 +133,10 @@ services:
       # The configurations of redis cache connection.
       REDIS_HOST: redis
       REDIS_PORT: 6379
+      REDIS_USERNAME: ''
       REDIS_PASSWORD: difyai123456
       REDIS_DB: 0
+      REDIS_USE_SSL: 'false'
       # The configurations of celery broker.
       CELERY_BROKER_URL: redis://:difyai123456@redis:6379/1
       # The type of storage to use for storing user files. Supported values are `local` and `s3`, Default: `local`