From 29471467e3a753276e9321658721c026596b8d59 Mon Sep 17 00:00:00 2001
From: Max Inden <mail@max-inden.de>
Date: Wed, 18 Mar 2020 10:15:33 +0100
Subject: [PATCH] protocols/kad: Fix right shift overflow panic in
 record_received  (#1492)

* protocols/kad: Add test to reproduce right shift overflow panic

* protocols/kad: Fix right shift overflow panic in record_received

Within `Behaviour::record_received` the exponentially decreasing
expiration based on the distance to the target for a record is
calculated as following:

1. Calculate the amount of nodes between us and the record key beyond
the k replication constant as `n`.

2. Shift the configured record time-to-live `n` times to the right to
calculate an exponentially decreasing expiration.

The configured record time-to-live is a u64. If `n` is larger or equal
to 64 the right shift will lead to an overflow which panics in debug
mode.

This patch uses a checked right shift instead, defaulting to 0 (`now +
0`) for the expiration on overflow.

* protocols/kad: Put attribute below comment

* protocols/kad: Extract shifting logic and rework test

Extract right shift into isolated function and replace complex
regression test with small isolated one.

* protocols/kad/src/behaviour: Refactor exp_decr_expiration

Co-authored-by: Roman Borschel <romanb@users.noreply.github.com>
---
 protocols/kad/src/behaviour.rs      | 10 ++++++----
 protocols/kad/src/behaviour/test.rs | 12 ++++++++++++
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/protocols/kad/src/behaviour.rs b/protocols/kad/src/behaviour.rs
index 0c04e5da224..be51c3dcf09 100644
--- a/protocols/kad/src/behaviour.rs
+++ b/protocols/kad/src/behaviour.rs
@@ -956,9 +956,7 @@ where
         let num_between = self.kbuckets.count_nodes_between(&target);
         let k = self.queries.config().replication_factor.get();
         let num_beyond_k = (usize::max(k, num_between) - k) as u32;
-        let expiration = self.record_ttl.map(|ttl|
-            now + Duration::from_secs(ttl.as_secs() >> num_beyond_k)
-        );
+        let expiration = self.record_ttl.map(|ttl| now + exp_decrease(ttl, num_beyond_k));
         // The smaller TTL prevails. Only if neither TTL is set is the record
         // stored "forever".
         record.expires = record.expires.or(expiration).min(expiration);
@@ -1030,6 +1028,11 @@ where
     }
 }
 
+/// Exponentially decrease the given duration (base 2).
+fn exp_decrease(ttl: Duration, exp: u32) -> Duration {
+    Duration::from_secs(ttl.as_secs().checked_shr(exp).unwrap_or(0))
+}
+
 impl<TStore> NetworkBehaviour for Kademlia<TStore>
 where
     for<'a> TStore: RecordStore<'a>,
@@ -1911,4 +1914,3 @@ impl QueryInfo {
         }
     }
 }
-
diff --git a/protocols/kad/src/behaviour/test.rs b/protocols/kad/src/behaviour/test.rs
index 274e5e17af5..224368118c0 100644
--- a/protocols/kad/src/behaviour/test.rs
+++ b/protocols/kad/src/behaviour/test.rs
@@ -682,3 +682,15 @@ fn exceed_jobs_max_queries() {
         })
     )
 }
+
+#[test]
+fn exp_decr_expiration_overflow() {
+    fn prop_no_panic(ttl: Duration, factor: u32) {
+        exp_decrease(ttl, factor);
+    }
+
+    // Right shifting a u64 by >63 results in a panic.
+    prop_no_panic(KademliaConfig::default().record_ttl.unwrap(), 64);
+
+    quickcheck(prop_no_panic as fn(_, _))
+}