优化重发逻辑

runner/retry.go

@@ -10,37 +10,20 @@ import (
 func (r *runner) retry(ctx context.Context) {
 	for {
 		// 循环检测超时的队列
-		currentTime := time.Now().Unix()
-		var preSend []string
+		now := time.Now()
 		r.hm.Scan(func(key string, v statusdb.Item) error {
-			// Scan自带锁，不要调用其他r.hm下的函数。。
 			if v.Retry > r.maxRetry {
 				r.hm.Del(key)
 				atomic.AddUint64(&r.faildIndex, 1)
 				return nil
 			}
-			if currentTime-v.Time >= r.timeout {
+			if int64(now.Sub(v.Time)) >= r.timeout {
 				// 重新发送
-				preSend = append(preSend, key)
+				r.sender <- key
 			}
 			return nil
 		})
-		for _, d := range preSend {
-			r.sender <- d
-		}
-		// 延时，map越多延时越大
-		length := r.hm.Length()
-		if length < 100 {
-			length = 500
-		} else if length < 1000 {
-			length = 1500
-		} else if length < 5000 {
-			length = 3500
-		} else if length < 10000 {
-			length = 4500
-		} else {
-			length = 5500
-		}
+		length := 1000
 		time.Sleep(time.Millisecond * time.Duration(length))
 	}
 }

runner/runner.go

@@ -11,6 +11,7 @@ import (
 	"ksubdomain/core/gologger"
 	options2 "ksubdomain/core/options"
 	"ksubdomain/runner/statusdb"
+	"math"
 	"math/rand"
 	"os"
 	"strings"
@@ -36,6 +37,7 @@ type runner struct {
 	ctx             context.Context
 	fisrtloadChanel chan string // 数据加载完毕的chanel
 	startTime       time.Time
+	domains         []string
 }
 
 func GetDeviceConfig() *device.EtherTable {
@@ -66,23 +68,32 @@ func New(options *options2.Options) (*runner, error) {
 	version := pcap.Version()
 	r := new(runner)
 	gologger.Infof(version + "\n")
-	//if options.ListNetwork {
-	//	device.GetIpv4Devices()
-	//	os.Exit(0)
-	//}
+
 	r.options = options
 	r.ether = GetDeviceConfig()
 	r.hm = statusdb.CreateMemoryDB()
 
-	gologger.Infof("Rate:%dpps\n", options.Rate)
 	gologger.Infof("DNS:%s\n", options.Resolvers)
 	r.handle, err = device.PcapInit(r.ether.Device)
 	if err != nil {
 		return nil, err
 	}
-	r.limit = ratelimit.New(int(options.Rate)) // per second
-	r.sender = make(chan string, 999)          // 多个协程发送
-	r.recver = make(chan core.RecvResult, 99)  // 多个协程接收
+	// 根据发包总数和timeout时间来合理分配每秒速度
+	// limit2 = 发包总数/timeout
+	// limit = min(limit2,options.Rate)
+
+	allPacket := r.loadTargets()
+	calcLimit := float64(allPacket/options.TimeOut) * 0.85
+	limit := int(math.Min(calcLimit, float64(options.Rate)))
+	if limit == 0 {
+		panic("预估长度失败")
+		//limit = int(options.Rate)
+	}
+	r.limit = ratelimit.New(limit) // per second
+	gologger.Infof("Rate:%dpps\n", limit)
+
+	r.sender = make(chan string, 99)          // 多个协程发送
+	r.recver = make(chan core.RecvResult, 99) // 多个协程接收
 
 	freePort, err := freeport.GetFreePort()
 	if err != nil {
@@ -97,21 +108,28 @@ func New(options *options2.Options) (*runner, error) {
 	r.fisrtloadChanel = make(chan string)
 	r.startTime = time.Now()
 
-	go r.loadTargets()
+	go func() {
+		for _, msg := range r.domains {
+			r.sender <- msg
+		}
+		r.domains = nil
+		r.fisrtloadChanel <- "ok"
+	}()
 	return r, nil
 }
 func (r *runner) choseDns() string {
 	dns := r.options.Resolvers
 	return dns[rand.Intn(len(dns))]
 }
 
-func (r *runner) loadTargets() {
+func (r *runner) loadTargets() int {
 	// get targets
 	var reader *bufio.Reader
 	options := r.options
 	if options.Method == "verify" {
 		if options.Stdin {
 			reader = bufio.NewReader(os.Stdin)
+
 		} else {
 			f2, err := os.Open(options.FileName)
 			if err != nil {
@@ -167,15 +185,15 @@ func (r *runner) loadTargets() {
 		msg := string(line)
 		if r.options.Method == "verify" {
 			// send msg
-			r.sender <- msg
+			r.domains = append(r.domains, msg)
 		} else {
 			for _, tmpDomain := range r.options.Domain {
 				newDomain := msg + "." + tmpDomain
-				r.sender <- newDomain
+				r.domains = append(r.domains, newDomain)
 			}
 		}
 	}
-	r.fisrtloadChanel <- "ok"
+	return len(r.domains)
 }
 func (r *runner) PrintStatus() {
 	queue := r.hm.Length()
@@ -185,9 +203,10 @@ func (r *runner) PrintStatus() {
 func (r *runner) RunEnumeration() {
 	ctx, cancel := context.WithCancel(r.ctx)
 	defer cancel()
-
-	go r.recvChanel(ctx)   // 启动接收线程
-	go r.sendCycle(ctx)    // 发送线程
+	go r.recvChanel(ctx) // 启动接收线程
+	for i := 0; i < 3; i++ {
+		go r.sendCycle(ctx) // 发送线程
+	}
 	go r.handleResult(ctx) // 处理结果，打印输出
 
 	var isLoadOver bool = false // 是否加载文件完毕

runner/send.go

@@ -21,14 +21,14 @@ func (r *runner) sendCycle(ctx context.Context) {
 			v = statusdb.Item{
 				Domain:      domain,
 				Dns:         r.choseDns(),
-				Time:        time.Now().Unix(),
+				Time:        time.Now(),
 				Retry:       0,
 				DomainLevel: 0,
 			}
 			r.hm.Add(domain, v)
 		} else {
 			v.Retry += 1
-			v.Time = time.Now().Unix()
+			v.Time = time.Now()
 			v.Dns = r.choseDns()
 			r.hm.Set(domain, v)
 		}

runner/statusdb/db.go

@@ -3,14 +3,15 @@ package statusdb
 import (
 	"sync"
 	"sync/atomic"
+	"time"
 )
 
 type Item struct {
-	Domain      string // 查询域名
-	Dns         string // 查询dns
-	Time        int64  // 发送时间
-	Retry       int    // 重试次数
-	DomainLevel int    // 域名层级
+	Domain      string    // 查询域名
+	Dns         string    // 查询dns
+	Time        time.Time // 发送时间
+	Retry       int       // 重试次数
+	DomainLevel int       // 域名层级
 }
 
 type StatusDb struct {