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

140 lines
2.1 KiB
Go
Raw Normal View History

2021-06-18 16:48:36 +08:00
package streamserver
import (
2021-06-21 19:27:26 +08:00
"bufio"
"fmt"
2021-06-21 19:27:26 +08:00
"io"
"net"
libnet "github.com/fatedier/frp/pkg/util/net"
"github.com/fatedier/frp/test/e2e/pkg/rpc"
)
2021-06-18 16:48:36 +08:00
type Type string
const (
2021-06-18 16:48:36 +08:00
TCP Type = "tcp"
UDP Type = "udp"
Unix Type = "unix"
)
type Server struct {
2021-06-18 16:48:36 +08:00
netType Type
bindAddr string
bindPort int
respContent []byte
2021-06-21 19:27:26 +08:00
handler func(net.Conn)
l net.Listener
}
type Option func(*Server) *Server
2021-06-18 16:48:36 +08:00
func New(netType Type, options ...Option) *Server {
s := &Server{
netType: netType,
bindAddr: "127.0.0.1",
}
2021-06-21 19:27:26 +08:00
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
}
}
2021-06-21 19:27:26 +08:00
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
}
2021-06-21 19:27:26 +08:00
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()
2021-06-21 19:27:26 +08:00
var reader io.Reader = c
if s.netType == UDP {
reader = bufio.NewReader(c)
}
for {
2021-06-21 19:27:26 +08:00
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
}