From 46fe2d83691ca38b1b97533415d78f54829f8697 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E7=A5=A5=E8=B6=85?= <iwind.liu@gmail.com>
Date: Fri, 12 Apr 2024 11:35:52 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E5=A5=97=E9=A4=90=E5=8F=98?=
 =?UTF-8?q?=E6=9B=B4=E5=90=8E=E7=BD=91=E7=AB=99=E9=99=90=E6=B5=81=E7=8A=B6?=
 =?UTF-8?q?=E6=80=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/db/models/server_dao.go              | 42 +++++++++++++++----
 internal/db/models/server_dao_ext.go          | 12 ++++++
 .../db/models/user_plan_stat_dao_community.go |  6 +++
 3 files changed, 53 insertions(+), 7 deletions(-)
 create mode 100644 internal/db/models/server_dao_ext.go

diff --git a/internal/db/models/server_dao.go b/internal/db/models/server_dao.go
index 9c725d4a..be66a5a2 100644
--- a/internal/db/models/server_dao.go
+++ b/internal/db/models/server_dao.go
@@ -2504,7 +2504,7 @@ func (this *ServerDAO) UpdateServerTrafficLimitStatus(tx *dbs.Tx, serverId int64
 		}
 		if len(oldStatus.UntilDay) > 0 &&
 			oldStatus.UntilDay >= day /** 如果已经限制，且比当前日期长，则无需重复 **/ &&
-			oldStatus.PlanId == planId /** 套餐无变化 **/ {
+			oldStatus.PlanId == planId {
 			// no need to change
 			return nil
 		}
@@ -2555,7 +2555,7 @@ func (this *ServerDAO) UpdateServersTrafficLimitStatusWithUserPlanId(tx *dbs.Tx,
 	return nil
 }
 
-// ResetServersTrafficLimitStatusWithPlanId 重置网站限流状态
+// ResetServersTrafficLimitStatusWithPlanId 重置某个套餐相关网站限流状态
 func (this *ServerDAO) ResetServersTrafficLimitStatusWithPlanId(tx *dbs.Tx, planId int64) error {
 	return this.Query(tx).
 		Where("JSON_EXTRACT(trafficLimitStatus, '$.planId')=:planId").
@@ -2632,13 +2632,17 @@ func (this *ServerDAO) UpdateServerUserPlanId(tx *dbs.Tx, serverId int64, userPl
 		return errors.New("serverId should not be smaller than 0")
 	}
 
-	oldClusterId, err := this.Query(tx).
+	oldServerOne, queryErr := SharedServerDAO.
+		Query(tx).
 		Pk(serverId).
-		Result("clusterId").
-		FindInt64Col(0)
-	if err != nil {
-		return err
+		Result("clusterId", "userPlanId").
+		Find()
+	if queryErr != nil || oldServerOne == nil {
+		return queryErr
 	}
+	var oldServer = oldServerOne.(*Server)
+	var oldClusterId = int64(oldServer.ClusterId)
+	var oldUserPlanId = int64(oldServer.UserPlanId)
 
 	// 取消套餐
 	if userPlanId <= 0 {
@@ -2670,6 +2674,15 @@ func (this *ServerDAO) UpdateServerUserPlanId(tx *dbs.Tx, serverId int64, userPl
 		if err != nil {
 			return err
 		}
+
+		// 重置以往的用户套餐状态
+		if oldUserPlanId > 0 {
+			err = SharedUserPlanStatDAO.ResetUserPlanStatsWithUserPlanId(tx, oldUserPlanId)
+			if err != nil {
+				return err
+			}
+		}
+
 		err = this.NotifyUpdate(tx, serverId)
 		if err != nil {
 			return err
@@ -2717,6 +2730,21 @@ func (this *ServerDAO) UpdateServerUserPlanId(tx *dbs.Tx, serverId int64, userPl
 	if err != nil {
 		return err
 	}
+
+	// 重置以往的用户套餐统计状态
+	if oldUserPlanId > 0 {
+		err = SharedUserPlanStatDAO.ResetUserPlanStatsWithUserPlanId(tx, oldUserPlanId)
+		if err != nil {
+			return err
+		}
+	}
+
+	// 重置当前用户套餐统计状态
+	err = SharedUserPlanStatDAO.ResetUserPlanStatsWithUserPlanId(tx, userPlanId)
+	if err != nil {
+		return err
+	}
+
 	err = this.NotifyUpdate(tx, serverId)
 	if err != nil {
 		return err
diff --git a/internal/db/models/server_dao_ext.go b/internal/db/models/server_dao_ext.go
new file mode 100644
index 00000000..95d5e30f
--- /dev/null
+++ b/internal/db/models/server_dao_ext.go
@@ -0,0 +1,12 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+//go:build !plus
+
+package models
+
+import "github.com/iwind/TeaGo/dbs"
+
+// ResetServersTrafficLimitStatusWithUserPlanId 重置用户套餐相关网站限流状态
+func (this *ServerDAO) ResetServersTrafficLimitStatusWithUserPlanId(tx *dbs.Tx, userPlanId int64) error {
+	// stub
+	return nil
+}
diff --git a/internal/db/models/user_plan_stat_dao_community.go b/internal/db/models/user_plan_stat_dao_community.go
index c50df458..3ebae037 100644
--- a/internal/db/models/user_plan_stat_dao_community.go
+++ b/internal/db/models/user_plan_stat_dao_community.go
@@ -6,5 +6,11 @@ package models
 import "github.com/iwind/TeaGo/dbs"
 
 func (this *UserPlanStatDAO) IncreaseUserPlanStat(tx *dbs.Tx, userPlanId int64, trafficBytes int64, countRequests int64, countWebsocketConnections int64) error {
+	// stub
+	return nil
+}
+
+func (this *UserPlanStatDAO) ResetUserPlanStatsWithUserPlanId(tx *dbs.Tx, userPlanId int64) error {
+	// stub
 	return nil
 }