重新实现套餐

internal/nodes/client_conn.go

@@ -198,9 +198,9 @@ func (this *ClientConn) Write(b []byte) (n int, err error) {
 
 				var cost = time.Since(before).Seconds()
 				if cost > 1 {
-					stats.SharedBandwidthStatManager.AddBandwidth(this.userId, this.serverId, int64(float64(n)/cost), int64(n))
+					stats.SharedBandwidthStatManager.AddBandwidth(this.userId, this.userPlanId, this.serverId, int64(float64(n)/cost), int64(n))
 				} else {
-					stats.SharedBandwidthStatManager.AddBandwidth(this.userId, this.serverId, int64(n), int64(n))
+					stats.SharedBandwidthStatManager.AddBandwidth(this.userId, this.userPlanId, this.serverId, int64(n), int64(n))
 				}
 			}
 		}

internal/nodes/client_conn_base.go

@@ -16,6 +16,7 @@ type BaseClientConn struct {
 
 	isBound    bool
 	userId     int64
+	userPlanId int64
 	serverId   int64
 	remoteAddr string
 	hasLimit   bool
@@ -106,11 +107,31 @@ func (this *BaseClientConn) SetUserId(userId int64) {
 	}
 }
 
+func (this *BaseClientConn) SetUserPlanId(userPlanId int64) {
+	this.userPlanId = userPlanId
+
+	// 设置包装前连接
+	switch conn := this.rawConn.(type) {
+	case *tls.Conn:
+		nativeConn, ok := conn.NetConn().(ClientConnInterface)
+		if ok {
+			nativeConn.SetUserPlanId(userPlanId)
+		}
+	case *ClientConn:
+		conn.SetUserPlanId(userPlanId)
+	}
+}
+
 // UserId 获取当前连接所属服务的用户ID
 func (this *BaseClientConn) UserId() int64 {
 	return this.userId
 }
 
