From bacdc374664c1f2731285e601e20b7bed7e5d486 Mon Sep 17 00:00:00 2001
From: click33 <2393584716@qq.com>
Date: Sat, 20 May 2023 18:43:59 +0800
Subject: [PATCH] =?UTF-8?q?StpLogic=20=E7=9A=84=E5=88=9B=E5=BB=BA=E8=BD=AC?=
 =?UTF-8?q?=E7=A7=BB=E5=88=B0=20SaStrategy=20=E5=85=A8=E5=B1=80=E7=AD=96?=
 =?UTF-8?q?=E7=95=A5=E4=B8=AD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../main/java/cn/dev33/satoken/SaManager.java |  7 ++++++
 .../java/cn/dev33/satoken/stp/StpLogic.java   | 14 +++++++----
 .../java/cn/dev33/satoken/stp/StpUtil.java    |  2 +-
 .../cn/dev33/satoken/strategy/SaStrategy.java | 23 +++++++++++++++++++
 4 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/sa-token-core/src/main/java/cn/dev33/satoken/SaManager.java b/sa-token-core/src/main/java/cn/dev33/satoken/SaManager.java
index dc4c75e6..5c95e69a 100644
--- a/sa-token-core/src/main/java/cn/dev33/satoken/SaManager.java
+++ b/sa-token-core/src/main/java/cn/dev33/satoken/SaManager.java
@@ -285,6 +285,13 @@ public class SaManager {
 		stpLogicMap.put(stpLogic.getLoginType(), stpLogic);
 	}
 
+	/**
+	 * 在全局集合中 移除 一个 StpLogic
+	 */
+	public static void removeStpLogic(String loginType) {
+		stpLogicMap.remove(loginType);
+	}
+
 	/**
 	 * 根据 LoginType 获取对应的StpLogic，如果不存在则新建并返回 
 	 * @param loginType 对应的账号类型 
diff --git a/sa-token-core/src/main/java/cn/dev33/satoken/stp/StpLogic.java b/sa-token-core/src/main/java/cn/dev33/satoken/stp/StpLogic.java
index c66c1043..76baae4f 100644
--- a/sa-token-core/src/main/java/cn/dev33/satoken/stp/StpLogic.java
+++ b/sa-token-core/src/main/java/cn/dev33/satoken/stp/StpLogic.java
@@ -67,10 +67,7 @@ public class StpLogic {
 	 * @param loginType 账号类型标识
 	 */
 	public StpLogic(String loginType) {
-		this.loginType = loginType;
-
-		// 在 SaManager 中记录下此 StpLogic，以便后续根据 LoginType 进行查找此对象
-		SaManager.putStpLogic(this);
+		setLoginType(loginType);
 	}
 
 	/**
@@ -89,9 +86,16 @@ public class StpLogic {
 	 * @return 对象自身
 	 */
 	public StpLogic setLoginType(String loginType){
+
+		// 先清除此 StpLogic 在全局 SaManager 中的记录
+		if(SaFoxUtil.isNotEmpty(this.loginType)) {
+			SaManager.removeStpLogic(this.loginType);
+		}
+
+		// 赋值
 		this.loginType = loginType;
 
-		// 因为 loginType 变了，所以需要往 SaManager 中重新 put 一下
+		// 将新的 loginType -> StpLogic 映射关系 put 到 SaManager 全局集合中，以便后续根据 LoginType 进行查找此对象
 		SaManager.putStpLogic(this);
 
 		return this;
diff --git a/sa-token-core/src/main/java/cn/dev33/satoken/stp/StpUtil.java b/sa-token-core/src/main/java/cn/dev33/satoken/stp/StpUtil.java
index 2c9d398a..e0dd2310 100644
--- a/sa-token-core/src/main/java/cn/dev33/satoken/stp/StpUtil.java
+++ b/sa-token-core/src/main/java/cn/dev33/satoken/stp/StpUtil.java
@@ -40,7 +40,7 @@ public class StpUtil {
 	/**
 	 * 底层使用的 StpLogic 对象
 	 */
-	public static StpLogic stpLogic = new StpLogic(TYPE); 
+	public static StpLogic stpLogic = SaManager.getStpLogic(TYPE);
 
 	/**
 	 * 获取当前 StpLogic 的账号类型
diff --git a/sa-token-core/src/main/java/cn/dev33/satoken/strategy/SaStrategy.java b/sa-token-core/src/main/java/cn/dev33/satoken/strategy/SaStrategy.java
index 708c0c8e..8e89b64a 100644
--- a/sa-token-core/src/main/java/cn/dev33/satoken/strategy/SaStrategy.java
+++ b/sa-token-core/src/main/java/cn/dev33/satoken/strategy/SaStrategy.java
@@ -21,6 +21,7 @@ import cn.dev33.satoken.basic.SaBasicUtil;
 import cn.dev33.satoken.exception.SaTokenException;
 import cn.dev33.satoken.fun.SaGenerateUniqueTokenFunction;
 import cn.dev33.satoken.session.SaSession;
+import cn.dev33.satoken.stp.StpLogic;
 import cn.dev33.satoken.util.SaFoxUtil;
 import cn.dev33.satoken.util.SaTokenConsts;
 
@@ -240,6 +241,15 @@ public final class SaStrategy {
 		}
 	};
 
+	/**
+	 * 创建 StpLogic 的算法
+	 *
+	 * <p> 参数 [ 账号体系标识 ]
+	 */
+	public Function<String, StpLogic> createStpLogic = (loginType) -> {
+		return new StpLogic(loginType);
+	};
+
 
 	// ----------------------- 重写策略 set连缀风格
 
@@ -332,4 +342,17 @@ public final class SaStrategy {
 		return this;
 	}
 
+	/**
+	 * 创建 StpLogic 的算法
+	 *
+	 * <p> 参数 [ 账号体系标识 ]
+	 *
+	 * @param createStpLogic /
+	 * @return 对象自身
+	 */
+	public SaStrategy setCreateStpLogic(Function<String, StpLogic> createStpLogic) {
+		this.createStpLogic = createStpLogic;
+		return this;
+	}
+
 }