使用trie优化效率

requirements.txt

@@ -16,6 +16,7 @@ MarkupSafe==2.1.5
 multidict==6.0.5
 pydantic==2.8.2
 pydantic_core==2.20.1
+pygtrie==2.5.0
 PyMySQL==1.1.1
 python-multipart==0.0.9
 sniffio==1.3.1

server/email.py

@@ -15,6 +15,9 @@ async def sendemail(to_addr, links):
     print(links)
     html = "验证链接：[<a href=\""+links+"\">"+links+"</a>]"
 
+    if (SMTP_SRV == ''):
+        return 0
+
     msg = MIMEText(html, 'html', 'utf-8')
     msg['From'] = ADDR
     msg['To'] = to_addr

server/main.py

@@ -14,7 +14,7 @@ import re
 import asyncio
 import uvicorn
 from .cfg import config
-
+import pygtrie
 
 ALGORITHM = config["common"]["algorithm"]
 ACCESS_TOKEN_EXPIRE_MINUTES = config["common"]["access_token_expire_minutes"]
@@ -25,23 +25,31 @@ MANAGE_KEY = config["common"]["manage_key"]
 REDIRECT_URL_WHITELIST = config["common"]["redirect_url_whitelist"]
 
 
-tokens = {}
-apikeys = {}
-emails = {}
+tokens = pygtrie.StringTrie()
+apikeys = pygtrie.StringTrie()
+emails = pygtrie.StringTrie()
 email_send_lst = []
 
 
+def prep_uuid(uuid: str):
+    return '/'.join(list(uuid))
+
+
+def clean_uuid(uuid: str):
+    return uuid.replace("/", "")
+
+
 async def clean_sys():
     while 1:
         await asyncio.sleep(CLEAN_TIMEOUT)
-        for k in list(tokens.keys()):
-            if (tokens[k][1] < datetime.now()):
+        for k, v in tokens.items():
+            if (v[1] < datetime.now()):
                 del tokens[k]
-        for k in list(apikeys.keys()):
-            if (apikeys[k][1] < datetime.now()):
+        for k, v in apikeys.items():
+            if (v[1] < datetime.now()):
                 del apikeys[k]
-        for k in list(emails.keys()):
-            if (emails[k][2] < datetime.now()):
+        for k, v in emails.items():
+            if (v[2] < datetime.now()):
                 del emails[k]
 
 
@@ -89,7 +97,7 @@ async def authenticate_user(username: str, password: str):
 
 
 async def create_token(username: str):
-    tkn = uuid.uuid4().hex
+    tkn = prep_uuid(uuid.uuid4().hex)
     tokens[tkn] = (username, datetime.now() +
                    timedelta(minutes=float(ACCESS_TOKEN_EXPIRE_MINUTES)))
     return tkn
@@ -119,9 +127,9 @@ async def check_apikey(tkn: str):
 async def login_callback(response: Response, username: str = Form(), password: str = Form()):
     if (await authenticate_user(username, password)):
         tokennow = await create_token(username)
-        tkn = uuid.uuid4().hex
+        tkn = prep_uuid(uuid.uuid4().hex)
         apikeys[tkn] = tokens[tokennow]
-        response.set_cookie("session", tokennow)
+        response.set_cookie("session", clean_uuid(tokennow))
         return {"msg": "", "key": tkn}
     else:
         return {"msg": "用户名或密码错误", "key": ""}
