frp/test/e2e/pkg/request/request.go

140 lines
2.6 KiB
Go
Raw Normal View History

2020-06-02 22:48:55 +08:00
package request
import (
"fmt"
"net"
"time"
2021-03-31 16:57:39 +08:00
libnet "github.com/fatedier/golib/net"
2020-06-02 22:48:55 +08:00
)
2021-03-31 16:57:39 +08:00
type Request struct {
protocol string
addr string
port int
body []byte
timeout time.Duration
proxyURL string
proxyHost string
}
func New() *Request {
return &Request{
protocol: "tcp",
}
}
func (r *Request) Protocol(protocol string) *Request {
r.protocol = protocol
return r
}
func (r *Request) TCP() *Request {
r.protocol = "tcp"
return r
}
func (r *Request) UDP() *Request {
r.protocol = "udp"
return r
}
func (r *Request) Proxy(url, host string) *Request {
r.proxyURL = url
r.proxyHost = host
return r
}
func (r *Request) Addr(addr string) *Request {
r.addr = addr
return r
}
func (r *Request) Port(port int) *Request {
r.port = port
return r
}
func (r *Request) Timeout(timeout time.Duration) *Request {
r.timeout = timeout
return r
}
func (r *Request) Body(content []byte) *Request {
r.body = content
return r
}
func (r *Request) Do() ([]byte, error) {
var (
conn net.Conn
err error
)
if len(r.proxyURL) > 0 {
if r.protocol != "tcp" {
return nil, fmt.Errorf("only tcp protocol is allowed for proxy")
}
conn, err = libnet.DialTcpByProxy(r.proxyURL, r.proxyHost)
if err != nil {
return nil, err
}
} else {
if r.addr == "" {
r.addr = fmt.Sprintf("127.0.0.1:%d", r.port)
}
switch r.protocol {
case "tcp":
conn, err = net.Dial("tcp", r.addr)
case "udp":
conn, err = net.Dial("udp", r.addr)
default:
return nil, fmt.Errorf("invalid protocol")
}
if err != nil {
return nil, err
}
}
defer conn.Close()
if r.timeout > 0 {
conn.SetDeadline(time.Now().Add(r.timeout))
}
return sendRequestByConn(conn, r.body)
}
func SendTCPRequest(port int, content []byte, timeout time.Duration) ([]byte, error) {
2020-06-02 22:48:55 +08:00
c, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", port))
if err != nil {
2021-03-31 16:57:39 +08:00
return nil, fmt.Errorf("connect to tcp server error: %v", err)
2020-06-02 22:48:55 +08:00
}
defer c.Close()
c.SetDeadline(time.Now().Add(timeout))
2020-09-07 14:57:23 +08:00
return sendRequestByConn(c, content)
2020-06-02 22:48:55 +08:00
}
2021-03-31 16:57:39 +08:00
func SendUDPRequest(port int, content []byte, timeout time.Duration) ([]byte, error) {
2020-09-07 14:57:23 +08:00
c, err := net.Dial("udp", fmt.Sprintf("127.0.0.1:%d", port))
if err != nil {
2021-03-31 16:57:39 +08:00
return nil, fmt.Errorf("connect to udp server error: %v", err)
2020-09-07 14:57:23 +08:00
}
defer c.Close()
c.SetDeadline(time.Now().Add(timeout))
return sendRequestByConn(c, content)
}
2021-03-31 16:57:39 +08:00
func sendRequestByConn(c net.Conn, content []byte) ([]byte, error) {
_, err := c.Write(content)
if err != nil {
return nil, fmt.Errorf("write error: %v", err)
}
2020-06-02 22:48:55 +08:00
buf := make([]byte, 2048)
2020-09-07 14:57:23 +08:00
n, err := c.Read(buf)
if err != nil {
2021-03-31 16:57:39 +08:00
return nil, fmt.Errorf("read error: %v", err)
2020-06-02 22:48:55 +08:00
}
2021-03-31 16:57:39 +08:00
return buf[:n], nil
2020-06-02 22:48:55 +08:00
}