From 785e83b221aadd1d4be84d80686d96b82d50e3c1 Mon Sep 17 00:00:00 2001 From: xxlaila Date: Tue, 11 Mar 2025 14:48:25 +0800 Subject: [PATCH 1/4] Add enterprise level LDAP authentication --- api/.env.example | 16 +++- api/app_factory.py | 2 + api/configs/app_config.py | 3 + api/configs/ldap/__init__.py | 74 +++++++++++++++++ api/extensions/ext_ldap.py | 121 ++++++++++++++++++++++++++++ api/extensions/ext_login.py | 4 + api/poetry.lock | 25 ++++++ api/services/account_service.py | 138 +++++++++++++++++++++++++++++++- 8 files changed, 379 insertions(+), 4 deletions(-) create mode 100644 api/configs/ldap/__init__.py create mode 100644 api/extensions/ext_ldap.py diff --git a/api/.env.example b/api/.env.example index 880453161e..e898974803 100644 --- a/api/.env.example +++ b/api/.env.example @@ -444,4 +444,18 @@ CREATE_TIDB_SERVICE_JOB_ENABLED=false # Maximum number of submitted thread count in a ThreadPool for parallel node execution MAX_SUBMIT_COUNT=100 # Lockout duration in seconds -LOGIN_LOCKOUT_DURATION=86400 \ No newline at end of file +LOGIN_LOCKOUT_DURATION=86400 + +# Enable ldap authentication login +LDAP_ENABLED=false +AUTH_LDAP_SERVER_URI=ldap://127.0.0.1:389 +AUTH_LDAP_BIND_DN=CN=ArcherMaster,CN=Users,DC=dify,DC=com +AUTH_LDAP_BIND_PASSWORD=dify@123456 +AUTH_LDAP_SEARCH_BASE_DN=OU=dify技术有限公司,DC=dify,DC=com +AUTH_LDAP_USER_FILTER=(mail=%(user)s) +AUTH_LDAP_USER_ATTR_MAP={"first_name": "uid", "last_name": "cn", "email": "mail"} +CACHE_LDAP_USER_KEY=LDAP_USER_INFO_DATA +CACHE_LDAP_USER_EX=86400 +LDAP_DEFAULT_ROLE=normal +LDAP_CONN_TIMEOUT=10 +LDAP_POOL_SIZE=10 \ No newline at end of file diff --git a/api/app_factory.py b/api/app_factory.py index 52ae05583a..f9815eabd4 100644 --- a/api/app_factory.py +++ b/api/app_factory.py @@ -38,6 +38,7 @@ def create_app() -> DifyApp: def initialize_extensions(app: DifyApp): from extensions import ( + ext_ldap, ext_app_metrics, ext_blueprints, ext_celery, @@ -61,6 +62,7 @@ def initialize_extensions(app: DifyApp): ) extensions = [ + ext_ldap, ext_timezone, ext_logging, ext_warnings, diff --git a/api/configs/app_config.py b/api/configs/app_config.py index ac1ce9db10..52eb441d9a 100644 --- a/api/configs/app_config.py +++ b/api/configs/app_config.py @@ -5,6 +5,7 @@ from pydantic.fields import FieldInfo from pydantic_settings import BaseSettings, PydanticBaseSettingsSource, SettingsConfigDict from .deploy import DeploymentConfig +from .ldap import AuthenticationConfig from .enterprise import EnterpriseFeatureConfig from .extra import ExtraServiceConfig from .feature import FeatureConfig @@ -49,6 +50,8 @@ class RemoteSettingsSourceFactory(PydanticBaseSettingsSource): class DifyConfig( + # Authentication configs + AuthenticationConfig, # Packaging info PackagingInfo, # Deployment configs diff --git a/api/configs/ldap/__init__.py b/api/configs/ldap/__init__.py new file mode 100644 index 0000000000..48b4155560 --- /dev/null +++ b/api/configs/ldap/__init__.py @@ -0,0 +1,74 @@ +# -*- coding: utf-8 -*- +""" +@File : __init__.py.py +@Time : 2025/3/6 {TIME} +@Author : xxlaila +@Software: dify +""" +from pydantic import Field +from pydantic_settings import BaseSettings + +class AuthenticationConfig(BaseSettings): + """LDAP authentication related configuration""" + # LDAP Authentication + LDAP_ENABLED: bool = Field( + False, + description="Whether to enable LDAP authentication", + validation_alias="LDAP_ENABLED" + ) + AUTH_LDAP_SERVER_URI: str = Field( + "ldap://localhost", + description="LDAP server address", + validation_alias="AUTH_LDAP_SERVER_URI" + ) + AUTH_LDAP_BIND_DN: str = Field( + "cn=admin,dc=example,dc=com", + description="LDAP Bind DN", + validation_alias="AUTH_LDAP_BIND_DN" + ) + AUTH_LDAP_BIND_PASSWORD: str = Field( + "", + description="LDAP bind password", + validation_alias="AUTH_LDAP_BIND_PASSWORD" + ) + AUTH_LDAP_SEARCH_BASE_DN: str = Field( + default="OU=Limited Company,DC=example,DC=com", + description="LDAP search base DN", + validation_alias="AUTH_LDAP_SEARCH_BASE_DN" + ) + AUTH_LDAP_USER_FILTER: str = Field( + "(objectClass=person)", + description="LDAP User Filter", + validation_alias="AUTH_LDAP_USER_FILTER" + ) + AUTH_LDAP_USER_ATTR_MAP: dict = Field( + {"first_name": "uid", "last_name": "cn", "email": "mail"}, + description="Mapping of LDAP attributes to user models", + validation_alias="AUTH_LDAP_USER_ATTR_MAP" + ) + CACHE_LDAP_USER_KEY: str = Field( + # Cache LDAP user information + "ldap_user", + description="Key for caching LDAP user information", + validation_alias="CACHE_LDAP_USER_KEY" + ) + CACHE_LDAP_USER_EX: int = Field( + 86400, + description="Expiration time of cached LDAP user information", + validation_alias="CACHE_LDAP_USER_EX" + ) + LDAP_DEFAULT_ROLE: str = Field( + default="normal", + description="Default Roles", + validation_alias="LDAP_DEFAULT_ROLE" + ) + LDAP_CONN_TIMEOUT: int = Field( + default=10, + description="LDAP connection timeout", + validation_alias="LDAP_CONN_TIMEOUT" + ) + LDAP_POOL_SIZE: int = Field( + default=10, + description="LDAP connection pool size", + validation_alias="LDAP_POOL_SIZE" + ) \ No newline at end of file diff --git a/api/extensions/ext_ldap.py b/api/extensions/ext_ldap.py new file mode 100644 index 0000000000..29fc9fb1f9 --- /dev/null +++ b/api/extensions/ext_ldap.py @@ -0,0 +1,121 @@ +# -*- coding: utf-8 -*- +""" +@File : ext_ldap.py +@Time : 2025/3/5 {TIME} +@Author : xxlaila +@Software: dify +""" +from flask_ldap3_login import LDAP3LoginManager +from configs import dify_config +from dify_app import DifyApp +import json +import logging +from queue import Queue +from ldap3 import Server, Connection, ALL + +def is_enabled(): + return getattr(dify_config, 'LDAP_ENABLED', False) + +def create_ldap_pool(): + pool = Queue(maxsize=dify_config.LDAP_POOL_SIZE) + for i in range(dify_config.LDAP_POOL_SIZE): + server = Server(dify_config.AUTH_LDAP_SERVER_URI, get_info=ALL) + conn = Connection( + server, + user=dify_config.AUTH_LDAP_BIND_DN, + password=dify_config.AUTH_LDAP_BIND_PASSWORD, + receive_timeout=dify_config.LDAP_CONN_TIMEOUT + ) + if conn.bind(): + logging.info(f"LDAP connection {i} bound successfully") + pool.put(conn) + else: + logging.error(f"LDAP connection {i} failed to bind") + return pool + +def get_ldap_connection(): + """Get a connection from the connection pool""" + conn = LDAP_POOL.get() + + # Add connection validity check + if not conn.bound or not conn.server: + try: + if not conn.bind(): + logging.error("LDAP connection is unbound, rebinding...") + conn.unbind() + server = Server(dify_config.AUTH_LDAP_SERVER_URI, get_info=ALL) + new_conn = Connection( + server, + user=dify_config.AUTH_LDAP_BIND_DN, + password=dify_config.AUTH_LDAP_BIND_PASSWORD, + receive_timeout=dify_config.LDAP_CONN_TIMEOUT + ) + if new_conn.bind(): + return new_conn + raise Exception("LDAP connection reconstruction failed") + except Exception as e: + logging.error(f"LDAP connection recovery failed: {str(e)}") + raise + return conn + +def release_ldap_connection(conn): + """Return the connection to the connection pool""" + try: + # Reset connection status + if conn.bound: + conn.unbind() + conn.open() # Reopen connection without binding + LDAP_POOL.put(conn) + except Exception as e: + logging.error(f"Failed to recycle LDAP connection: {str(e)}") + conn.unbind() + +def init_app(app: DifyApp): + """Initialize LDAP authentication integration""" + if not is_enabled(): + app.ldap_manager = None # Explicitly set the manager to None + logging.info("LDAP authentication is disabled") + return + + # Global LDAP connection pool + global LDAP_POOL + LDAP_POOL = create_ldap_pool() + + # Parsing User Attribute Mapping + if isinstance(dify_config.AUTH_LDAP_USER_ATTR_MAP, str): + ldap_user_attr_map = json.loads(dify_config.AUTH_LDAP_USER_ATTR_MAP) + else: + ldap_user_attr_map = dify_config.AUTH_LDAP_USER_ATTR_MAP + + # Setting up LDAP configuration + app.config.update({ + "LDAP_HOST": dify_config.AUTH_LDAP_SERVER_URI, + "LDAP_BASE_DN": dify_config.AUTH_LDAP_SEARCH_BASE_DN, + "LDAP_BIND_DN": dify_config.AUTH_LDAP_BIND_DN, + "LDAP_BIND_PASSWORD": dify_config.AUTH_LDAP_BIND_PASSWORD, + "LDAP_USER_FILTER": dify_config.AUTH_LDAP_USER_FILTER, + "LDAP_USER_RDN_ATTR": "uid", + "LDAP_USER_LOGIN_ATTR": "uid", + "LDAP_USER_SEARCH_SCOPE": "SUBTREE", + "LDAP_USER_MAPPING": ldap_user_attr_map, + "LDAP_DEFAULT_ROLE": dify_config.LDAP_DEFAULT_ROLE, + }) + + + # Initializing the LDAP Manager + ldap_manager = LDAP3LoginManager() + ldap_manager.init_app(app) + + # Mount the LDAP manager into the app + app.ldap_manager = ldap_manager + # Confirm that the mount was successful + logging.info(f"LDAP manager mounted: {hasattr(app, 'ldap_manager')}") + + + # Configuring Logging + if app.debug: + app.logger.info("LDAP configuration loaded：") + app.logger.info(f"Server: {app.config['LDAP_HOST']}") + app.logger.info(f"Base DN: {app.config['LDAP_BASE_DN']}") + + return ldap_manager diff --git a/api/extensions/ext_login.py b/api/extensions/ext_login.py index 10fb89eb73..6a29ae84dd 100644 --- a/api/extensions/ext_login.py +++ b/api/extensions/ext_login.py @@ -12,6 +12,10 @@ from services.account_service import AccountService login_manager = flask_login.LoginManager() +@login_manager.user_loader +def load_user(user_id): + """Support loading LDAP users from session""" + return AccountService.load_user(user_id) # Flask-Login configuration @login_manager.request_loader diff --git a/api/poetry.lock b/api/poetry.lock index ffa5810982..cb8b8664ff 100644 --- a/api/poetry.lock +++ b/api/poetry.lock @@ -2210,6 +2210,7 @@ files = [ brotli = {version = "*", markers = "platform_python_implementation != \"PyPy\""} brotlicffi = {version = "*", markers = "platform_python_implementation == \"PyPy\""} flask = "*" +ldap3 = ">=2.0.7" zstandard = [ {version = "*", markers = "platform_python_implementation != \"PyPy\""}, {version = "*", extras = ["cffi"], markers = "platform_python_implementation == \"PyPy\""}, @@ -2231,6 +2232,18 @@ files = [ [package.dependencies] Flask = ">=0.9" +[[package]] +name = "flask-ldap3-login" +version = "1.0.2" +description = "LDAP Support for Flask" +optional = false +python-versions = ">=3.5" +groups = ["main"] +files = [ + {file = "flask_ldap3_login-1.0.2-py3-none-any.whl", hash = "sha256:72368ea58b5faac4935d1f3e5ec1fcfdce51cc16adbec41a74c834247422ac61"}, + {file = "flask_ldap3_login-1.0.2.tar.gz", hash = "sha256:969d2e4263aa0908383174d9700e03198a5eb3db483b26a18f9ea634032cbdd4"}, +] + [[package]] name = "flask-login" version = "0.6.3" @@ -4053,6 +4066,18 @@ requests-toolbelt = ">=1.0.0,<2.0.0" [package.extras] langsmith-pyo3 = ["langsmith-pyo3 (>=0.1.0rc2,<0.2.0)"] +[[package]] +name = "ldap3" +version = "2.9.1" +description = "A strictly RFC 4510 conforming LDAP V3 pure Python client library" +optional = false +python-versions = "*" +groups = ["main"] +files = [ + {file = "ldap3-2.9.1-py2.py3-none-any.whl", hash = "sha256:5869596fc4948797020d3f03b7939da938778a0f9e2009f7a072ccf92b8e8d70"}, + {file = "ldap3-2.9.1.tar.gz", hash = "sha256:f3e7fc4718e3f09dda568b57100095e0ce58633bcabbed8667ce3f8fbaa4229f"}, +] + [[package]] name = "levenshtein" version = "0.27.1" diff --git a/api/services/account_service.py b/api/services/account_service.py index 2923750298..5e747cd6ea 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -57,7 +57,10 @@ from tasks.mail_account_deletion_task import send_account_deletion_verification_ from tasks.mail_email_code_login import send_email_code_login_mail_task from tasks.mail_invite_member_task import send_invite_member_mail_task from tasks.mail_reset_password_task import send_reset_password_mail_task +from extensions.ext_ldap import get_ldap_connection, release_ldap_connection +from ldap3 import Server, Connection, ALL +from flask import current_app as app class TokenPair(BaseModel): access_token: str @@ -146,10 +149,89 @@ class AccountService: @staticmethod def authenticate(email: str, password: str, invite_token: Optional[str] = None) -> Account: """authenticate account with email and password""" + account = None + is_ldap_success = False - account = db.session.query(Account).filter_by(email=email).first() - if not account: - raise AccountNotFoundError() + # Initialize LDAP connection parameters + def _get_ldap_connection(): + return get_ldap_connection() + + def _handle_ldap_user(conn: Connection, email: str, password: str) -> Account: + """Handle LDAP user authentication""" + conn.search( + search_base=dify_config.AUTH_LDAP_SEARCH_BASE_DN, + search_filter=f"(mail={email})", + attributes=list(dify_config.AUTH_LDAP_USER_ATTR_MAP.values()) + ) + if not conn.entries: + logging.warning(f"No LDAP entry found for: {email}") + raise ValueError("LDAP entry not found") + + entry = conn.entries[0] + logging.info(f"LDAP entry found for: {entry}") + user_dn = entry.entry_dn # Obtain the complete DN of the user + logging.info(f"LDAP user DN: {user_dn}") + # Add password verification steps + user_conn = Connection( + Server(dify_config.AUTH_LDAP_SERVER_URI), + user=user_dn, + password=password + ) + if not user_conn.bind(): + logging.error(f"LDAP password verification failed: {email}") + raise AccountPasswordError("Password error") + + display_name = entry.cn.value if entry.cn else email.split('@')[0] + account = db.session.query(Account).filter_by(email=email).with_for_update().first() + if not account: + logging.info(f"Creating new LDAP account for: {email}") + return AccountService.create_ldap_user( + email=email, + name=display_name, + ldap_attrs=entry + ) + + if account.name != display_name: + logging.info(f"Updating LDAP account name for: {email}") + account.name = display_name + db.session.commit() + + return account + + # Main certification process + try: + # Try LDAP authentication first + if hasattr(app, 'ldap_manager') and app.ldap_manager: + ldap_user = app.ldap_manager.authenticate(email, password) + logging.info(f"LDAP authentication successful for user: {email}") + if ldap_user and ldap_user.status: + conn = _get_ldap_connection() + try: + if not conn.bind(): + logging.error("Failed to bind to LDAP server.") + raise ValueError("Failed to bind to LDAP server.") + account = _handle_ldap_user(conn, email, password) # 传入password参数 + logging.info(f"Returning account {email} after successful LDAP authentication.") + is_ldap_success = True + return account + finally: + release_ldap_connection(conn) + except Exception as e: + logging.error(f"LDAP authentication error: {str(e)}") + + # Perform local authentication only if LDAP authentication explicitly fails + if not is_ldap_success: + logging.info(f"Falling back to local authentication for: {email}") + account = db.session.query(Account).filter_by(email=email).first() + if not account: + logging.error(f"Account not found for: {email}") + + raise AccountNotFoundError() + + # Key modification: If it is an LDAP user, they must pass LDAP authentication + if account.password is None: + logging.error(f"LDAP user {email} attempted local authentication") + raise AccountPasswordError("LDAP users must authenticate via LDAP") if account.status == AccountStatus.BANNED.value: raise AccountLoginError("Account is banned.") @@ -174,6 +256,56 @@ class AccountService: return cast(Account, account) + @staticmethod + def create_ldap_user(email: str, name: str, ldap_attrs: dict) -> Account: + """ + Create LDAP users and assign default permissions + Note: LDAP users do not require a local password, so password and password_stalt are set to None. + """ + try: + # First check if the account already exists + account = db.session.query(Account).filter_by(email=email).first() + if account: + logging.info(f"Account {email} already exists, skipping creation.") + else: + account = Account( + email=email, + name=name, + status=AccountStatus.ACTIVE.value, + password=None, # LDAP user has no local password + password_salt=None, + interface_language=languages[0] + ) + db.session.add(account) + db.session.commit() # Submit the transaction to ensure that the account is available + logging.info(f"Local account created successfully for: {email}") + + # Get the default tenant (remove the specific name and only take the first one) + default_tenant = db.session.query(Tenant).first() + if not default_tenant: + default_tenant = TenantService.create_tenant(name="Default Workspace") + db.session.commit() + logging.info(f"Created default tenant: {default_tenant.id}") + + # Check if the user has joined the tenant + existing_member = db.session.query(TenantAccountJoin).filter_by( + tenant_id=default_tenant.id, account_id=account.id + ).first() + + if not existing_member: + TenantService.create_tenant_member(tenant=default_tenant, account=account, + role="normal") + logging.info(f"User {email} added to tenant: {default_tenant.id}") + else: + logging.info(f"User {email} is already a member of tenant {default_tenant.id}") + + return account # Make sure to return a valid account + + except Exception as e: + db.session.rollback() # Transaction rollback to prevent database pollution + logging.error(f"Failed to create LDAP user or add to tenant: {str(e)}") + raise AccountRegisterError("Failed to create LDAP user.") + @staticmethod def update_account_password(account, password, new_password): """update account password""" From d1ed6ddd9dfa1295cad7bd04d0912b31b6ee21cd Mon Sep 17 00:00:00 2001 From: xxlaila Date: Tue, 11 Mar 2025 14:51:22 +0800 Subject: [PATCH 2/4] Add enterprise level LDAP authentication --- api/app_factory.py | 2 +- api/configs/app_config.py | 2 +- api/configs/ldap/__init__.py | 2 +- api/extensions/ext_ldap.py | 16 +++++++++------- api/services/account_service.py | 10 +++++----- 5 files changed, 17 insertions(+), 15 deletions(-) diff --git a/api/app_factory.py b/api/app_factory.py index f9815eabd4..8b081446fe 100644 --- a/api/app_factory.py +++ b/api/app_factory.py @@ -38,7 +38,6 @@ def create_app() -> DifyApp: def initialize_extensions(app: DifyApp): from extensions import ( - ext_ldap, ext_app_metrics, ext_blueprints, ext_celery, @@ -48,6 +47,7 @@ def initialize_extensions(app: DifyApp): ext_database, ext_hosting_provider, ext_import_modules, + ext_ldap, ext_logging, ext_login, ext_mail, diff --git a/api/configs/app_config.py b/api/configs/app_config.py index 52eb441d9a..028053086b 100644 --- a/api/configs/app_config.py +++ b/api/configs/app_config.py @@ -5,10 +5,10 @@ from pydantic.fields import FieldInfo from pydantic_settings import BaseSettings, PydanticBaseSettingsSource, SettingsConfigDict from .deploy import DeploymentConfig -from .ldap import AuthenticationConfig from .enterprise import EnterpriseFeatureConfig from .extra import ExtraServiceConfig from .feature import FeatureConfig +from .ldap import AuthenticationConfig from .middleware import MiddlewareConfig from .packaging import PackagingInfo from .remote_settings_sources import RemoteSettingsSource, RemoteSettingsSourceConfig, RemoteSettingsSourceName diff --git a/api/configs/ldap/__init__.py b/api/configs/ldap/__init__.py index 48b4155560..e74fdf5f95 100644 --- a/api/configs/ldap/__init__.py +++ b/api/configs/ldap/__init__.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- """ @File : __init__.py.py @Time : 2025/3/6 {TIME} @@ -8,6 +7,7 @@ from pydantic import Field from pydantic_settings import BaseSettings + class AuthenticationConfig(BaseSettings): """LDAP authentication related configuration""" # LDAP Authentication diff --git a/api/extensions/ext_ldap.py b/api/extensions/ext_ldap.py index 29fc9fb1f9..ef380de57c 100644 --- a/api/extensions/ext_ldap.py +++ b/api/extensions/ext_ldap.py @@ -1,17 +1,19 @@ -# -*- coding: utf-8 -*- """ @File : ext_ldap.py @Time : 2025/3/5 {TIME} @Author : xxlaila @Software: dify """ -from flask_ldap3_login import LDAP3LoginManager -from configs import dify_config -from dify_app import DifyApp import json import logging from queue import Queue -from ldap3 import Server, Connection, ALL + +from flask_ldap3_login import LDAP3LoginManager +from ldap3 import ALL, Connection, Server + +from configs import dify_config +from dify_app import DifyApp + def is_enabled(): return getattr(dify_config, 'LDAP_ENABLED', False) @@ -54,7 +56,7 @@ def get_ldap_connection(): return new_conn raise Exception("LDAP connection reconstruction failed") except Exception as e: - logging.error(f"LDAP connection recovery failed: {str(e)}") + logging.exception(f"LDAP connection recovery failed: {e}") raise return conn @@ -67,7 +69,7 @@ def release_ldap_connection(conn): conn.open() # Reopen connection without binding LDAP_POOL.put(conn) except Exception as e: - logging.error(f"Failed to recycle LDAP connection: {str(e)}") + logging.exception(f"Failed to recycle LDAP connection: {str(e)}") conn.unbind() def init_app(app: DifyApp): diff --git a/api/services/account_service.py b/api/services/account_service.py index 5e747cd6ea..a548335de6 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -8,6 +8,8 @@ from datetime import UTC, datetime, timedelta from hashlib import sha256 from typing import Any, Optional, cast +from flask import current_app as app +from ldap3 import Connection, Server from pydantic import BaseModel from sqlalchemy import func from sqlalchemy.orm import Session @@ -17,6 +19,7 @@ from configs import dify_config from constants.languages import language_timezone_mapping, languages from events.tenant_event import tenant_was_created from extensions.ext_database import db +from extensions.ext_ldap import get_ldap_connection, release_ldap_connection from extensions.ext_redis import redis_client from libs.helper import RateLimiter, TokenManager from libs.passport import PassportService @@ -57,10 +60,7 @@ from tasks.mail_account_deletion_task import send_account_deletion_verification_ from tasks.mail_email_code_login import send_email_code_login_mail_task from tasks.mail_invite_member_task import send_invite_member_mail_task from tasks.mail_reset_password_task import send_reset_password_mail_task -from extensions.ext_ldap import get_ldap_connection, release_ldap_connection -from ldap3 import Server, Connection, ALL -from flask import current_app as app class TokenPair(BaseModel): access_token: str @@ -217,7 +217,7 @@ class AccountService: finally: release_ldap_connection(conn) except Exception as e: - logging.error(f"LDAP authentication error: {str(e)}") + logging.exception(f"LDAP authentication error: {str(e)}") # Perform local authentication only if LDAP authentication explicitly fails if not is_ldap_success: @@ -303,7 +303,7 @@ class AccountService: except Exception as e: db.session.rollback() # Transaction rollback to prevent database pollution - logging.error(f"Failed to create LDAP user or add to tenant: {str(e)}") + logging.exception(f"Failed to create LDAP user or add to tenant: {str(e)}") raise AccountRegisterError("Failed to create LDAP user.") @staticmethod From c13f3da5022ace87fd21c3b6b958d5ad6191342d Mon Sep 17 00:00:00 2001 From: xxlaila Date: Tue, 11 Mar 2025 14:54:53 +0800 Subject: [PATCH 3/4] add pyproject ldap --- api/pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/api/pyproject.toml b/api/pyproject.toml index 493bfa240b..965152bd8c 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -83,6 +83,8 @@ yarl = "~1.18.3" # Related transparent dependencies with pinned version # required by main implementations ############################################################ +flask-ldap3-login = "^1.0.2" +ldap3 = "^2.9.1" [tool.poetry.group.indirect.dependencies] kaleido = "0.2.1" rank-bm25 = "~0.2.2" From a1b074c684c16b913f70550aec94a67a0af452c0 Mon Sep 17 00:00:00 2001 From: xxlaila Date: Tue, 11 Mar 2025 18:05:32 +0800 Subject: [PATCH 4/4] Delete LDAP connection pool issue, use direct --- api/extensions/ext_ldap.py | 60 ---------------------- api/services/account_service.py | 91 ++++++++++++++++----------------- 2 files changed, 45 insertions(+), 106 deletions(-) diff --git a/api/extensions/ext_ldap.py b/api/extensions/ext_ldap.py index ef380de57c..34c002871d 100644 --- a/api/extensions/ext_ldap.py +++ b/api/extensions/ext_ldap.py @@ -6,10 +6,8 @@ """ import json import logging -from queue import Queue from flask_ldap3_login import LDAP3LoginManager -from ldap3 import ALL, Connection, Server from configs import dify_config from dify_app import DifyApp @@ -18,60 +16,6 @@ from dify_app import DifyApp def is_enabled(): return getattr(dify_config, 'LDAP_ENABLED', False) -def create_ldap_pool(): - pool = Queue(maxsize=dify_config.LDAP_POOL_SIZE) - for i in range(dify_config.LDAP_POOL_SIZE): - server = Server(dify_config.AUTH_LDAP_SERVER_URI, get_info=ALL) - conn = Connection( - server, - user=dify_config.AUTH_LDAP_BIND_DN, - password=dify_config.AUTH_LDAP_BIND_PASSWORD, - receive_timeout=dify_config.LDAP_CONN_TIMEOUT - ) - if conn.bind(): - logging.info(f"LDAP connection {i} bound successfully") - pool.put(conn) - else: - logging.error(f"LDAP connection {i} failed to bind") - return pool - -def get_ldap_connection(): - """Get a connection from the connection pool""" - conn = LDAP_POOL.get() - - # Add connection validity check - if not conn.bound or not conn.server: - try: - if not conn.bind(): - logging.error("LDAP connection is unbound, rebinding...") - conn.unbind() - server = Server(dify_config.AUTH_LDAP_SERVER_URI, get_info=ALL) - new_conn = Connection( - server, - user=dify_config.AUTH_LDAP_BIND_DN, - password=dify_config.AUTH_LDAP_BIND_PASSWORD, - receive_timeout=dify_config.LDAP_CONN_TIMEOUT - ) - if new_conn.bind(): - return new_conn - raise Exception("LDAP connection reconstruction failed") - except Exception as e: - logging.exception(f"LDAP connection recovery failed: {e}") - raise - return conn - -def release_ldap_connection(conn): - """Return the connection to the connection pool""" - try: - # Reset connection status - if conn.bound: - conn.unbind() - conn.open() # Reopen connection without binding - LDAP_POOL.put(conn) - except Exception as e: - logging.exception(f"Failed to recycle LDAP connection: {str(e)}") - conn.unbind() - def init_app(app: DifyApp): """Initialize LDAP authentication integration""" if not is_enabled(): @@ -79,10 +23,6 @@ def init_app(app: DifyApp): logging.info("LDAP authentication is disabled") return - # Global LDAP connection pool - global LDAP_POOL - LDAP_POOL = create_ldap_pool() - # Parsing User Attribute Mapping if isinstance(dify_config.AUTH_LDAP_USER_ATTR_MAP, str): ldap_user_attr_map = json.loads(dify_config.AUTH_LDAP_USER_ATTR_MAP) diff --git a/api/services/account_service.py b/api/services/account_service.py index a548335de6..5f725d6ab4 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -8,8 +8,7 @@ from datetime import UTC, datetime, timedelta from hashlib import sha256 from typing import Any, Optional, cast -from flask import current_app as app -from ldap3 import Connection, Server +from ldap3 import ALL, Connection, Server from pydantic import BaseModel from sqlalchemy import func from sqlalchemy.orm import Session @@ -19,7 +18,7 @@ from configs import dify_config from constants.languages import language_timezone_mapping, languages from events.tenant_event import tenant_was_created from extensions.ext_database import db -from extensions.ext_ldap import get_ldap_connection, release_ldap_connection +from extensions.ext_ldap import is_enabled from extensions.ext_redis import redis_client from libs.helper import RateLimiter, TokenManager from libs.passport import PassportService @@ -150,14 +149,17 @@ class AccountService: def authenticate(email: str, password: str, invite_token: Optional[str] = None) -> Account: """authenticate account with email and password""" account = None - is_ldap_success = False - # Initialize LDAP connection parameters - def _get_ldap_connection(): - return get_ldap_connection() - - def _handle_ldap_user(conn: Connection, email: str, password: str) -> Account: + def _handle_ldap_user(email: str, password: str) -> Optional[Account]: """Handle LDAP user authentication""" + server = Server(dify_config.AUTH_LDAP_SERVER_URI, get_info=ALL) + conn = Connection( + server, + user=dify_config.AUTH_LDAP_BIND_DN, + password=dify_config.AUTH_LDAP_BIND_PASSWORD, + receive_timeout=dify_config.LDAP_CONN_TIMEOUT, + auto_bind=True + ) conn.search( search_base=dify_config.AUTH_LDAP_SEARCH_BASE_DN, search_filter=f"(mail={email})", @@ -165,15 +167,13 @@ class AccountService: ) if not conn.entries: logging.warning(f"No LDAP entry found for: {email}") - raise ValueError("LDAP entry not found") + return None entry = conn.entries[0] - logging.info(f"LDAP entry found for: {entry}") user_dn = entry.entry_dn # Obtain the complete DN of the user - logging.info(f"LDAP user DN: {user_dn}") # Add password verification steps user_conn = Connection( - Server(dify_config.AUTH_LDAP_SERVER_URI), + server, user=user_dn, password=password ) @@ -199,39 +199,38 @@ class AccountService: return account # Main certification process - try: - # Try LDAP authentication first - if hasattr(app, 'ldap_manager') and app.ldap_manager: - ldap_user = app.ldap_manager.authenticate(email, password) - logging.info(f"LDAP authentication successful for user: {email}") - if ldap_user and ldap_user.status: - conn = _get_ldap_connection() - try: - if not conn.bind(): - logging.error("Failed to bind to LDAP server.") - raise ValueError("Failed to bind to LDAP server.") - account = _handle_ldap_user(conn, email, password) # 传入password参数 - logging.info(f"Returning account {email} after successful LDAP authentication.") - is_ldap_success = True - return account - finally: - release_ldap_connection(conn) - except Exception as e: - logging.exception(f"LDAP authentication error: {str(e)}") + if is_enabled(): + try: + # Try LDAP authentication first + server = Server(dify_config.AUTH_LDAP_SERVER_URI, get_info=ALL) + conn = Connection( + server, + user=dify_config.AUTH_LDAP_BIND_DN, + password=dify_config.AUTH_LDAP_BIND_PASSWORD, + receive_timeout=dify_config.LDAP_CONN_TIMEOUT, + auto_bind=True + ) + if not conn.bind(): + logging.error("Failed to bind to LDAP server.") + raise ValueError("Failed to bind to LDAP server.") + account = _handle_ldap_user(email, password) + if account: + logging.info(f"Returning account {email} after successful LDAP authentication.") + + return account + except Exception as e: + logging.info(f"LDAP authentication error: {e}") # Perform local authentication only if LDAP authentication explicitly fails - if not is_ldap_success: - logging.info(f"Falling back to local authentication for: {email}") - account = db.session.query(Account).filter_by(email=email).first() - if not account: - logging.error(f"Account not found for: {email}") + account = db.session.query(Account).filter_by(email=email).first() + if not account: + logging.error(f"Account not found for: {email}") + raise AccountNotFoundError() - raise AccountNotFoundError() - - # Key modification: If it is an LDAP user, they must pass LDAP authentication - if account.password is None: - logging.error(f"LDAP user {email} attempted local authentication") - raise AccountPasswordError("LDAP users must authenticate via LDAP") + # LDAP users must authenticate via LDAP + if account.password is None and account.password_salt is None and not invite_token: + logging.error(f"LDAP user {email} attempted local authentication") + raise AccountPasswordError("LDAP users must authenticate via LDAP") if account.status == AccountStatus.BANNED.value: raise AccountLoginError("Account is banned.") @@ -289,8 +288,8 @@ class AccountService: # Check if the user has joined the tenant existing_member = db.session.query(TenantAccountJoin).filter_by( - tenant_id=default_tenant.id, account_id=account.id - ).first() + tenant_id=default_tenant.id, account_id=account.id + ).first() if not existing_member: TenantService.create_tenant_member(tenant=default_tenant, account=account, @@ -303,7 +302,7 @@ class AccountService: except Exception as e: db.session.rollback() # Transaction rollback to prevent database pollution - logging.exception(f"Failed to create LDAP user or add to tenant: {str(e)}") + logging.info(f"Failed to create LDAP user or add to tenant: {e}") raise AccountRegisterError("Failed to create LDAP user.") @staticmethod