From 3f18f6cb365186df0b6a0989d6e3e4f0ee94c850 Mon Sep 17 00:00:00 2001
From: David Renshaw <dwrenshaw@gmail.com>
Date: Fri, 12 Jan 2024 08:36:46 -0500
Subject: [PATCH] implement this_cap() as a method of Params

---
 capnp-rpc/src/local.rs          | 14 ++++++++++++--
 capnp-rpc/src/rpc.rs            | 14 ++++++++++++--
 capnp/src/capability.rs         |  3 +++
 capnp/src/private/capability.rs |  1 +
 4 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/capnp-rpc/src/local.rs b/capnp-rpc/src/local.rs
index 50ac0d128..21e8533da 100644
--- a/capnp-rpc/src/local.rs
+++ b/capnp-rpc/src/local.rs
@@ -62,14 +62,20 @@ impl ResponseHook for Response {
 struct Params {
     request: message::Builder<message::HeapAllocator>,
     cap_table: Vec<Option<Box<dyn ClientHook>>>,
+    this_cap: Box<dyn ClientHook>,
 }
 
 impl Params {
     fn new(
         request: message::Builder<message::HeapAllocator>,
         cap_table: Vec<Option<Box<dyn ClientHook>>>,
+        this_cap: Box<dyn ClientHook>,
     ) -> Self {
-        Self { request, cap_table }
+        Self {
+            request,
+            cap_table,
+            this_cap,
+        }
     }
 }
 
@@ -79,6 +85,10 @@ impl ParamsHook for Params {
         result.imbue(&self.cap_table);
         Ok(result)
     }
+
+    fn this_cap(&self) -> Box<dyn ClientHook> {
+        self.this_cap.add_ref()
+    }
 }
 
 struct Results {
@@ -218,7 +228,7 @@ impl RequestHook for Request {
             method_id,
             client,
         } = tmp;
-        let params = Params::new(message, cap_table);
+        let params = Params::new(message, cap_table, client.add_ref());
 
         let (results_done_fulfiller, results_done_promise) =
             oneshot::channel::<Box<dyn ResultsDoneHook>>();
diff --git a/capnp-rpc/src/rpc.rs b/capnp-rpc/src/rpc.rs
index 729d2395b..3b39b7518 100644
--- a/capnp-rpc/src/rpc.rs
+++ b/capnp-rpc/src/rpc.rs
@@ -946,7 +946,7 @@ impl<VatId> ConnectionState<VatId> {
                     )));
                 }
 
-                let params = Params::new(message, cap_table_array);
+                let params = Params::new(message, cap_table_array, capability.add_ref());
 
                 let answer = Answer::new();
 
@@ -2135,14 +2135,20 @@ impl<VatId> PipelineHook for Pipeline<VatId> {
 pub struct Params {
     request: Box<dyn crate::IncomingMessage>,
     cap_table: Vec<Option<Box<dyn ClientHook>>>,
+    this_cap: Box<dyn ClientHook>,
 }
 
 impl Params {
     fn new(
         request: Box<dyn crate::IncomingMessage>,
         cap_table: Vec<Option<Box<dyn ClientHook>>>,
+        this_cap: Box<dyn ClientHook>,
     ) -> Self {
-        Self { request, cap_table }
+        Self {
+            request,
+            cap_table,
+            this_cap,
+        }
     }
 }
 
@@ -2161,6 +2167,10 @@ impl ParamsHook for Params {
             }
         }
     }
+
+    fn this_cap(&self) -> Box<dyn ClientHook> {
+        self.this_cap.add_ref()
+    }
 }
 
 enum ResultsVariant {
diff --git a/capnp/src/capability.rs b/capnp/src/capability.rs
index 7cff20bbb..950e92bb9 100644
--- a/capnp/src/capability.rs
+++ b/capnp/src/capability.rs
@@ -237,6 +237,9 @@ impl<T> Params<T> {
     {
         self.hook.get()?.get_as()
     }
+    pub fn this_cap<C: FromClientHook>(&self) -> C {
+        FromClientHook::new(self.hook.this_cap())
+    }
 }
 
 /// The return values of a method, written in-place by the method body.
diff --git a/capnp/src/private/capability.rs b/capnp/src/private/capability.rs
index ba725fb60..67568581e 100644
--- a/capnp/src/private/capability.rs
+++ b/capnp/src/private/capability.rs
@@ -107,6 +107,7 @@ pub trait ResultsHook {
 
 pub trait ParamsHook {
     fn get(&self) -> crate::Result<crate::any_pointer::Reader<'_>>;
+    fn this_cap(&self) -> Box<dyn ClientHook>;
 }
 
 // Where should this live?