frp/test/e2e/framework/mockservers.go

package framework

import (
	"fmt"
	"os"

	"github.com/fatedier/frp/test/e2e/mock/echoserver"
	"github.com/fatedier/frp/test/e2e/pkg/port"
)

const (
	TCPEchoServerPort = "TCPEchoServerPort"
	UDPEchoServerPort = "UDPEchoServerPort"
	UDSEchoServerAddr = "UDSEchoServerAddr"
)

type MockServers struct {
	tcpEchoServer *echoserver.Server
	udpEchoServer *echoserver.Server
	udsEchoServer *echoserver.Server
}

func NewMockServers(portAllocator *port.Allocator) *MockServers {
	s := &MockServers{}
	tcpPort := portAllocator.Get()
	udpPort := portAllocator.Get()
	s.tcpEchoServer = echoserver.New(echoserver.Options{
		Type:      echoserver.TCP,
		BindAddr:  "127.0.0.1",
		BindPort:  int32(tcpPort),
		RepeatNum: 1,
	})
	s.udpEchoServer = echoserver.New(echoserver.Options{
		Type:      echoserver.UDP,
		BindAddr:  "127.0.0.1",
		BindPort:  int32(udpPort),
		RepeatNum: 1,
	})

	udsIndex := portAllocator.Get()
	udsAddr := fmt.Sprintf("%s/frp_echo_server_%d.sock", os.TempDir(), udsIndex)
	os.Remove(udsAddr)
	s.udsEchoServer = echoserver.New(echoserver.Options{
		Type:      echoserver.Unix,
		BindAddr:  udsAddr,
		RepeatNum: 1,
	})
	return s
}

func (m *MockServers) Run() error {
	if err := m.tcpEchoServer.Run(); err != nil {
		return err
	}
	if err := m.udpEchoServer.Run(); err != nil {
		return err
	}
	if err := m.udsEchoServer.Run(); err != nil {
		return err
	}
	return nil
}

func (m *MockServers) Close() {
	m.tcpEchoServer.Close()
	m.udpEchoServer.Close()
	m.udsEchoServer.Close()
	os.Remove(m.udsEchoServer.GetOptions().BindAddr)
}

func (m *MockServers) GetTemplateParams() map[string]interface{} {
	ret := make(map[string]interface{})
	ret[TCPEchoServerPort] = m.tcpEchoServer.GetOptions().BindPort
	ret[UDPEchoServerPort] = m.udpEchoServer.GetOptions().BindPort
	ret[UDSEchoServerAddr] = m.udsEchoServer.GetOptions().BindAddr
	return ret
}

func (m *MockServers) GetParam(key string) interface{} {
	params := m.GetTemplateParams()
	if v, ok := params[key]; ok {
		return v
	}
	return nil
}