From d1fd1bc04051e0dd4dcc01f24f2df30efab747da Mon Sep 17 00:00:00 2001 From: boy-hack <34109680@qq.com> Date: Tue, 22 Sep 2020 11:43:34 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=9C=AC=E5=9C=B0=E7=8A=B6?= =?UTF-8?q?=E6=80=81=E8=A1=A8=E6=95=B0=E6=8D=AE=E7=BB=93=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core/localStruct.go | 114 ++++++++++++++++++++++++++++++++++++++++++++ core/recv.go | 16 +++---- core/send.go | 7 +-- core/stack.go | 4 +- core/start.go | 49 +++++++++---------- core/struct.go | 4 +- 6 files changed, 152 insertions(+), 42 deletions(-) create mode 100644 core/localStruct.go diff --git a/core/localStruct.go b/core/localStruct.go new file mode 100644 index 0000000..4707d80 --- /dev/null +++ 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 index 37d4748..a61d3c1 100644 --- 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, ð, &ipv4, &udp, &dns) + layers.LayerTypeEthernet, ð, &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 index 09153e0..bd765f9 100644 --- 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 index 8c0842b..8efa4a9 100644 --- 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 index c0aab1f..3be97bd 100644 --- 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 index e9a1592..6611629 100644 --- 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