From e3d72954aef8cffdb9a4ed7498320d54309c8e94 Mon Sep 17 00:00:00 2001 From: fit2cloud-chenyw Date: Tue, 23 Dec 2025 16:15:35 +0800 Subject: [PATCH] perf: Optimizing API Key Authentication --- backend/alembic/versions/056_api_key_ddl.py | 39 +++++++++++ backend/apps/api.py | 3 +- backend/apps/system/api/apikey.py | 60 +++++++++++++++++ backend/apps/system/crud/apikey_manage.py | 17 +++++ backend/apps/system/middleware/auth.py | 55 ++++++++++++++- backend/apps/system/models/system_model.py | 13 +++- backend/apps/system/schemas/auth.py | 1 + backend/apps/system/schemas/system_schema.py | 10 +++ frontend/src/components/layout/Apikey.vue | 70 ++++++++------------ 9 files changed, 223 insertions(+), 45 deletions(-) create mode 100644 backend/alembic/versions/056_api_key_ddl.py create mode 100644 backend/apps/system/api/apikey.py create mode 100644 backend/apps/system/crud/apikey_manage.py diff --git a/backend/alembic/versions/056_api_key_ddl.py b/backend/alembic/versions/056_api_key_ddl.py new file mode 100644 index 000000000..e2657927d --- /dev/null +++ b/backend/alembic/versions/056_api_key_ddl.py @@ -0,0 +1,39 @@ +"""056_api_key_ddl + +Revision ID: d9a5589fc00b +Revises: 3d4bd2d673dc +Create Date: 2025-12-23 13:41:26.705947 + +""" +from alembic import op +import sqlalchemy as sa +import sqlmodel.sql.sqltypes +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = 'd9a5589fc00b' +down_revision = '3d4bd2d673dc' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('sys_apikey', + sa.Column('access_key', sqlmodel.sql.sqltypes.AutoString(length=255), nullable=False), + sa.Column('secret_key', sqlmodel.sql.sqltypes.AutoString(length=255), nullable=False), + sa.Column('create_time', sa.BigInteger(), nullable=False), + sa.Column('uid', sa.BigInteger(), nullable=False), + sa.Column('status', sa.Boolean(), nullable=False), + sa.Column('id', sa.BigInteger(), nullable=False), + sa.PrimaryKeyConstraint('id') + ) + op.create_index(op.f('ix_sys_apikey_id'), 'sys_apikey', ['id'], unique=False) + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index(op.f('ix_sys_apikey_id'), table_name='sys_apikey') + op.drop_table('sys_apikey') + # ### end Alembic commands ### diff --git a/backend/apps/api.py b/backend/apps/api.py index f53f7fa48..06dd11b59 100644 --- a/backend/apps/api.py +++ b/backend/apps/api.py @@ -5,7 +5,7 @@ from apps.data_training.api import data_training from apps.datasource.api import datasource, table_relation, recommended_problem from apps.mcp import mcp -from apps.system.api import login, user, aimodel, workspace, assistant, parameter +from apps.system.api import login, user, aimodel, workspace, assistant, parameter, apikey from apps.terminology.api import terminology from apps.settings.api import base @@ -24,5 +24,6 @@ api_router.include_router(mcp.router) api_router.include_router(table_relation.router) api_router.include_router(parameter.router) +api_router.include_router(apikey.router) api_router.include_router(recommended_problem.router) diff --git a/backend/apps/system/api/apikey.py b/backend/apps/system/api/apikey.py new file mode 100644 index 000000000..cfed5ae74 --- /dev/null +++ b/backend/apps/system/api/apikey.py @@ -0,0 +1,60 @@ + +from fastapi import APIRouter +from sqlmodel import func, select +from apps.system.crud.apikey_manage import clear_api_key_cache +from apps.system.models.system_model import ApiKeyModel +from apps.system.schemas.system_schema import ApikeyGridItem, ApikeyStatus +from common.core.deps import CurrentUser, SessionDep +from common.utils.time import get_timestamp +import secrets + +router = APIRouter(tags=["system_apikey"], prefix="/system/apikey") + +@router.get("") +async def grid(session: SessionDep, current_user: CurrentUser) -> list[ApikeyGridItem]: + query = select(ApiKeyModel).where(ApiKeyModel.uid == current_user.id).order_by(ApiKeyModel.create_time.desc()) + return session.exec(query).all() + +@router.post("") +async def create(session: SessionDep, current_user: CurrentUser): + count = session.exec(select(func.count()).select_from(ApiKeyModel)).one() + if count >= 5: + raise ValueError("Maximum of 5 API keys allowed") + access_key = secrets.token_urlsafe(16) + secret_key = secrets.token_urlsafe(32) + api_key = ApiKeyModel( + access_key=access_key, + secret_key=secret_key, + create_time=get_timestamp(), + uid=current_user.id, + status=True + ) + session.add(api_key) + session.commit() + return api_key.id + +@router.put("/status") +async def status(session: SessionDep, current_user: CurrentUser, dto: ApikeyStatus): + api_key = session.get(ApiKeyModel, dto.id) + if not api_key: + raise ValueError("API Key not found") + if api_key.uid != current_user.id: + raise PermissionError("No permission to modify this API Key") + if dto.status == api_key.status: + return + api_key.status = dto.status + await clear_api_key_cache(api_key.access_key) + session.add(api_key) + session.commit() + +@router.delete("/{id}") +async def delete(session: SessionDep, current_user: CurrentUser, id: int): + api_key = session.get(ApiKeyModel, id) + if not api_key: + raise ValueError("API Key not found") + if api_key.uid != current_user.id: + raise PermissionError("No permission to delete this API Key") + await clear_api_key_cache(api_key.access_key) + session.delete(api_key) + session.commit() + \ No newline at end of file diff --git a/backend/apps/system/crud/apikey_manage.py b/backend/apps/system/crud/apikey_manage.py new file mode 100644 index 000000000..ce6b43fcc --- /dev/null +++ b/backend/apps/system/crud/apikey_manage.py @@ -0,0 +1,17 @@ + +from sqlmodel import select + +from apps.system.models.system_model import ApiKeyModel +from apps.system.schemas.auth import CacheName, CacheNamespace +from common.core.deps import SessionDep +from common.core.sqlbot_cache import cache, clear_cache +from common.utils.utils import SQLBotLogUtil + +@cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.ASK_INFO, keyExpression="access_key") +async def get_api_key(session: SessionDep, access_key: str) -> ApiKeyModel | None: + query = select(ApiKeyModel).where(ApiKeyModel.access_key == access_key) + return session.exec(query).first() + +@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.ASK_INFO, keyExpression="access_key") +async def clear_api_key_cache(access_key: str): + SQLBotLogUtil.info(f"Api key cache for [{access_key}] has been cleaned") \ No newline at end of file diff --git a/backend/apps/system/middleware/auth.py b/backend/apps/system/middleware/auth.py index 41167624c..bdf887118 100644 --- a/backend/apps/system/middleware/auth.py +++ b/backend/apps/system/middleware/auth.py @@ -7,7 +7,8 @@ import jwt from sqlmodel import Session from starlette.middleware.base import BaseHTTPMiddleware -from apps.system.models.system_model import AssistantModel +from apps.system.crud.apikey_manage import get_api_key +from apps.system.models.system_model import ApiKeyModel, AssistantModel from common.core.db import engine from apps.system.crud.assistant import get_assistant_info, get_assistant_user from apps.system.crud.user import get_user_by_account, get_user_info @@ -33,7 +34,15 @@ async def dispatch(self, request, call_next): return await call_next(request) assistantTokenKey = settings.ASSISTANT_TOKEN_KEY assistantToken = request.headers.get(assistantTokenKey) + askToken = request.headers.get("X-SQLBOT-ASK-TOKEN") trans = await get_i18n(request) + if askToken: + validate_pass, data = await self.validateAskToken(askToken, trans) + if validate_pass: + request.state.current_user = data + return await call_next(request) + message = trans('i18n_permission.authenticate_invalid', msg = data) + return JSONResponse(message, status_code=401, headers={"Access-Control-Allow-Origin": "*"}) #if assistantToken and assistantToken.lower().startswith("assistant "): if assistantToken: validator: tuple[any] = await self.validateAssistant(assistantToken, trans) @@ -62,6 +71,50 @@ async def dispatch(self, request, call_next): def is_options(self, request: Request): return request.method == "OPTIONS" + async def validateAskToken(self, askToken: Optional[str], trans: I18n): + if not askToken: + return False, f"Miss Token[X-SQLBOT-ASK-TOKEN]!" + schema, param = get_authorization_scheme_param(askToken) + if schema.lower() != "sk": + return False, f"Token schema error!" + try: + payload = jwt.decode( + param, options={"verify_signature": False, "verify_exp": False}, algorithms=[security.ALGORITHM] + ) + access_key = payload.get('access_key', None) + + if not access_key: + return False, f"Miss access_key payload error!" + with Session(engine) as session: + api_key_model = await get_api_key(session, access_key) + api_key_model = ApiKeyModel.model_validate(api_key_model) if api_key_model else None + if not api_key_model: + return False, f"Invalid access_key!" + if not api_key_model.status: + return False, f"Disabled access_key!" + payload = jwt.decode( + param, api_key_model.secret_key, algorithms=[security.ALGORITHM] + ) + uid = api_key_model.uid + session_user = await get_user_info(session = session, user_id = uid) + if not session_user: + message = trans('i18n_not_exist', msg = trans('i18n_user.account')) + raise Exception(message) + session_user = UserInfoDTO.model_validate(session_user) + if session_user.status != 1: + message = trans('i18n_login.user_disable', msg = trans('i18n_concat_admin')) + raise Exception(message) + if not session_user.oid or session_user.oid == 0: + message = trans('i18n_login.no_associated_ws', msg = trans('i18n_concat_admin')) + raise Exception(message) + return True, session_user + except Exception as e: + msg = str(e) + SQLBotLogUtil.exception(f"Token validation error: {msg}") + if 'expired' in msg: + return False, jwt.ExpiredSignatureError(trans('i18n_permission.token_expired')) + return False, e + async def validateToken(self, token: Optional[str], trans: I18n): if not token: return False, f"Miss Token[{settings.TOKEN_KEY}]!" diff --git a/backend/apps/system/models/system_model.py b/backend/apps/system/models/system_model.py index 281b737c7..681c6c73d 100644 --- a/backend/apps/system/models/system_model.py +++ b/backend/apps/system/models/system_model.py @@ -67,4 +67,15 @@ class AuthenticationModel(SnowflakeBase, AuthenticationBaseModel, table=True): __tablename__ = "sys_authentication" create_time: Optional[int] = Field(default=0, sa_type=BigInteger()) enable: bool = Field(default=False, nullable=False) - valid: bool = Field(default=False, nullable=False) \ No newline at end of file + valid: bool = Field(default=False, nullable=False) + + +class ApiKeyBaseModel(SQLModel): + access_key: str = Field(max_length=255, nullable=False) + secret_key: str = Field(max_length=255, nullable=False) + create_time: int = Field(default=0, sa_type=BigInteger()) + uid: int = Field(default=0,nullable=False, sa_type=BigInteger()) + status: bool = Field(default=True, nullable=False) + +class ApiKeyModel(SnowflakeBase, ApiKeyBaseModel, table=True): + __tablename__ = "sys_apikey" \ No newline at end of file diff --git a/backend/apps/system/schemas/auth.py b/backend/apps/system/schemas/auth.py index 3b3816ab0..14db17d97 100644 --- a/backend/apps/system/schemas/auth.py +++ b/backend/apps/system/schemas/auth.py @@ -16,6 +16,7 @@ class CacheName(Enum): USER_INFO = "user:info" ASSISTANT_INFO = "assistant:info" ASSISTANT_DS = "assistant:ds" + ASK_INFO = "ask:info" def __str__(self): return self.value diff --git a/backend/apps/system/schemas/system_schema.py b/backend/apps/system/schemas/system_schema.py index 67d66e1da..9c67f92f0 100644 --- a/backend/apps/system/schemas/system_schema.py +++ b/backend/apps/system/schemas/system_schema.py @@ -207,3 +207,13 @@ class AssistantUiSchema(BaseCreatorDTO): name: Optional[str] = None welcome: Optional[str] = None welcome_desc: Optional[str] = None + +class ApikeyStatus(BaseModel): + id: int = Field(description=f"{PLACEHOLDER_PREFIX}id") + status: bool = Field(description=f"{PLACEHOLDER_PREFIX}status") + +class ApikeyGridItem(BaseCreatorDTO): + access_key: str = Field(description=f"Access Key") + secret_key: str = Field(description=f"Secret Key") + status: bool = Field(description=f"{PLACEHOLDER_PREFIX}status") + create_time: int = Field(description=f"{PLACEHOLDER_PREFIX}create_time") \ No newline at end of file diff --git a/frontend/src/components/layout/Apikey.vue b/frontend/src/components/layout/Apikey.vue index 36382e73f..e56d9d5a3 100644 --- a/frontend/src/components/layout/Apikey.vue +++ b/frontend/src/components/layout/Apikey.vue @@ -1,5 +1,5 @@