@@ -138,12 +146,13 @@ async def login_callback(username: str = Form(), password: str = Form(), email:
     if (not re.fullmatch(regex, email)):
         return {"msg": "邮箱不合法", "code": 1}
     if not (await db.check_user(username)):
-        tkn = uuid.uuid4().hex
+        tkn = prep_uuid(uuid.uuid4().hex)
         emails[tkn] = (username, hashlib.sha256(
             password.encode("utf-8")).hexdigest(), datetime.now() +
             timedelta(minutes=float(ACCESS_EMAIL_EXPIRE_MINUTES)), email)
 
-        email_send_lst.append((email, ROOT+"/api/checkemail?uid="+tkn))
+        email_send_lst.append(
+            (email, ROOT+"/api/checkemail?uid="+clean_uuid(tkn)))
         return {"msg": "验证邮件已发送到邮箱，请在10分钟内完成验证", "code": 0}
     else:
         return {"msg": "用户名重复", "code": 1}
@@ -151,6 +160,7 @@ async def login_callback(username: str = Form(), password: str = Form(), email:
 
 @app.get("/api/checkemail")
 async def checkemail(uid: str):
+    uid = prep_uuid(uid)
     if (uid not in emails):
         return templates.TemplateResponse("checkemail.html", {"request": {}, "msg": "不存在的注册id"})
     if (emails[uid][2] < datetime.now()):
@@ -168,6 +178,7 @@ async def checkemail(uid: str):
 
 @app.get("/api/resetpasswd", response_class=HTMLResponse)
 async def resetpasswd(uid: str, response: Response):
+    uid = prep_uuid(uid)
     if (uid not in emails):
         return templates.TemplateResponse("checkemail.html", {"request": {}, "msg": "不存在的验证id"})
     if (emails[uid][2] < datetime.now()):
@@ -176,9 +187,9 @@ async def resetpasswd(uid: str, response: Response):
     if (emails[uid][1] != ""):
         return templates.TemplateResponse("checkemail.html", {"request": {}, "msg": "不存在的注册id"})
     tokennow = await create_token(emails[uid][0])
-    tkn = uuid.uuid4().hex
+    tkn = prep_uuid(uuid.uuid4().hex)
     apikeys[tkn] = tokens[tokennow]
-    response.set_cookie("session", tokennow)
+    response.set_cookie("session", clean_uuid(tokennow))
     del emails[uid]
     return '<html><head><meta http-equiv="refresh" content="0;url=/user"><title>正在跳转</title></head><body>正在跳转</body></html>'
 
@@ -187,11 +198,12 @@ async def resetpasswd(uid: str, response: Response):
 async def resetpasswd(username: str = Form()):
     if (await db.check_user(username)):
         email = await db.get_email(username)
-        tkn = uuid.uuid4().hex
+        tkn = prep_uuid(uuid.uuid4().hex)
         emails[tkn] = (username, "", datetime.now() +
                        timedelta(minutes=float(ACCESS_EMAIL_EXPIRE_MINUTES)), email)
 
-        email_send_lst.append((email, ROOT+"/api/resetpasswd?uid="+tkn))
+        email_send_lst.append(
+            (email, ROOT+"/api/resetpasswd?uid="+clean_uuid(tkn)))
         return {"msg": "验证邮件已发送到邮箱，请在10分钟内完成验证", "code": 0}
     else:
         return {"msg": "用户名不存在", "code": 1}
@@ -199,6 +211,7 @@ async def resetpasswd(username: str = Form()):
 
 @app.get("/api/getinfo")
 async def get_user_info(uid: str):
+    uid = prep_uuid(uid)
     username = await check_apikey(uid)
     if (username == ""):
         return {"code": 1, "msg": "token无效", "data": {}}
@@ -210,6 +223,7 @@ async def changepasswd(password: str = Form(), session: Annotated[str | None, Co
     if (check_passwd(password)):
         return {"msg": "密码强度弱，请至少包含一个大小写字符与数字且长度>8"}
     if (session is not None):
+        session = prep_uuid(session)
         username = await check_token(session)
         if (username != ""):
             await db.update_passwd(username, hashlib.sha256(
@@ -229,9 +243,10 @@ async def login(state: str = "", client_id: str = "", redirect_url: str = "/user
     if (now_redirect_url not in REDIRECT_URL_WHITELIST):
         return templates.TemplateResponse("checkemail.html", {"request": {}, "msg": "无效的重定向URL"})
     if (session is not None):
+        session = prep_uuid(session)
         username = await check_token(session)
         if (username != ""):
-            tkn = uuid.uuid4().hex
+            tkn = prep_uuid(uuid.uuid4().hex)
             apikeys[tkn] = tokens[session]
             return RedirectResponse(url=redirect_url+f"#access_token={tkn}&token_type=Bearer&state={state}")
     return templates.TemplateResponse("login.html", {"request": {}, "redirect_url": redirect_url, "state": state})
@@ -246,6 +261,7 @@ async def login(redirect_url: str):
 async def login(session: Annotated[str | None, Cookie()] = None):
 
     if (session is not None):
+        session = prep_uuid(session)
         username = await check_token(session)
         if (username == ""):
             return RedirectResponse(url="/login?redirect_url=/user")