优化本地状态表数据结构

knownsec · Sep 22, 2020 · d1fd1bc · d1fd1bc
1 parent 24b9bb3
commit d1fd1bc
Show file tree

Hide file tree

Showing 6 changed files with 152 additions and 42 deletions.
diff --git a/core/localStruct.go b/core/localStruct.go
@@ -0,0 +1,114 @@
+package core
+
+import (
+	"errors"
+	"sync"
+	"time"
+)
+
+type (
+	LocalStruct struct {
+		header *localNode
+		length int
+		lock   sync.RWMutex
+	}
+	localNode struct {
+		v     *StatusTable
+		index uint32
+		next  *localNode
+	}
+	LocalRetryStruct struct {
+		v     StatusTable
+		index uint32
+	}
+)
+
+// Create a new stack
+func NewLocalStruct() *LocalStruct {
+	return &LocalStruct{nil, 0, sync.RWMutex{}}
+}
+
+// Return the number of items in the stack
+func (this *LocalStruct) Len() int {
+	this.lock.RLock()
+	defer this.lock.RUnlock()
+	return this.length
+}
+
+func (this *LocalStruct) Empty() bool {
+	return this.Len() == 0
+}
+
+func (this *LocalStruct) Append(node *StatusTable, index uint32) {
+	newNode := &localNode{node, index, nil}
+	if this.Empty() {
+		this.lock.Lock()
+		this.header = newNode
+		this.lock.Unlock()
+	} else {
+		this.lock.Lock()
+		current := this.header
+		for current.next != nil {
+			current = current.next
+		}
+		current.next = newNode
+		this.lock.Unlock()
+	}
+	this.lock.Lock()
+	this.length++
+	this.lock.Unlock()
+}
+
+// 搜索并删除节点
+func (this *LocalStruct) SearchFromIndexAndDelete(index uint32) (LocalRetryStruct, error) {
+	this.lock.Lock()
+	defer this.lock.Unlock()
+	prev := this.header
+	if prev == nil {
+		return LocalRetryStruct{}, errors.New("data length is null")
+	}
+	if prev.index == index {
+		v := *prev.v
+		this.header = prev.next
+		this.length--
+		return LocalRetryStruct{v, index}, nil
+	}
+	for prev.next != nil {
+		if prev.next.index == index {
+			v := *prev.next.v
+			after := prev.next.next
+			prev.next = after
+			this.length--
+			return LocalRetryStruct{v, index}, nil
+		}
+		prev = prev.next
+	}
+	return LocalRetryStruct{}, errors.New("data not found")
+}
+
+// 从链表中取出超时的数据，可选每次取出最多多少数据
+func (this *LocalStruct) GetTimeoutData(max int) []LocalRetryStruct {
+	this.lock.Lock()
+	defer this.lock.Unlock()
+	current := this.header
+	currentTime := time.Now().Unix()
+	index := 0
+	var tables []LocalRetryStruct
+	for current != nil {
+		if currentTime-current.v.Time < 5 {
+			break
+		}
+		if index > max {
+			break
+		}
+		index++
+		tables = append(tables, LocalRetryStruct{*current.v, current.index})
+		current = current.next
+	}
+	// 删除掉这些选择的链表数据
+	if index > 0 {
+		this.header = current
+		this.length -= index
+	}
+	return tables
+}
diff --git a/core/recv.go b/core/recv.go
@@ -35,9 +35,10 @@ func Recv(device string, options *Options, flagID uint16, retryChan chan RetrySt
 	var dns layers.DNS
 	var eth layers.Ethernet
 	var ipv4 layers.IPv4
+	var ipv6 layers.IPv6
 
 	parser := gopacket.NewDecodingLayerParser(
-		layers.LayerTypeEthernet, &eth, &ipv4, &udp, &dns)
+		layers.LayerTypeEthernet, &eth, &ipv4, &ipv6, &udp, &dns)
 	var isWrite bool = false
 	var isttl bool = options.TTL
 	var isSummary bool = options.Summary
