Implement caching dns resolver

antireplay/noop.go

@@ -5,7 +5,6 @@ import "github.com/9seconds/mtg/v2/mtglib"
 type noop struct{}
 
 func (n noop) SeenBefore(_ []byte) bool { return false }
-func (n noop) Shutdown()                {}
 
 func NewNoop() mtglib.AntiReplayCache {
 	return noop{}

antireplay/noop_test.go

@@ -18,8 +18,6 @@ func (suite *NoopTestSuite) TestOp() {
 	suite.False(filter.SeenBefore([]byte{4, 5, 6}))
 	suite.False(filter.SeenBefore([]byte{1, 2, 3}))
 	suite.False(filter.SeenBefore([]byte{4, 5, 6}))
-
-	filter.Shutdown()
 }
 
 func TestNoop(t *testing.T) {

antireplay/stable_bloom_filter.go

@@ -20,8 +20,6 @@ func (s *stableBloomFilter) SeenBefore(digest []byte) bool {
 	return s.filter.TestAndAdd(digest)
 }
 
-func (s *stableBloomFilter) Shutdown() {}
-
 func NewStableBloomFilter(byteSize uint, errorRate float64) mtglib.AntiReplayCache {
 	sf := boom.NewDefaultStableBloomFilter(byteSize*8, errorRate) // nolint: gomnd
 	sf.SetHash(xxhash.New64())

antireplay/stable_bloom_filter_test.go

@@ -18,8 +18,6 @@ func (suite *StableBloomFilterTestSuite) TestOp() {
 	suite.False(filter.SeenBefore([]byte{4, 5, 6}))
 	suite.True(filter.SeenBefore([]byte{1, 2, 3}))
 	suite.True(filter.SeenBefore([]byte{4, 5, 6}))
-
-	filter.Shutdown()
 }
 
 func TestStableBloomFilter(t *testing.T) {

cli/proxy.go

@@ -59,12 +59,6 @@ func (c *Proxy) Execute() error { // nolint: funlen
 		PreferIP:           c.Config.PreferIP.Value(mtglib.DefaultPreferIP),
 	}
 
-	defer func() {
-		opts.AntiReplayCache.Shutdown()
-		opts.IPBlocklist.Shutdown()
-		opts.EventStream.Shutdown()
-	}()
-
 	if opts.Concurrency == 0 {
 		opts.Concurrency = mtglib.DefaultConcurrency
 	}

events/event_stream.go

@@ -9,13 +9,13 @@ import (
 	"github.com/OneOfOne/xxhash"
 )
 
-type eventStream struct {
+type EventStream struct {
 	ctx       context.Context
 	ctxCancel context.CancelFunc
 	chans     []chan mtglib.Event
 }
 
-func (e eventStream) Send(ctx context.Context, evt mtglib.Event) {
+func (e EventStream) Send(ctx context.Context, evt mtglib.Event) {
 	var chanNo uint32
 
 	if streamID := evt.StreamID(); streamID != "" {
@@ -31,17 +31,17 @@ func (e eventStream) Send(ctx context.Context, evt mtglib.Event) {
 	}
 }
 
-func (e eventStream) Shutdown() {
+func (e EventStream) Shutdown() {
 	e.ctxCancel()
 }
 
-func NewEventStream(observerFactories []ObserverFactory) mtglib.EventStream {
+func NewEventStream(observerFactories []ObserverFactory) EventStream {
 	if len(observerFactories) == 0 {
 		observerFactories = append(observerFactories, NewNoopObserver)
 	}
 
 	ctx, cancel := context.WithCancel(context.Background())
-	rv := eventStream{
+	rv := EventStream{
 		ctx:       ctx,
 		ctxCancel: cancel,
 		chans:     make([]chan mtglib.Event, runtime.NumCPU()),

events/event_stream_test.go

@@ -19,7 +19,7 @@ type EventStreamTestSuite struct {
 	ctxCancel     context.CancelFunc
 	observerMock1 *ObserverMock
 	observerMock2 *ObserverMock
-	stream        mtglib.EventStream
+	stream        events.EventStream
 }
 
 func (suite *EventStreamTestSuite) SetupTest() {

events/noop.go

@@ -9,7 +9,6 @@ import (
 type noop struct{}
 
 func (n noop) Send(ctx context.Context, evt mtglib.Event) {}
-func (n noop) Shutdown()                                  {}
 
 func NewNoopStream() mtglib.EventStream {
 	return noop{}

events/noop_test.go

@@ -70,8 +70,6 @@ func (suite *NoopTestSuite) TestStream() {
 			stream.Send(suite.ctx, value)
 		})
 	}
-
-	stream.Shutdown()
 }
 
 func (suite *NoopTestSuite) TestObserver() {

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5
 	github.com/babolivier/go-doh-client v0.0.0-20201028162107-a76cff4cb8b6
 	github.com/d4l3k/messagediff v1.2.1 // indirect
+	github.com/dgraph-io/ristretto v0.0.3 // indirect
 	github.com/jarcoal/httpmock v1.0.8
 	github.com/kentik/patricia v0.0.0-20201202224819-f9447a6e25f1
 	github.com/libp2p/go-reuseport v0.0.2

go.sum

@@ -2,6 +2,7 @@ cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMT
 cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/Knetic/govaluate v3.0.1-0.20171022003610-9aa49832a739+incompatible/go.mod h1:r7JcOSlj0wfOMncg0iLm8Leh48TZaKVeNIfJntJ2wa0=
+github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
 github.com/OneOfOne/xxhash v1.2.8 h1:31czK/TI9sNkxIKfaUfGlU47BAxQ0ztGgd9vPyqimf8=
 github.com/OneOfOne/xxhash v1.2.8/go.mod h1:eZbhyaAYD41SGSSsnmcpxVoRiQ/MPUTjUdIIOT9Um7Q=
 github.com/Shopify/sarama v1.19.0/go.mod h1:FVkBWblsNy7DGZRfXLU0O9RCGt5g3g3yEuWXgklEdEo=
@@ -38,6 +39,8 @@ github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kB
 github.com/casbin/casbin/v2 v2.1.2/go.mod h1:YcPU1XXisHhLzuxH9coDNf2FbKpjGlbCg3n9yuLkIJQ=
 github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
+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 h1:6MnRN8NT7+YBpUIWxHtefFZOKTAPgGjpQSxqLNn0+qY=
 github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/clbanning/x2j v0.0.0-20191024224557-825249438eec/go.mod h1:jMjuTZXRI4dUb/I5gc9Hdhagfvm9+RyrPryS/auMzxE=
@@ -55,7 +58,10 @@ github.com/d4l3k/messagediff v1.2.1/go.mod h1:Oozbb1TVXFac9FtSIxHBMnBCq2qeH/2KkE
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dgraph-io/ristretto v0.0.3 h1:jh22xisGBjrEVnRZ1DVTpBVQm0Xndu8sMl0CWDzSIBI=
+github.com/dgraph-io/ristretto v0.0.3/go.mod h1:KPxhHT9ZxKefz+PCeOGsrHpl1qZ7i70dGTu2u+Ahh6E=
 github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
+github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
 github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/eapache/go-resiliency v1.1.0/go.mod h1:kFI+JgMyC7bLPUVY133qvEBtVayf5mFgVsvEsIPBvNs=
 github.com/eapache/go-xerial-snappy v0.0.0-20180814174437-776d5712da21/go.mod h1:+020luEh2TKB4/GOp8oxxtq0Daoen/Cii55CzbTV6DU=
@@ -277,6 +283,7 @@ github.com/smira/go-statsd v1.3.2 h1:1EeuzxNZ/TD9apbTOFSM9nulqfcsQFmT4u1A2DREabI
 github.com/smira/go-statsd v1.3.2/go.mod h1:1srXJ9/pbnN04G8f4F1jUzsGOnwkPKXciyqpewGlkC4=
 github.com/soheilhy/cmux v0.1.4/go.mod h1:IM3LyeVVIOuxMH7sFAkER9+bJ4dT7Ms6E4xg4kGIyLM=
 github.com/sony/gobreaker v0.4.1/go.mod h1:ZKptC7FHNvhBz7dN2LGjPVBz2sZJmc0/PkyDJOjmxWY=
+github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
 github.com/spf13/cobra v0.0.3/go.mod h1:1l0Ry5zgKvJasoi3XT1TypsSe7PqH0Sj9dhYf7v3XqQ=
 github.com/spf13/pflag v1.0.1/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnInEg4=
 github.com/streadway/amqp v0.0.0-20190404075320-75d898a42a94/go.mod h1:AZpEONHx3DKn8O/DFsRAY58/XVQiIPMTMB1SddzLXVw=

ipblocklist/noop.go

@@ -9,7 +9,6 @@ import (
 type noop struct{}
 
 func (n noop) Contains(ip net.IP) bool { return false }
-func (n noop) Shutdown()               {}
 
 func NewNoop() mtglib.IPBlocklist {
 	return noop{}

mtglib/init.go

@@ -35,12 +35,10 @@ type Network interface {
 
 type AntiReplayCache interface {
 	SeenBefore(data []byte) bool
-	Shutdown()
 }
 
 type IPBlocklist interface {
 	Contains(net.IP) bool
-	Shutdown()
 }
 
 type Event interface {
@@ -50,7 +48,6 @@ type Event interface {
 
 type EventStream interface {
 	Send(context.Context, Event)
-	Shutdown()
 }
 
 type TimeAttackDetector interface {

network/dns_resolver.go

@@ -0,0 +1,88 @@
+package network
+
+import (
+	"net/http"
+	"time"
+
+	doh "github.com/babolivier/go-doh-client"
+	"github.com/dgraph-io/ristretto"
+)
+
+const (
+	dnsResolverSize     = 1024 * 1024 // 1mb
+	dnsResolverKeepTime = 10 * time.Minute
+)
+
+type dnsResolver struct {
+	resolver doh.Resolver
+	cache    *ristretto.Cache
+}
+
+func (d dnsResolver) LookupA(hostname string) []string {
+	key := "\x00." + hostname
+
+	if value, ok := d.cache.Get(key); ok {
+		return value.([]string)
+	}
+
+	var ips []string
+
+	if recs, _, err := d.resolver.LookupA(hostname); err == nil {
+		for _, v := range recs {
+			ips = append(ips, v.IP4)
+		}
+
+		d.cache.SetWithTTL(key, ips, 0, dnsResolverKeepTime)
+	}
+
+	return ips
+}
+
+func (d dnsResolver) LookupAAAA(hostname string) []string {
+	key := "\x01." + hostname
+
+	if value, ok := d.cache.Get(key); ok {
+		return value.([]string)
+	}
+
+	var ips []string
+
+	if recs, _, err := d.resolver.LookupAAAA(hostname); err == nil {
+		for _, v := range recs {
+			ips = append(ips, v.IP6)
+		}
+
+		d.cache.SetWithTTL(key, ips, 0, dnsResolverKeepTime)
+	}
+
+	return ips
+}
+
+func newDNSResolver(hostname string, httpClient *http.Client) dnsResolver {
+	cache, err := ristretto.NewCache(&ristretto.Config{
+		NumCounters: 10 * dnsResolverSize, // nolint: gomnd // taken from official doc as a best practice value
+		MaxCost:     dnsResolverSize,
+		BufferItems: 64, // nolint: gomnd // taken from official doc as a best practice value
+		Cost: func(value interface{}) int64 {
+			var cost int64
+
+			for _, v := range value.([]string) {
+				cost += int64(len([]byte(v)))
+			}
+
+			return cost
+		},
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	return dnsResolver{
+		resolver: doh.Resolver{
+			Host:       hostname,
+			Class:      doh.IN,
+			HTTPClient: httpClient,
+		},
+		cache: cache,
+	}
+}

network/dns_resolver_internal_test.go

@@ -0,0 +1,50 @@
+package network
+
+import (
+	"net"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/suite"
+)
+
+type DNSResolverTestSuite struct {
+	suite.Suite
+
+	d dnsResolver
+}
+
+func (suite *DNSResolverTestSuite) TestLookupA() {
+	suite.d.LookupA("google.com")
+	time.Sleep(10 * time.Millisecond)
+
+	addrs := suite.d.LookupA("google.com")
+
+	for _, v := range addrs {
+		suite.NotEmpty(v)
+		suite.NotNil(net.ParseIP(v).To4())
+	}
+}
+
+func (suite *DNSResolverTestSuite) TestLookupAAAA() {
+	suite.d.LookupAAAA("google.com")
+	time.Sleep(10 * time.Millisecond)
+
+	addrs := suite.d.LookupAAAA("google.com")
+
+	for _, v := range addrs {
+		suite.NotEmpty(v)
+		suite.Nil(net.ParseIP(v).To4())
+		suite.NotNil(net.ParseIP(v).To16())
+	}
+}
+
+func (suite *DNSResolverTestSuite) SetupTest() {
+	suite.d = newDNSResolver("1.1.1.1", &http.Client{})
+}
+
+func TestDNSResolver(t *testing.T) {
+	t.Parallel()
+	suite.Run(t, &DNSResolverTestSuite{})
+}

network/network.go

@@ -10,7 +10,6 @@ import (
 	"time"
 
 	"github.com/9seconds/mtg/v2/mtglib"
-	doh "github.com/babolivier/go-doh-client"
 )
 
 type networkHTTPTransport struct {
@@ -26,7 +25,7 @@ func (n networkHTTPTransport) RoundTrip(req *http.Request) (*http.Response, erro
 
 type network struct {
 	dialer      Dialer
-	dns         doh.Resolver
+	dns         dnsResolver
 	httpTimeout time.Duration
 	userAgent   string
 }
@@ -84,14 +83,11 @@ func (n *network) dnsResolve(protocol, address string) ([]string, error) {
 		go func() {
 			defer wg.Done()
 
-			if recs, _, err := n.dns.LookupA(address); err == nil {
+			resolved := n.dns.LookupA(address)
-				mutex.Lock()
-				defer mutex.Unlock()
 
-				for _, v := range recs {
+			mutex.Lock()
-					ips = append(ips, v.IP4)
+			ips = append(ips, resolved...)
-				}
+			mutex.Unlock()
-			}
 		}()
 	}
 
@@ -102,14 +98,11 @@ func (n *network) dnsResolve(protocol, address string) ([]string, error) {
 		go func() {
 			defer wg.Done()
 
-			if recs, _, err := n.dns.LookupAAAA(address); err == nil {
+			resolved := n.dns.LookupAAAA(address)
-				mutex.Lock()
-				defer mutex.Unlock()
 
-				for _, v := range recs {
+			mutex.Lock()
-					ips = append(ips, v.IP6)
+			ips = append(ips, resolved...)
-				}
+			mutex.Unlock()
-			}
 		}()
 	}
 
@@ -140,11 +133,8 @@ func NewNetwork(dialer Dialer,
 		dialer:      dialer,
 		httpTimeout: httpTimeout,
 		userAgent:   userAgent,
-		dns: doh.Resolver{
+		dns: newDNSResolver(dohHostname,
-			Host:       dohHostname,
+			makeHTTPClient(userAgent, DNSTimeout, dialer.DialContext)),
-			Class:      doh.IN,
-			HTTPClient: makeHTTPClient(userAgent, DNSTimeout, dialer.DialContext),
-		},
 	}, nil
 }