From 7f29acd964dfca9c9e90a87af607a7897eb8ea67 Mon Sep 17 00:00:00 2001
From: Felix Giese <felix.giese@icloud.com>
Date: Sun, 26 Jul 2020 09:40:29 +0200
Subject: [PATCH] time: fix resetting expired timers causing panics (#2587)

* Add Unit Test demonstrating the issue

This test demonstrates a panic that occurs when the user inserts an
item with an instant in the past, and then tries to reset the timeout
using the returned key

* Guard reset_at against removals of expired items

Trying to remove an already expired Timer Wheel entry (called by
DelayQueue.reset()) causes panics in some cases as described in (#2573)

This prevents this panic by removing the item from the expired queue and
not the wheel in these cases

Fixes: #2473
---
 tokio/src/time/delay_queue.rs   | 28 ++++++++++++++++++----------
 tokio/tests/time_delay_queue.rs | 21 +++++++++++++++++++++
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/tokio/src/time/delay_queue.rs b/tokio/src/time/delay_queue.rs
index 55ec7cd68..a947cc6fe 100644
--- a/tokio/src/time/delay_queue.rs
+++ b/tokio/src/time/delay_queue.rs
@@ -427,6 +427,22 @@ impl<T> DelayQueue<T> {
         }
     }
 
+    /// Removes the key fom the expired queue or the timer wheel
+    /// depending on its expiration status
+    ///
+    /// # Panics
+    /// Panics if the key is not contained in the expired queue or the wheel
+    fn remove_key(&mut self, key: &Key) {
+        use crate::time::wheel::Stack;
+
+        // Special case the `expired` queue
+        if self.slab[key.index].expired {
+            self.expired.remove(&key.index, &mut self.slab);
+        } else {
+            self.wheel.remove(&key.index, &mut self.slab);
+        }
+    }
+
     /// Removes the item associated with `key` from the queue.
     ///
     /// There must be an item associated with `key`. The function returns the
@@ -456,15 +472,7 @@ impl<T> DelayQueue<T> {
     /// # }
     /// ```
     pub fn remove(&mut self, key: &Key) -> Expired<T> {
-        use crate::time::wheel::Stack;
-
-        // Special case the `expired` queue
-        if self.slab[key.index].expired {
-            self.expired.remove(&key.index, &mut self.slab);
-        } else {
-            self.wheel.remove(&key.index, &mut self.slab);
-        }
-
+        self.remove_key(key);
         let data = self.slab.remove(key.index);
 
         Expired {
@@ -508,7 +516,7 @@ impl<T> DelayQueue<T> {
     /// # }
     /// ```
     pub fn reset_at(&mut self, key: &Key, when: Instant) {
-        self.wheel.remove(&key.index, &mut self.slab);
+        self.remove_key(key);
 
         // Normalize the deadline. Values cannot be set to expire in the past.
         let when = self.normalize_deadline(when);
diff --git a/tokio/tests/time_delay_queue.rs b/tokio/tests/time_delay_queue.rs
index f04576d5a..ba3c98581 100644
--- a/tokio/tests/time_delay_queue.rs
+++ b/tokio/tests/time_delay_queue.rs
@@ -481,6 +481,27 @@ async fn reset_later_after_slot_starts() {
     assert_eq!(entry, "foo");
 }
 
+#[tokio::test]
+async fn reset_inserted_expired() {
+    time::pause();
+    let mut queue = task::spawn(DelayQueue::new());
+    let now = Instant::now();
+
+    let key = queue.insert_at("foo", now - ms(100));
+
+    // this causes the panic described in #2473
+    queue.reset_at(&key, now + ms(100));
+
+    assert_eq!(1, queue.len());
+
+    delay_for(ms(200)).await;
+
+    let entry = assert_ready_ok!(poll!(queue)).into_inner();
+    assert_eq!(entry, "foo");
+
+    assert_eq!(queue.len(), 0);
+}
+
 #[tokio::test]
 async fn reset_earlier_after_slot_starts() {
     time::pause();