@@ -68,27 +69,26 @@ func Recv(device string, options *Options, flagID uint16, retryChan chan RetrySt
 			atomic.AddUint64(&RecvIndex, 1)
 			udp, _ := packet.Layer(layers.LayerTypeUDP).(*layers.UDP)
 			index := GenerateMapIndex(dns.ID%100, uint16(udp.DstPort))
-			if _data, ok := LocalStauts.Load(uint32(index)); ok {
-				data := _data.(StatusTable)
+			data, err := LocalStauts.SearchFromIndexAndDelete(uint32(index))
+			if err == nil {
 				// 处理多级域名
-				if dns.ANCount > 0 && data.DomainLevel < options.DomainLevel {
+				if dns.ANCount > 0 && data.v.DomainLevel < options.DomainLevel {
 					running := true
 					if options.SkipWildCard {
-						if IsWildCard(data.Domain) {
+						if IsWildCard(data.v.Domain) {
 							running = false
 						}
 					}
 					if running {
 						for _, sub := range GetSubNextData() {
-							subdomain := sub + "." + data.Domain
-							retryChan <- RetryStruct{Domain: subdomain, Dns: data.Dns, SrcPort: 0, FlagId: 0, DomainLevel: data.DomainLevel + 1}
+							subdomain := sub + "." + data.v.Domain
+							retryChan <- RetryStruct{Domain: subdomain, Dns: data.v.Dns, SrcPort: 0, FlagId: 0, DomainLevel: data.v.DomainLevel + 1}
 						}
 					}
 				}
 				if LocalStack.Len() <= 50000 {
 					LocalStack.Push(uint32(index))
 				}
-				LocalStauts.Delete(uint32(index))
 			}
 			if dns.ANCount > 0 {
 				atomic.AddUint64(&SuccessIndex, 1)

diff --git a/core/send.go b/core/send.go
@@ -79,11 +79,8 @@ func (d *SendDog) BuildStatusTable(domain string, dns string, domainlevel int) (
 		}
 	}
 	index := GenerateMapIndex(d.flagID2, d.index)
-	if _, ok := LocalStauts.Load(index); !ok {
-		LocalStauts.Store(uint32(index), StatusTable{Domain: domain, Dns: dns, Time: time.Now().Unix(), Retry: 0, DomainLevel: domainlevel})
-	} else {
-		gologger.Warningf("LocalStatus 状态重复")
-	}
+	s := StatusTable{Domain: domain, Dns: dns, Time: time.Now().Unix(), Retry: 0, DomainLevel: domainlevel}
+	LocalStauts.Append(&s, uint32(index))
 	return d.flagID2, d.index
 }
 

diff --git a/core/stack.go b/core/stack.go
@@ -9,7 +9,7 @@ type (
 	Stack struct {
 		top    *node
 		length int
-		lock   *sync.RWMutex
+		lock   sync.RWMutex
 	}
 	node struct {
 		value uint32
@@ -21,7 +21,7 @@ var LocalStack *Stack
 
 // Create a new stack
 func NewStack() *Stack {
-	return &Stack{nil, 0, &sync.RWMutex{}}
+	return &Stack{nil, 0, sync.RWMutex{}}
 }
 
 // Return the number of items in the stack

diff --git a/core/start.go b/core/start.go
@@ -98,27 +98,33 @@ func Start(options *Options) {
 	go func() {
 		for {
 			// 循环检测超时的队列
-			//遍历该map，参数是个函数，该函数参的两个参数是遍历获得的key和value，返回一个bool值，当返回false时，遍历立刻结束。
-			LocalStauts.Range(func(k, v interface{}) bool {
-				index := k.(uint32)
-				value := v.(StatusTable)
-				if value.Retry >= 25 {
+			maxLength := int(options.Rate / 10)
+			datas := LocalStauts.GetTimeoutData(maxLength)
+			isdelay := true
+			if len(datas) <= 100 {
+				isdelay = false
+			}
+			for _, localdata := range datas {
+				index := localdata.index
+				value := localdata.v
+				if value.Retry >= 15 {
 					atomic.AddUint64(&FaildIndex, 1)
-					LocalStauts.Delete(index)
-					return true
+					LocalStauts.SearchFromIndexAndDelete(index)
+					continue
 				}
-				if time.Now().Unix()-value.Time >= 5 {
-					_ = limiter.Wait(ctx)
-					value.Retry++
-					value.Time = time.Now().Unix()
-					value.Dns = sendog.ChoseDns()
-					LocalStauts.Store(index, value)
-					flag2, srcport := GenerateFlagIndexFromMap(index)
-					retryChan <- RetryStruct{Domain: value.Domain, Dns: value.Dns, SrcPort: srcport, FlagId: flag2, DomainLevel: value.DomainLevel}
+				_ = limiter.Wait(ctx)
+				value.Retry++
+				value.Time = time.Now().Unix()
+				value.Dns = sendog.ChoseDns()
+				// 先删除，再重新创建
+				LocalStauts.SearchFromIndexAndDelete(index)
+				LocalStauts.Append(&value, index)
+				flag2, srcport := GenerateFlagIndexFromMap(index)
+				retryChan <- RetryStruct{Domain: value.Domain, Dns: value.Dns, SrcPort: srcport, FlagId: flag2, DomainLevel: value.DomainLevel}
+				if isdelay {
+					time.Sleep(time.Microsecond * time.Duration(rand.Intn(300)+100))
 				}
-				time.Sleep(time.Microsecond * time.Duration(rand.Intn(300)+100))
-				return true
-			})
+			}
 		}
 	}()
 	// 多级域名检测
@@ -155,12 +161,7 @@ func Start(options *Options) {
 		}
 	}
 	for {
-		var isbreak bool = true
-		LocalStauts.Range(func(k, v interface{}) bool {
-			isbreak = false
-			return false
-		})
-		if isbreak {
+		if LocalStauts.Empty() {
 			break
 		}
 		time.Sleep(time.Millisecond * 723)

diff --git a/core/struct.go b/core/struct.go
@@ -3,7 +3,6 @@ package core
 import (
 	"github.com/google/gopacket/layers"
 	"net"
-	"sync"
 )
 
 // 本地状态表
@@ -45,8 +44,7 @@ type EthTable struct {
 }
 
 var wait_chain *Stack
-var LocalStauts sync.Map
-var DnsChoice sync.Map
+var LocalStauts LocalStruct
 var RecvIndex uint64 = 0
 var FaildIndex uint64 = 0
 var SentIndex uint64 = 0