From 89fff7d11de98d48b36283a5508aa6f5078e246e Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Thu, 9 Feb 2023 01:11:00 +0800
Subject: [PATCH] e2e: add test case for bandwidth_limit_mode server (#3295)

---
 Makefile                             |  3 ++
 test/e2e/features/bandwidth_limit.go | 64 ++++++++++++++++++++++++++--
 test/e2e/framework/process.go        |  4 +-
 3 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/Makefile b/Makefile
index d3f09bd0..4600abcb 100644
--- a/Makefile
+++ b/Makefile
@@ -16,6 +16,9 @@ file:
 fmt:
 	go fmt ./...
 
+fmt-more:
+	gofumpt -l -w .
+
 vet:
 	go vet ./...
 
diff --git a/test/e2e/features/bandwidth_limit.go b/test/e2e/features/bandwidth_limit.go
index a912b06a..238055a8 100644
--- a/test/e2e/features/bandwidth_limit.go
+++ b/test/e2e/features/bandwidth_limit.go
@@ -7,16 +7,18 @@ import (
 
 	"github.com/onsi/ginkgo"
 
+	plugin "github.com/fatedier/frp/pkg/plugin/server"
 	"github.com/fatedier/frp/test/e2e/framework"
 	"github.com/fatedier/frp/test/e2e/framework/consts"
 	"github.com/fatedier/frp/test/e2e/mock/server/streamserver"
 	"github.com/fatedier/frp/test/e2e/pkg/request"
+	plugintest "github.com/fatedier/frp/test/e2e/plugin"
 )
 
 var _ = ginkgo.Describe("[Feature: Bandwidth Limit]", func() {
 	f := framework.NewDefaultFramework()
 
-	ginkgo.It("Proxy Bandwidth Limit", func() {
+	ginkgo.It("Proxy Bandwidth Limit by Client", func() {
 		serverConf := consts.DefaultServerConfig
 		clientConf := consts.DefaultClientConfig
 
@@ -40,8 +42,64 @@ var _ = ginkgo.Describe("[Feature: Bandwidth Limit]", func() {
 		framework.NewRequestExpect(f).Port(remotePort).RequestModify(func(r *request.Request) {
 			r.Body([]byte(content)).Timeout(30 * time.Second)
 		}).ExpectResp([]byte(content)).Ensure()
-		duration := time.Since(start)
 
-		framework.ExpectTrue(duration.Seconds() > 7, "100Kb with 10KB limit, want > 7 seconds, but got %d seconds", duration.Seconds())
+		duration := time.Since(start)
+		framework.Logf("request duration: %s", duration.String())
+
+		framework.ExpectTrue(duration.Seconds() > 8, "100Kb with 10KB limit, want > 8 seconds, but got %s", duration.String())
+	})
+
+	ginkgo.It("Proxy Bandwidth Limit by Server", func() {
+		// new test plugin server
+		newFunc := func() *plugin.Request {
+			var r plugin.Request
+			r.Content = &plugin.NewProxyContent{}
+			return &r
+		}
+		pluginPort := f.AllocPort()
+		handler := func(req *plugin.Request) *plugin.Response {
+			var ret plugin.Response
+			content := req.Content.(*plugin.NewProxyContent)
+			content.BandwidthLimit = "10KB"
+			content.BandwidthLimitMode = "server"
+			ret.Content = content
+			return &ret
+		}
+		pluginServer := plugintest.NewHTTPPluginServer(pluginPort, newFunc, handler, nil)
+
+		f.RunServer("", pluginServer)
+
+		serverConf := consts.DefaultServerConfig + fmt.Sprintf(`
+		[plugin.test]
+		addr = 127.0.0.1:%d
+		path = /handler
+		ops = NewProxy
+		`, pluginPort)
+		clientConf := consts.DefaultClientConfig
+
+		localPort := f.AllocPort()
+		localServer := streamserver.New(streamserver.TCP, streamserver.WithBindPort(localPort))
+		f.RunServer("", localServer)
+
+		remotePort := f.AllocPort()
+		clientConf += fmt.Sprintf(`
+			[tcp]
+			type = tcp
+			local_port = %d
+			remote_port = %d
+			`, localPort, remotePort)
+
+		f.RunProcesses([]string{serverConf}, []string{clientConf})
+
+		content := strings.Repeat("a", 50*1024) // 5KB
+		start := time.Now()
+		framework.NewRequestExpect(f).Port(remotePort).RequestModify(func(r *request.Request) {
+			r.Body([]byte(content)).Timeout(30 * time.Second)
+		}).ExpectResp([]byte(content)).Ensure()
+
+		duration := time.Since(start)
+		framework.Logf("request duration: %s", duration.String())
+
+		framework.ExpectTrue(duration.Seconds() > 8, "100Kb with 10KB limit, want > 8 seconds, but got %s", duration.String())
 	})
 })
diff --git a/test/e2e/framework/process.go b/test/e2e/framework/process.go
index 7c328dd4..6761d8f4 100644
--- a/test/e2e/framework/process.go
+++ b/test/e2e/framework/process.go
@@ -69,7 +69,7 @@ func (f *Framework) RunFrps(args ...string) (*process.Process, string, error) {
 		return p, p.StdOutput(), err
 	}
 	// sleep for a while to get std output
-	time.Sleep(500 * time.Millisecond)
+	time.Sleep(time.Second)
 	return p, p.StdOutput(), nil
 }
 
@@ -80,7 +80,7 @@ func (f *Framework) RunFrpc(args ...string) (*process.Process, string, error) {
 	if err != nil {
 		return p, p.StdOutput(), err
 	}
-	time.Sleep(500 * time.Millisecond)
+	time.Sleep(time.Second)
 	return p, p.StdOutput(), nil
 }