frp/test/e2e/framework/framework.go

package framework

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
	"text/template"

	"github.com/fatedier/frp/test/e2e/pkg/port"
	"github.com/fatedier/frp/test/e2e/pkg/process"

	"github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/config"
)

type Options struct {
	TotalParallelNode int
	CurrentNodeIndex  int
	FromPortIndex     int
	ToPortIndex       int
}

type Framework struct {
	TempDirectory string
	UsedPorts     map[string]int

	portAllocator *port.Allocator

	// Multiple mock servers used for e2e testing.
	mockServers *MockServers

	// To make sure that this framework cleans up after itself, no matter what,
	// we install a Cleanup action before each test and clear it after.  If we
	// should abort, the AfterSuite hook should run all Cleanup actions.
	cleanupHandle CleanupActionHandle

	// beforeEachStarted indicates that BeforeEach has started
	beforeEachStarted bool

	serverConfPaths []string
	serverProcesses []*process.Process
	clientConfPaths []string
	clientProcesses []*process.Process
}

func NewDefaultFramework() *Framework {
	options := Options{
		TotalParallelNode: config.GinkgoConfig.ParallelTotal,
		CurrentNodeIndex:  config.GinkgoConfig.ParallelNode,
		FromPortIndex:     20000,
		ToPortIndex:       50000,
	}
	return NewFramework(options)
}

func NewFramework(opt Options) *Framework {
	f := &Framework{
		portAllocator: port.NewAllocator(opt.FromPortIndex, opt.ToPortIndex, opt.TotalParallelNode, opt.CurrentNodeIndex-1),
	}
	f.mockServers = NewMockServers(f.portAllocator)
	if err := f.mockServers.Run(); err != nil {
		Failf("%v", err)
	}

	ginkgo.BeforeEach(f.BeforeEach)
	ginkgo.AfterEach(f.AfterEach)
	return f
}

// BeforeEach create a temp directory.
func (f *Framework) BeforeEach() {
	f.beforeEachStarted = true

	f.cleanupHandle = AddCleanupAction(f.AfterEach)

	dir, err := ioutil.TempDir(os.TempDir(), "frpe2e-test-*")
	ExpectNoError(err)
	f.TempDirectory = dir
}

func (f *Framework) AfterEach() {
	if !f.beforeEachStarted {
		return
	}

	RemoveCleanupAction(f.cleanupHandle)

	os.RemoveAll(f.TempDirectory)
	f.TempDirectory = ""
	f.UsedPorts = nil
	f.serverConfPaths = nil
	f.clientConfPaths = nil
	for _, p := range f.serverProcesses {
		p.Stop()
	}
	for _, p := range f.clientProcesses {
		p.Stop()
	}
	f.serverProcesses = nil
	f.clientProcesses = nil
}

var portRegex = regexp.MustCompile(`{{ \.Port.*? }}`)

// RenderPortsTemplate render templates with ports.
//
// Local: {{ .Port1 }}
// Target: {{ .Port2 }}
//
// return rendered content and all allocated ports.
func (f *Framework) genPortsFromTemplates(templates []string) (ports map[string]int, err error) {
	ports = make(map[string]int)
	for _, t := range templates {
		arrs := portRegex.FindAllString(t, -1)
		for _, str := range arrs {
			str = strings.TrimPrefix(str, "{{ .")
			str = strings.TrimSuffix(str, " }}")
			str = strings.TrimSpace(str)
			ports[str] = 0
		}
	}
	defer func() {
		if err != nil {
			for _, port := range ports {
				f.portAllocator.Release(port)
			}
		}
	}()

	for name := range ports {
		port := f.portAllocator.Get()
		if port <= 0 {
			return nil, fmt.Errorf("can't allocate port")
		}
		ports[name] = port
	}
	return

}

func (f *Framework) RenderTemplates(templates []string) (outs []string, ports map[string]int, err error) {
	ports, err = f.genPortsFromTemplates(templates)
	if err != nil {
		return
	}

	params := f.mockServers.GetTemplateParams()
	for name, port := range ports {
		params[name] = port
	}

	for _, t := range templates {
		tmpl, err := template.New("").Parse(t)
		if err != nil {
			return nil, nil, err
		}
		buffer := bytes.NewBuffer(nil)
		if err = tmpl.Execute(buffer, params); err != nil {
			return nil, nil, err
		}
		outs = append(outs, buffer.String())
	}
	return
}