feat: 增加超出流量后立即断开用户连接的逻辑

cedar2025 · Mar 6, 2024 · 19e695a · 19e695a
1 parent 161c1a4
commit 19e695a
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 38 deletions.
diff --git a/app/cmd/server.go b/app/cmd/server.go
@@ -608,18 +608,8 @@ func (c *serverConfig) fillAuthenticator(hyConfig *server.Config) error {
 		if v2boardConfig.ApiHost == "" || v2boardConfig.ApiKey == "" || v2boardConfig.NodeID == 0 {
 			return configError{Field: "auth.v2board", Err: errors.New("v2board config error")}
 		}
-		// 创建一个url.Values来存储查询参数
-		queryParams := url.Values{}
-		queryParams.Add("token", v2boardConfig.ApiKey)
-		queryParams.Add("node_id", strconv.Itoa(int(v2boardConfig.NodeID)))
-		queryParams.Add("node_type", "hysteria")
-		// 创建完整的URL，包括查询参数
-		url := v2boardConfig.ApiHost + "/api/v1/server/UniProxy/user?" + queryParams.Encode()
-
 		// 创建定时更新用户UUID协程
-		go auth.UpdateUsers(url, time.Second*5)
-
-		hyConfig.Authenticator = &auth.V2boardApiProvider{URL: url}
+		hyConfig.Authenticator = &auth.V2boardApiProvider{URL: fmt.Sprintf("%s?token=%s&node_id=%d&node_type=hysteria", c.V2board.ApiHost+"/api/v1/server/UniProxy/user", c.V2board.ApiKey, c.V2board.NodeID)}
 
 		return nil
 
@@ -639,14 +629,12 @@ func (c *serverConfig) fillTrafficLogger(hyConfig *server.Config) error {
 		hyConfig.TrafficLogger = tss
 		// 添加定时更新用户使用流量协程
 		if c.V2board != nil && c.V2board.ApiHost != "" {
-			// 创建一个url.Values来存储查询参数
-			queryParams := url.Values{}
-			queryParams.Add("token", c.V2board.ApiKey)
-			queryParams.Add("node_id", strconv.Itoa(int(c.V2board.NodeID)))
-			queryParams.Add("node_type", "hysteria")
-			go hyConfig.TrafficLogger.PushTrafficToV2boardInterval(c.V2board.ApiHost+"/api/v1/server/UniProxy/push?"+queryParams.Encode(), time.Second*60)
+			go auth.UpdateUsers(fmt.Sprintf("%s?token=%s&node_id=%d&node_type=hysteria", c.V2board.ApiHost+"/api/v1/server/UniProxy/user", c.V2board.ApiKey, c.V2board.NodeID), time.Second*5, hyConfig.TrafficLogger)
+			go hyConfig.TrafficLogger.PushTrafficToV2boardInterval(fmt.Sprintf("%s?token=%s&node_id=%d&node_type=hysteria", c.V2board.ApiHost+"/api/v1/server/UniProxy/push", c.V2board.ApiKey, c.V2board.NodeID), time.Second*60)
 		}
 		go runTrafficStatsServer(c.TrafficStats.Listen, tss)
+	} else {
+		go auth.UpdateUsers(fmt.Sprintf("%s?token=%s&node_id=%d&node_type=hysteria", c.V2board.ApiHost+"/api/v1/server/UniProxy/user", c.V2board.ApiKey, c.V2board.NodeID), time.Second*5, nil)
 	}
 	return nil
 }

diff --git a/core/server/config.go b/core/server/config.go
@@ -197,4 +197,5 @@ type EventLogger interface {
 type TrafficLogger interface {
 	Log(id string, tx, rx uint64) (ok bool)
 	PushTrafficToV2boardInterval(url string, interval time.Duration)
+	NewKick(id string) (ok bool)
 }
diff --git a/extras/auth/v2board.go b/extras/auth/v2board.go
@@ -50,29 +50,36 @@ func getUserList(url string) ([]User, error) {
 	return responseData.Users, nil
 }
 
-func UpdateUsers(url string, interval time.Duration) {
+func UpdateUsers(url string, interval time.Duration, trafficlogger server.TrafficLogger) {
 
 	fmt.Println("用户列表自动更新服务已激活")
-
 	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
 
-	for {
-		select {
-		case <-ticker.C:
-			userList, err := getUserList(url)
-			if err != nil {
-				fmt.Println("Error:", err)
-				continue
-			}
-			lock.Lock()
-			usersMap = make(map[string]User)
-			for _, user := range userList {
-				usersMap[user.UUID] = user
+	for range ticker.C {
+		userList, err := getUserList(url)
+		if err != nil {
+			fmt.Println("Error:", err)
+			continue
+		}
+		lock.Lock()
+		newUsersMap := make(map[string]User)
+		for _, user := range userList {
+			newUsersMap[user.UUID] = user
+		}
+		if trafficlogger != nil {
+			for uuid := range usersMap {
+				if _, exists := newUsersMap[uuid]; !exists {
+					fmt.Println(usersMap[uuid].ID)
+					trafficlogger.NewKick(strconv.Itoa(usersMap[uuid].ID))
+				}
 			}
-			lock.Unlock()
 		}
+
+		usersMap = newUsersMap
+		lock.Unlock()
 	}
+
 }
 
 // 验证代码

diff --git a/extras/trafficlogger/http.go b/extras/trafficlogger/http.go
@@ -43,14 +43,12 @@ func (s *trafficStatsServerImpl) PushTrafficToV2boardInterval(url string, interv
 	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
 
-	for {
-		select {
-		case <-ticker.C:
-			if err := s.PushTrafficToV2board(url); err != nil {
-				fmt.Println("用户流量信息提交失败:", err)
-			}
+	for range ticker.C {
+		if err := s.PushTrafficToV2board(url); err != nil {
+			fmt.Println("用户流量信息提交失败:", err)
 		}
 	}
+
 }
 
 // 向v2board 提交用户流量使用情况
@@ -185,3 +183,11 @@ func (s *trafficStatsServerImpl) kick(w http.ResponseWriter, r *http.Request) {
 
 	w.WriteHeader(http.StatusOK)
 }
+
+// 踢出用户名单
+func (s *trafficStatsServerImpl) NewKick(id string) bool {
+	s.Mutex.Lock()
+	s.KickMap[id] = struct{}{}
+	s.Mutex.Unlock()
+	return true
+}