From 10b78322a3d216b28d421b60cc213f3c2fdb3a38 Mon Sep 17 00:00:00 2001 From: 9seconds Date: Tue, 30 Mar 2021 14:13:42 +0300 Subject: [PATCH] Implement caching dns resolver --- antireplay/noop.go | 1 - antireplay/noop_test.go | 2 - antireplay/stable_bloom_filter.go | 2 - antireplay/stable_bloom_filter_test.go | 2 - cli/proxy.go | 6 -- events/event_stream.go | 10 +-- events/event_stream_test.go | 2 +- events/noop.go | 1 - events/noop_test.go | 2 - go.mod | 1 + go.sum | 7 ++ ipblocklist/noop.go | 1 - mtglib/init.go | 3 - network/dns_resolver.go | 88 ++++++++++++++++++++++++++ network/dns_resolver_internal_test.go | 50 +++++++++++++++ network/network.go | 32 ++++------ 16 files changed, 163 insertions(+), 47 deletions(-) create mode 100644 network/dns_resolver.go create mode 100644 network/dns_resolver_internal_test.go diff --git a/antireplay/noop.go b/antireplay/noop.go index b07d34f..8613e40 100644 --- a/antireplay/noop.go +++ b/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{} diff --git a/antireplay/noop_test.go b/antireplay/noop_test.go index 42f39eb..b78c48d 100644 --- a/antireplay/noop_test.go +++ b/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) { diff --git a/antireplay/stable_bloom_filter.go b/antireplay/stable_bloom_filter.go index ea4f5a7..b52631e 100644 --- a/antireplay/stable_bloom_filter.go +++ b/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()) diff --git a/antireplay/stable_bloom_filter_test.go b/antireplay/stable_bloom_filter_test.go index 9b2517e..accb590 100644 --- a/antireplay/stable_bloom_filter_test.go +++ b/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) { diff --git a/cli/proxy.go b/cli/proxy.go index 22ebc38..94b228f 100644 --- a/cli/proxy.go +++ b/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 } diff --git a/events/event_stream.go b/events/event_stream.go index 891382c..9995eb0 100644 --- a/events/event_stream.go +++ b/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()), diff --git a/events/event_stream_test.go b/events/event_stream_test.go index 179e29f..599d839 100644 --- a/events/event_stream_test.go +++ b/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() { diff --git a/events/noop.go b/events/noop.go index afc2c66..2fdb238 100644 --- a/events/noop.go +++ b/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{} diff --git a/events/noop_test.go b/events/noop_test.go index 024febc..4a7b20b 100644 --- a/events/noop_test.go +++ b/events/noop_test.go @@ -70,8 +70,6 @@ func (suite *NoopTestSuite) TestStream() { stream.Send(suite.ctx, value) }) } - - stream.Shutdown() } func (suite *NoopTestSuite) TestObserver() { diff --git a/go.mod b/go.mod index ad324f1..9112aba 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 0cf083c..12d7451 100644 --- a/go.sum +++ b/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= diff --git a/ipblocklist/noop.go b/ipblocklist/noop.go index 4b29449..dcbed19 100644 --- a/ipblocklist/noop.go +++ b/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{} diff --git a/mtglib/init.go b/mtglib/init.go index 0660e9e..010a78c 100644 --- a/mtglib/init.go +++ b/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 { diff --git a/network/dns_resolver.go b/network/dns_resolver.go new file mode 100644 index 0000000..3205b8f --- /dev/null +++ b/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, + } +} diff --git a/network/dns_resolver_internal_test.go b/network/dns_resolver_internal_test.go new file mode 100644 index 0000000..a5934e8 --- /dev/null +++ b/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{}) +} diff --git a/network/network.go b/network/network.go index 9df70a0..c65c282 100644 --- a/network/network.go +++ b/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 { - mutex.Lock() - defer mutex.Unlock() + resolved := n.dns.LookupA(address) - for _, v := range recs { - ips = append(ips, v.IP4) - } - } + mutex.Lock() + 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 { - mutex.Lock() - defer mutex.Unlock() + resolved := n.dns.LookupAAAA(address) - for _, v := range recs { - ips = append(ips, v.IP6) - } - } + mutex.Lock() + ips = append(ips, resolved...) + mutex.Unlock() }() } @@ -140,11 +133,8 @@ func NewNetwork(dialer Dialer, dialer: dialer, httpTimeout: httpTimeout, userAgent: userAgent, - dns: doh.Resolver{ - Host: dohHostname, - Class: doh.IN, - HTTPClient: makeHTTPClient(userAgent, DNSTimeout, dialer.DialContext), - }, + dns: newDNSResolver(dohHostname, + makeHTTPClient(userAgent, DNSTimeout, dialer.DialContext)), }, nil }