From ab8fd4d005bd8d3fdf677f12c3b59f636c06f1a8 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Wed, 24 Mar 2021 13:37:20 +0100
Subject: [PATCH] p2p/dnsdisc: rate limit resolving before checking cache
 (#22566)

This makes the rate limit apply regardless of whether the node is
already cached.
---
 p2p/dnsdisc/client.go | 35 +++++++++++++++++------------------
 1 file changed, 17 insertions(+), 18 deletions(-)

diff --git a/p2p/dnsdisc/client.go b/p2p/dnsdisc/client.go
index d3e8111ab53b..3e4b50aaddb0 100644
--- a/p2p/dnsdisc/client.go
+++ b/p2p/dnsdisc/client.go
@@ -37,9 +37,10 @@ import (
 
 // Client discovers nodes by querying DNS servers.
 type Client struct {
-	cfg     Config
-	clock   mclock.Clock
-	entries *lru.Cache
+	cfg       Config
+	clock     mclock.Clock
+	entries   *lru.Cache
+	ratelimit *rate.Limiter
 }
 
 // Config holds configuration options for the client.
@@ -97,8 +98,12 @@ func NewClient(cfg Config) *Client {
 		panic(err)
 	}
 	rlimit := rate.NewLimiter(rate.Limit(cfg.RateLimit), 10)
-	cfg.Resolver = &rateLimitResolver{cfg.Resolver, rlimit}
-	return &Client{cfg: cfg, entries: cache, clock: mclock.System{}}
+	return &Client{
+		cfg:       cfg,
+		entries:   cache,
+		clock:     mclock.System{},
+		ratelimit: rlimit,
+	}
 }
 
 // SyncTree downloads the entire node tree at the given URL.
@@ -157,6 +162,13 @@ func parseAndVerifyRoot(txt string, loc *linkEntry) (rootEntry, error) {
 // resolveEntry retrieves an entry from the cache or fetches it from the network
 // if it isn't cached.
 func (c *Client) resolveEntry(ctx context.Context, domain, hash string) (entry, error) {
+	// The rate limit always applies, even when the result might be cached. This is
+	// important because it avoids hot-spinning in consumers of node iterators created on
+	// this client.
+	if err := c.ratelimit.Wait(ctx); err != nil {
+		return nil, err
+	}
+
 	cacheKey := truncateHash(hash)
 	if e, ok := c.entries.Get(cacheKey); ok {
 		return e.(entry), nil
@@ -196,19 +208,6 @@ func (c *Client) doResolveEntry(ctx context.Context, domain, hash string) (entry
 	return nil, nameError{name, errNoEntry}
 }
 
-// rateLimitResolver applies a rate limit to a Resolver.
-type rateLimitResolver struct {
-	r       Resolver
-	limiter *rate.Limiter
-}
-
-func (r *rateLimitResolver) LookupTXT(ctx context.Context, domain string) ([]string, error) {
-	if err := r.limiter.Wait(ctx); err != nil {
-		return nil, err
-	}
-	return r.r.LookupTXT(ctx, domain)
-}
-
 // randomIterator traverses a set of trees and returns nodes found in them.
 type randomIterator struct {
 	cur      *enode.Node