From e2626aa374cbeea1e5eb4288a95669fed064dda4 Mon Sep 17 00:00:00 2001 From: cnlh Date: Wed, 1 Apr 2020 02:39:36 +0800 Subject: [PATCH] add proxy --- benchmark.go | 10 ++++-- connection.go | 57 ++------------------------------ proxy.go | 90 +++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 100 insertions(+), 57 deletions(-) create mode 100644 proxy.go diff --git a/benchmark.go b/benchmark.go index 2ba5471..0ca0c14 100644 --- a/benchmark.go +++ b/benchmark.go @@ -31,7 +31,12 @@ type benchmark struct { // Start benchmark with the param has setting func (pf *benchmark) Run() { fmt.Printf("Running %d test @ %s by %d connections\n", pf.reqNum, pf.target, pf.connectionNum) - var err error + dialer, err := NewProxyConn(pf.proxy) + if err != nil { + fmt.Println(err) + os.Exit(0) + } + pf.startTime = time.Now() pf.wg.Add(pf.connectionNum) for i := 0; i < pf.connectionNum; i++ { @@ -44,7 +49,7 @@ func (pf *benchmark) Run() { remoteAddr: pf.target, schema: pf.schema, buf: make([]byte, 65535), - proxy: pf.proxy, + dialer: dialer, } go func() { if err = rc.Start(); err != nil { @@ -57,6 +62,7 @@ func (pf *benchmark) Run() { } pf.wg.Wait() pf.endTime = time.Now() + return } diff --git a/connection.go b/connection.go index 8fbf264..9883357 100644 --- a/connection.go +++ b/connection.go @@ -5,15 +5,10 @@ package main import ( - "bufio" "bytes" "crypto/tls" - "errors" - "golang.org/x/net/proxy" "io" "net" - "net/http" - "net/url" "strconv" "sync/atomic" "time" @@ -41,7 +36,7 @@ type ReqConn struct { remoteAddr string schema string buf []byte - proxy string + dialer ProxyConn } // Connect to the server, http and socks5 proxy support @@ -50,29 +45,7 @@ func (rc *ReqConn) dial() error { if rc.conn != nil { rc.conn.Close() } - var err error - var conn net.Conn - if rc.proxy != "" { - var u *url.URL - var d proxy.Dialer - u, err = url.Parse(rc.proxy) - if err != nil { - return err - } - switch u.Scheme { - case "socks5": - d, err = proxy.FromURL(u, nil) - if err != nil { - return err - } - conn, err = d.Dial("tcp", rc.remoteAddr) - default: - conn, err = NewHttpProxyConn(u, rc.remoteAddr) - } - - } else { - conn, err = net.DialTimeout("tcp", rc.remoteAddr, time.Millisecond*time.Duration(rc.timeout)) - } + conn, err := rc.dialer.Dial("tcp", rc.remoteAddr, time.Millisecond*time.Duration(rc.timeout)) if err != nil { return err } @@ -157,29 +130,3 @@ re: func Bytes2str(b []byte) string { return *(*string)(unsafe.Pointer(&b)) } - -// Create a connection by http proxy server -func NewHttpProxyConn(url *url.URL, remoteAddr string) (net.Conn, error) { - req, err := http.NewRequest("CONNECT", "http://"+remoteAddr, nil) - if err != nil { - return nil, err - } - password, _ := url.User.Password() - req.SetBasicAuth(url.User.Username(), password) - proxyConn, err := net.Dial("tcp", url.Host) - if err != nil { - return nil, err - } - if err := req.Write(proxyConn); err != nil { - return nil, err - } - res, err := http.ReadResponse(bufio.NewReader(proxyConn), req) - if err != nil { - return nil, err - } - _ = res.Body.Close() - if res.StatusCode != 200 { - return nil, errors.New("Proxy error " + res.Status) - } - return proxyConn, nil -} diff --git a/proxy.go b/proxy.go new file mode 100644 index 0000000..b5c02f6 --- /dev/null +++ b/proxy.go @@ -0,0 +1,90 @@ +// Copyright 2020 The benchmark. All rights reserved. +// Use of this source code is governed by a MIT-style +// license that can be found in the LICENSE file. + +package main + +import ( + "bufio" + "errors" + "golang.org/x/net/proxy" + "net" + "net/http" + "net/url" + "time" +) + +// Return based on proxy url +func NewProxyConn(proxyUrl string) (ProxyConn, error) { + u, err := url.Parse(proxyUrl) + if err != nil { + return nil, err + } + switch u.Scheme { + case "socks5": + return &Socks5Client{u}, nil + case "http": + return &HttpClient{u}, nil + default: + return DefaultClient{}, nil + } +} + +// ProxyConn is used to define the proxy +type ProxyConn interface { + Dial(network string, address string, timeout time.Duration) (net.Conn, error) +} + +// DefaultClient is used to implement a proxy in default +type DefaultClient struct { +} + +// socks5 implementation of ProxyConn +func (s5 DefaultClient) Dial(network string, address string, timeout time.Duration) (net.Conn, error) { + return net.DialTimeout(network, address, timeout) +} + +// Socks5Client is used to implement a proxy in socks5 +type Socks5Client struct { + proxyUrl *url.URL +} + +// Socks5 implementation of ProxyConn +func (s5 *Socks5Client) Dial(network string, address string, timeout time.Duration) (net.Conn, error) { + d, err := proxy.FromURL(s5.proxyUrl, nil) + if err != nil { + return nil, err + } + return d.Dial(network, address) +} + +// Socks5Client is used to implement a proxy in http +type HttpClient struct { + proxyUrl *url.URL +} + +// Http implementation of ProxyConn +func (hc *HttpClient) Dial(network string, address string, timeout time.Duration) (net.Conn, error) { + req, err := http.NewRequest("CONNECT", "http://"+address, nil) + if err != nil { + return nil, err + } + password, _ := hc.proxyUrl.User.Password() + req.SetBasicAuth(hc.proxyUrl.User.Username(), password) + proxyConn, err := net.Dial("tcp", hc.proxyUrl.Host) + if err != nil { + return nil, err + } + if err := req.Write(proxyConn); err != nil { + return nil, err + } + res, err := http.ReadResponse(bufio.NewReader(proxyConn), req) + if err != nil { + return nil, err + } + _ = res.Body.Close() + if res.StatusCode != 200 { + return nil, errors.New("Proxy error " + res.Status) + } + return proxyConn, nil +}