diff --git a/Makefile b/Makefile index d3f09bd..4600abc 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 a912b06..238055a 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 7c328dd..6761d8f 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 }