frp/test/e2e/features/group.go

package features

import (
	"fmt"
	"strconv"
	"sync"
	"time"

	"github.com/onsi/ginkgo/v2"

	"github.com/fatedier/frp/test/e2e/framework"
	"github.com/fatedier/frp/test/e2e/framework/consts"
	"github.com/fatedier/frp/test/e2e/mock/server/httpserver"
	"github.com/fatedier/frp/test/e2e/mock/server/streamserver"
	"github.com/fatedier/frp/test/e2e/pkg/request"
)

var _ = ginkgo.Describe("[Feature: Group]", func() {
	f := framework.NewDefaultFramework()

	newHTTPServer := func(port int, respContent string) *httpserver.Server {
		return httpserver.New(
			httpserver.WithBindPort(port),
			httpserver.WithHandler(framework.SpecifiedHTTPBodyHandler([]byte(respContent))),
		)
	}

	validateFooBarResponse := func(resp *request.Response) bool {
		if string(resp.Content) == "foo" || string(resp.Content) == "bar" {
			return true
		}
		return false
	}

	doFooBarHTTPRequest := func(vhostPort int, host string) []string {
		results := []string{}
		var wait sync.WaitGroup
		var mu sync.Mutex
		expectFn := func() {
			framework.NewRequestExpect(f).Port(vhostPort).
				RequestModify(func(r *request.Request) {
					r.HTTP().HTTPHost(host)
				}).
				Ensure(validateFooBarResponse, func(resp *request.Response) bool {
					mu.Lock()
					defer mu.Unlock()
					results = append(results, string(resp.Content))
					return true
				})
		}
		for i := 0; i < 10; i++ {
			wait.Add(1)
			go func() {
				defer wait.Done()
				expectFn()
			}()
		}

		wait.Wait()
		return results
	}

	ginkgo.Describe("Load Balancing", func() {
		ginkgo.It("TCP", func() {
			serverConf := consts.DefaultServerConfig
			clientConf := consts.DefaultClientConfig

			fooPort := f.AllocPort()
			fooServer := streamserver.New(streamserver.TCP, streamserver.WithBindPort(fooPort), streamserver.WithRespContent([]byte("foo")))
			f.RunServer("", fooServer)

			barPort := f.AllocPort()
			barServer := streamserver.New(streamserver.TCP, streamserver.WithBindPort(barPort), streamserver.WithRespContent([]byte("bar")))
			f.RunServer("", barServer)

			remotePort := f.AllocPort()
			clientConf += fmt.Sprintf(`
			[foo]
			type = tcp
			local_port = %d
			remote_port = %d
			group = test
			group_key = 123

			[bar]
			type = tcp
			local_port = %d
			remote_port = %d
			group = test
			group_key = 123
			`, fooPort, remotePort, barPort, remotePort)

			f.RunProcesses([]string{serverConf}, []string{clientConf})

			fooCount := 0
			barCount := 0
			for i := 0; i < 10; i++ {
				framework.NewRequestExpect(f).Explain("times " + strconv.Itoa(i)).Port(remotePort).Ensure(func(resp *request.Response) bool {
					switch string(resp.Content) {
					case "foo":
						fooCount++
					case "bar":
						barCount++
					default:
						return false
					}
					return true
				})
			}

			framework.ExpectTrue(fooCount > 1 && barCount > 1, "fooCount: %d, barCount: %d", fooCount, barCount)
		})
	})

	ginkgo.Describe("Health Check", func() {
		ginkgo.It("TCP", func() {
			serverConf := consts.DefaultServerConfig
			clientConf := consts.DefaultClientConfig

			fooPort := f.AllocPort()
			fooServer := streamserver.New(streamserver.TCP, streamserver.WithBindPort(fooPort), streamserver.WithRespContent([]byte("foo")))
			f.RunServer("", fooServer)

			barPort := f.AllocPort()
			barServer := streamserver.New(streamserver.TCP, streamserver.WithBindPort(barPort), streamserver.WithRespContent([]byte("bar")))
			f.RunServer("", barServer)

			remotePort := f.AllocPort()
			clientConf += fmt.Sprintf(`
			[foo]
			type = tcp
			local_port = %d
			remote_port = %d
			group = test
			group_key = 123
			health_check_type = tcp
			health_check_interval_s = 1

			[bar]
			type = tcp
			local_port = %d
			remote_port = %d
			group = test
			group_key = 123
			health_check_type = tcp
			health_check_interval_s = 1
			`, fooPort, remotePort, barPort, remotePort)

			f.RunProcesses([]string{serverConf}, []string{clientConf})

			// check foo and bar is ok
			results := []string{}
			for i := 0; i < 10; i++ {
				framework.NewRequestExpect(f).Port(remotePort).Ensure(validateFooBarResponse, func(resp *request.Response) bool {
					results = append(results, string(resp.Content))
					return true
				})
			}
			framework.ExpectContainElements(results, []string{"foo", "bar"})

			// close bar server, check foo is ok
			barServer.Close()
			time.Sleep(2 * time.Second)
			for i := 0; i < 10; i++ {
				framework.NewRequestExpect(f).Port(remotePort).ExpectResp([]byte("foo")).Ensure()
			}

			// resume bar server, check foo and bar is ok
			f.RunServer("", barServer)
			time.Sleep(2 * time.Second)
			results = []string{}
			for i := 0; i < 10; i++ {
				framework.NewRequestExpect(f).Port(remotePort).Ensure(validateFooBarResponse, func(resp *request.Response) bool {
					results = append(results, string(resp.Content))
					return true
				})
			}
			framework.ExpectContainElements(results, []string{"foo", "bar"})
		})

		ginkgo.It("HTTP", func() {
			vhostPort := f.AllocPort()
			serverConf := consts.DefaultServerConfig + fmt.Sprintf(`
			vhost_http_port = %d
			`, vhostPort)
			clientConf := consts.DefaultClientConfig

			fooPort := f.AllocPort()
			fooServer := newHTTPServer(fooPort, "foo")
			f.RunServer("", fooServer)

			barPort := f.AllocPort()
			barServer := newHTTPServer(barPort, "bar")
			f.RunServer("", barServer)

			clientConf += fmt.Sprintf(`
			[foo]
			type = http
			local_port = %d
			custom_domains = example.com
			group = test
			group_key = 123
			health_check_type = http
			health_check_interval_s = 1
			health_check_url = /healthz

			[bar]
			type = http
			local_port = %d
			custom_domains = example.com
			group = test
			group_key = 123
			health_check_type = http
			health_check_interval_s = 1
			health_check_url = /healthz
			`, fooPort, barPort)

			f.RunProcesses([]string{serverConf}, []string{clientConf})

			// send first HTTP request
			var contents []string
			framework.NewRequestExpect(f).Port(vhostPort).
				RequestModify(func(r *request.Request) {
					r.HTTP().HTTPHost("example.com")
				}).
				Ensure(func(resp *request.Response) bool {
					contents = append(contents, string(resp.Content))
					return true
				})

			// send second HTTP request, should be forwarded to another service
			framework.NewRequestExpect(f).Port(vhostPort).
				RequestModify(func(r *request.Request) {
					r.HTTP().HTTPHost("example.com")
				}).
				Ensure(func(resp *request.Response) bool {
					contents = append(contents, string(resp.Content))
					return true
				})

			framework.ExpectContainElements(contents, []string{"foo", "bar"})

			// check foo and bar is ok
			results := doFooBarHTTPRequest(vhostPort, "example.com")
			framework.ExpectContainElements(results, []string{"foo", "bar"})

			// close bar server, check foo is ok
			barServer.Close()
			time.Sleep(2 * time.Second)
			results = doFooBarHTTPRequest(vhostPort, "example.com")
			framework.ExpectContainElements(results, []string{"foo"})
			framework.ExpectNotContainElements(results, []string{"bar"})

			// resume bar server, check foo and bar is ok
			f.RunServer("", barServer)
			time.Sleep(2 * time.Second)
			results = doFooBarHTTPRequest(vhostPort, "example.com")
			framework.ExpectContainElements(results, []string{"foo", "bar"})
		})
	})
})