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

246 lines
4.5 KiB
Go
Raw Normal View History

2020-06-02 22:48:55 +08:00
package request
import (
2021-06-21 19:27:26 +08:00
"bufio"
2021-06-18 16:48:36 +08:00
"bytes"
2021-08-02 13:07:28 +08:00
"crypto/tls"
2020-06-02 22:48:55 +08:00
"fmt"
2021-06-18 16:48:36 +08:00
"io"
2020-06-02 22:48:55 +08:00
"net"
2021-06-18 16:48:36 +08:00
"net/http"
"net/url"
"strconv"
2020-06-02 22:48:55 +08:00
"time"
2021-03-31 16:57:39 +08:00
"github.com/fatedier/frp/test/e2e/pkg/rpc"
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 {
2021-06-18 16:48:36 +08:00
protocol string
// for all protocol
addr string
port int
body []byte
timeout time.Duration
2021-08-02 13:07:28 +08:00
// for http or https
method string
host string
path string
headers map[string]string
tlsConfig *tls.Config
2021-06-18 16:48:36 +08:00
proxyURL string
2021-03-31 16:57:39 +08:00
}
func New() *Request {
return &Request{
protocol: "tcp",
2021-06-18 16:48:36 +08:00
addr: "127.0.0.1",
method: "GET",
path: "/",
2021-03-31 16:57:39 +08:00
}
}
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
}
2021-06-18 16:48:36 +08:00
func (r *Request) HTTP() *Request {
r.protocol = "http"
return r
}
2021-08-02 13:07:28 +08:00
func (r *Request) HTTPS() *Request {
r.protocol = "https"
return r
}
2021-06-18 16:48:36 +08:00
func (r *Request) Proxy(url string) *Request {
2021-03-31 16:57:39 +08:00
r.proxyURL = url
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
}
2021-06-18 16:48:36 +08:00
func (r *Request) HTTPParams(method, host, path string, headers map[string]string) *Request {
r.method = method
r.host = host
r.path = path
r.headers = headers
return r
}
func (r *Request) HTTPHost(host string) *Request {
r.host = host
return r
}
func (r *Request) HTTPPath(path string) *Request {
r.path = path
return r
}
func (r *Request) HTTPHeaders(headers map[string]string) *Request {
r.headers = headers
return r
}
2021-08-02 13:07:28 +08:00
func (r *Request) TLSConfig(tlsConfig *tls.Config) *Request {
r.tlsConfig = tlsConfig
return r
}
2021-03-31 16:57:39 +08:00
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
}
2021-06-18 16:48:36 +08:00
func (r *Request) Do() (*Response, error) {
2021-03-31 16:57:39 +08:00
var (
conn net.Conn
err error
)
2021-06-18 16:48:36 +08:00
addr := net.JoinHostPort(r.addr, strconv.Itoa(r.port))
2021-08-02 13:07:28 +08:00
// for protocol http and https
if r.protocol == "http" || r.protocol == "https" {
return r.sendHTTPRequest(r.method, fmt.Sprintf("%s://%s%s", r.protocol, addr, r.path),
r.host, r.headers, r.proxyURL, r.body, r.tlsConfig)
2021-06-18 16:48:36 +08:00
}
// for protocol tcp and udp
2021-03-31 16:57:39 +08:00
if len(r.proxyURL) > 0 {
if r.protocol != "tcp" {
return nil, fmt.Errorf("only tcp protocol is allowed for proxy")
}
2021-06-18 16:48:36 +08:00
conn, err = libnet.DialTcpByProxy(r.proxyURL, addr)
2021-03-31 16:57:39 +08:00
if err != nil {
return nil, err
}
} else {
switch r.protocol {
case "tcp":
2021-06-18 16:48:36 +08:00
conn, err = net.Dial("tcp", addr)
2021-03-31 16:57:39 +08:00
case "udp":
2021-06-18 16:48:36 +08:00
conn, err = net.Dial("udp", addr)
2021-03-31 16:57:39 +08:00
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))
}
2021-06-21 19:27:26 +08:00
buf, err := r.sendRequestByConn(conn, r.body)
2020-06-02 22:48:55 +08:00
if err != nil {
2021-06-18 16:48:36 +08:00
return nil, err
2020-06-02 22:48:55 +08:00
}
2021-06-18 16:48:36 +08:00
return &Response{Content: buf}, nil
}
2020-06-02 22:48:55 +08:00
2021-06-18 16:48:36 +08:00
type Response struct {
Code int
Header http.Header
Content []byte
2020-06-02 22:48:55 +08:00
}
2021-08-02 13:07:28 +08:00
func (r *Request) sendHTTPRequest(method, urlstr string, host string, headers map[string]string,
proxy string, body []byte, tlsConfig *tls.Config,
) (*Response, error) {
2021-06-18 16:48:36 +08:00
var inBody io.Reader
if len(body) != 0 {
inBody = bytes.NewReader(body)
}
req, err := http.NewRequest(method, urlstr, inBody)
2020-09-07 14:57:23 +08:00
if err != nil {
2021-06-18 16:48:36 +08:00
return nil, err
}
if host != "" {
req.Host = host
}
for k, v := range headers {
req.Header.Set(k, v)
}
tr := &http.Transport{
DialContext: (&net.Dialer{
Timeout: time.Second,
KeepAlive: 30 * time.Second,
DualStack: true,
}).DialContext,
MaxIdleConns: 100,
IdleConnTimeout: 90 * time.Second,
TLSHandshakeTimeout: 10 * time.Second,
ExpectContinueTimeout: 1 * time.Second,
2021-08-02 13:07:28 +08:00
TLSClientConfig: tlsConfig,
2021-06-18 16:48:36 +08:00
}
if len(proxy) != 0 {
tr.Proxy = func(req *http.Request) (*url.URL, error) {
return url.Parse(proxy)
}
}
client := http.Client{Transport: tr}
resp, err := client.Do(req)
if err != nil {
return nil, err
2020-09-07 14:57:23 +08:00
}
2021-06-18 16:48:36 +08:00
ret := &Response{Code: resp.StatusCode, Header: resp.Header}
buf, err := io.ReadAll(resp.Body)
2021-06-18 16:48:36 +08:00
if err != nil {
return nil, err
}
ret.Content = buf
return ret, nil
2020-09-07 14:57:23 +08:00
}
2021-06-21 19:27:26 +08:00
func (r *Request) sendRequestByConn(c net.Conn, content []byte) ([]byte, error) {
_, err := rpc.WriteBytes(c, content)
2021-03-31 16:57:39 +08:00
if err != nil {
return nil, fmt.Errorf("write error: %v", err)
}
2020-06-02 22:48:55 +08:00
2021-06-21 19:27:26 +08:00
var reader io.Reader = c
if r.protocol == "udp" {
reader = bufio.NewReader(c)
}
buf, err := rpc.ReadBytes(reader)
2020-09-07 14:57:23 +08:00
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
}
return buf, nil
2020-06-02 22:48:55 +08:00
}