From 4761bd5fc83dc9a10cbebd6af30070d765bf86d4 Mon Sep 17 00:00:00 2001
From: David Renshaw <dwrenshaw@gmail.com>
Date: Mon, 9 Sep 2024 11:17:58 -0400
Subject: [PATCH] avoid panic when set_pipeline() is called twice

---
 capnp-rpc/src/local.rs | 8 ++++----
 capnp-rpc/src/rpc.rs   | 9 ++++-----
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/capnp-rpc/src/local.rs b/capnp-rpc/src/local.rs
index 6c99971c9..bd8dbe8b3 100644
--- a/capnp-rpc/src/local.rs
+++ b/capnp-rpc/src/local.rs
@@ -143,10 +143,10 @@ impl ResultsHook for Results {
         root2.imbue_mut(&mut cap_table2);
         root2.set_as(root.into_reader())?;
         let hook = Box::new(ResultsDone::new(message2, cap_table2)) as Box<dyn ResultsDoneHook>;
-        self.pipeline_sender
-            .take()
-            .unwrap()
-            .complete(Box::new(Pipeline::new(hook)));
+        let Some(sender) = self.pipeline_sender.take() else {
+            return Err(Error::failed("set_pipeline() called twice".into()));
+        };
+        sender.complete(Box::new(Pipeline::new(hook)));
         Ok(())
     }
 
diff --git a/capnp-rpc/src/rpc.rs b/capnp-rpc/src/rpc.rs
index 6aff2dbf2..6d8081076 100644
--- a/capnp-rpc/src/rpc.rs
+++ b/capnp-rpc/src/rpc.rs
@@ -2271,11 +2271,10 @@ impl<VatId> ResultsHook for Results<VatId> {
         let Some(ref mut inner) = self.inner else {
             unreachable!();
         };
-        inner
-            .pipeline_sender
-            .take()
-            .unwrap()
-            .complete(Box::new(local::Pipeline::new(hook)));
+        let Some(sender) = inner.pipeline_sender.take() else {
+            return Err(Error::failed("set_pipeline() called twice".into()));
+        };
+        sender.complete(Box::new(local::Pipeline::new(hook)));
         Ok(())
     }