并发优化

Author: SummerSec

cmd/commons/core/options.go

@@ -56,7 +56,7 @@ func (o Options) toString() interface{} {
 func ParseOptions() *Options {
 	options := &Options{}
 	flag.IntVar(&options.Mode, "m", 0, "debug mode off (debug mode = 1) default mode = 0")
-	flag.IntVar(&options.Thread, "t", 1, "threads number ")
+	flag.IntVar(&options.Thread, "t", 50, "threads number ")
 	flag.StringVar(&options.File, "f", "", "file to read example: -file=test.txt  http(s)://host:port/path/ (notes: The last line must be empty)")
 	flag.StringVar(&options.Url, "u", "", "url to read example: -url=http://www.baidu.com:80/")
 	flag.StringVar(&options.Proxy, "proxy", "", "proxy example: -proxy=http(socks5)://127.0.0.1:8080 ")

cmd/commons/core/runner.go

@@ -7,6 +7,7 @@ import (
 	"github.com/fatih/structs"
 	log "github.com/sirupsen/logrus"
 	"net/url"
+	"sync"
 )
 
 type Runner struct {
@@ -47,6 +48,7 @@ func (r *Runner) Run() {
 	log.Debugln("URLs: ", urls)
 	var i = 0
 	k := r.options.Thread
+	var wg sync.WaitGroup
 	hashmap := structs.Map(&r.options)
 	for i < len(urls) {
 		for t := 0; t < k; t++ {
@@ -57,9 +59,15 @@ func (r *Runner) Run() {
 				log.Debugln("Running attack on: ", urls[i])
 				// 通道通信 发送url 并且 i++
 				c := make(chan int)
-				go Start(urls[i], hashmap, i, c) // Start k goroutines
+				wg.Add(1)
+				go func() {
+					log.Debugf("Running go func() %d", t)
+					Start(urls[i], hashmap, i, c) // Start k goroutines
+					wg.Done()
+				}()
 				i = <-c
 			} else {
+				wg.Wait()
 				i++
 				break
 			}

cmd/test/threads.go

@@ -0,0 +1,27 @@
+package main
+
+import (
+	"fmt"
+	"strconv"
+	"sync"
+)
+
+func hello(str int) {
+	i := strconv.Itoa(str)
+	fmt.Println("Hello, world." + i)
+
+}
+
+func main() {
+	var wg sync.WaitGroup
+
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		fmt.Println("Starting goroutine", i)
+		go func() {
+			hello(i)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+}