套餐可以设置带宽限制

internal/nodes/http_request.go

@@ -201,9 +201,10 @@ func (this *HTTPRequest) Do() {
 
 		// 流量限制
 		if this.ReqServer.TrafficLimitStatus != nil && this.ReqServer.TrafficLimitStatus.IsValid() {
-			this.doTrafficLimit()
-			this.doEnd()
-			return
+			if this.doTrafficLimit(this.ReqServer.TrafficLimitStatus) {
+				this.doEnd()
+				return
+			}
 		}
 
 		// WAF

internal/nodes/http_request_traffic_limit.go

@@ -7,7 +7,19 @@ import (
 )
 
 // 流量限制
-func (this *HTTPRequest) doTrafficLimit() {
+func (this *HTTPRequest) doTrafficLimit(status *serverconfigs.TrafficLimitStatus) (blocked bool) {
+	if status == nil {
+		return false
+	}
+
+	// 如果是网站单独设置的流量限制，则检查是否已关闭
+	var config = this.ReqServer.TrafficLimit
+	if (config == nil || !config.IsOn) && status.PlanId == 0 {
+		return false
+	}
+
+	// 如果是套餐设置的流量限制，即使套餐变更了（变更套餐或者变更套餐的限制），仍然会提示流量超限
+
 	this.tags = append(this.tags, "trafficLimit")
 
 	var statusCode = 509
@@ -17,10 +29,19 @@ func (this *HTTPRequest) doTrafficLimit() {
 	this.writer.Header().Set("Content-Type", "text/html; charset=utf-8")
 	this.writer.WriteHeader(statusCode)
 
-	var config = this.ReqServer.TrafficLimit
+	// check plan traffic limit
+	if (config == nil || !config.IsOn) && this.ReqServer.PlanId() > 0 && this.nodeConfig != nil {
+		var planConfig = this.nodeConfig.FindPlan(this.ReqServer.PlanId())
+		if planConfig != nil && planConfig.TrafficLimit != nil && planConfig.TrafficLimit.IsOn {
+			config = planConfig.TrafficLimit
+		}
+	}
+
 	if config != nil && len(config.NoticePageBody) != 0 {
 		_, _ = this.writer.WriteString(this.Format(config.NoticePageBody))
 	} else {
 		_, _ = this.writer.WriteString(this.Format(serverconfigs.DefaultTrafficLimitNoticePageBody))
 	}
+
+	return true
 }

internal/nodes/http_writer.go

@@ -807,6 +807,8 @@ func (this *HTTPWriter) Write(data []byte) (n int, err error) {
 	}
 	n, err = this.writer.Write(data)
 
+	this.checkPlanBandwidth(n)
+
 	return
 }
 

internal/nodes/http_writer_ext.go

@@ -11,3 +11,7 @@ import (
 func (this *HTTPWriter) canSendfile() (*os.File, bool) {
 	return nil, false
 }
+
+func (this *HTTPWriter) checkPlanBandwidth(n int) {
+	// stub
+}

internal/nodes/node_tasks.go

@@ -102,6 +102,8 @@ func (this *Node) execTask(rpcClient *rpc.RPCClient, task *pb.NodeTask) error {
 		err = this.execNetworkSecurityPolicyChangedTask(rpcClient)
 	case "webPPolicyChanged":
 		err = this.execWebPPolicyChangedTask(rpcClient)
+	case "planChanged":
+		err = this.execPlanChangedTask(rpcClient)
 	default:
 		// 特殊任务
 		if strings.HasPrefix(task.Type, "ipListDeleted") { // 删除IP名单

internal/nodes/node_tasks_ext.go

@@ -34,3 +34,7 @@ func (this *Node) execNetworkSecurityPolicyChangedTask(rpcClient *rpc.RPCClient)
 	// stub
 	return nil
 }
+
+func (this *Node) execPlanChangedTask(rpcClient *rpc.RPCClient) error {
+	return nil
+}

internal/rpc/rpc_client.go

@@ -55,6 +55,7 @@ type RPCClient struct {
 	ClientAgentIPRPC       pb.ClientAgentIPServiceClient
 	AuthorityKeyRPC        pb.AuthorityKeyServiceClient
 	UpdatingServerListRPC  pb.UpdatingServerListServiceClient
+	PlanRPC                pb.PlanServiceClient
 }
 
 func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) {
@@ -91,6 +92,7 @@ func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) {
 	client.ClientAgentIPRPC = pb.NewClientAgentIPServiceClient(client)
 	client.AuthorityKeyRPC = pb.NewAuthorityKeyServiceClient(client)
 	client.UpdatingServerListRPC = pb.NewUpdatingServerListServiceClient(client)
+	client.PlanRPC = pb.NewPlanServiceClient(client)
 
 	err := client.init()
 	if err != nil {

internal/utils/ratelimit/bandwidth.go

@@ -0,0 +1,61 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package ratelimit
+
+import (
+	"context"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
+	"sync/atomic"
+	"time"
+)
+
+// Bandwidth lossy bandwidth limiter
+type Bandwidth struct {
+	totalBytes int64
+
+	currentTimestamp int64
+	currentBytes     int64
+}
+
+// NewBandwidth create new bandwidth limiter
+func NewBandwidth(totalBytes int64) *Bandwidth {
+	return &Bandwidth{totalBytes: totalBytes}
+}
+
+// Ack acquire next chance to send data
+func (this *Bandwidth) Ack(ctx context.Context, newBytes int) {
+	if newBytes <= 0 {
+		return
+	}
+	if this.totalBytes <= 0 {
+		return
+	}
+
+	var timestamp = fasttime.Now().Unix()
+	if this.currentTimestamp != 0 && this.currentTimestamp != timestamp {
+		this.currentTimestamp = timestamp
+		this.currentBytes = int64(newBytes)
+
+		// 第一次发送直接放行，不需要判断
+		return
+	}
+
+	if this.currentTimestamp == 0 {
+		this.currentTimestamp = timestamp
+	}
+	if atomic.AddInt64(&this.currentBytes, int64(newBytes)) <= this.totalBytes {
+		return
+	}
+
+	var timeout = time.NewTimer(1 * time.Second)
+	if ctx != nil {
+		select {
+		case <-timeout.C:
+		case <-ctx.Done():
+		}
+	} else {
+		select {
+		case <-timeout.C:
+		}
+	}
+}

internal/utils/ratelimit/bandwidth_test.go

@@ -0,0 +1,27 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package ratelimit_test
+
+import (
+	"context"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/ratelimit"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/testutils"
+	"testing"
+)
+
+func TestBandwidth(t *testing.T) {
+	if !testutils.IsSingleTesting() {
+		return
+	}
+
+	var bandwidth = ratelimit.NewBandwidth(32 << 10)
+	bandwidth.Ack(context.Background(), 123)
+	bandwidth.Ack(context.Background(), 16 << 10)
+	bandwidth.Ack(context.Background(), 32 << 10)
+}
+
+func TestBandwidth_0(t *testing.T) {
+	var bandwidth = ratelimit.NewBandwidth(0)
+	bandwidth.Ack(context.Background(), 123)
+	bandwidth.Ack(context.Background(), 123456)
+}

internal/utils/ratelimit/counter_test.go

@@ -1,14 +1,20 @@
 // Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
 
-package ratelimit
+package ratelimit_test
 
 import (
+	"github.com/TeaOSLab/EdgeNode/internal/utils/ratelimit"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/testutils"
 	"testing"
 	"time"
 )
 
 func TestCounter_ACK(t *testing.T) {
-	var counter = NewCounter(10)
+	if !testutils.IsSingleTesting() {
+		return
+	}
+
+	var counter = ratelimit.NewCounter(10)
 
 	go func() {
 		for i := 0; i < 10; i++ {
@@ -26,7 +32,7 @@ func TestCounter_ACK(t *testing.T) {
 }
 
 func TestCounter_Release(t *testing.T) {
-	var counter = NewCounter(10)
+	var counter = ratelimit.NewCounter(10)
 
 	for i := 0; i < 10; i++ {
 		counter.Ack()
@@ -34,5 +40,5 @@ func TestCounter_Release(t *testing.T) {
 	for i := 0; i < 10; i++ {
 		counter.Release()
 	}
-	t.Log(len(counter.sem))
+	t.Log(counter.Len())
 }