From ab60c985a6fcc94d40c4f3b635f35e718598fcf1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BD=B1=E8=88=9E=E8=80=85?= <shadow1ng@qq.com>
Date: Mon, 9 May 2022 12:08:29 +0800
Subject: [PATCH] Update client.go

---
 WebScan/lib/client.go | 33 ++++++++++++++++-----------------
 1 file changed, 16 insertions(+), 17 deletions(-)

diff --git a/WebScan/lib/client.go b/WebScan/lib/client.go
index f34aa31..f9ed985 100644
--- a/WebScan/lib/client.go
+++ b/WebScan/lib/client.go
@@ -31,28 +31,13 @@ func Inithttp(PocInfo common.PocInfo) {
 
 func InitHttpClient(ThreadsNum int, DownProxy string, Timeout time.Duration) error {
 	type DialContext = func(ctx context.Context, network, addr string) (net.Conn, error)
-	var dialContext DialContext
 	dialer := &net.Dialer{
 		Timeout:   dialTimout,
 		KeepAlive: keepAlive,
 	}
-	if common.Socks5Proxy != "" {
-		dialSocksProxy, err := common.Socks5Dailer(dialer)
-		if err != nil {
-			return err
-		}
-		if contextDialer, ok := dialSocksProxy.(proxy.ContextDialer); ok {
-			dialContext = contextDialer.DialContext
-		} else {
-			return errors.New("Failed type assertion to DialContext")
-		}
-	}else {
-		dialContext = dialer.DialContext
-	}
-
 
 	tr := &http.Transport{
-		DialContext:         dialContext,
+		DialContext:         dialer.DialContext,
 		MaxConnsPerHost:     5,
 		MaxIdleConns:        0,
 		MaxIdleConnsPerHost: ThreadsNum * 2,
@@ -61,7 +46,18 @@ func InitHttpClient(ThreadsNum int, DownProxy string, Timeout time.Duration) err
 		TLSHandshakeTimeout: 5 * time.Second,
 		DisableKeepAlives:   false,
 	}
-	if DownProxy != "" {
+
+	if common.Socks5Proxy != "" {
+		dialSocksProxy, err := common.Socks5Dailer(dialer)
+		if err != nil {
+			return err
+		}
+		if contextDialer, ok := dialSocksProxy.(proxy.ContextDialer); ok {
+			tr.DialContext = contextDialer.DialContext
+		} else {
+			return errors.New("Failed type assertion to DialContext")
+		}
+	}else if DownProxy != "" {
 		if DownProxy == "1" {
 			DownProxy = "http://127.0.0.1:8080"
 		} else if DownProxy == "2" {
@@ -69,6 +65,9 @@ func InitHttpClient(ThreadsNum int, DownProxy string, Timeout time.Duration) err
 		} else if !strings.Contains(DownProxy, "://") {
 			DownProxy = "http://127.0.0.1:" + DownProxy
 		}
+		if !strings.HasPrefix(DownProxy,"socks") && !strings.HasPrefix(DownProxy,"http") {
+			return errors.New("no support this proxy")
+		}
 		u, err := url.Parse(DownProxy)
 		if err != nil {
 			return err