重新实现套餐

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, "bandwidth")
+	this.tags = append(this.tags, "trafficLimit")
 
 	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 {