mirror of https://gitee.com/IrisVega/frp.git
112 lines
1.7 KiB
Go
112 lines
1.7 KiB
Go
package echoserver
|
|
|
|
import (
|
|
"fmt"
|
|
"net"
|
|
"strings"
|
|
|
|
fnet "github.com/fatedier/frp/pkg/util/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))
|
|
}
|
|
}
|