frp/tests/mock/echo_server.go

121 lines
1.9 KiB
Go
Raw Normal View History

2018-07-11 23:27:47 +08:00
package mock
2016-05-17 19:13:37 +08:00
import (
"fmt"
"io"
"net"
"os"
"syscall"
2016-05-17 19:13:37 +08:00
frpNet "github.com/fatedier/frp/utils/net"
2016-05-17 19:13:37 +08:00
)
2018-12-09 21:56:46 +08:00
type EchoServer struct {
2019-10-12 20:13:12 +08:00
l net.Listener
2016-05-17 19:13:37 +08:00
2018-12-09 21:56:46 +08:00
port int
repeatedNum int
specifyStr string
}
2016-05-17 19:13:37 +08:00
2018-12-09 21:56:46 +08:00
func NewEchoServer(port int, repeatedNum int, specifyStr string) *EchoServer {
if repeatedNum <= 0 {
repeatedNum = 1
}
return &EchoServer{
port: port,
repeatedNum: repeatedNum,
specifyStr: specifyStr,
2016-05-17 19:13:37 +08:00
}
}
2018-12-09 21:56:46 +08:00
func (es *EchoServer) Start() error {
2019-10-12 20:13:12 +08:00
l, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", es.port))
2018-05-23 14:39:12 +08:00
if err != nil {
2018-12-09 21:56:46 +08:00
fmt.Printf("echo server listen error: %v\n", err)
return err
2018-05-23 14:39:12 +08:00
}
2018-12-09 21:56:46 +08:00
es.l = l
2018-05-23 14:39:12 +08:00
2018-12-09 21:56:46 +08:00
go func() {
for {
c, err := l.Accept()
if err != nil {
return
}
go echoWorker(c, es.repeatedNum, es.specifyStr)
2018-05-23 14:39:12 +08:00
}
2018-12-09 21:56:46 +08:00
}()
return nil
}
2018-05-23 14:39:12 +08:00
2018-12-09 21:56:46 +08:00
func (es *EchoServer) Stop() {
es.l.Close()
2018-05-23 14:39:12 +08:00
}
2018-07-11 23:27:47 +08:00
func StartUdpEchoServer(port int) {
l, err := frpNet.ListenUDP("127.0.0.1", port)
2017-03-28 00:27:30 +08:00
if err != nil {
fmt.Printf("udp echo server listen error: %v\n", err)
return
}
for {
c, err := l.Accept()
if err != nil {
fmt.Printf("udp echo server accept error: %v\n", err)
return
}
2018-12-09 21:56:46 +08:00
go echoWorker(c, 1, "")
2017-03-28 00:27:30 +08:00
}
}
2018-07-11 23:27:47 +08:00
func StartUnixDomainServer(unixPath string) {
os.Remove(unixPath)
syscall.Umask(0)
l, err := net.Listen("unix", unixPath)
if err != nil {
fmt.Printf("unix domain server listen error: %v\n", err)
return
}
for {
c, err := l.Accept()
if err != nil {
fmt.Printf("unix domain server accept error: %v\n", err)
return
}
2018-12-09 21:56:46 +08:00
go echoWorker(c, 1, "")
}
}
2018-12-09 21:56:46 +08:00
func echoWorker(c net.Conn, repeatedNum int, specifyStr string) {
2017-12-16 23:59:46 +08:00
buf := make([]byte, 2048)
2016-05-17 19:13:37 +08:00
for {
2017-12-16 23:59:46 +08:00
n, err := c.Read(buf)
2016-05-17 19:13:37 +08:00
if err != nil {
2017-12-16 23:59:46 +08:00
if err == io.EOF {
c.Close()
break
} else {
fmt.Printf("echo server read error: %v\n", err)
return
}
2016-05-17 19:13:37 +08:00
}
2018-12-09 21:56:46 +08:00
if specifyStr != "" {
c.Write([]byte(specifyStr))
} else {
var w []byte
for i := 0; i < repeatedNum; i++ {
w = append(w, buf[:n]...)
2018-05-23 14:39:12 +08:00
}
2018-12-09 21:56:46 +08:00
c.Write(w)
2018-05-23 14:39:12 +08:00
}
}
}