From 7644795950481432f6ba230421d41349b69774e8 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Mon, 29 Mar 2021 14:17:35 +0200
Subject: [PATCH] eth/protocols/snap: try to prevent requests timing out

---
 eth/protocols/snap/handler.go | 13 +++++++++----
 eth/protocols/snap/sync.go    |  2 +-
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/eth/protocols/snap/handler.go b/eth/protocols/snap/handler.go
index 4169306c2..f7939964f 100644
--- a/eth/protocols/snap/handler.go
+++ b/eth/protocols/snap/handler.go
@@ -50,6 +50,11 @@ const (
 	// maxTrieNodeLookups is the maximum number of state trie nodes to serve. This
 	// number is there to limit the number of disk lookups.
 	maxTrieNodeLookups = 1024
+
+	// maxTrieNodeTimeSpent is the maximum time we should spend on looking up trie nodes.
+	// If we spend too much time, then it's a fairly high chance of timing out
+	// at the remote side, which means all the work is in vain.
+	maxTrieNodeTimeSpent = 5 * time.Second
 )
 
 // Handler is a callback to invoke from an outside runner after the boilerplate
@@ -129,7 +134,7 @@ func handleMessage(backend Backend, peer *Peer) error {
 		return fmt.Errorf("%w: %v > %v", errMsgTooLarge, msg.Size, maxMessageSize)
 	}
 	defer msg.Discard()
-
+	start := time.Now()
 	// Track the emount of time it takes to serve the request and run the handler
 	if metrics.Enabled {
 		h := fmt.Sprintf("%s/%s/%d/%#02x", p2p.HandleHistName, ProtocolName, peer.Version(), msg.Code)
@@ -140,7 +145,7 @@ func handleMessage(backend Backend, peer *Peer) error {
 				)
 			}
 			metrics.GetOrRegisterHistogramLazy(h, nil, sampler).Update(time.Since(start).Microseconds())
-		}(time.Now())
+		}(start)
 	}
 	// Handle the message depending on its contents
 	switch {
@@ -470,13 +475,13 @@ func handleMessage(backend Backend, peer *Peer) error {
 					bytes += uint64(len(blob))
 
 					// Sanity check limits to avoid DoS on the store trie loads
-					if bytes > req.Bytes || loads > maxTrieNodeLookups {
+					if bytes > req.Bytes || loads > maxTrieNodeLookups || time.Since(start) > maxTrieNodeTimeSpent {
 						break
 					}
 				}
 			}
 			// Abort request processing if we've exceeded our limits
-			if bytes > req.Bytes || loads > maxTrieNodeLookups {
+			if bytes > req.Bytes || loads > maxTrieNodeLookups || time.Since(start) > maxTrieNodeTimeSpent {
 				break
 			}
 		}
diff --git a/eth/protocols/snap/sync.go b/eth/protocols/snap/sync.go
index af581df07..2924fa080 100644
--- a/eth/protocols/snap/sync.go
+++ b/eth/protocols/snap/sync.go
@@ -85,7 +85,7 @@ const (
 var (
 	// requestTimeout is the maximum time a peer is allowed to spend on serving
 	// a single network request.
-	requestTimeout = 10 * time.Second // TODO(karalabe): Make it dynamic ala fast-sync?
+	requestTimeout = 15 * time.Second // TODO(karalabe): Make it dynamic ala fast-sync?
 )
 
 // ErrCancelled is returned from snap syncing if the operation was prematurely