From 35ae56aae717365a4d7762d9f1c6beb544b9f5aa Mon Sep 17 00:00:00 2001 From: Jeff Foley Date: Thu, 1 Feb 2024 11:47:31 -0500 Subject: [PATCH] major update to the package --- backoff.go | 22 +++++----- backoff_test.go | 2 +- conn.go | 106 ++++++++++++++++++++------------------------- conn_test.go | 5 ++- ecs.go | 43 ++++++++++-------- go.mod | 7 +-- go.sum | 27 ++---------- msgs.go | 2 +- msgs_test.go | 2 +- rate.go | 9 ++-- rate_test.go | 12 ++--- resolvers.go | 102 +++++++++++++++++++++++++------------------ selector.go | 41 ++++++++++-------- selector_test.go | 2 +- sub.go | 2 +- sub_test.go | 2 +- thresholds_test.go | 2 +- traversal.go | 2 +- traversal_test.go | 2 +- walk.go | 2 +- walk_test.go | 2 +- wildcards.go | 4 +- wildcards_test.go | 2 +- 23 files changed, 204 insertions(+), 198 deletions(-) diff --git a/backoff.go b/backoff.go index 51f556c..81388e9 100644 --- a/backoff.go +++ b/backoff.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2022-2023. All rights reserved. +// Copyright © by Jeff Foley 2022-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 @@ -10,6 +10,8 @@ import ( "time" ) +const numOfUnits int = 100 + // ExponentialBackoff returns a Duration equal to 2^events multiplied by the provided delay // and jitter added equal to [0,delay). func ExponentialBackoff(events int, delay time.Duration) time.Duration { @@ -19,21 +21,19 @@ func ExponentialBackoff(events int, delay time.Duration) time.Duration { // TruncatedExponentialBackoff returns a Duration equal to ExponentialBackoff with a provided // maximum Duration used to truncate the result. func TruncatedExponentialBackoff(events int, delay, max time.Duration) time.Duration { - backoff := ExponentialBackoff(events, delay) - - if backoff > max { - backoff = max + if backoff := ExponentialBackoff(events, delay); backoff < max { + return backoff } - return backoff + return max } // BackoffJitter returns a random Duration between the provided min and max parameters. func BackoffJitter(min, max time.Duration) time.Duration { - delta := max - min - if delta <= 0 { + if max < min { return time.Duration(0) } - - one := delta / 1000 - return min + (time.Duration(rand.Intn(1000)) * one) + if period := max - min; period > time.Duration(numOfUnits) { + return min + (time.Duration(rand.Intn(numOfUnits)) * (period / time.Duration(numOfUnits))) + } + return min } diff --git a/backoff_test.go b/backoff_test.go index c88aebd..fb46683 100644 --- a/backoff_test.go +++ b/backoff_test.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2022-2023. All rights reserved. +// Copyright © by Jeff Foley 2022-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 diff --git a/conn.go b/conn.go index 5d7cf24..4b5171c 100644 --- a/conn.go +++ b/conn.go @@ -1,33 +1,35 @@ -// Copyright © by Jeff Foley 2022-2023. All rights reserved. +// Copyright © by Jeff Foley 2022-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 package resolve import ( + "context" "fmt" "net" + "runtime" "sync" + "syscall" "time" "github.com/caffix/queue" "github.com/miekg/dns" + "golang.org/x/sys/unix" ) -const maxUDPBufferSize = 64 * 1024 * 1024 +const headerSize = 12 type resp struct { Msg *dns.Msg - Addr *net.UDPAddr + Addr net.Addr } type connections struct { sync.Mutex done chan struct{} - conns []*net.UDPConn + conns []net.PacketConn resps queue.Queue - rbufSize int - wbufSize int nextWrite int } @@ -58,7 +60,7 @@ func (c *connections) Close() { } } -func (c *connections) Next() *net.UDPConn { +func (c *connections) Next() net.PacketConn { c.Lock() defer c.Unlock() @@ -69,22 +71,46 @@ func (c *connections) Next() *net.UDPConn { func (c *connections) Add() error { var err error - var addr *net.UDPAddr - var conn *net.UDPConn - - if addr, err = net.ResolveUDPAddr("udp", ":0"); err == nil { - if conn, err = net.ListenUDP("udp", addr); err == nil { - _ = conn.SetDeadline(time.Time{}) - c.setMaxReadBufSize(conn) - c.setMaxWriteBufSize(conn) - c.conns = append(c.conns, conn) - go c.responses(conn) - } + var conn net.PacketConn + + if runtime.GOOS == "linux" { + conn, err = c.linuxListenPacket() + } else { + conn, err = net.ListenPacket("udp", ":0") + } + + if err == nil { + _ = conn.SetDeadline(time.Time{}) + c.conns = append(c.conns, conn) + go c.responses(conn) } return err } -func (c *connections) WriteMsg(msg *dns.Msg, addr *net.UDPAddr) error { +func (c *connections) linuxListenPacket() (net.PacketConn, error) { + lc := net.ListenConfig{ + Control: func(network, address string, c syscall.RawConn) error { + var operr error + + if err := c.Control(func(fd uintptr) { + operr = unix.SetsockoptInt(int(fd), unix.SOL_SOCKET, unix.SO_REUSEPORT, 1) + }); err != nil { + return err + } + + return operr + }, + } + + laddr := ":0" + if len(c.conns) > 0 { + laddr = c.conns[0].LocalAddr().String() + } + + return lc.ListenPacket(context.Background(), "udp", laddr) +} + +func (c *connections) WriteMsg(msg *dns.Msg, addr net.Addr) error { var n int var err error var out []byte @@ -93,14 +119,14 @@ func (c *connections) WriteMsg(msg *dns.Msg, addr *net.UDPAddr) error { conn := c.Next() _ = conn.SetWriteDeadline(time.Now().Add(500 * time.Millisecond)) - if n, err = conn.WriteToUDP(out, addr); err == nil && n < len(out) { + if n, err = conn.WriteTo(out, addr); err == nil && n < len(out) { err = fmt.Errorf("only wrote %d bytes of the %d byte message", n, len(out)) } } return err } -func (c *connections) responses(conn *net.UDPConn) { +func (c *connections) responses(conn net.PacketConn) { b := make([]byte, dns.DefaultMsgSize) for { @@ -109,7 +135,7 @@ func (c *connections) responses(conn *net.UDPConn) { return default: } - if n, addr, err := conn.ReadFromUDP(b); err == nil && n >= headerSize { + if n, addr, err := conn.ReadFrom(b); err == nil && n >= headerSize { m := new(dns.Msg) if err := m.Unpack(b[:n]); err == nil && len(m.Question) > 0 { @@ -121,39 +147,3 @@ func (c *connections) responses(conn *net.UDPConn) { } } } - -func (c *connections) setMaxReadBufSize(conn *net.UDPConn) { - c.Lock() - defer c.Unlock() - - if c.rbufSize != 0 { - _ = conn.SetReadBuffer(c.rbufSize) - return - } - - min := 1024 - for size := maxUDPBufferSize; size > min; size /= 2 { - if err := conn.SetReadBuffer(size); err == nil { - c.rbufSize = size - return - } - } -} - -func (c *connections) setMaxWriteBufSize(conn *net.UDPConn) { - c.Lock() - defer c.Unlock() - - if c.wbufSize != 0 { - _ = conn.SetWriteBuffer(c.wbufSize) - return - } - - min := 1024 - for size := maxUDPBufferSize; size > min; size /= 2 { - if err := conn.SetWriteBuffer(size); err == nil { - c.wbufSize = size - return - } - } -} diff --git a/conn_test.go b/conn_test.go index 1bb369b..821fd32 100644 --- a/conn_test.go +++ b/conn_test.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2022-2023. All rights reserved. +// Copyright © by Jeff Foley 2022-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 @@ -6,6 +6,7 @@ package resolve import ( "net" + "runtime" "testing" "time" @@ -25,7 +26,7 @@ func TestConnections(t *testing.T) { defer func() { _ = s.Shutdown() }() resps := queue.NewQueue() - conn := newConnections(1, resps) + conn := newConnections(runtime.NumCPU(), resps) defer conn.Close() for i := 0; i < 100; i++ { diff --git a/ecs.go b/ecs.go index fa82964..4714b14 100644 --- a/ecs.go +++ b/ecs.go @@ -5,6 +5,7 @@ package resolve import ( + "sync" "time" "github.com/miekg/dns" @@ -16,37 +17,45 @@ func (r *Resolvers) ClientSubnetCheck() { all := r.pool.AllResolvers() alen := len(all) ch := make(chan *dns.Msg, alen) + var msglock sync.Mutex msgsToRes := make(map[string]*resolver) - send := func(res *resolver) { - msg := QueryMsg("o-o.myaddr.l.google.com", dns.TypeTXT) - msgsToRes[xchgKey(msg.Id, msg.Question[0].Name)] = res - r.writeReq(&request{ - Res: res, - Msg: msg, - Result: ch, - }) - } + go func() { + var count int + + for _, res := range all { + msg := QueryMsg("o-o.myaddr.l.google.com", dns.TypeTXT) + key := xchgKey(msg.Id, msg.Question[0].Name) + msglock.Lock() + msgsToRes[key] = res + msglock.Unlock() + res.writeReq(&request{ + Res: res, + Msg: msg, + Result: ch, + }) - var count int - for _, res := range all { - send(res) - count++ - if count == 100 { - count = 0 - time.Sleep(100 * time.Millisecond) + count++ + if count == 250 { + count = 0 + time.Sleep(100 * time.Millisecond) + } } - } + }() for i := 0; i < alen; i++ { resp := <-ch // pull the resolver associated with this message key := xchgKey(resp.Id, resp.Question[0].Name) + + msglock.Lock() res, found := msgsToRes[key] if !found { + msglock.Unlock() continue } delete(msgsToRes, key) + msglock.Unlock() // check if the resolver responded, but did not return a successful response if resp.Rcode != dns.RcodeSuccess || (!resp.Authoritative && !resp.RecursionAvailable) { if res != nil { diff --git a/go.mod b/go.mod index 7766113..c293b18 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,8 @@ module github.com/owasp-amass/resolve -go 1.19 +go 1.21 + +toolchain go1.21.4 require ( github.com/caffix/queue v0.1.5 @@ -8,6 +10,7 @@ require ( github.com/miekg/dns v1.1.58 go.uber.org/ratelimit v0.3.0 golang.org/x/net v0.20.0 + golang.org/x/sys v0.16.0 ) require ( @@ -22,8 +25,6 @@ require ( github.com/pkg/errors v0.9.1 // indirect github.com/stretchr/testify v1.7.0 // indirect golang.org/x/mod v0.14.0 // indirect - golang.org/x/sys v0.16.0 // indirect golang.org/x/tools v0.17.0 // indirect - golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect google.golang.org/protobuf v1.32.0 // indirect ) diff --git a/go.sum b/go.sum index 1e1173e..fd661a6 100644 --- a/go.sum +++ b/go.sum @@ -5,15 +5,10 @@ github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAE github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8= github.com/benbjohnson/clock v1.3.5 h1:VvXlSJBzZpA/zum6Sj74hxwYI2DIxRWuNIoXAzHZz5o= github.com/benbjohnson/clock v1.3.5/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA= -github.com/caffix/queue v0.1.4 h1:sQbFzwGaPM1tRnQHWCgHOwj7hLuhDQ3BhY1/1TFbBiE= -github.com/caffix/queue v0.1.4/go.mod h1:l8Eg7UTUHTRlc5aQ37mRVjzLN6eC7hgwimN0pA4UHe8= github.com/caffix/queue v0.1.5 h1:uFx535UrEMkKIhaHeixz8C+ZJgvn1lZ07dLsuW/Na3I= github.com/caffix/queue v0.1.5/go.mod h1:ReDjIxSFHpMhVnomlQlZj86BzF9kTFDMyM33jXyDJGg= -github.com/caffix/stringset v0.1.1 h1:Tm4b7SBFAsRTBbBX90eP8xBv6BxSuU2w+6G/JNXtNpg= -github.com/caffix/stringset v0.1.1/go.mod h1:9Ztc521vlcp8IWdtIowZyWbbddMKR9Rdr+d0pgnjcvk= github.com/caffix/stringset v0.1.2 h1:AnBiZ5dH8AqOtDsUPdFt7ZzHk5RqmGixmfZFlxzZh4U= github.com/caffix/stringset v0.1.2/go.mod h1:eWeJ1l/1Tc3SO5eybwwMIltkoPNkej2y5d4sHQlHOxw= -github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko= github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc= github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= @@ -37,25 +32,21 @@ github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkp github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= -github.com/golang/glog v1.1.2 h1:DVjP2PbBOzHyzA+dn3WhHIq4NdVu3Q+pvivFICf/7fo= -github.com/golang/glog v1.1.2/go.mod h1:zR+okUeTbrL6EL3xHUDxZuEtGv04p5shwip1+mL/rLQ= github.com/golang/glog v1.2.0 h1:uCdmnmatrKCgMBlM4rMuJZWOkPDqdbZPnrMXDY4gI68= github.com/golang/glog v1.2.0/go.mod h1:6AhwSGph0fcJtXVM/PEHPqZlFeoLxhs7/t5UDAwmO+w= github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg= github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY= -github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ= github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8= github.com/kr/pretty v0.2.0/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ= -github.com/miekg/dns v1.1.55 h1:GoQ4hpsj0nFLYe+bWiCToyrBEJXkQfOOIvFGFy0lEgo= -github.com/miekg/dns v1.1.55/go.mod h1:uInx36IzPl7FYnDcMeVWxj9byh7DutNykX4G9Sj60FY= github.com/miekg/dns v1.1.58 h1:ca2Hdkz+cDg/7eNF6V56jjzuZ4aCAE+DbVkILdQWG/4= github.com/miekg/dns v1.1.58/go.mod h1:Ypv+3b/KadlvW9vJfXOTf300O4UqaHFzFCuHz+rPkBY= github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0= @@ -83,40 +74,30 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/ github.com/ugorji/go/codec v0.0.0-20181204163529-d75b2dcb6bc8/go.mod h1:VFNgLljTbGfSG7qAOspJ7OScBnGdDN/yBr0sguwnwf0= github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q= go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw= +go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= go.uber.org/ratelimit v0.3.0 h1:IdZd9wqvFXnvLvSEBo0KPcGfkoBGNkpTHlrE3Rcjkjw= go.uber.org/ratelimit v0.3.0/go.mod h1:So5LG7CV1zWpY1sHe+DXTJqQvOx+FFPFaAs2SnoyBaI= golang.org/x/crypto v0.0.0-20181203042331-505ab145d0a9/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= -golang.org/x/mod v0.12.0 h1:rmsUpXtvNzj340zd98LZ4KntptpfRHwpFOHG188oHXc= -golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= golang.org/x/mod v0.14.0 h1:dGoOF9QVLYng8IHTm7BAyWqCqSheQ5pYWGhzW00YJr0= golang.org/x/mod v0.14.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= -golang.org/x/net v0.15.0 h1:ugBLEUaxABaB5AJqW9enI0ACdci2RUd4eP51NTBvuJ8= -golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk= golang.org/x/net v0.20.0 h1:aCL9BSgETF1k+blQaYUBx9hJ9LOGP3gAVemcZlf1Kpo= golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY= -golang.org/x/sync v0.3.0 h1:ftCYgMx6zT/asHUrPw8BLLscYtGznsLAnjq5RH9P66E= +golang.org/x/sync v0.6.0 h1:5BMeUDZ7vkXGfEr1x9B4bRcTH4lpkTkpdh0T/J+qjbQ= +golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= golang.org/x/sys v0.0.0-20181205085412-a5c9d58dba9a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190626221950-04f50cda93cb/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20221010170243-090e33056c14/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.12.0 h1:CM0HF96J0hcLAwsHPJZjfdNzs0gftsLfgKt57wWHJ0o= -golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU= golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= -golang.org/x/tools v0.13.0 h1:Iey4qkscZuv0VvIt8E0neZjtPVQFSc870HQ448QgEmQ= -golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58= golang.org/x/tools v0.17.0 h1:FvmRgNOcs3kOa+T20R1uhfP9F6HgG2mfxDv1vrx1Htc= golang.org/x/tools v0.17.0/go.mod h1:xsh6VxdV005rRVaS6SSAf9oiAqljS7UZUacMZ8Bnsps= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= -golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 h1:H2TDz8ibqkAF6YGhCdN3jS9O0/s90v0rJh3X/OLHEUk= -golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2/go.mod h1:K8+ghG5WaK9qNqU5K3HdILfMLy1f3aNYFI/wnl100a8= google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= -google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8= -google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= google.golang.org/protobuf v1.32.0 h1:pPC6BG5ex8PDFnkbrGU3EixyhKcQ2aDuBS36lqK/C7I= google.golang.org/protobuf v1.32.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/msgs.go b/msgs.go index b5d5ac2..5464194 100644 --- a/msgs.go +++ b/msgs.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2021-2023. All rights reserved. +// Copyright © by Jeff Foley 2021-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 diff --git a/msgs_test.go b/msgs_test.go index 0992bec..4e7df25 100644 --- a/msgs_test.go +++ b/msgs_test.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2021-2023. All rights reserved. +// Copyright © by Jeff Foley 2021-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 diff --git a/rate.go b/rate.go index 629c2f7..fcc3af7 100644 --- a/rate.go +++ b/rate.go @@ -16,9 +16,9 @@ import ( const ( maxQPSPerNameserver = 100 - numIntervalSeconds = 2 + numIntervalSeconds = 10 rateUpdateInterval = numIntervalSeconds * time.Second - maxTimeoutPercentage = 0.1 + maxTimeoutPercentage = 0.2 ) type rateTrack struct { @@ -129,12 +129,13 @@ func (rt *rateTrack) update() { } // timeouts in excess of maxTimeoutPercentage indicate a need to slow down if float64(rt.timeout)/float64(rt.success+rt.timeout) > maxTimeoutPercentage { - rt.qps-- + p := float64(rt.success) / float64(rt.success+rt.timeout) + rt.qps = int(float64(rt.qps) * p) if rt.qps <= 0 { rt.qps = 1 } } else { - rt.qps++ + rt.qps = rt.qps + 10 } // update the QPS rate limiter and reset counters rt.rate = ratelimit.New(rt.qps) diff --git a/rate_test.go b/rate_test.go index d0abc92..67ed029 100644 --- a/rate_test.go +++ b/rate_test.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2022-2023. All rights reserved. +// Copyright © by Jeff Foley 2022-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 @@ -29,14 +29,14 @@ func TestUpdateRateLimiters(t *testing.T) { for i := 0; i < num; i++ { rt.Timeout(domain) } - time.Sleep(2 * rateUpdateInterval) + time.Sleep(rateUpdateInterval + (rateUpdateInterval / 2)) tracker.Lock() qps2 := tracker.qps tracker.Unlock() // the QPS should now be lower if qps2 >= qps { - t.Errorf("Unexpected QPS, expected %d, got %d", qps-1, qps2) + t.Errorf("Unexpected QPS, expected QPS lower than %d, got %d", qps, qps2) } tracker.Lock() @@ -55,13 +55,13 @@ func TestUpdateRateLimiters(t *testing.T) { for i := 0; i < qps; i++ { rt.Success(domain) } - time.Sleep(2 * rateUpdateInterval) + time.Sleep(rateUpdateInterval + (rateUpdateInterval / 2)) tracker.Lock() qps2 = tracker.qps tracker.Unlock() - // the QPS should now be lower + // the QPS should now be higher if qps2 <= qps { - t.Errorf("Unexpected QPS, expected %d, got %d", qps+1, qps2) + t.Errorf("Unexpected QPS, expected QPS higher than %d, got %d", qps, qps2) } } diff --git a/resolvers.go b/resolvers.go index 739ccd3..fefe3c5 100644 --- a/resolvers.go +++ b/resolvers.go @@ -19,8 +19,6 @@ import ( "go.uber.org/ratelimit" ) -const headerSize = 12 - // Resolvers is a pool of DNS resolvers managed for brute forcing using random selection. type Resolvers struct { sync.Mutex @@ -43,9 +41,12 @@ type Resolvers struct { type resolver struct { done chan struct{} + pool *Resolvers + queue queue.Queue xchgs *xchgMgr address *net.UDPAddr qps int + rate ratelimit.Limiter stats *stats } @@ -59,11 +60,15 @@ func (r *Resolvers) initializeResolver(qps int, addr string) *resolver { if uaddr, err := net.ResolveUDPAddr("udp", addr); err == nil { res = &resolver{ done: make(chan struct{}, 1), + pool: r, + queue: queue.NewQueue(), xchgs: newXchgMgr(r.timeout), address: uaddr, qps: qps, + rate: ratelimit.New(qps), stats: new(stats), } + go res.processRequests() } return res } @@ -179,8 +184,8 @@ func (r *Resolvers) AddResolvers(qps int, addrs ...string) error { addr = net.JoinHostPort(addr, "53") } // check that this address will not create a duplicate resolver - if uaddr, err := net.ResolveUDPAddr("udp", addr); err == nil { - if _, found := r.rmap[uaddr.IP.String()]; !found { + if host, _, err := net.SplitHostPort(addr); err == nil { + if _, found := r.rmap[host]; !found { if res := r.initializeResolver(qps, addr); res != nil { r.rmap[res.address.IP.String()] = struct{}{} r.pool.AddResolver(res) @@ -266,18 +271,12 @@ func (r *Resolvers) QueryBlocking(ctx context.Context, msg *dns.Msg) (*dns.Msg, default: } - ch := r.QueryChan(ctx, msg) - - select { - case <-ctx.Done(): - return msg, errors.New("the context expired") - case resp := <-ch: - var err error - if resp == nil { - err = errors.New("query failed") - } - return resp, err + var err error + resp := <-r.QueryChan(ctx, msg) + if resp == nil { + err = errors.New("query failed") } + return resp, err } func (r *Resolvers) enforceMaxQPS() { @@ -287,17 +286,20 @@ loop: case <-r.done: break loop case <-r.queue.Signal(): + element, found := r.queue.Next() + if !found { + continue loop + } + if r.rate != nil { - r.rate.Take() + _ = r.rate.Take() } - if e, yes := r.queue.Next(); yes { - if req, ok := e.(*request); ok { - if res := r.pool.GetResolver(); res != nil { - req.Res = res - r.writeReq(req) - continue loop - } + if req, ok := element.(*request); ok { + if res := r.pool.GetResolver(); res != nil { + req.Res = res + res.queue.Append(req) + } else { req.errNoResponse() req.release() } @@ -322,8 +324,8 @@ func (r *Resolvers) processResponses() { } r.resps.Process(func(element interface{}) { - if response, ok := element.(*resp); ok && r != nil { - r.processSingleResp(response) + if response, ok := element.(*resp); ok && response != nil { + go r.processSingleResp(response) } }) } @@ -331,13 +333,13 @@ func (r *Resolvers) processResponses() { func (r *Resolvers) processSingleResp(response *resp) { var res *resolver - addr := response.Addr.IP.String() + addr, _, _ := net.SplitHostPort(response.Addr.String()) if res = r.pool.LookupResolver(addr); res == nil { - detector := r.getDetectionResolver() - - if detector != nil && addr == detector.address.IP.String() { - res = detector + if detector := r.getDetectionResolver(); detector != nil { + if detector.address.IP.String() == addr { + res = detector + } } } if res == nil { @@ -362,7 +364,14 @@ func (r *Resolvers) processSingleResp(response *resp) { } func (r *Resolvers) timeouts() { - for { + r.Lock() + d := r.timeout / 2 + r.Unlock() + + t := time.NewTicker(d) + defer t.Stop() + + for range t.C { select { case <-r.done: return @@ -389,24 +398,33 @@ func (r *Resolvers) timeouts() { } } } - // wait a bit before checking again - r.Lock() - d := r.timeout / 2 - r.Unlock() - if d > 0 { - time.Sleep(d) + } +} + +func (r *resolver) processRequests() { + for { + select { + case <-r.done: + return + case <-r.queue.Signal(): } + + r.queue.Process(func(element interface{}) { + if req, ok := element.(*request); ok && req != nil { + _ = r.rate.Take() + go r.writeReq(req) + } + }) } } -func (r *Resolvers) writeReq(req *request) { - res := req.Res +func (r *resolver) writeReq(req *request) { msg := req.Msg.Copy() req.Timestamp = time.Now() - if res.xchgs.add(req) == nil { - if err := r.conns.WriteMsg(msg, res.address); err != nil { - _ = res.xchgs.remove(msg.Id, msg.Question[0].Name) + if r.xchgs.add(req) == nil { + if err := r.pool.conns.WriteMsg(msg, r.address); err != nil { + _ = r.xchgs.remove(msg.Id, msg.Question[0].Name) req.errNoResponse() req.release() } diff --git a/selector.go b/selector.go index 37e67ac..860cbc9 100644 --- a/selector.go +++ b/selector.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2022-2023. All rights reserved. +// Copyright © by Jeff Foley 2022-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 @@ -41,46 +41,51 @@ func newRandomSelector() *randomSelector { // GetResolver performs random selection on the pool of resolvers. func (r *randomSelector) GetResolver() *resolver { - var low int + max := r.maxQPS() + if max == -1 { + return nil + } + sel := rand.Intn(max) + + r.Lock() + defer r.Unlock() + + var cur int var chosen *resolver loop: - for _, res := range r.randList() { + for _, res := range r.list { select { case <-res.done: continue loop default: } - if cur := res.xchgs.len(); chosen == nil || cur < low { + + cur += res.qps + if sel < cur { chosen = res - low = cur - } - if low == 0 { break } } return chosen } -func (r *randomSelector) randList() []*resolver { +func (r *randomSelector) maxQPS() int { r.Lock() defer r.Unlock() - rlen := len(r.list) - if rlen == 0 { - return nil + if len(r.list) == 0 { + return -1 } - slen := min(rlen, 25) - var list []*resolver - for a, i, j := 0, 0, rand.Intn(rlen); i < rlen && a < slen; i, j = i+1, (j+1)%rlen { + var max int + for _, res := range r.list { select { - case <-r.list[j].done: + case <-res.done: default: - list = append(list, r.list[j]) - a++ + max += res.qps } } - return list + return max } func (r *randomSelector) LookupResolver(addr string) *resolver { diff --git a/selector_test.go b/selector_test.go index 76d8c95..194ccf5 100644 --- a/selector_test.go +++ b/selector_test.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2022-2023. All rights reserved. +// Copyright © by Jeff Foley 2022-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 diff --git a/sub.go b/sub.go index 2fb2fc2..e355936 100644 --- a/sub.go +++ b/sub.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2021-2023. All rights reserved. +// Copyright © by Jeff Foley 2021-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 diff --git a/sub_test.go b/sub_test.go index b7b07ac..c5a6fa7 100644 --- a/sub_test.go +++ b/sub_test.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2021-2023. All rights reserved. +// Copyright © by Jeff Foley 2021-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 diff --git a/thresholds_test.go b/thresholds_test.go index 6d32b2a..a6af4eb 100644 --- a/thresholds_test.go +++ b/thresholds_test.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2017-2023. All rights reserved. +// Copyright © by Jeff Foley 2017-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 diff --git a/traversal.go b/traversal.go index 8a1bc2c..47f7751 100644 --- a/traversal.go +++ b/traversal.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2022-2023. All rights reserved. +// Copyright © by Jeff Foley 2022-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 diff --git a/traversal_test.go b/traversal_test.go index b4aac75..c77a306 100644 --- a/traversal_test.go +++ b/traversal_test.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2022-2023. All rights reserved. +// Copyright © by Jeff Foley 2022-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 diff --git a/walk.go b/walk.go index 3d0a859..65d79c3 100644 --- a/walk.go +++ b/walk.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2017-2023. All rights reserved. +// Copyright © by Jeff Foley 2017-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 diff --git a/walk_test.go b/walk_test.go index 4c1ecfe..8c2a4c8 100644 --- a/walk_test.go +++ b/walk_test.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2021-2023. All rights reserved. +// Copyright © by Jeff Foley 2021-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 diff --git a/wildcards.go b/wildcards.go index 81cd187..3bfc7a9 100644 --- a/wildcards.go +++ b/wildcards.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2017-2023. All rights reserved. +// Copyright © by Jeff Foley 2017-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0 @@ -244,7 +244,7 @@ loop: Result: ch, } - r.writeReq(req) + detector.writeReq(req) select { case <-ctx.Done(): break loop diff --git a/wildcards_test.go b/wildcards_test.go index 1043ada..be4c65c 100644 --- a/wildcards_test.go +++ b/wildcards_test.go @@ -1,4 +1,4 @@ -// Copyright © by Jeff Foley 2021-2023. All rights reserved. +// Copyright © by Jeff Foley 2021-2024. All rights reserved. // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file. // SPDX-License-Identifier: Apache-2.0