From a0ece197fd07e6a9837906e29df6a77b266f8807 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Mon, 10 Jan 2022 13:13:01 +1100
Subject: [PATCH] Avoid race condition in `address_map`

The check for being connected and sending a message is not atomic. To
not trigger a panic, we need to attempt sending the message and react to
the error accordingly.

Fixes #1049.
---
 daemon/src/address_map.rs | 39 ++++++++++++++++++++-------------------
 1 file changed, 20 insertions(+), 19 deletions(-)

diff --git a/daemon/src/address_map.rs b/daemon/src/address_map.rs
index 7cfc9615a..129ce3cc8 100644
--- a/daemon/src/address_map.rs
+++ b/daemon/src/address_map.rs
@@ -54,15 +54,12 @@ where
         M: Message<Result = ()>,
         A: Handler<M> + ActorName,
     {
-        match self.inner.get(key) {
-            Some(addr) if addr.is_connected() => {
-                addr.send(msg)
-                    .await
-                    .expect("we checked that we are connected");
-                Ok(())
-            }
-            _ => Err(NotConnected::new::<A>()),
-        }
+        self.get(key)?
+            .send(msg)
+            .await
+            .map_err(|_| NotConnected::new::<A>())?;
+
+        Ok(())
     }
 
     /// Sends a message to the actor stored with the given key.
@@ -71,16 +68,20 @@ where
         M: Message<Result = anyhow::Result<()>>,
         A: Handler<M> + ActorName,
     {
-        match self.inner.get(key) {
-            Some(addr) if addr.is_connected() => {
-                let res = addr
-                    .send(msg)
-                    .await
-                    .expect("we checked that we are connected");
-                Ok(res)
-            }
-            _ => Err(NotConnected::new::<A>()),
-        }
+        let result = self
+            .get(key)?
+            .send(msg)
+            .await
+            .map_err(|_| NotConnected::new::<A>())?;
+
+        Ok(result)
+    }
+
+    fn get(&self, key: &K) -> Result<&Address<A>, NotConnected>
+    where
+        A: ActorName,
+    {
+        self.inner.get(key).ok_or_else(|| NotConnected::new::<A>())
     }
 }