add proxy

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
 }
 

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
-}

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
+}