avoid panic when set_pipeline() is called twice

capnproto · Sep 9, 2024 · 4761bd5 · 4761bd5
1 parent 6647fb3
commit 4761bd5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 9 deletions.
diff --git 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
@@ -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(())
     }