+// UserPlanId 用户套餐ID
+func (this *BaseClientConn) UserPlanId() int64 {
+	return this.userPlanId
+}
+
 // RawIP 原本IP
 func (this *BaseClientConn) RawIP() string {
 	if len(this.rawIP) > 0 {

internal/nodes/client_conn_interface.go

@@ -18,9 +18,12 @@ type ClientConnInterface interface {
 	// SetServerId 设置服务ID
 	SetServerId(serverId int64) (goNext bool)
 
-	// SetUserId 设置所属服务的用户ID
+	// SetUserId 设置所属网站的用户ID
 	SetUserId(userId int64)
 
+	// SetUserPlanId 设置
+	SetUserPlanId(userPlanId int64)
+
 	// UserId 获取当前连接所属服务的用户ID
 	UserId() int64
 

internal/nodes/http_request.go

@@ -198,7 +198,7 @@ func (this *HTTPRequest) Do() {
 		}
 
 		// 流量限制
-		if this.ReqServer.TrafficLimit != nil && this.ReqServer.TrafficLimit.IsOn && !this.ReqServer.TrafficLimit.IsEmpty() && this.ReqServer.TrafficLimitStatus != nil && this.ReqServer.TrafficLimitStatus.IsValid() {
+		if this.ReqServer.TrafficLimitStatus != nil && this.ReqServer.TrafficLimitStatus.IsValid() {
 			this.doTrafficLimit()
 			this.doEnd()
 			return

internal/nodes/http_request_traffic_limit.go

@@ -8,15 +8,17 @@ import (
 
 // 流量限制
 func (this *HTTPRequest) doTrafficLimit() {
-	var config = this.ReqServer.TrafficLimit
+	this.tags = append(this.tags, "trafficLimit")
-
-	this.tags = append(this.tags, "bandwidth")
 
 	var statusCode = 509
+	this.writer.statusCode = statusCode
 	this.ProcessResponseHeaders(this.writer.Header(), statusCode)
 
+	this.writer.Header().Set("Content-Type", "text/html; charset=utf-8")
 	this.writer.WriteHeader(statusCode)
-	if len(config.NoticePageBody) != 0 {
+
+	var config = this.ReqServer.TrafficLimit
+	if config != nil && len(config.NoticePageBody) != 0 {
 		_, _ = this.writer.WriteString(this.Format(config.NoticePageBody))
 	} else {
 		_, _ = this.writer.WriteString(this.Format(serverconfigs.DefaultTrafficLimitNoticePageBody))

internal/nodes/listener_http.go

@@ -177,6 +177,12 @@ func (this *HTTPListener) ServeHTTP(rawWriter http.ResponseWriter, rawReq *http.
 					return
 				}
 				clientConn.SetUserId(server.UserId)
+
+				var userPlanId int64
+				if server.UserPlan != nil && server.UserPlan.Id > 0 {
+					userPlanId = server.UserPlan.Id
+				}
+				clientConn.SetUserPlanId(userPlanId)
 			}
 		}
 	}

internal/nodes/listener_tcp.go

@@ -80,6 +80,12 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
 			return nil
 		}
 		clientConn.SetUserId(server.UserId)
+
+		var userPlanId int64
+		if server.UserPlan != nil && server.UserPlan.Id > 0 {
+			userPlanId = server.UserPlan.Id
+		}
+		clientConn.SetUserPlanId(userPlanId)
 	} else {
 		tlsConn, ok := conn.(*tls.Conn)
 		if ok {
@@ -92,6 +98,12 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
 						return nil
 					}
 					clientConn.SetUserId(server.UserId)
+
+					var userPlanId int64
+					if server.UserPlan != nil && server.UserPlan.Id > 0 {
+						userPlanId = server.UserPlan.Id
+					}
+					clientConn.SetUserPlanId(userPlanId)
 				}
 			}
 		}

internal/nodes/listener_udp.go

@@ -404,7 +404,11 @@ func NewUDPConn(server *serverconfigs.ServerConfig, addr net.Addr, proxyListener
 					stats.SharedTrafficStatManager.Add(server.UserId, server.Id, "", int64(n), 0, 0, 0, 0, 0, server.ShouldCheckTrafficLimit(), server.PlanId())
 
 					// 带宽
-					stats.SharedBandwidthStatManager.AddBandwidth(server.UserId, server.Id, int64(n), int64(n))
+					var userPlanId int64
+					if server.UserPlan != nil && server.UserPlan.Id > 0 {
+						userPlanId = server.UserPlan.Id
+					}
+					stats.SharedBandwidthStatManager.AddBandwidth(server.UserId, userPlanId, server.Id, int64(n), int64(n))
 				}
 			}
 			if err != nil {

internal/stats/bandwidth_stat_manager.go

@@ -62,6 +62,7 @@ type BandwidthStat struct {
 	CountRequests       int64 `json:"countRequests"`
 	CountCachedRequests int64 `json:"countCachedRequests"`
 	CountAttackRequests int64 `json:"countAttackRequests"`
+	UserPlanId          int64 `json:"userPlanId"`
 }
 
 // BandwidthStatManager 服务带宽统计
@@ -153,6 +154,7 @@ func (this *BandwidthStatManager) Loop() error {
 				CountRequests:       stat.CountRequests,
 				CountCachedRequests: stat.CountCachedRequests,
 				CountAttackRequests: stat.CountAttackRequests,
+				UserPlanId:          stat.UserPlanId,
 				NodeRegionId:        regionId,
 			})
 			delete(this.m, key)
@@ -178,7 +180,7 @@ func (this *BandwidthStatManager) Loop() error {
 }
 
 // AddBandwidth 添加带宽数据
-func (this *BandwidthStatManager) AddBandwidth(userId int64, serverId int64, peekBytes int64, totalBytes int64) {
+func (this *BandwidthStatManager) AddBandwidth(userId int64, userPlanId int64, serverId int64, peekBytes int64, totalBytes int64) {
 	if serverId <= 0 || (peekBytes == 0 && totalBytes == 0) {
 		return
 	}
@@ -217,6 +219,7 @@ func (this *BandwidthStatManager) AddBandwidth(userId int64, serverId int64, pee
 			Day:              day,
 			TimeAt:           timeAt,
 			UserId:           userId,
+			UserPlanId:       userPlanId,
 			ServerId:         serverId,
 			CurrentBytes:     peekBytes,
 			MaxBytes:         peekBytes,

internal/stats/bandwidth_stat_manager_test.go

@@ -12,22 +12,22 @@ import (
 
 func TestBandwidthStatManager_Add(t *testing.T) {
 	var manager = stats.NewBandwidthStatManager()
-	manager.AddBandwidth(1, 1, 10, 10)
+	manager.AddBandwidth(1, 0, 1, 10, 10)
-	manager.AddBandwidth(1, 1, 10, 10)
+	manager.AddBandwidth(1, 0, 1, 10, 10)
-	manager.AddBandwidth(1, 1, 10, 10)
+	manager.AddBandwidth(1, 0, 1, 10, 10)
 	time.Sleep(1 * time.Second)
-	manager.AddBandwidth(1, 1, 85, 85)
+	manager.AddBandwidth(1, 0, 1, 85, 85)
 	time.Sleep(1 * time.Second)
-	manager.AddBandwidth(1, 1, 25, 25)
+	manager.AddBandwidth(1, 0, 1, 25, 25)
-	manager.AddBandwidth(1, 1, 75, 75)
+	manager.AddBandwidth(1, 0, 1, 75, 75)
 	manager.Inspect()
 }
 
 func TestBandwidthStatManager_Loop(t *testing.T) {
 	var manager = stats.NewBandwidthStatManager()
-	manager.AddBandwidth(1, 1, 10, 10)
+	manager.AddBandwidth(1, 0, 1, 10, 10)
-	manager.AddBandwidth(1, 1, 10, 10)
+	manager.AddBandwidth(1, 0, 1, 10, 10)
-	manager.AddBandwidth(1, 1, 10, 10)
+	manager.AddBandwidth(1, 0, 1, 10, 10)
 	err := manager.Loop()
 	if err != nil {
 		t.Fatal(err)
@@ -40,7 +40,7 @@ func BenchmarkBandwidthStatManager_Add(b *testing.B) {
 		var i int
 		for pb.Next() {
 			i++
-			manager.AddBandwidth(1, int64(i%100), 10, 10)
+			manager.AddBandwidth(1, 0, int64(i%100), 10, 10)
 		}
 	})
 }