feat: add multiple thread

README.md

@@ -21,7 +21,13 @@ git clone https://github.com/chaitin/blazehttp.git && cd blazehttp
 ## 🚀 一键运行
 
 ``` bash
-bash build.sh && ./build/blazehttp -t http://127.0.0.1:8008
+bash build.sh && ./build/blazehttp -t http://127.0.0.1:8080
+```
+
+也可配置 `thread` 参数提高测试速度，默认 2
+
+```
+bash build.sh && ./build/blazehttp -t http://127.0.0.1:8080 -thread 4
 ```
 
 ## 🕹️ 靶机服务

cmd/blazehttp/main.go

@@ -12,6 +12,7 @@ import (
 	"regexp"
 	"sort"
 	"strings"
+	"sync"
 	"time"
 
 	blazehttp "github.com/chaitin/blazehttp/http"
@@ -27,6 +28,7 @@ var (
 	glob              string // use glob expression to select multi files
 	timeout           = 1000 // default 1000 ms
 	mHost             string // modify host header
+	thread            = 2    // send request thread
 	requestPerSession bool   // send request per session
 )
 
@@ -38,6 +40,7 @@ func init() {
 	flag.StringVar(&target, "t", "", "target website, example: http://192.168.0.1:8080")
 	flag.StringVar(&glob, "g", "./testcases/", "glob expression, example: *.http")
 	flag.IntVar(&timeout, "timeout", 1000, "connection timeout, default 1000 ms")
+	flag.IntVar(&thread, "thread", 2, "request thread, default 2")
 	flag.StringVar(&mHost, "H", "", "modify host header")
 	flag.BoolVar(&requestPerSession, "rps", true, "send request per session")
 	flag.Parse()
@@ -141,8 +144,56 @@ func getAllFiles(path string) ([]string, error) {
 	return files, nil
 }
 
+func work(addr string, isHttps bool, blockStatusCode int, f string) (bool, bool, int64, error) {
+	req := new(blazehttp.Request)
+
+	if err := req.ReadFile(f); err != nil {
+		return false, false, 0, fmt.Errorf("read request file: %s error: %s", f, err)
+	}
+	if mHost != "" {
+		// 修改host header
+		req.SetHost(mHost)
+	} else {
+		// 不修改会导致域名备案拦截
+		req.SetHost(addr)
+	}
+
+	if requestPerSession {
+		// one http request one connection
+		req.SetHeader("Connection", "close")
+	}
+
+	req.CalculateContentLength()
+
+	start := time.Now()
+	conn := connect(addr, isHttps, timeout)
+	if conn == nil {
+		return false, false, 0, fmt.Errorf("connect to %s failed", addr)
+	}
+	nWrite, err := req.WriteTo(*conn)
+	if err != nil {
+		return false, false, 0, fmt.Errorf("send request poc: %s length: %d error: %s", f, nWrite, err)
+	}
+
+	rsp := new(blazehttp.Response)
+	if err = rsp.ReadConn(*conn); err != nil {
+		return false, false, 0, fmt.Errorf("read poc file: %s response, error: %s", f, err)
+	}
+	(*conn).Close()
+	isWhite := false // black case
+	if strings.HasSuffix(f, "white") {
+		isWhite = true // white case
+	}
+	isPass := true
+	code := rsp.GetStatusCode()
+	if code == blockStatusCode {
+		isPass = false
+	}
+	return isWhite, isPass, time.Since(start).Nanoseconds(), nil
+}
+
 func main() {
-	// mcl := true
+
 	isHttps := false
 	addr := target
 
@@ -193,74 +244,49 @@ func main() {
 		fmt.Println("目标网站未开启waf")
 		os.Exit(1)
 	}
-	for _, f := range fileList {
-		_ = bar.Add(1)
-		req := new(blazehttp.Request)
-		if err = req.ReadFile(f); err != nil {
-			fmt.Printf("read request file: %s error: %s\n", f, err)
-			continue
-		}
-		if mHost != "" {
-			// 修改host header
-			req.SetHost(mHost)
-		} else {
-			// 不修改会导致域名备案拦截
-			req.SetHost(addr)
-		}
-
-		if requestPerSession {
-			// one http request one connection
-			req.SetHeader("Connection", "close")
-		}
-
-		req.CalculateContentLength()
-
-		start := time.Now()
-		conn := connect(addr, isHttps, timeout)
-		if conn == nil {
-			fmt.Printf("connect to %s failed!\n", addr)
-			continue
-		}
-		nWrite, err := req.WriteTo(*conn)
-		if err != nil {
-			fmt.Printf("send request poc: %s length: %d error: %s", f, nWrite, err)
-			continue
-		}
-
-		rsp := new(blazehttp.Response)
-		if err = rsp.ReadConn(*conn); err != nil {
-			fmt.Printf("read poc file: %s response, error: %s", f, err)
-			continue
-		}
-		elap = append(elap, time.Since(start).Nanoseconds())
-		(*conn).Close()
-		success++
 
-		isWhite := false // black case
-		if strings.HasSuffix(f, "white") {
-			isWhite = true // white case
-		}
+	var mutex sync.Mutex
+	concurrency := thread
+	wg := &sync.WaitGroup{}
+	jobs := make(chan string, len(fileList))
 
-		isPass := true
-		code := rsp.GetStatusCode()
-		if code == blockStatusCode {
-			isPass = false
-		}
-
-		if isWhite {
-			if isPass {
-				TN += 1
-			} else {
-				FP += 1
-			}
-		} else {
-			if isPass {
-				FN += 1
-			} else {
-				TP += 1
+	for i := 0; i < concurrency; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for f := range jobs {
+				isWhite, isPass, time, err := work(addr, isHttps, blockStatusCode, f)
+				mutex.Lock()
+				bar.Add(1)
+				if err != nil {
+					fmt.Printf("%s\n", err)
+					continue
+				}
+				elap = append(elap, time)
+				success++
+				if isWhite {
+					if isPass {
+						TN += 1
+					} else {
+						FP += 1
+					}
+				} else {
+					if isPass {
+						FN += 1
+					} else {
+						TP += 1
+					}
+				}
+				mutex.Unlock()
 			}
-		}
+		}()
+	}
+
+	for _, f := range fileList {
+		jobs <- f
 	}
+	close(jobs)
+	wg.Wait()
 
 	fmt.Printf("总样本数量: %d    成功: %d    错误: %d\n", len(fileList), success, (len(fileList) - success))
 	fmt.Printf("检出率: %.2f%% (恶意样本总数: %d , 正确拦截: %d , 漏报放行: %d)\n", float64(TP)*100/float64(TP+FN), TP+FN, TP, FN)
@@ -273,5 +299,5 @@ func main() {
 	for _, v := range elap {
 		sum += v
 	}
-	fmt.Printf("平均耗时: %.2f毫秒\n", float64(sum)/float64(all)/1000000)
+	fmt.Printf("平均耗时: %.2f 毫秒\n", float64(sum)/float64(all)/1000000)
 }