feat: add Redis SSL/TLS certificate authentication support (#23624)
This commit is contained in:
-LAN-
@@ -1,5 +1,6 @@
|
||||
import functools
|
||||
import logging
|
||||
import ssl
|
||||
from collections.abc import Callable
|
||||
from datetime import timedelta
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
@@ -116,76 +117,132 @@ class RedisClientWrapper:
|
||||
redis_client: RedisClientWrapper = RedisClientWrapper()
|
||||
|
||||
|
||||
def init_app(app: DifyApp):
|
||||
global redis_client
|
||||
connection_class: type[Union[Connection, SSLConnection]] = Connection
|
||||
if dify_config.REDIS_USE_SSL:
|
||||
connection_class = SSLConnection
|
||||
resp_protocol = dify_config.REDIS_SERIALIZATION_PROTOCOL
|
||||
if dify_config.REDIS_ENABLE_CLIENT_SIDE_CACHE:
|
||||
if resp_protocol >= 3:
|
||||
clientside_cache_config = CacheConfig()
|
||||
else:
|
||||
raise ValueError("Client side cache is only supported in RESP3")
|
||||
else:
|
||||
clientside_cache_config = None
|
||||
def _get_ssl_configuration() -> tuple[type[Union[Connection, SSLConnection]], dict[str, Any]]:
|
||||
"""Get SSL configuration for Redis connection."""
|
||||
if not dify_config.REDIS_USE_SSL:
|
||||
return Connection, {}
|
||||
|
||||
redis_params: dict[str, Any] = {
|
||||
cert_reqs_map = {
|
||||
"CERT_NONE": ssl.CERT_NONE,
|
||||
"CERT_OPTIONAL": ssl.CERT_OPTIONAL,
|
||||
"CERT_REQUIRED": ssl.CERT_REQUIRED,
|
||||
}
|
||||
ssl_cert_reqs = cert_reqs_map.get(dify_config.REDIS_SSL_CERT_REQS, ssl.CERT_NONE)
|
||||
|
||||
ssl_kwargs = {
|
||||
"ssl_cert_reqs": ssl_cert_reqs,
|
||||
"ssl_ca_certs": dify_config.REDIS_SSL_CA_CERTS,
|
||||
"ssl_certfile": dify_config.REDIS_SSL_CERTFILE,
|
||||
"ssl_keyfile": dify_config.REDIS_SSL_KEYFILE,
|
||||
}
|
||||
|
||||
return SSLConnection, ssl_kwargs
|
||||
|
||||
|
||||
def _get_cache_configuration() -> CacheConfig | None:
|
||||
"""Get client-side cache configuration if enabled."""
|
||||
if not dify_config.REDIS_ENABLE_CLIENT_SIDE_CACHE:
|
||||
return None
|
||||
|
||||
resp_protocol = dify_config.REDIS_SERIALIZATION_PROTOCOL
|
||||
if resp_protocol < 3:
|
||||
raise ValueError("Client side cache is only supported in RESP3")
|
||||
|
||||
return CacheConfig()
|
||||
|
||||
|
||||
def _get_base_redis_params() -> dict[str, Any]:
|
||||
"""Get base Redis connection parameters."""
|
||||
return {
|
||||
"username": dify_config.REDIS_USERNAME,
|
||||
"password": dify_config.REDIS_PASSWORD or None, # Temporary fix for empty password
|
||||
"password": dify_config.REDIS_PASSWORD or None,
|
||||
"db": dify_config.REDIS_DB,
|
||||
"encoding": "utf-8",
|
||||
"encoding_errors": "strict",
|
||||
"decode_responses": False,
|
||||
"protocol": resp_protocol,
|
||||
"cache_config": clientside_cache_config,
|
||||
"protocol": dify_config.REDIS_SERIALIZATION_PROTOCOL,
|
||||
"cache_config": _get_cache_configuration(),
|
||||
}
|
||||
|
||||
if dify_config.REDIS_USE_SENTINEL:
|
||||
assert dify_config.REDIS_SENTINELS is not None, "REDIS_SENTINELS must be set when REDIS_USE_SENTINEL is True"
|
||||
assert dify_config.REDIS_SENTINEL_SERVICE_NAME is not None, (
|
||||
"REDIS_SENTINEL_SERVICE_NAME must be set when REDIS_USE_SENTINEL is True"
|
||||
)
|
||||
sentinel_hosts = [
|
||||
(node.split(":")[0], int(node.split(":")[1])) for node in dify_config.REDIS_SENTINELS.split(",")
|
||||
]
|
||||
sentinel = Sentinel(
|
||||
sentinel_hosts,
|
||||
sentinel_kwargs={
|
||||
"socket_timeout": dify_config.REDIS_SENTINEL_SOCKET_TIMEOUT,
|
||||
"username": dify_config.REDIS_SENTINEL_USERNAME,
|
||||
"password": dify_config.REDIS_SENTINEL_PASSWORD,
|
||||
},
|
||||
)
|
||||
master = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **redis_params)
|
||||
redis_client.initialize(master)
|
||||
elif dify_config.REDIS_USE_CLUSTERS:
|
||||
assert dify_config.REDIS_CLUSTERS is not None, "REDIS_CLUSTERS must be set when REDIS_USE_CLUSTERS is True"
|
||||
nodes = [
|
||||
ClusterNode(host=node.split(":")[0], port=int(node.split(":")[1]))
|
||||
for node in dify_config.REDIS_CLUSTERS.split(",")
|
||||
]
|
||||
redis_client.initialize(
|
||||
RedisCluster(
|
||||
startup_nodes=nodes,
|
||||
password=dify_config.REDIS_CLUSTERS_PASSWORD,
|
||||
protocol=resp_protocol,
|
||||
cache_config=clientside_cache_config,
|
||||
)
|
||||
)
|
||||
else:
|
||||
redis_params.update(
|
||||
{
|
||||
"host": dify_config.REDIS_HOST,
|
||||
"port": dify_config.REDIS_PORT,
|
||||
"connection_class": connection_class,
|
||||
"protocol": resp_protocol,
|
||||
"cache_config": clientside_cache_config,
|
||||
}
|
||||
)
|
||||
pool = redis.ConnectionPool(**redis_params)
|
||||
redis_client.initialize(redis.Redis(connection_pool=pool))
|
||||
|
||||
def _create_sentinel_client(redis_params: dict[str, Any]) -> Union[redis.Redis, RedisCluster]:
|
||||
"""Create Redis client using Sentinel configuration."""
|
||||
if not dify_config.REDIS_SENTINELS:
|
||||
raise ValueError("REDIS_SENTINELS must be set when REDIS_USE_SENTINEL is True")
|
||||
|
||||
if not dify_config.REDIS_SENTINEL_SERVICE_NAME:
|
||||
raise ValueError("REDIS_SENTINEL_SERVICE_NAME must be set when REDIS_USE_SENTINEL is True")
|
||||
|
||||
sentinel_hosts = [(node.split(":")[0], int(node.split(":")[1])) for node in dify_config.REDIS_SENTINELS.split(",")]
|
||||
|
||||
sentinel = Sentinel(
|
||||
sentinel_hosts,
|
||||
sentinel_kwargs={
|
||||
"socket_timeout": dify_config.REDIS_SENTINEL_SOCKET_TIMEOUT,
|
||||
"username": dify_config.REDIS_SENTINEL_USERNAME,
|
||||
"password": dify_config.REDIS_SENTINEL_PASSWORD,
|
||||
},
|
||||
)
|
||||
|
||||
master: redis.Redis = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **redis_params)
|
||||
return master
|
||||
|
||||
|
||||
def _create_cluster_client() -> Union[redis.Redis, RedisCluster]:
|
||||
"""Create Redis cluster client."""
|
||||
if not dify_config.REDIS_CLUSTERS:
|
||||
raise ValueError("REDIS_CLUSTERS must be set when REDIS_USE_CLUSTERS is True")
|
||||
|
||||
nodes = [
|
||||
ClusterNode(host=node.split(":")[0], port=int(node.split(":")[1]))
|
||||
for node in dify_config.REDIS_CLUSTERS.split(",")
|
||||
]
|
||||
|
||||
cluster: RedisCluster = RedisCluster(
|
||||
startup_nodes=nodes,
|
||||
password=dify_config.REDIS_CLUSTERS_PASSWORD,
|
||||
protocol=dify_config.REDIS_SERIALIZATION_PROTOCOL,
|
||||
cache_config=_get_cache_configuration(),
|
||||
)
|
||||
return cluster
|
||||
|
||||
|
||||
def _create_standalone_client(redis_params: dict[str, Any]) -> Union[redis.Redis, RedisCluster]:
|
||||
"""Create standalone Redis client."""
|
||||
connection_class, ssl_kwargs = _get_ssl_configuration()
|
||||
|
||||
redis_params.update(
|
||||
{
|
||||
"host": dify_config.REDIS_HOST,
|
||||
"port": dify_config.REDIS_PORT,
|
||||
"connection_class": connection_class,
|
||||
}
|
||||
)
|
||||
|
||||
if ssl_kwargs:
|
||||
redis_params.update(ssl_kwargs)
|
||||
|
||||
pool = redis.ConnectionPool(**redis_params)
|
||||
client: redis.Redis = redis.Redis(connection_pool=pool)
|
||||
return client
|
||||
|
||||
|
||||
def init_app(app: DifyApp):
|
||||
"""Initialize Redis client and attach it to the app."""
|
||||
global redis_client
|
||||
|
||||
# Determine Redis mode and create appropriate client
|
||||
if dify_config.REDIS_USE_SENTINEL:
|
||||
redis_params = _get_base_redis_params()
|
||||
client = _create_sentinel_client(redis_params)
|
||||
elif dify_config.REDIS_USE_CLUSTERS:
|
||||
client = _create_cluster_client()
|
||||
else:
|
||||
redis_params = _get_base_redis_params()
|
||||
client = _create_standalone_client(redis_params)
|
||||
|
||||
# Initialize the wrapper and attach to app
|
||||
redis_client.initialize(client)
|
||||
app.extensions["redis"] = redis_client
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user