frp/test/e2e/mock/server/streamserver/server.go

package streamserver

import (
	"bufio"
	"fmt"
	"io"
	"net"

	libnet "github.com/fatedier/frp/pkg/util/net"
	"github.com/fatedier/frp/test/e2e/pkg/rpc"
)

type Type string

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

type Server struct {
	netType     Type
	bindAddr    string
	bindPort    int
	respContent []byte

	handler func(net.Conn)

	l net.Listener
}

type Option func(*Server) *Server

func New(netType Type, options ...Option) *Server {
	s := &Server{
		netType:  netType,
		bindAddr: "127.0.0.1",
	}
	s.handler = s.handle

	for _, option := range options {
		s = option(s)
	}
	return s
}

func WithBindAddr(addr string) Option {
	return func(s *Server) *Server {
		s.bindAddr = addr
		return s
	}
}

func WithBindPort(port int) Option {
	return func(s *Server) *Server {
		s.bindPort = port
		return s
	}
}

func WithRespContent(content []byte) Option {
	return func(s *Server) *Server {
		s.respContent = content
		return s
	}
}

func WithCustomHandler(handler func(net.Conn)) Option {
	return func(s *Server) *Server {
		s.handler = handler
		return s
	}
}

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.handler(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.netType {
	case TCP:
		s.l, err = net.Listen("tcp", fmt.Sprintf("%s:%d", s.bindAddr, s.bindPort))
	case UDP:
		s.l, err = libnet.ListenUDP(s.bindAddr, s.bindPort)
	case Unix:
		s.l, err = net.Listen("unix", s.bindAddr)
	default:
		return fmt.Errorf("unknown server type: %s", s.netType)
	}
	return err
}

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

	var reader io.Reader = c
	if s.netType == UDP {
		reader = bufio.NewReader(c)
	}
	for {
		buf, err := rpc.ReadBytes(reader)
		if err != nil {
			return
		}

		if len(s.respContent) > 0 {
			buf = s.respContent
		}
		rpc.WriteBytes(c, buf)
	}
}

func (s *Server) BindAddr() string {
	return s.bindAddr
}

func (s *Server) BindPort() int {
	return s.bindPort
}