frp/test/e2e/mock/echoserver/echoserver.go

package echoserver

import (
	"fmt"
	"net"
	"strings"

	fnet "github.com/fatedier/frp/utils/net"
)

type ServerType string

const (
	TCP  ServerType = "tcp"
	UDP  ServerType = "udp"
	Unix ServerType = "unix"
)

type Options struct {
	Type              ServerType
	BindAddr          string
	BindPort          int32
	RepeatNum         int
	SpecifiedResponse string
}

type Server struct {
	opt Options

	l net.Listener
}

func New(opt Options) *Server {
	if opt.Type == "" {
		opt.Type = TCP
	}
	if opt.BindAddr == "" {
		opt.BindAddr = "127.0.0.1"
	}
	if opt.RepeatNum <= 0 {
		opt.RepeatNum = 1
	}
	return &Server{
		opt: opt,
	}
}

func (s *Server) GetOptions() Options {
	return s.opt
}

func (s *Server) Run() error {
	if err := s.initListener(); err != nil {
		return err
	}

	go func() {
		for {
			c, err := s.l.Accept()
			if err != nil {
				return
			}
			go s.handle(c)
		}
	}()
	return nil
}

func (s *Server) Close() error {
	if s.l != nil {
		return s.l.Close()
	}
	return nil
}

func (s *Server) initListener() (err error) {
	switch s.opt.Type {
	case TCP:
		s.l, err = net.Listen("tcp", fmt.Sprintf("%s:%d", s.opt.BindAddr, s.opt.BindPort))
	case UDP:
		s.l, err = fnet.ListenUDP(s.opt.BindAddr, int(s.opt.BindPort))
	case Unix:
		s.l, err = net.Listen("unix", s.opt.BindAddr)
	default:
		return fmt.Errorf("unknown server type: %s", s.opt.Type)
	}
	if err != nil {
		return
	}
	return nil
}

func (s *Server) handle(c net.Conn) {
	defer c.Close()

	buf := make([]byte, 2048)
	for {
		n, err := c.Read(buf)
		if err != nil {
			return
		}

		var response string
		if len(s.opt.SpecifiedResponse) > 0 {
			response = s.opt.SpecifiedResponse
		} else {
			response = strings.Repeat(string(buf[:n]), s.opt.RepeatNum)
		}
		c.Write([]byte(response))
	}
}