major update to the package

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
 }

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
 

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
-		}
-	}
-}

conn_test.go

@@ -1,11 +1,12 @@
-// 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 (
 	"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++ {

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 {

go.mod

@@ -1,13 +1,16 @@
 module github.com/owasp-amass/resolve
 
-go 1.19
+go 1.21
+
+toolchain go1.21.4
 
 require (
 	github.com/caffix/queue v0.1.5
 	github.com/caffix/stringset v0.1.2
 	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
 )