Skip to main content

capnp_rpc/
rpc.rs

1// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
2// Licensed under the MIT License:
3//
4// Permission is hereby granted, free of charge, to any person obtaining a copy
5// of this software and associated documentation files (the "Software"), to deal
6// in the Software without restriction, including without limitation the rights
7// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
8// copies of the Software, and to permit persons to whom the Software is
9// furnished to do so, subject to the following conditions:
10//
11// The above copyright notice and this permission notice shall be included in
12// all copies or substantial portions of the Software.
13//
14// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
15// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
16// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
17// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
18// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
19// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
20// THE SOFTWARE.
21
22use std::pin::Pin;
23use std::task::{Context, Poll};
24
25use capnp::any_pointer;
26use capnp::capability::Promise;
27use capnp::private::capability::{
28    ClientHook, ParamsHook, PipelineHook, PipelineOp, RequestHook, ResponseHook, ResultsHook,
29};
30use capnp::Error;
31
32use futures::channel::oneshot;
33use futures::{future, Future, FutureExt, TryFutureExt};
34
35use std::cell::{Cell, RefCell};
36use std::cmp::Reverse;
37use std::collections::binary_heap::BinaryHeap;
38use std::collections::hash_map::{self, HashMap};
39use std::mem;
40use std::rc::{Rc, Weak};
41
42use crate::attach::Attach;
43use crate::local::ResultsDoneHook;
44use crate::rpc_capnp::{
45    bootstrap, call, cap_descriptor, disembargo, exception, finish, message, message_target,
46    payload, promised_answer, resolve, return_,
47};
48use crate::task_set::TaskSet;
49use crate::{broken, local, queued};
50
51pub(crate) type QuestionId = u32;
52pub(crate) type AnswerId = QuestionId;
53pub(crate) type ExportId = u32;
54pub(crate) type ImportId = ExportId;
55
56pub(crate) struct ImportTable<T> {
57    slots: HashMap<u32, T>,
58}
59
60impl<T> ImportTable<T> {
61    pub(crate) fn new() -> Self {
62        Self {
63            slots: HashMap::new(),
64        }
65    }
66}
67
68struct ExportTable<T> {
69    slots: Vec<Option<T>>,
70
71    // prioritize lower values
72    free_ids: BinaryHeap<Reverse<u32>>,
73}
74
75struct ExportTableIter<'a, T>
76where
77    T: 'a,
78{
79    table: &'a ExportTable<T>,
80    idx: usize,
81}
82
83impl<'a, T> ::std::iter::Iterator for ExportTableIter<'a, T>
84where
85    T: 'a,
86{
87    type Item = &'a T;
88    fn next(&mut self) -> Option<&'a T> {
89        while self.idx < self.table.slots.len() {
90            let idx = self.idx;
91            self.idx += 1;
92            if let Some(v) = &self.table.slots[idx] {
93                return Some(v);
94            }
95        }
96        None
97    }
98}
99
100impl<T> ExportTable<T> {
101    pub(crate) fn new() -> Self {
102        Self {
103            slots: Vec::new(),
104            free_ids: BinaryHeap::new(),
105        }
106    }
107
108    pub(crate) fn erase(&mut self, id: u32) {
109        self.slots[id as usize] = None;
110        self.free_ids.push(Reverse(id));
111    }
112
113    pub(crate) fn push(&mut self, val: T) -> u32 {
114        match self.free_ids.pop() {
115            Some(Reverse(id)) => {
116                self.slots[id as usize] = Some(val);
117                id
118            }
119            None => {
120                self.slots.push(Some(val));
121                self.slots.len() as u32 - 1
122            }
123        }
124    }
125
126    pub(crate) fn find(&mut self, id: u32) -> Option<&mut T> {
127        let idx = id as usize;
128        if idx < self.slots.len() {
129            self.slots[idx].as_mut()
130        } else {
131            None
132        }
133    }
134
135    pub(crate) fn iter(&self) -> ExportTableIter<'_, T> {
136        ExportTableIter {
137            table: self,
138            idx: 0,
139        }
140    }
141}
142
143struct Question<VatId>
144where
145    VatId: 'static,
146{
147    is_awaiting_return: bool,
148
149    #[allow(dead_code)]
150    param_exports: Vec<ExportId>,
151
152    #[allow(dead_code)]
153    is_tail_call: bool,
154
155    /// The local QuestionRef, set to None when it is destroyed.
156    self_ref: Option<Weak<RefCell<QuestionRef<VatId>>>>,
157
158    /// If true, don't send a Finish message.
159    skip_finish: bool,
160}
161
162impl<VatId> Question<VatId> {
163    fn new() -> Self {
164        Self {
165            is_awaiting_return: true,
166            param_exports: Vec::new(),
167            is_tail_call: false,
168            self_ref: None,
169            skip_finish: false,
170        }
171    }
172}
173
174/// A reference to an entry on the question table.  Used to detect when the `Finish` message
175/// can be sent.
176struct QuestionRef<VatId>
177where
178    VatId: 'static,
179{
180    connection_state: Rc<ConnectionState<VatId>>,
181    id: QuestionId,
182    fulfiller: Option<oneshot::Sender<Promise<Response<VatId>, Error>>>,
183}
184
185impl<VatId> QuestionRef<VatId> {
186    fn new(
187        state: Rc<ConnectionState<VatId>>,
188        id: QuestionId,
189        fulfiller: oneshot::Sender<Promise<Response<VatId>, Error>>,
190    ) -> Self {
191        Self {
192            connection_state: state,
193            id,
194            fulfiller: Some(fulfiller),
195        }
196    }
197    fn fulfill(&mut self, response: Promise<Response<VatId>, Error>) {
198        if let Some(fulfiller) = self.fulfiller.take() {
199            let _ = fulfiller.send(response);
200        }
201    }
202
203    fn reject(&mut self, err: Error) {
204        if let Some(fulfiller) = self.fulfiller.take() {
205            let _ = fulfiller.send(Promise::err(err));
206        }
207    }
208}
209
210impl<VatId> Drop for QuestionRef<VatId> {
211    fn drop(&mut self) {
212        let mut questions = self.connection_state.questions.borrow_mut();
213        let Some(q) = &mut questions.slots[self.id as usize] else {
214            unreachable!()
215        };
216        if let Ok(ref mut c) = *self.connection_state.connection.borrow_mut() {
217            if !q.skip_finish {
218                let mut message = c.new_outgoing_message(5);
219                {
220                    let root: message::Builder = message.get_body().unwrap().init_as();
221                    let mut builder = root.init_finish();
222                    builder.set_question_id(self.id);
223
224                    // If we're still awaiting a return, then this request is being
225                    // canceled, and we're going to ignore any capabilities in the return
226                    // message, so set releaseResultCaps true. If we already received the
227                    // return, then we've already built local proxies for the caps and will
228                    // send Release messages when those are destroyed.
229                    builder.set_release_result_caps(q.is_awaiting_return);
230                }
231                let _ = message.send();
232            }
233        }
234
235        if q.is_awaiting_return {
236            // Still waiting for return, so just remove the QuestionRef pointer from the table.
237            q.self_ref = None;
238        } else {
239            // Call has already returned, so we can now remove it from the table.
240            questions.erase(self.id)
241        }
242    }
243}
244
245struct Answer<VatId>
246where
247    VatId: 'static,
248{
249    return_has_been_sent: bool,
250
251    // Send pipelined calls here.  Becomes null as soon as a `Finish` is received.
252    pipeline: Option<Box<dyn PipelineHook>>,
253
254    // For locally-redirected calls (Call.sendResultsTo.yourself), this is a promise for the call
255    // result, to be picked up by a subsequent `Return`.
256    redirected_results: Option<Promise<Response<VatId>, Error>>,
257
258    received_finish: Rc<Cell<bool>>,
259    call_completion_promise: Option<Promise<(), Error>>,
260
261    // List of exports that were sent in the results.  If the finish has `releaseResultCaps` these
262    // will need to be released.
263    result_exports: Vec<ExportId>,
264}
265
266impl<VatId> Answer<VatId> {
267    fn new() -> Self {
268        Self {
269            return_has_been_sent: false,
270            pipeline: None,
271            redirected_results: None,
272            received_finish: Rc::new(Cell::new(false)),
273            call_completion_promise: None,
274            result_exports: Vec::new(),
275        }
276    }
277}
278
279pub(crate) struct Export {
280    refcount: u32,
281
282    /// If true, this is the canonical export entry for this clientHook, that is,
283    /// `exports_by_cap[clientHook]` points to this entry.
284    canonical: bool,
285
286    client_hook: Box<dyn ClientHook>,
287
288    // If this export is a promise (not a settled capability), the `resolve_op` represents the
289    // ongoing operation to wait for that promise to resolve and then send a `Resolve` message.
290    resolve_op: Promise<(), Error>,
291}
292
293impl Export {
294    fn new(client_hook: Box<dyn ClientHook>) -> Self {
295        Self {
296            refcount: 1,
297            canonical: false,
298            client_hook,
299            resolve_op: Promise::err(Error::failed("no resolve op".to_string())),
300        }
301    }
302}
303
304pub(crate) struct Import<VatId>
305where
306    VatId: 'static,
307{
308    import_client: Weak<RefCell<ImportClient<VatId>>>,
309
310    // Either a copy of importClient, or, in the case of promises, the wrapping PromiseClient.
311    // Becomes null when it is discarded *or* when the import is destroyed (e.g. the promise is
312    // resolved and the import is no longer needed).
313    app_client: Option<WeakClient<VatId>>,
314
315    // If non-null, the import is a promise.
316    promise_client_to_resolve: Option<Weak<RefCell<PromiseClient<VatId>>>>,
317}
318
319impl<VatId> Import<VatId> {
320    fn new(import_client: &Rc<RefCell<ImportClient<VatId>>>) -> Self {
321        Self {
322            import_client: Rc::downgrade(import_client),
323            app_client: None,
324            promise_client_to_resolve: None,
325        }
326    }
327}
328
329struct Embargo {
330    fulfiller: Option<oneshot::Sender<Result<(), Error>>>,
331}
332
333impl Embargo {
334    fn new(fulfiller: oneshot::Sender<Result<(), Error>>) -> Self {
335        Self {
336            fulfiller: Some(fulfiller),
337        }
338    }
339}
340
341fn to_pipeline_ops(
342    ops: ::capnp::struct_list::Reader<promised_answer::op::Owned>,
343) -> ::capnp::Result<Vec<PipelineOp>> {
344    let mut result = Vec::new();
345    for op in ops {
346        match op.which()? {
347            promised_answer::op::Noop(()) => {
348                result.push(PipelineOp::Noop);
349            }
350            promised_answer::op::GetPointerField(idx) => {
351                result.push(PipelineOp::GetPointerField(idx));
352            }
353        }
354    }
355    Ok(result)
356}
357
358fn from_error(error: &Error, mut builder: exception::Builder) {
359    let typ = match error.kind {
360        ::capnp::ErrorKind::Failed => exception::Type::Failed,
361        ::capnp::ErrorKind::Overloaded => exception::Type::Overloaded,
362        ::capnp::ErrorKind::Disconnected => exception::Type::Disconnected,
363        ::capnp::ErrorKind::Unimplemented => exception::Type::Unimplemented,
364        ::capnp::ErrorKind::SettingDynamicCapabilitiesIsUnsupported => {
365            exception::Type::Unimplemented
366        }
367        _ => exception::Type::Failed,
368    };
369    builder.set_type(typ);
370    match error.kind {
371        ::capnp::ErrorKind::Failed
372        | ::capnp::ErrorKind::Overloaded
373        | ::capnp::ErrorKind::Disconnected
374        | ::capnp::ErrorKind::Unimplemented => {
375            builder.set_reason(&error.extra);
376        }
377        _ => {
378            // There is extra information in `error.kind` that is not
379            // captured by `typ`. We call `error.to_string()` to allow that
380            // information to be recorded in the `reason` field.
381            builder.set_reason(error.to_string());
382        }
383    }
384}
385
386fn remote_exception_to_error(exception: exception::Reader) -> Error {
387    let (kind, reason) = match (exception.get_type(), exception.get_reason()) {
388        (Ok(exception::Type::Failed), Ok(reason)) => (::capnp::ErrorKind::Failed, reason),
389        (Ok(exception::Type::Overloaded), Ok(reason)) => (::capnp::ErrorKind::Overloaded, reason),
390        (Ok(exception::Type::Disconnected), Ok(reason)) => {
391            (::capnp::ErrorKind::Disconnected, reason)
392        }
393        (Ok(exception::Type::Unimplemented), Ok(reason)) => {
394            (::capnp::ErrorKind::Unimplemented, reason)
395        }
396        _ => (::capnp::ErrorKind::Failed, "(malformed error)".into()),
397    };
398    let reason_str = reason
399        .to_str()
400        .unwrap_or("<malformed utf-8 in error reason>");
401    Error {
402        extra: format!("remote exception: {reason_str}"),
403        kind,
404    }
405}
406
407pub(crate) struct ConnectionErrorHandler<VatId>
408where
409    VatId: 'static,
410{
411    weak_state: Weak<ConnectionState<VatId>>,
412}
413
414impl<VatId> ConnectionErrorHandler<VatId> {
415    fn new(weak_state: Weak<ConnectionState<VatId>>) -> Self {
416        Self { weak_state }
417    }
418}
419
420impl<VatId> crate::task_set::TaskReaper<capnp::Error> for ConnectionErrorHandler<VatId> {
421    fn task_failed(&mut self, error: ::capnp::Error) {
422        if let Some(state) = self.weak_state.upgrade() {
423            state.disconnect(error)
424        }
425    }
426}
427
428pub struct ConnectionState<VatId>
429where
430    VatId: 'static,
431{
432    bootstrap_cap: Box<dyn ClientHook>,
433    exports: RefCell<ExportTable<Export>>,
434    questions: RefCell<ExportTable<Question<VatId>>>,
435    answers: RefCell<ImportTable<Answer<VatId>>>,
436    imports: RefCell<ImportTable<Import<VatId>>>,
437
438    /// Exports keyed by ClientHook::get_ptr().
439    exports_by_cap: RefCell<HashMap<usize, ExportId>>,
440
441    embargoes: RefCell<ExportTable<Embargo>>,
442
443    tasks: RefCell<Option<crate::task_set::TaskSetHandle<capnp::Error>>>,
444    connection: RefCell<::std::result::Result<Box<dyn crate::Connection<VatId>>, ::capnp::Error>>,
445    disconnect_fulfiller: RefCell<Option<oneshot::Sender<Promise<(), Error>>>>,
446
447    // Set when `disconnect()` is called. Resolves once the connection's `shutdown()`
448    // has completed. `Disconnector` futures wait on a clone of this.
449    disconnect_promise: RefCell<Option<future::Shared<Promise<(), Error>>>>,
450
451    client_downcast_map: RefCell<HashMap<usize, WeakClient<VatId>>>,
452}
453
454impl<VatId> ConnectionState<VatId> {
455    pub fn new(
456        bootstrap_cap: Box<dyn ClientHook>,
457        connection: Box<dyn crate::Connection<VatId>>,
458        disconnect_fulfiller: oneshot::Sender<Promise<(), Error>>,
459    ) -> (TaskSet<Error>, Rc<Self>) {
460        let state = Rc::new(Self {
461            bootstrap_cap,
462            exports: RefCell::new(ExportTable::new()),
463            questions: RefCell::new(ExportTable::new()),
464            answers: RefCell::new(ImportTable::new()),
465            imports: RefCell::new(ImportTable::new()),
466            exports_by_cap: RefCell::new(HashMap::new()),
467            embargoes: RefCell::new(ExportTable::new()),
468            tasks: RefCell::new(None),
469            connection: RefCell::new(Ok(connection)),
470            disconnect_fulfiller: RefCell::new(Some(disconnect_fulfiller)),
471            disconnect_promise: RefCell::new(None),
472            client_downcast_map: RefCell::new(HashMap::new()),
473        });
474        let (mut handle, tasks) =
475            TaskSet::new(Box::new(ConnectionErrorHandler::new(Rc::downgrade(&state))));
476
477        handle.add(Self::message_loop(Rc::downgrade(&state)));
478        *state.tasks.borrow_mut() = Some(handle);
479        (tasks, state)
480    }
481
482    fn new_outgoing_message(
483        &self,
484        first_segment_words: u32,
485    ) -> capnp::Result<Box<dyn crate::OutgoingMessage>> {
486        match self.connection.borrow_mut().as_mut() {
487            Err(e) => Err(e.clone()),
488            Ok(c) => Ok(c.new_outgoing_message(first_segment_words)),
489        }
490    }
491
492    fn disconnect(&self, error: ::capnp::Error) {
493        if self.connection.borrow().is_err() {
494            // Already disconnected.
495            return;
496        }
497
498        // Carefully pull all the objects out of the tables prior to releasing them because their
499        // destructors could come back and mess with the tables.
500        let mut pipelines_to_release = Vec::new();
501        let mut clients_to_release = Vec::new();
502        //let mut tail_calls_to_release = Vec::new();
503        let mut resolve_ops_to_release = Vec::new();
504
505        for q in self.questions.borrow().iter() {
506            if let Some(ref weak_question_ref) = q.self_ref {
507                if let Some(question_ref) = weak_question_ref.upgrade() {
508                    question_ref.borrow_mut().reject(error.clone());
509                }
510            }
511        }
512
513        {
514            let answer_slots = &mut self.answers.borrow_mut().slots;
515            for (_, ref mut answer) in answer_slots.iter_mut() {
516                // TODO tail call
517                pipelines_to_release.push(answer.pipeline.take())
518            }
519        }
520
521        let len = self.exports.borrow().slots.len();
522        for idx in 0..len {
523            if let Some(exp) = self.exports.borrow_mut().slots[idx].take() {
524                let Export {
525                    client_hook,
526                    resolve_op,
527                    ..
528                } = exp;
529                clients_to_release.push(client_hook);
530                resolve_ops_to_release.push(resolve_op);
531            }
532        }
533        *self.exports.borrow_mut() = ExportTable::new();
534
535        {
536            let import_slots = &mut self.imports.borrow_mut().slots;
537            for (_, ref mut import) in import_slots.iter_mut() {
538                if let Some(f) = import.promise_client_to_resolve.take() {
539                    if let Some(promise_client) = f.upgrade() {
540                        promise_client.borrow_mut().resolve(Err(error.clone()));
541                    }
542                }
543            }
544        }
545
546        let len = self.embargoes.borrow().slots.len();
547        for idx in 0..len {
548            if let Some(ref mut emb) = self.embargoes.borrow_mut().slots[idx] {
549                if let Some(f) = emb.fulfiller.take() {
550                    let _ = f.send(Err(error.clone()));
551                }
552            }
553        }
554        *self.embargoes.borrow_mut() = ExportTable::new();
555
556        drop(pipelines_to_release);
557        drop(clients_to_release);
558        drop(resolve_ops_to_release);
559        // TODO drop tail calls
560
561        match *self.connection.borrow_mut() {
562            Ok(ref mut c) => {
563                let mut message = c.new_outgoing_message(100); // TODO estimate size
564                {
565                    let builder = message
566                        .get_body()
567                        .unwrap()
568                        .init_as::<message::Builder>()
569                        .init_abort();
570                    from_error(&error, builder);
571                }
572                let _ = message.send();
573            }
574            Err(_) => unreachable!(),
575        }
576
577        let connection = mem::replace(&mut *self.connection.borrow_mut(), Err(error.clone()));
578
579        let Ok(mut c) = connection else {
580            unreachable!()
581        };
582        let promise = c.shutdown(Err(error)).then(|r| match r {
583            Ok(()) => Promise::ok(()),
584            Err(e) => {
585                if e.kind != ::capnp::ErrorKind::Disconnected {
586                    // Don't report disconnects as an error.
587                    Promise::err(e)
588                } else {
589                    Promise::ok(())
590                }
591            }
592        });
593        let shutdown_promise = Promise::from_future(promise.attach(c)).shared();
594        *self.disconnect_promise.borrow_mut() = Some(shutdown_promise.clone());
595        let Some(fulfiller) = self.disconnect_fulfiller.borrow_mut().take() else {
596            unreachable!()
597        };
598        let _ = fulfiller.send(Promise::from_future(shutdown_promise));
599    }
600
601    // Transform a future into a promise that gets executed even if it is never polled.
602    // Dropping the returned promise cancels the computation.
603    fn eagerly_evaluate<T, F>(&self, task: F) -> Promise<T, Error>
604    where
605        F: Future<Output = Result<T, Error>> + 'static + Unpin,
606        T: 'static,
607    {
608        let (tx, rx) = oneshot::channel::<Result<T, Error>>();
609        let (tx2, rx2) = oneshot::channel::<()>();
610        let f1 = Box::pin(task.map(move |r| {
611            let _ = tx.send(r);
612        })) as Pin<Box<dyn Future<Output = ()> + Unpin>>;
613        let f2 = Box::pin(rx2.map(drop)) as Pin<Box<dyn Future<Output = ()> + Unpin>>;
614
615        self.add_task(future::select(f1, f2).map(|_| Ok(())));
616        Promise::from_future(rx.map_err(crate::canceled_to_error).map(|r| {
617            drop(tx2);
618            r?
619        }))
620    }
621
622    fn add_task<F>(&self, task: F)
623    where
624        F: Future<Output = Result<(), Error>> + 'static,
625    {
626        if let Some(ref mut tasks) = *self.tasks.borrow_mut() {
627            tasks.add(task);
628        }
629    }
630
631    pub fn bootstrap(state: &Rc<Self>) -> Box<dyn ClientHook> {
632        let question_id = state.questions.borrow_mut().push(Question::new());
633
634        let (fulfiller, promise) = oneshot::channel();
635        let promise = promise.map_err(crate::canceled_to_error);
636        let promise = promise.and_then(|response_promise| response_promise);
637        let question_ref = Rc::new(RefCell::new(QuestionRef::new(
638            state.clone(),
639            question_id,
640            fulfiller,
641        )));
642        let promise = promise.attach(question_ref.clone());
643        match state.questions.borrow_mut().slots[question_id as usize] {
644            Some(ref mut q) => {
645                q.self_ref = Some(Rc::downgrade(&question_ref));
646            }
647            None => unreachable!(),
648        }
649        match *state.connection.borrow_mut() {
650            Ok(ref mut c) => {
651                let mut message = c.new_outgoing_message(5);
652                {
653                    let mut builder = message
654                        .get_body()
655                        .unwrap()
656                        .init_as::<message::Builder>()
657                        .init_bootstrap();
658                    builder.set_question_id(question_id);
659                }
660                let _ = message.send();
661            }
662            Err(_) => panic!(),
663        }
664
665        let pipeline = Pipeline::new(state, question_ref, Some(Promise::from_future(promise)));
666        pipeline.get_pipelined_cap_move(Vec::new())
667    }
668
669    fn message_loop(weak_state: Weak<Self>) -> Promise<(), capnp::Error> {
670        let Some(state) = weak_state.upgrade() else {
671            return Promise::err(Error::disconnected(
672                "message loop cannot continue without a connection".into(),
673            ));
674        };
675
676        let promise = match *state.connection.borrow_mut() {
677            Err(_) => return Promise::ok(()),
678            Ok(ref mut connection) => connection.receive_incoming_message(),
679        };
680
681        Promise::from_future(async move {
682            match promise.await? {
683                Some(m) => {
684                    Self::handle_message(&weak_state, m)?;
685                    weak_state
686                        .upgrade()
687                        .expect("message loop outlived connection state?")
688                        .add_task(Self::message_loop(weak_state));
689                }
690                None => {
691                    weak_state
692                        .upgrade()
693                        .expect("message loop outlived connection state?")
694                        .disconnect(Error::disconnected("Peer disconnected.".to_string()));
695                }
696            }
697            Ok(())
698        })
699    }
700
701    fn send_unimplemented(
702        connection_state: &Rc<Self>,
703        message: &dyn crate::IncomingMessage,
704    ) -> capnp::Result<()> {
705        let mut out_message = connection_state.new_outgoing_message(50)?; // XXX size hint
706        {
707            let mut root: message::Builder = out_message.get_body()?.get_as()?;
708            root.set_unimplemented(message.get_body()?.get_as()?)?;
709        }
710        let _ = out_message.send();
711        Ok(())
712    }
713
714    fn handle_unimplemented(
715        connection_state: &Rc<Self>,
716        message: message::Reader,
717    ) -> capnp::Result<()> {
718        match message.which()? {
719            message::Resolve(resolve) => {
720                let resolve = resolve?;
721                match resolve.which()? {
722                    resolve::Cap(c) => match c?.which()? {
723                        cap_descriptor::None(()) => (),
724                        cap_descriptor::SenderHosted(export_id) => {
725                            connection_state.release_export(export_id, 1)?;
726                        }
727                        cap_descriptor::SenderPromise(export_id) => {
728                            connection_state.release_export(export_id, 1)?;
729                        }
730                        cap_descriptor::ReceiverAnswer(_) | cap_descriptor::ReceiverHosted(_) => (),
731                        cap_descriptor::ThirdPartyHosted(_) => {
732                            return Err(Error::failed(
733                                "Peer claims we resolved a ThirdPartyHosted cap.".to_string(),
734                            ));
735                        }
736                    },
737                    resolve::Exception(_) => (),
738                }
739            }
740            _ => {
741                return Err(Error::failed(
742                    "Peer did not implement required RPC message type.".to_string(),
743                ));
744            }
745        }
746        Ok(())
747    }
748
749    fn handle_bootstrap(
750        connection_state: &Rc<Self>,
751        bootstrap: bootstrap::Reader,
752    ) -> capnp::Result<()> {
753        use ::capnp::traits::ImbueMut;
754
755        let answer_id = bootstrap.get_question_id();
756        if connection_state.connection.borrow().is_err() {
757            // Disconnected; ignore.
758            return Ok(());
759        }
760
761        let mut response = connection_state.new_outgoing_message(10)?;
762
763        let result_exports = {
764            let mut ret = response
765                .get_body()?
766                .init_as::<message::Builder>()
767                .init_return();
768            ret.set_answer_id(answer_id);
769
770            let cap = connection_state.bootstrap_cap.clone();
771            let mut cap_table = Vec::new();
772            let mut payload = ret.init_results();
773            {
774                let mut content = payload.reborrow().get_content();
775                content.imbue_mut(&mut cap_table);
776                content.set_as_capability(cap);
777            }
778            assert_eq!(cap_table.len(), 1);
779
780            Self::write_descriptors(connection_state, &cap_table, payload)
781        };
782
783        let slots = &mut connection_state.answers.borrow_mut().slots;
784        let hash_map::Entry::Vacant(slot) = slots.entry(answer_id) else {
785            connection_state.release_exports(&result_exports)?;
786            return Err(Error::failed("questionId is already in use".to_string()));
787        };
788        let mut answer = Answer::new();
789        answer.return_has_been_sent = true;
790        answer.result_exports = result_exports;
791        answer.pipeline = Some(Box::new(SingleCapPipeline::new(
792            connection_state.bootstrap_cap.clone(),
793        )));
794        slot.insert(answer);
795
796        let _ = response.send();
797        Ok(())
798    }
799
800    fn handle_finish(connection_state: &Rc<Self>, finish: finish::Reader) -> capnp::Result<()> {
801        let mut exports_to_release = Vec::new();
802        let answer_id = finish.get_question_id();
803
804        let answers_slots = &mut connection_state.answers.borrow_mut().slots;
805        match answers_slots.entry(answer_id) {
806            hash_map::Entry::Vacant(_) => {
807                // The `Finish` message targets a question ID that isn't present in our answer table.
808                // Probably, we sent a `Return` with `noFinishNeeded = true`, but the other side didn't
809                // recognize this hint and sent a `Finish` anyway, or the `Finish` was already in-flight at
810                // the time we sent the `Return`. We can silently ignore this.
811            }
812            hash_map::Entry::Occupied(mut entry) => {
813                let answer = entry.get_mut();
814                answer.received_finish.set(true);
815
816                if finish.get_release_result_caps() {
817                    exports_to_release = ::std::mem::take(&mut answer.result_exports);
818                }
819
820                // If the pipeline has not been cloned, the following two lines cancel the call.
821                answer.pipeline.take();
822                answer.call_completion_promise.take();
823
824                if answer.return_has_been_sent {
825                    entry.remove();
826                }
827            }
828        }
829
830        connection_state.release_exports(&exports_to_release)?;
831        Ok(())
832    }
833
834    fn handle_resolve(connection_state: &Rc<Self>, resolve: resolve::Reader) -> capnp::Result<()> {
835        let replacement_or_error = match resolve.which()? {
836            resolve::Cap(c) => match Self::receive_cap(connection_state, c?)? {
837                Some(cap) => Ok(cap),
838                None => {
839                    return Err(Error::failed(
840                        "'Resolve' contained 'CapDescriptor.none'.".to_string(),
841                    ));
842                }
843            },
844            resolve::Exception(e) => {
845                // We can't set `replacement` to a new broken cap here because this will
846                // confuse PromiseClient::Resolve() into thinking that the remote
847                // promise resolved to a local capability and therefore a Disembargo is
848                // needed. We must actually reject the promise.
849                Err(remote_exception_to_error(e?))
850            }
851        };
852
853        // If the import is in the table, fulfill it.
854        let slots = &mut connection_state.imports.borrow_mut().slots;
855        if let Some(import) = slots.get_mut(&resolve.get_promise_id()) {
856            match import.promise_client_to_resolve.take() {
857                Some(weak_promise_client) => {
858                    if let Some(promise_client) = weak_promise_client.upgrade() {
859                        promise_client.borrow_mut().resolve(replacement_or_error);
860                    }
861                }
862                None => {
863                    return Err(Error::failed(
864                        "Got 'Resolve' for a non-promise import.".to_string(),
865                    ));
866                }
867            }
868        }
869        Ok(())
870    }
871
872    fn handle_disembargo(
873        connection_state: &Rc<Self>,
874        disembargo: disembargo::Reader,
875    ) -> capnp::Result<()> {
876        let context = disembargo.get_context();
877        match context.which()? {
878            disembargo::context::SenderLoopback(embargo_id) => {
879                let mut target = connection_state.get_message_target(disembargo.get_target()?)?;
880                while let Some(resolved) = target.get_resolved() {
881                    target = resolved;
882                }
883
884                if target.get_brand() != connection_state.get_brand() {
885                    return Err(Error::failed(
886                        "'Disembargo' of type 'senderLoopback' sent to an object that does not point \
887                         back to the sender.".to_string()));
888                }
889
890                let connection_state_ref = connection_state.clone();
891                let connection_state_ref1 = connection_state.clone();
892                let task = async move {
893                    if let Ok(ref mut c) = *connection_state_ref.connection.borrow_mut() {
894                        let mut message = c.new_outgoing_message(100); // TODO estimate size
895                        {
896                            let root: message::Builder = message.get_body()?.init_as();
897                            let mut disembargo = root.init_disembargo();
898                            disembargo
899                                .reborrow()
900                                .init_context()
901                                .set_receiver_loopback(embargo_id);
902
903                            let redirect =
904                                match Client::from_ptr(target.get_ptr(), &connection_state_ref1) {
905                                    Some(c) => c.write_target(disembargo.init_target()),
906                                    None => unreachable!(),
907                                };
908                            if redirect.is_some() {
909                                return Err(Error::failed(
910                                    "'Disembargo' of type 'senderLoopback' sent to an object that \
911                                     does not appear to have been the subject of a previous \
912                                     'Resolve' message."
913                                        .to_string(),
914                                ));
915                            }
916                        }
917                        let _ = message.send();
918                    }
919                    Ok(())
920                };
921                connection_state.add_task(task);
922            }
923            disembargo::context::ReceiverLoopback(embargo_id) => {
924                if let Some(embargo) = connection_state.embargoes.borrow_mut().find(embargo_id) {
925                    let fulfiller = embargo.fulfiller.take().unwrap();
926                    let _ = fulfiller.send(Ok(()));
927                } else {
928                    return Err(Error::failed(
929                        "Invalid embargo ID in `Disembargo.context.receiverLoopback".to_string(),
930                    ));
931                }
932                connection_state.embargoes.borrow_mut().erase(embargo_id);
933            }
934            disembargo::context::Accept(_) | disembargo::context::Provide(_) => {
935                return Err(Error::unimplemented(
936                    "Disembargo::Context::Provide/Accept not implemented".to_string(),
937                ));
938            }
939        }
940        Ok(())
941    }
942
943    fn handle_message(
944        weak_state: &Weak<Self>,
945        message: Box<dyn crate::IncomingMessage>,
946    ) -> ::capnp::Result<()> {
947        let Some(connection_state) = weak_state.upgrade() else {
948            return Err(Error::disconnected(
949                "handle_message() cannot continue without a connection".into(),
950            ));
951        };
952
953        let reader = message.get_body()?.get_as::<message::Reader>()?;
954        match reader.which() {
955            Ok(message::Unimplemented(message)) => {
956                Self::handle_unimplemented(&connection_state, message?)?
957            }
958            Ok(message::Abort(abort)) => return Err(remote_exception_to_error(abort?)),
959            Ok(message::Bootstrap(bootstrap)) => {
960                Self::handle_bootstrap(&connection_state, bootstrap?)?
961            }
962            Ok(message::Call(call)) => {
963                let call = call?;
964                let capability = connection_state.get_message_target(call.get_target()?)?;
965                let (interface_id, method_id, question_id, cap_table_array, redirect_results) = {
966                    let redirect_results = match call.get_send_results_to().which()? {
967                        call::send_results_to::Caller(()) => false,
968                        call::send_results_to::Yourself(()) => true,
969                        call::send_results_to::ThirdParty(_) => {
970                            return Err(Error::failed(
971                                "Unsupported `Call.sendResultsTo`.".to_string(),
972                            ))
973                        }
974                    };
975                    let payload = call.get_params()?;
976
977                    (
978                        call.get_interface_id(),
979                        call.get_method_id(),
980                        call.get_question_id(),
981                        Self::receive_caps(&connection_state, payload.get_cap_table()?)?,
982                        redirect_results,
983                    )
984                };
985
986                if connection_state
987                    .answers
988                    .borrow()
989                    .slots
990                    .contains_key(&question_id)
991                {
992                    return Err(Error::failed(format!(
993                        "Received a new call on in-use question id {question_id}"
994                    )));
995                }
996
997                let params = Params::new(message, cap_table_array);
998
999                let answer = Answer::new();
1000
1001                let (results_inner_fulfiller, results_inner_promise) = oneshot::channel();
1002                let results_inner_promise = results_inner_promise.map_err(crate::canceled_to_error);
1003
1004                let (pipeline_sender, mut pipeline) = queued::Pipeline::new();
1005                let results = Results::new(
1006                    &connection_state,
1007                    question_id,
1008                    redirect_results,
1009                    results_inner_fulfiller,
1010                    answer.received_finish.clone(),
1011                    Some(pipeline_sender.weak_clone()),
1012                );
1013
1014                let (redirected_results_done_promise, redirected_results_done_fulfiller) =
1015                    if redirect_results {
1016                        let (f, p) = oneshot::channel::<Result<Response<VatId>, Error>>();
1017                        let p = p.map_err(crate::canceled_to_error).and_then(future::ready);
1018                        (Some(Promise::from_future(p)), Some(f))
1019                    } else {
1020                        (None, None)
1021                    };
1022
1023                {
1024                    let slots = &mut connection_state.answers.borrow_mut().slots;
1025                    let hash_map::Entry::Vacant(slot) = slots.entry(question_id) else {
1026                        return Err(Error::failed("questionId is already in use".to_string()));
1027                    };
1028                    slot.insert(answer);
1029                }
1030
1031                let call_promise =
1032                    capability.call(interface_id, method_id, Box::new(params), Box::new(results));
1033
1034                let promise = call_promise
1035                    .then(move |call_result| {
1036                        results_inner_promise.then(move |result| {
1037                            future::ready(ResultsDone::from_results_inner(
1038                                result,
1039                                call_result,
1040                                pipeline_sender,
1041                            ))
1042                        })
1043                    })
1044                    .then(move |v| {
1045                        if let Some(f) = redirected_results_done_fulfiller {
1046                            match v {
1047                                Ok(r) => drop(f.send(Ok(Response::redirected(r.clone())))),
1048                                Err(e) => drop(f.send(Err(e))),
1049                            }
1050                        }
1051                        Promise::ok(())
1052                    });
1053
1054                let fork = promise.shared();
1055                pipeline.drive(fork.clone());
1056
1057                {
1058                    let slots = &mut connection_state.answers.borrow_mut().slots;
1059                    let Some(answer) = slots.get_mut(&question_id) else {
1060                        unreachable!()
1061                    };
1062                    answer.pipeline = Some(Box::new(pipeline));
1063                    if redirect_results {
1064                        answer.redirected_results = redirected_results_done_promise;
1065                        // More to do here?
1066                    } else {
1067                        answer.call_completion_promise =
1068                            Some(connection_state.eagerly_evaluate(fork));
1069                    }
1070                }
1071            }
1072            Ok(message::Return(oret)) => {
1073                let ret = oret?;
1074                let question_id = ret.get_answer_id();
1075
1076                let mut questions = connection_state.questions.borrow_mut();
1077                match questions.find(question_id) {
1078                    Some(ref mut question) => {
1079                        question.is_awaiting_return = false;
1080                        if ret.get_no_finish_needed() {
1081                            question.skip_finish = true;
1082                        }
1083                        match question.self_ref {
1084                            Some(ref question_ref) => match ret.which()? {
1085                                return_::Results(results) => {
1086                                    let cap_table = Self::receive_caps(
1087                                        &connection_state,
1088                                        results?.get_cap_table()?,
1089                                    )?;
1090
1091                                    let question_ref =
1092                                        question_ref.upgrade().expect("dangling question ref?");
1093                                    let response = Response::new(
1094                                        connection_state.clone(),
1095                                        question_ref.clone(),
1096                                        message,
1097                                        cap_table,
1098                                    );
1099                                    question_ref.borrow_mut().fulfill(Promise::ok(response));
1100                                }
1101                                return_::Exception(e) => {
1102                                    let tmp =
1103                                        question_ref.upgrade().expect("dangling question ref?");
1104                                    tmp.borrow_mut().reject(remote_exception_to_error(e?));
1105                                }
1106                                return_::Canceled(_) => {
1107                                    Self::send_unimplemented(&connection_state, message.as_ref())?;
1108                                }
1109                                return_::ResultsSentElsewhere(_) => {
1110                                    Self::send_unimplemented(&connection_state, message.as_ref())?;
1111                                }
1112                                return_::TakeFromOtherQuestion(id) => {
1113                                    if let Some(answer) =
1114                                        connection_state.answers.borrow_mut().slots.get_mut(&id)
1115                                    {
1116                                        if let Some(res) = answer.redirected_results.take() {
1117                                            let tmp = question_ref
1118                                                .upgrade()
1119                                                .expect("dangling question ref?");
1120                                            tmp.borrow_mut().fulfill(res);
1121                                        } else {
1122                                            return Err(Error::failed("return.takeFromOtherQuestion referenced a call that \
1123                                                     did not use sendResultsTo.yourself.".to_string()));
1124                                        }
1125                                    } else {
1126                                        return Err(Error::failed(
1127                                            "return.takeFromOtherQuestion had invalid answer ID."
1128                                                .to_string(),
1129                                        ));
1130                                    }
1131                                }
1132                                return_::AcceptFromThirdParty(_) => {
1133                                    drop(questions);
1134                                    Self::send_unimplemented(&connection_state, message.as_ref())?;
1135                                }
1136                            },
1137                            None => {
1138                                if let return_::TakeFromOtherQuestion(_) = ret.which()? {
1139                                    return Self::send_unimplemented(
1140                                        &connection_state,
1141                                        message.as_ref(),
1142                                    );
1143                                }
1144                                // Looks like this question was canceled earlier, so `Finish`
1145                                // was already sent, with `releaseResultCaps` set true so that
1146                                // we don't have to release them here. We can go ahead and
1147                                // delete it from the table.
1148                                questions.erase(question_id);
1149                            }
1150                        }
1151                    }
1152                    None => {
1153                        return Err(Error::failed(format!(
1154                            "Invalid question ID in Return message: {question_id}"
1155                        )));
1156                    }
1157                }
1158            }
1159            Ok(message::Finish(finish)) => Self::handle_finish(&connection_state, finish?)?,
1160            Ok(message::Resolve(resolve)) => Self::handle_resolve(&connection_state, resolve?)?,
1161            Ok(message::Release(release)) => {
1162                let release = release?;
1163                connection_state.release_export(release.get_id(), release.get_reference_count())?;
1164            }
1165            Ok(message::Disembargo(disembargo)) => {
1166                Self::handle_disembargo(&connection_state, disembargo?)?
1167            }
1168            Ok(
1169                message::Provide(_)
1170                | message::Accept(_)
1171                | message::Join(_)
1172                | message::ObsoleteSave(_)
1173                | message::ObsoleteDelete(_),
1174            )
1175            | Err(::capnp::NotInSchema(_)) => {
1176                Self::send_unimplemented(&connection_state, message.as_ref())?;
1177            }
1178        }
1179        Ok(())
1180    }
1181
1182    fn answer_has_sent_return(&self, id: AnswerId, result_exports: Vec<ExportId>) {
1183        let answers_slots = &mut self.answers.borrow_mut().slots;
1184        let hash_map::Entry::Occupied(mut entry) = answers_slots.entry(id) else {
1185            unreachable!()
1186        };
1187        let a = entry.get_mut();
1188        a.return_has_been_sent = true;
1189        if a.received_finish.get() {
1190            entry.remove();
1191        } else {
1192            a.result_exports = result_exports;
1193        }
1194    }
1195
1196    fn release_export(&self, id: ExportId, refcount: u32) -> ::capnp::Result<()> {
1197        let mut exports = self.exports.borrow_mut();
1198        let Some(e) = exports.find(id) else {
1199            return Err(Error::failed(
1200                "Tried to release invalid export ID.".to_string(),
1201            ));
1202        };
1203        if refcount > e.refcount {
1204            return Err(Error::failed(
1205                "Tried to drop export's refcount below zero.".to_string(),
1206            ));
1207        }
1208        e.refcount -= refcount;
1209        if e.refcount == 0 {
1210            let client_ptr = e.client_hook.get_ptr();
1211            if e.canonical {
1212                self.exports_by_cap.borrow_mut().remove(&client_ptr);
1213            }
1214            exports.erase(id);
1215        }
1216        Ok(())
1217    }
1218
1219    fn release_exports(&self, exports: &[ExportId]) -> ::capnp::Result<()> {
1220        for &export_id in exports {
1221            self.release_export(export_id, 1)?;
1222        }
1223        Ok(())
1224    }
1225
1226    fn get_brand(&self) -> usize {
1227        self as *const _ as usize
1228    }
1229
1230    fn get_message_target(
1231        &self,
1232        target: message_target::Reader,
1233    ) -> ::capnp::Result<Box<dyn ClientHook>> {
1234        match target.which()? {
1235            message_target::ImportedCap(export_id) => {
1236                match self.exports.borrow().slots.get(export_id as usize) {
1237                    Some(Some(exp)) => Ok(exp.client_hook.clone()),
1238                    _ => Err(Error::failed(
1239                        "Message target is not a current export ID.".to_string(),
1240                    )),
1241                }
1242            }
1243            message_target::PromisedAnswer(promised_answer) => {
1244                let promised_answer = promised_answer?;
1245                let question_id = promised_answer.get_question_id();
1246
1247                let pipeline = match self.answers.borrow().slots.get(&question_id) {
1248                    None => Box::new(broken::Pipeline::new(Error::failed(
1249                        "Pipeline call on a request that returned no capabilities or was already closed.".to_string(),
1250                    ))) as Box<dyn PipelineHook>,
1251                    Some(base) => {
1252                        match base.pipeline {
1253                            Some(ref pipeline) => pipeline.add_ref(),
1254                            None => Box::new(broken::Pipeline::new(Error::failed(
1255                                "Pipeline call on a request that returned not capabilities or was \
1256                                 already closed."
1257                                    .to_string(),
1258                            ))) as Box<dyn PipelineHook>,
1259                        }
1260                    }
1261                };
1262                let ops = to_pipeline_ops(promised_answer.get_transform()?)?;
1263                Ok(pipeline.get_pipelined_cap(&ops))
1264            }
1265        }
1266    }
1267
1268    /// If calls to the given capability should pass over this connection, fill in `target`
1269    /// appropriately for such a call and return None. Otherwise, return a `ClientHook` to which
1270    /// the call should be forwarded; the caller should then delegate the call to that `ClientHook`.
1271    ///
1272    /// The main case where this ends up returning Some(_) is if `cap` is a promise that has
1273    /// recently resolved. The application might have started building a request before the promise
1274    /// resolved, and so the request may have been built on the assumption that it would be sent over
1275    /// this network connection, but then the promise resolved to point somewhere else before the
1276    /// request was sent. Now the request has to be redirected to the new target instead.
1277    fn write_target(
1278        &self,
1279        cap: &dyn ClientHook,
1280        target: message_target::Builder,
1281    ) -> Option<Box<dyn ClientHook>> {
1282        if cap.get_brand() == self.get_brand() {
1283            match Client::from_ptr(cap.get_ptr(), self) {
1284                Some(c) => c.write_target(target),
1285                None => unreachable!(),
1286            }
1287        } else {
1288            Some(cap.add_ref())
1289        }
1290    }
1291
1292    /// If the given client just wraps some other client -- even if it is only *temporarily*
1293    /// wrapping that other client -- returns a reference to the other client, transitively.
1294    /// Otherwise, returns a new reference to *this.
1295    fn get_innermost_client(&self, mut client: Box<dyn ClientHook>) -> Box<dyn ClientHook> {
1296        while let Some(inner) = client.get_resolved() {
1297            client = inner;
1298        }
1299        if client.get_brand() == self.get_brand() {
1300            match self.client_downcast_map.borrow().get(&client.get_ptr()) {
1301                Some(c) => Box::new(c.upgrade().expect("dangling client?")),
1302                None => unreachable!(),
1303            }
1304        } else {
1305            client
1306        }
1307    }
1308
1309    /// Implements exporting of a promise.  The promise has been exported under the given ID, and is
1310    /// to eventually resolve to the ClientHook produced by `promise`.  This method waits for that
1311    /// resolve to happen and then sends the appropriate `Resolve` message to the peer.
1312    #[allow(clippy::await_holding_refcell_ref)] // https://github.com/rust-lang/rust-clippy/issues/6353
1313    fn resolve_exported_promise(
1314        state: &Rc<Self>,
1315        export_id: ExportId,
1316        promise: Promise<Box<dyn ClientHook>, Error>,
1317    ) -> Promise<(), Error> {
1318        let weak_connection_state = Rc::downgrade(state);
1319        state.eagerly_evaluate(Promise::from_future(async move {
1320            let resolution_result = promise.await;
1321            let connection_state = weak_connection_state
1322                .upgrade()
1323                .expect("dangling connection state?");
1324
1325            match resolution_result {
1326                Ok(resolution) => {
1327                    let resolution = connection_state.get_innermost_client(resolution.clone());
1328
1329                    let brand = resolution.get_brand();
1330
1331                    // Update the export table to point at this object instead. We know that our
1332                    // entry in the export table is still live because when it is destroyed the
1333                    // asynchronous resolution task (i.e. this code) is canceled.
1334                    let mut exports = connection_state.exports.borrow_mut();
1335                    let Some(exp) = exports.find(export_id) else {
1336                        return Err(Error::failed("export table entry not found".to_string()));
1337                    };
1338
1339                    if exp.canonical {
1340                        connection_state
1341                            .exports_by_cap
1342                            .borrow_mut()
1343                            .remove(&exp.client_hook.get_ptr());
1344                    }
1345                    exp.client_hook = resolution.clone();
1346
1347                    // The export now points to `resolution`, but it is not necessarily the
1348                    // canonical export for `resolution`. The export itself still represents
1349                    // the promise that ended up resolving to `resolution`, but `resolution`
1350                    // itself also needs to be exported under a separate export ID to
1351                    // distinguish from the promise. (Unless it's also a promise, see the next
1352                    // bit...)
1353                    exp.canonical = false;
1354
1355                    if brand != connection_state.get_brand() {
1356                        // We're resolving to a local capability. If we're resolving to a promise,
1357                        // we might be able to reuse our export table entry and avoid sending a
1358                        // message.
1359                        if let Some(promise) = resolution.when_more_resolved() {
1360                            // We're replacing a promise with another local promise. In this case,
1361                            // we might actually be able to just reuse the existing export table
1362                            // entry to represent the new promise -- unless it already has an entry.
1363                            // Let's check.
1364
1365                            let mut exports_by_cap = connection_state.exports_by_cap.borrow_mut();
1366
1367                            let replacement_export_id =
1368                                match exports_by_cap.entry(exp.client_hook.get_ptr()) {
1369                                    hash_map::Entry::Occupied(occ) => *occ.get(),
1370                                    hash_map::Entry::Vacant(vac) => {
1371                                        // The replacement capability isn't previously exported,
1372                                        // so assign it to the existing table entry.
1373                                        vac.insert(export_id);
1374                                        export_id
1375                                    }
1376                                };
1377                            if replacement_export_id == export_id {
1378                                // The new promise was not already in the table, therefore the existing
1379                                // export table entry has now been repurposed to represent it. There is
1380                                // no need to send a resolve message at all. We do, however, have to
1381                                // start resolving the next promise.
1382                                exp.canonical = true;
1383                                drop(exports);
1384                                drop(exports_by_cap);
1385                                return Self::resolve_exported_promise(
1386                                    &connection_state,
1387                                    export_id,
1388                                    promise,
1389                                )
1390                                .await;
1391                            }
1392                        }
1393                    }
1394                    // Prevent a double borrow in write_descriptor() below.
1395                    drop(exports);
1396
1397                    // OK, we have to send a `Resolve` message.
1398                    let mut message = connection_state.new_outgoing_message(15)?;
1399                    {
1400                        let root: message::Builder = message.get_body()?.get_as()?;
1401                        let mut resolve = root.init_resolve();
1402                        resolve.set_promise_id(export_id);
1403                        let _export = Self::write_descriptor(
1404                            &connection_state,
1405                            resolution,
1406                            resolve.init_cap(),
1407                        )?;
1408                    }
1409                    let _ = message.send();
1410                    Ok(())
1411                }
1412                Err(e) => {
1413                    // send error resolution
1414                    let mut message = connection_state.new_outgoing_message(15)?;
1415                    {
1416                        let root: message::Builder = message.get_body()?.get_as()?;
1417                        let mut resolve = root.init_resolve();
1418                        resolve.set_promise_id(export_id);
1419                        from_error(&e, resolve.init_exception());
1420                    }
1421                    let _ = message.send();
1422                    Ok(())
1423                }
1424            }
1425        }))
1426    }
1427
1428    fn write_descriptor(
1429        state: &Rc<Self>,
1430        mut inner: Box<dyn ClientHook>,
1431        mut descriptor: cap_descriptor::Builder,
1432    ) -> ::capnp::Result<Option<ExportId>> {
1433        // Find the innermost wrapped capability.
1434        while let Some(resolved) = inner.get_resolved() {
1435            inner = resolved;
1436        }
1437        if inner.get_brand() == state.get_brand() {
1438            if let Some(c) = Client::from_ptr(inner.get_ptr(), state) {
1439                return Ok(c.write_descriptor(descriptor));
1440            }
1441            // The hook claims to belong to this connection but the downcast
1442            // map has no live entry for it (e.g. a stale entry left by a
1443            // since-dropped duplicate wrapper — see the reuse logic in
1444            // `import()`). The hook itself still works for calls, so fall
1445            // through and export it as if it were foreign: the receiver gets
1446            // a functioning capability (at the cost of an extra round-trip)
1447            // instead of the event loop panicking.
1448        }
1449        {
1450            let ptr = inner.get_ptr();
1451            let contains_key = state.exports_by_cap.borrow().contains_key(&ptr);
1452            if contains_key {
1453                // We've already seen and exported this capability before.  Just up the refcount.
1454                let export_id = state.exports_by_cap.borrow()[&ptr];
1455                descriptor.set_sender_hosted(export_id);
1456                // Should never fail because exports_by_cap should match exports.
1457                state.exports.borrow_mut().find(export_id).unwrap().refcount += 1;
1458                Ok(Some(export_id))
1459            } else {
1460                // This is the first time we've seen this capability.
1461
1462                let mut exp = Export::new(inner.clone());
1463                exp.canonical = true;
1464                let export_id = state.exports.borrow_mut().push(exp);
1465                state.exports_by_cap.borrow_mut().insert(ptr, export_id);
1466                match inner.when_more_resolved() {
1467                    Some(wrapped) => {
1468                        // This is a promise.  Arrange for the `Resolve` message to be sent later.
1469                        if let Some(exp) = state.exports.borrow_mut().find(export_id) {
1470                            exp.resolve_op =
1471                                Self::resolve_exported_promise(state, export_id, wrapped);
1472                        }
1473                        descriptor.set_sender_promise(export_id);
1474                    }
1475                    None => {
1476                        descriptor.set_sender_hosted(export_id);
1477                    }
1478                }
1479                Ok(Some(export_id))
1480            }
1481        }
1482    }
1483
1484    fn write_descriptors(
1485        state: &Rc<Self>,
1486        cap_table: &[Option<Box<dyn ClientHook>>],
1487        payload: payload::Builder,
1488    ) -> Vec<ExportId> {
1489        let mut cap_table_builder = payload.init_cap_table(cap_table.len() as u32);
1490        let mut exports = Vec::new();
1491        for (idx, value) in cap_table.iter().enumerate() {
1492            match value {
1493                Some(cap) => {
1494                    if let Some(export_id) = Self::write_descriptor(
1495                        state,
1496                        cap.clone(),
1497                        cap_table_builder.reborrow().get(idx as u32),
1498                    )
1499                    .unwrap()
1500                    {
1501                        exports.push(export_id);
1502                    }
1503                }
1504                None => {
1505                    cap_table_builder.reborrow().get(idx as u32).set_none(());
1506                }
1507            }
1508        }
1509        exports
1510    }
1511
1512    fn import(state: &Rc<Self>, import_id: ImportId, is_promise: bool) -> Box<dyn ClientHook> {
1513        let import_client = {
1514            match state.imports.borrow_mut().slots.entry(import_id) {
1515                hash_map::Entry::Occupied(occ) => occ
1516                    .get()
1517                    .import_client
1518                    .upgrade()
1519                    .expect("dangling ref to import client?"),
1520                hash_map::Entry::Vacant(v) => {
1521                    let import_client = ImportClient::new(state, import_id);
1522                    v.insert(Import::new(&import_client));
1523                    import_client
1524                }
1525            }
1526        };
1527
1528        // We just received a copy of this import ID, so the remote refcount has gone up.
1529        import_client.borrow_mut().add_remote_ref();
1530
1531        let mut tmp = state.imports.borrow_mut();
1532        let Some(import) = tmp.slots.get_mut(&import_id) else {
1533            unreachable!()
1534        };
1535
1536        if is_promise {
1537            // We need to construct a PromiseClient around this import, if we haven't already.
1538            match &import.app_client {
1539                Some(c) => {
1540                    // Use the existing one.
1541                    Box::new(c.upgrade().expect("dangling client ref?"))
1542                }
1543                None => {
1544                    // Create a promise for this import's resolution.
1545
1546                    let client: Box<Client<VatId>> = Box::new(import_client.into());
1547                    let client: Box<dyn ClientHook> = client;
1548
1549                    // Here the C++ implementation does something like:
1550                    // ```
1551                    //   // Make sure the import is not destroyed while this promise exists.
1552                    //   let promise = promise.attach(client.add_ref());
1553                    // ```
1554                    // However, as far as I can tell that is unnecessary, because the
1555                    // PromiseClient holds `client` until it resolves, after which point
1556                    // there is no reason to keep the import alive.
1557
1558                    let client = PromiseClient::new(state, client, Some(import_id));
1559
1560                    import.promise_client_to_resolve = Some(Rc::downgrade(&client));
1561                    let client: Box<Client<VatId>> = Box::new(client.into());
1562                    import.app_client = Some(client.downgrade());
1563                    client
1564                }
1565            }
1566        } else {
1567            // Reuse the existing wrapper `Client` if one is still alive,
1568            // mirroring the promise branch above. Unconditionally creating a
1569            // new wrapper for an already-imported cap overwrites the
1570            // `client_downcast_map` entry (keyed by the shared inner
1571            // `ImportClient` pointer); when the newer wrapper is dropped
1572            // while an older one is still held by the application, the map's
1573            // weak reference dies and a later `write_descriptor` of the older
1574            // wrapper hits `Client::from_ptr() == None`.
1575            match import.app_client.as_ref().and_then(|c| c.upgrade()) {
1576                Some(c) => Box::new(c),
1577                None => {
1578                    let client: Box<Client<VatId>> = Box::new(import_client.into());
1579                    import.app_client = Some(client.downgrade());
1580                    client
1581                }
1582            }
1583        }
1584    }
1585
1586    fn receive_cap(
1587        state: &Rc<Self>,
1588        descriptor: cap_descriptor::Reader,
1589    ) -> ::capnp::Result<Option<Box<dyn ClientHook>>> {
1590        match descriptor.which()? {
1591            cap_descriptor::None(()) => Ok(None),
1592            cap_descriptor::SenderHosted(sender_hosted) => {
1593                Ok(Some(Self::import(state, sender_hosted, false)))
1594            }
1595            cap_descriptor::SenderPromise(sender_promise) => {
1596                Ok(Some(Self::import(state, sender_promise, true)))
1597            }
1598            cap_descriptor::ReceiverHosted(receiver_hosted) => {
1599                if let Some(exp) = state.exports.borrow_mut().find(receiver_hosted) {
1600                    Ok(Some(exp.client_hook.add_ref()))
1601                } else {
1602                    Ok(Some(broken::new_cap(Error::failed(
1603                        "invalid 'receiverHosted' export ID".to_string(),
1604                    ))))
1605                }
1606            }
1607            cap_descriptor::ReceiverAnswer(receiver_answer) => {
1608                let promised_answer = receiver_answer?;
1609                let question_id = promised_answer.get_question_id();
1610                if let Some(answer) = state.answers.borrow().slots.get(&question_id) {
1611                    if let Some(ref pipeline) = answer.pipeline {
1612                        let ops = to_pipeline_ops(promised_answer.get_transform()?)?;
1613                        return Ok(Some(pipeline.get_pipelined_cap(&ops)));
1614                    }
1615                }
1616                Ok(Some(broken::new_cap(Error::failed(
1617                    "invalid 'receiver answer'".to_string(),
1618                ))))
1619            }
1620            cap_descriptor::ThirdPartyHosted(_third_party_hosted) => Err(Error::unimplemented(
1621                "ThirdPartyHosted caps are not supported.".to_string(),
1622            )),
1623        }
1624    }
1625
1626    fn receive_caps(
1627        state: &Rc<Self>,
1628        cap_table: ::capnp::struct_list::Reader<cap_descriptor::Owned>,
1629    ) -> ::capnp::Result<Vec<Option<Box<dyn ClientHook>>>> {
1630        let mut result = Vec::new();
1631        for idx in 0..cap_table.len() {
1632            result.push(Self::receive_cap(state, cap_table.get(idx))?);
1633        }
1634        Ok(result)
1635    }
1636}
1637
1638enum DisconnectorState {
1639    New,
1640    Disconnecting(future::Shared<Promise<(), Error>>),
1641    Disconnected,
1642}
1643
1644/// A `Future` that can be run to disconnect an `RpcSystem`'s ConnectionState and wait for it to be closed.
1645pub struct Disconnector<VatId>
1646where
1647    VatId: 'static,
1648{
1649    connection_state: Rc<RefCell<Option<Rc<ConnectionState<VatId>>>>>,
1650    state: DisconnectorState,
1651}
1652
1653impl<VatId> Disconnector<VatId> {
1654    pub fn new(connection_state: Rc<RefCell<Option<Rc<ConnectionState<VatId>>>>>) -> Self {
1655        Self {
1656            connection_state,
1657            state: DisconnectorState::New,
1658        }
1659    }
1660    fn disconnect(&self) -> Option<future::Shared<Promise<(), Error>>> {
1661        if let Some(ref state) = *(self.connection_state.borrow()) {
1662            state.disconnect(::capnp::Error::disconnected(
1663                "client requested disconnect".to_owned(),
1664            ));
1665            state.disconnect_promise.borrow().clone()
1666        } else {
1667            None
1668        }
1669    }
1670}
1671
1672impl<VatId> Future for Disconnector<VatId>
1673where
1674    VatId: 'static,
1675{
1676    type Output = Result<(), capnp::Error>;
1677
1678    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
1679        let this = &mut *self;
1680        loop {
1681            match this.state {
1682                DisconnectorState::New => {
1683                    this.state = match this.disconnect() {
1684                        Some(shutdown_promise) => {
1685                            DisconnectorState::Disconnecting(shutdown_promise)
1686                        }
1687                        None => DisconnectorState::Disconnected,
1688                    };
1689                }
1690                DisconnectorState::Disconnecting(ref mut shutdown_promise) => {
1691                    let result = futures::ready!(Pin::new(shutdown_promise).poll(cx));
1692                    this.state = DisconnectorState::Disconnected;
1693                    return Poll::Ready(result);
1694                }
1695                DisconnectorState::Disconnected => return Poll::Ready(Ok(())),
1696            }
1697        }
1698    }
1699}
1700
1701struct ResponseState<VatId>
1702where
1703    VatId: 'static,
1704{
1705    _connection_state: Rc<ConnectionState<VatId>>,
1706    message: Box<dyn crate::IncomingMessage>,
1707    cap_table: Vec<Option<Box<dyn ClientHook>>>,
1708    _question_ref: Rc<RefCell<QuestionRef<VatId>>>,
1709}
1710
1711enum ResponseVariant<VatId>
1712where
1713    VatId: 'static,
1714{
1715    Rpc(ResponseState<VatId>),
1716    LocallyRedirected(Box<dyn ResultsDoneHook>),
1717}
1718
1719struct Response<VatId>
1720where
1721    VatId: 'static,
1722{
1723    variant: Rc<ResponseVariant<VatId>>,
1724}
1725
1726impl<VatId> Response<VatId> {
1727    fn new(
1728        connection_state: Rc<ConnectionState<VatId>>,
1729        question_ref: Rc<RefCell<QuestionRef<VatId>>>,
1730        message: Box<dyn crate::IncomingMessage>,
1731        cap_table_array: Vec<Option<Box<dyn ClientHook>>>,
1732    ) -> Self {
1733        Self {
1734            variant: Rc::new(ResponseVariant::Rpc(ResponseState {
1735                _connection_state: connection_state,
1736                message,
1737                cap_table: cap_table_array,
1738                _question_ref: question_ref,
1739            })),
1740        }
1741    }
1742    fn redirected(results_done: Box<dyn ResultsDoneHook>) -> Self {
1743        Self {
1744            variant: Rc::new(ResponseVariant::LocallyRedirected(results_done)),
1745        }
1746    }
1747}
1748
1749impl<VatId> Clone for Response<VatId> {
1750    fn clone(&self) -> Self {
1751        Self {
1752            variant: self.variant.clone(),
1753        }
1754    }
1755}
1756
1757impl<VatId> ResponseHook for Response<VatId> {
1758    fn get(&self) -> ::capnp::Result<any_pointer::Reader<'_>> {
1759        match *self.variant {
1760            ResponseVariant::Rpc(ref state) => {
1761                match state
1762                    .message
1763                    .get_body()?
1764                    .get_as::<message::Reader>()?
1765                    .which()?
1766                {
1767                    message::Return(Ok(ret)) => match ret.which()? {
1768                        return_::Results(Ok(mut payload)) => {
1769                            use ::capnp::traits::Imbue;
1770                            payload.imbue(&state.cap_table);
1771                            Ok(payload.get_content())
1772                        }
1773                        _ => unreachable!(),
1774                    },
1775                    _ => unreachable!(),
1776                }
1777            }
1778            ResponseVariant::LocallyRedirected(ref results_done) => results_done.get(),
1779        }
1780    }
1781}
1782
1783struct Request<VatId>
1784where
1785    VatId: 'static,
1786{
1787    connection_state: Rc<ConnectionState<VatId>>,
1788    target: Client<VatId>,
1789    message: Box<dyn crate::OutgoingMessage>,
1790    cap_table: Vec<Option<Box<dyn ClientHook>>>,
1791}
1792
1793fn get_call(message: &mut Box<dyn crate::OutgoingMessage>) -> ::capnp::Result<call::Builder<'_>> {
1794    let message_root: message::Builder = message.get_body()?.get_as()?;
1795    match message_root.which()? {
1796        message::Call(call) => call,
1797        _ => {
1798            unimplemented!()
1799        }
1800    }
1801}
1802
1803impl<VatId> Request<VatId>
1804where
1805    VatId: 'static,
1806{
1807    fn new(
1808        connection_state: Rc<ConnectionState<VatId>>,
1809        _size_hint: Option<::capnp::MessageSize>,
1810        target: Client<VatId>,
1811    ) -> ::capnp::Result<Self> {
1812        let message = connection_state.new_outgoing_message(1024)?;
1813        Ok(Self {
1814            connection_state,
1815            target,
1816            message,
1817            cap_table: Vec::new(),
1818        })
1819    }
1820
1821    fn init_call(&mut self) -> call::Builder<'_> {
1822        let message_root: message::Builder = self.message.get_body().unwrap().get_as().unwrap();
1823        message_root.init_call()
1824    }
1825
1826    fn send_internal(
1827        connection_state: &Rc<ConnectionState<VatId>>,
1828        mut message: Box<dyn crate::OutgoingMessage>,
1829        cap_table: &[Option<Box<dyn ClientHook>>],
1830        is_tail_call: bool,
1831    ) -> (
1832        Rc<RefCell<QuestionRef<VatId>>>,
1833        Promise<Response<VatId>, Error>,
1834    ) {
1835        // Build the cap table.
1836        let exports = ConnectionState::write_descriptors(
1837            connection_state,
1838            cap_table,
1839            get_call(&mut message).unwrap().get_params().unwrap(),
1840        );
1841
1842        // Init the question table.  Do this after writing descriptors to avoid interference.
1843        let mut question = Question::<VatId>::new();
1844        question.is_awaiting_return = true;
1845        question.param_exports = exports;
1846        question.is_tail_call = is_tail_call;
1847
1848        let question_id = connection_state.questions.borrow_mut().push(question);
1849        {
1850            let mut call_builder: call::Builder = get_call(&mut message).unwrap();
1851            // Finish and send.
1852            call_builder.reborrow().set_question_id(question_id);
1853            if is_tail_call {
1854                call_builder.get_send_results_to().set_yourself(());
1855            }
1856        }
1857        let _ = message.send();
1858        // Make the result promise.
1859        let (fulfiller, promise) = oneshot::channel::<Promise<Response<VatId>, Error>>();
1860        let promise = promise.map_err(crate::canceled_to_error).and_then(|x| x);
1861        let question_ref = Rc::new(RefCell::new(QuestionRef::new(
1862            connection_state.clone(),
1863            question_id,
1864            fulfiller,
1865        )));
1866
1867        match connection_state.questions.borrow_mut().slots[question_id as usize] {
1868            Some(ref mut q) => {
1869                q.self_ref = Some(Rc::downgrade(&question_ref));
1870            }
1871            None => unreachable!(),
1872        }
1873
1874        let promise = promise.attach(question_ref.clone());
1875        let promise2 = Promise::from_future(promise);
1876
1877        (question_ref, promise2)
1878    }
1879
1880    fn send_streaming_internal(
1881        connection_state: &Rc<ConnectionState<VatId>>,
1882        mut message: Box<dyn crate::OutgoingMessage>,
1883        cap_table: &[Option<Box<dyn ClientHook>>],
1884        flow: Rc<RefCell<Option<Box<dyn crate::FlowController>>>>,
1885    ) -> Promise<(), Error> {
1886        // Build the cap table.
1887        let exports = ConnectionState::write_descriptors(
1888            connection_state,
1889            cap_table,
1890            get_call(&mut message).unwrap().get_params().unwrap(),
1891        );
1892
1893        // Init the question table.  Do this after writing descriptors to avoid interference.
1894        let mut question = Question::<VatId>::new();
1895        question.is_awaiting_return = true;
1896        question.param_exports = exports;
1897        question.is_tail_call = false;
1898
1899        let question_id = connection_state.questions.borrow_mut().push(question);
1900        {
1901            let mut call_builder: call::Builder = get_call(&mut message).unwrap();
1902            call_builder.reborrow().set_question_id(question_id);
1903        }
1904
1905        // Make the result promise.
1906        let (fulfiller, promise) = oneshot::channel::<Promise<Response<VatId>, Error>>();
1907        let promise = promise.map_err(crate::canceled_to_error).and_then(|x| x);
1908        let question_ref = Rc::new(RefCell::new(QuestionRef::new(
1909            connection_state.clone(),
1910            question_id,
1911            fulfiller,
1912        )));
1913
1914        match connection_state.questions.borrow_mut().slots[question_id as usize] {
1915            Some(ref mut q) => {
1916                q.self_ref = Some(Rc::downgrade(&question_ref));
1917            }
1918            None => unreachable!(),
1919        }
1920        let promise = promise.attach(question_ref.clone());
1921
1922        let mut flow = flow.borrow_mut();
1923        if flow.is_none() {
1924            match connection_state.connection.borrow_mut().as_mut() {
1925                Err(_) => return Promise::err(Error::failed("no connection".into())),
1926                Ok(connection) => {
1927                    let (s, p) = connection.new_stream();
1928                    connection_state.add_task(p);
1929                    *flow = Some(s);
1930                }
1931            };
1932        }
1933        let Some(ref mut flow) = *flow else {
1934            unreachable!()
1935        };
1936        flow.send(
1937            message,
1938            Promise::from_future(async move {
1939                let _ = promise.await?;
1940                Ok(())
1941            }),
1942        )
1943    }
1944}
1945
1946impl<VatId> RequestHook for Request<VatId> {
1947    fn get(&mut self) -> any_pointer::Builder<'_> {
1948        use ::capnp::traits::ImbueMut;
1949        let mut builder = get_call(&mut self.message)
1950            .unwrap()
1951            .get_params()
1952            .unwrap()
1953            .get_content();
1954        builder.imbue_mut(&mut self.cap_table);
1955        builder
1956    }
1957    fn get_brand<'a>(&self) -> usize {
1958        self.connection_state.get_brand()
1959    }
1960    fn send(self: Box<Self>) -> ::capnp::capability::RemotePromise<any_pointer::Owned> {
1961        let tmp = *self;
1962        let Self {
1963            connection_state,
1964            target,
1965            mut message,
1966            cap_table,
1967        } = tmp;
1968        let write_target_result = {
1969            let call_builder: call::Builder = get_call(&mut message).unwrap();
1970            target.write_target(call_builder.get_target().unwrap())
1971        };
1972        if let Some(redirect) = write_target_result {
1973            // Whoops, this capability has been redirected while we were building the request!
1974            // We'll have to make a new request and do a copy.  Ick.
1975            let mut call_builder: call::Builder = get_call(&mut message).unwrap();
1976            let mut replacement = redirect.new_call(
1977                call_builder.reborrow().get_interface_id(),
1978                call_builder.reborrow().get_method_id(),
1979                None,
1980            );
1981
1982            replacement
1983                .set(
1984                    call_builder
1985                        .get_params()
1986                        .unwrap()
1987                        .get_content()
1988                        .into_reader(),
1989                )
1990                .unwrap();
1991            return replacement.send();
1992        }
1993        let (question_ref, promise) =
1994            Self::send_internal(&connection_state, message, &cap_table, false);
1995        let forked_promise1 = promise.shared();
1996        let forked_promise2 = forked_promise1.clone();
1997
1998        // The pipeline must get notified of resolution before the app does to maintain ordering.
1999        let pipeline = Pipeline::new(
2000            &connection_state,
2001            question_ref,
2002            Some(Promise::from_future(forked_promise1)),
2003        );
2004
2005        let resolved = pipeline.when_resolved();
2006
2007        let forked_promise2 = resolved.map(|_| Ok(())).and_then(|()| forked_promise2);
2008
2009        let app_promise = Promise::from_future(
2010            forked_promise2
2011                .map_ok(|response| ::capnp::capability::Response::new(Box::new(response))),
2012        );
2013
2014        ::capnp::capability::RemotePromise {
2015            promise: app_promise,
2016            pipeline: any_pointer::Pipeline::new(Box::new(pipeline)),
2017        }
2018    }
2019    fn send_streaming(self: Box<Self>) -> Promise<(), Error> {
2020        let tmp = *self;
2021        let Self {
2022            connection_state,
2023            target,
2024            mut message,
2025            cap_table,
2026        } = tmp;
2027        let write_target_result = {
2028            let call_builder: call::Builder = get_call(&mut message).unwrap();
2029            target.write_target(call_builder.get_target().unwrap())
2030        };
2031        if let Some(redirect) = write_target_result {
2032            // Whoops, this capability has been redirected while we were building the request!
2033            // We'll have to make a new request and do a copy.  Ick.
2034            let mut call_builder: call::Builder = get_call(&mut message).unwrap();
2035            let mut replacement = redirect.new_call(
2036                call_builder.reborrow().get_interface_id(),
2037                call_builder.reborrow().get_method_id(),
2038                None,
2039            );
2040
2041            replacement
2042                .set(
2043                    call_builder
2044                        .get_params()
2045                        .unwrap()
2046                        .get_content()
2047                        .into_reader(),
2048                )
2049                .unwrap();
2050            return replacement.hook.send_streaming();
2051        }
2052        Self::send_streaming_internal(
2053            &connection_state,
2054            message,
2055            &cap_table,
2056            target.flow_controller,
2057        )
2058    }
2059    fn tail_send(self: Box<Self>) -> Option<(u32, Promise<(), Error>, Box<dyn PipelineHook>)> {
2060        let tmp = *self;
2061        let Self {
2062            connection_state,
2063            target,
2064            mut message,
2065            cap_table,
2066        } = tmp;
2067
2068        if connection_state.connection.borrow().is_err() {
2069            // Disconnected; fall back to a regular send() which will fail appropriately.
2070            return None;
2071        }
2072
2073        let write_target_result = {
2074            let call_builder: crate::rpc_capnp::call::Builder = get_call(&mut message).unwrap();
2075            target.write_target(call_builder.get_target().unwrap())
2076        };
2077
2078        let (question_ref, promise) = match write_target_result {
2079            Some(_redirect) => {
2080                return None;
2081            }
2082            None => Self::send_internal(&connection_state, message, &cap_table, true),
2083        };
2084
2085        let promise = promise.map_ok(|_response| {
2086            // Response should be null if `Return` handling code is correct.
2087
2088            unimplemented!()
2089        });
2090
2091        let question_id = question_ref.borrow().id;
2092        let pipeline = Pipeline::never_done(connection_state, question_ref);
2093
2094        Some((
2095            question_id,
2096            Promise::from_future(promise),
2097            Box::new(pipeline),
2098        ))
2099    }
2100}
2101
2102enum PipelineVariant<VatId>
2103where
2104    VatId: 'static,
2105{
2106    Waiting(Rc<RefCell<QuestionRef<VatId>>>),
2107    Resolved(Response<VatId>),
2108    Broken(Error),
2109}
2110
2111struct PipelineState<VatId>
2112where
2113    VatId: 'static,
2114{
2115    variant: PipelineVariant<VatId>,
2116    redirect_later: Option<RefCell<futures::future::Shared<Promise<Response<VatId>, Error>>>>,
2117    connection_state: Rc<ConnectionState<VatId>>,
2118
2119    #[allow(dead_code)]
2120    resolve_self_promise: Promise<(), Error>,
2121
2122    promise_clients_to_resolve: RefCell<
2123        crate::sender_queue::SenderQueue<
2124            (Weak<RefCell<PromiseClient<VatId>>>, Vec<PipelineOp>),
2125            (),
2126        >,
2127    >,
2128    resolution_waiters: crate::sender_queue::SenderQueue<(), ()>,
2129}
2130
2131impl<VatId> PipelineState<VatId>
2132where
2133    VatId: 'static,
2134{
2135    fn resolve(state: &Rc<RefCell<Self>>, response: Result<Response<VatId>, Error>) {
2136        let to_resolve = {
2137            let tmp = state.borrow();
2138            let r = tmp.promise_clients_to_resolve.borrow_mut().drain();
2139            r
2140        };
2141        for ((c, ops), _) in to_resolve {
2142            let resolved = match response.clone() {
2143                Ok(v) => match v.get() {
2144                    Ok(x) => x.get_pipelined_cap(&ops),
2145                    Err(e) => Err(e),
2146                },
2147                Err(e) => Err(e),
2148            };
2149            if let Some(c) = c.upgrade() {
2150                c.borrow_mut().resolve(resolved);
2151            }
2152        }
2153
2154        let new_variant = match response {
2155            Ok(r) => PipelineVariant::Resolved(r),
2156            Err(e) => PipelineVariant::Broken(e),
2157        };
2158        let _old_variant = mem::replace(&mut state.borrow_mut().variant, new_variant);
2159
2160        let waiters = state.borrow_mut().resolution_waiters.drain();
2161        for (_, waiter) in waiters {
2162            let _ = waiter.send(());
2163        }
2164    }
2165}
2166
2167struct Pipeline<VatId>
2168where
2169    VatId: 'static,
2170{
2171    state: Rc<RefCell<PipelineState<VatId>>>,
2172}
2173
2174impl<VatId> Pipeline<VatId> {
2175    fn new(
2176        connection_state: &Rc<ConnectionState<VatId>>,
2177        question_ref: Rc<RefCell<QuestionRef<VatId>>>,
2178        redirect_later: Option<Promise<Response<VatId>, ::capnp::Error>>,
2179    ) -> Self {
2180        let state = Rc::new(RefCell::new(PipelineState {
2181            variant: PipelineVariant::Waiting(question_ref),
2182            connection_state: connection_state.clone(),
2183            redirect_later: None,
2184            resolve_self_promise: Promise::from_future(future::pending()),
2185            promise_clients_to_resolve: RefCell::new(crate::sender_queue::SenderQueue::new()),
2186            resolution_waiters: crate::sender_queue::SenderQueue::new(),
2187        }));
2188        if let Some(redirect_later_promise) = redirect_later {
2189            let fork = redirect_later_promise.shared();
2190            let this = Rc::downgrade(&state);
2191            let resolve_self_promise =
2192                connection_state.eagerly_evaluate(fork.clone().then(move |response| {
2193                    let Some(state) = this.upgrade() else {
2194                        return Promise::err(Error::failed("dangling reference to this".into()));
2195                    };
2196                    PipelineState::resolve(&state, response);
2197                    Promise::ok(())
2198                }));
2199
2200            state.borrow_mut().resolve_self_promise = resolve_self_promise;
2201            state.borrow_mut().redirect_later = Some(RefCell::new(fork));
2202        }
2203        Self { state }
2204    }
2205
2206    fn when_resolved(&self) -> Promise<(), Error> {
2207        self.state.borrow_mut().resolution_waiters.push(())
2208    }
2209
2210    fn never_done(
2211        connection_state: Rc<ConnectionState<VatId>>,
2212        question_ref: Rc<RefCell<QuestionRef<VatId>>>,
2213    ) -> Self {
2214        let state = Rc::new(RefCell::new(PipelineState {
2215            variant: PipelineVariant::Waiting(question_ref),
2216            connection_state,
2217            redirect_later: None,
2218            resolve_self_promise: Promise::from_future(future::pending()),
2219            promise_clients_to_resolve: RefCell::new(crate::sender_queue::SenderQueue::new()),
2220            resolution_waiters: crate::sender_queue::SenderQueue::new(),
2221        }));
2222
2223        Self { state }
2224    }
2225}
2226
2227impl<VatId> PipelineHook for Pipeline<VatId> {
2228    fn add_ref(&self) -> Box<dyn PipelineHook> {
2229        Box::new(Self {
2230            state: self.state.clone(),
2231        })
2232    }
2233    fn get_pipelined_cap(&self, ops: &[PipelineOp]) -> Box<dyn ClientHook> {
2234        self.get_pipelined_cap_move(ops.into())
2235    }
2236    fn get_pipelined_cap_move(&self, ops: Vec<PipelineOp>) -> Box<dyn ClientHook> {
2237        match *self.state.borrow() {
2238            PipelineState {
2239                variant: PipelineVariant::Waiting(ref question_ref),
2240                ref connection_state,
2241                ref redirect_later,
2242                ref promise_clients_to_resolve,
2243                ..
2244            } => {
2245                // Wrap a PipelineClient in a PromiseClient.
2246                let pipeline_client =
2247                    PipelineClient::new(connection_state, question_ref.clone(), ops.clone());
2248
2249                match redirect_later {
2250                    Some(_r) => {
2251                        let client: Client<VatId> = pipeline_client.into();
2252                        let promise_client =
2253                            PromiseClient::new(connection_state, Box::new(client), None);
2254                        promise_clients_to_resolve
2255                            .borrow_mut()
2256                            .push_detach((Rc::downgrade(&promise_client), ops));
2257                        let result: Client<VatId> = promise_client.into();
2258                        Box::new(result)
2259                    }
2260                    None => {
2261                        // Oh, this pipeline will never get redirected, so just return the PipelineClient.
2262                        let client: Client<VatId> = pipeline_client.into();
2263                        Box::new(client)
2264                    }
2265                }
2266            }
2267            PipelineState {
2268                variant: PipelineVariant::Resolved(ref response),
2269                ..
2270            } => response.get().unwrap().get_pipelined_cap(&ops[..]).unwrap(),
2271            PipelineState {
2272                variant: PipelineVariant::Broken(ref e),
2273                ..
2274            } => broken::new_cap(e.clone()),
2275        }
2276    }
2277}
2278
2279pub(crate) struct Params {
2280    request: Box<dyn crate::IncomingMessage>,
2281    cap_table: Vec<Option<Box<dyn ClientHook>>>,
2282}
2283
2284impl Params {
2285    fn new(
2286        request: Box<dyn crate::IncomingMessage>,
2287        cap_table: Vec<Option<Box<dyn ClientHook>>>,
2288    ) -> Self {
2289        Self { request, cap_table }
2290    }
2291}
2292
2293impl ParamsHook for Params {
2294    fn get(&self) -> ::capnp::Result<any_pointer::Reader<'_>> {
2295        let root: message::Reader = self.request.get_body()?.get_as()?;
2296        let message::Call(call) = root.which()? else {
2297            unreachable!()
2298        };
2299        use ::capnp::traits::Imbue;
2300        let mut content = call?.get_params()?.get_content();
2301        content.imbue(&self.cap_table);
2302        Ok(content)
2303    }
2304}
2305
2306enum ResultsVariant {
2307    Rpc(
2308        Box<dyn crate::OutgoingMessage>,
2309        Vec<Option<Box<dyn ClientHook>>>,
2310    ),
2311    LocallyRedirected(
2312        ::capnp::message::Builder<::capnp::message::HeapAllocator>,
2313        Vec<Option<Box<dyn ClientHook>>>,
2314    ),
2315}
2316
2317struct ResultsInner<VatId>
2318where
2319    VatId: 'static,
2320{
2321    connection_state: Rc<ConnectionState<VatId>>,
2322    variant: Option<ResultsVariant>,
2323    redirect_results: bool,
2324    answer_id: AnswerId,
2325    finish_received: Rc<Cell<bool>>,
2326    pipeline_sender: Option<queued::PipelineInnerSender>,
2327}
2328
2329impl<VatId> ResultsInner<VatId>
2330where
2331    VatId: 'static,
2332{
2333    fn ensure_initialized(&mut self) {
2334        let answer_id = self.answer_id;
2335        if self.variant.is_none() {
2336            match (
2337                self.redirect_results,
2338                self.connection_state.connection.borrow_mut().as_mut(),
2339            ) {
2340                (false, Ok(c)) => {
2341                    let mut message = c.new_outgoing_message(100); // size hint?
2342
2343                    {
2344                        let root: message::Builder = message.get_body().unwrap().init_as();
2345                        let mut ret = root.init_return();
2346                        ret.set_answer_id(answer_id);
2347                        ret.set_release_param_caps(false);
2348                    }
2349                    self.variant = Some(ResultsVariant::Rpc(message, Vec::new()));
2350                }
2351                _ => {
2352                    self.variant = Some(ResultsVariant::LocallyRedirected(
2353                        ::capnp::message::Builder::new_default(),
2354                        Vec::new(),
2355                    ));
2356                }
2357            }
2358        }
2359    }
2360}
2361
2362// This takes the place of both RpcCallContext and RpcServerResponse in capnproto-c++.
2363pub(crate) struct Results<VatId>
2364where
2365    VatId: 'static,
2366{
2367    inner: Option<ResultsInner<VatId>>,
2368    results_done_fulfiller: Option<oneshot::Sender<ResultsInner<VatId>>>,
2369}
2370
2371impl<VatId> Results<VatId>
2372where
2373    VatId: 'static,
2374{
2375    fn new(
2376        connection_state: &Rc<ConnectionState<VatId>>,
2377        answer_id: AnswerId,
2378        redirect_results: bool,
2379        fulfiller: oneshot::Sender<ResultsInner<VatId>>,
2380        finish_received: Rc<Cell<bool>>,
2381        pipeline_sender: Option<queued::PipelineInnerSender>,
2382    ) -> Self {
2383        Self {
2384            inner: Some(ResultsInner {
2385                variant: None,
2386                connection_state: connection_state.clone(),
2387                redirect_results,
2388                answer_id,
2389                finish_received,
2390                pipeline_sender,
2391            }),
2392            results_done_fulfiller: Some(fulfiller),
2393        }
2394    }
2395}
2396
2397impl<VatId> Drop for Results<VatId> {
2398    fn drop(&mut self) {
2399        match (self.inner.take(), self.results_done_fulfiller.take()) {
2400            (Some(inner), Some(fulfiller)) => {
2401                let _ = fulfiller.send(inner);
2402            }
2403            (None, None) => (),
2404            _ => unreachable!(),
2405        }
2406    }
2407}
2408
2409impl<VatId> ResultsHook for Results<VatId> {
2410    fn get(&mut self) -> ::capnp::Result<any_pointer::Builder<'_>> {
2411        use ::capnp::traits::ImbueMut;
2412        let Some(ref mut inner) = self.inner else {
2413            unreachable!();
2414        };
2415        inner.ensure_initialized();
2416        match inner.variant {
2417            None => unreachable!(),
2418            Some(ResultsVariant::Rpc(ref mut message, ref mut cap_table)) => {
2419                let root: message::Builder = message.get_body()?.get_as()?;
2420                let message::Return(ret) = root.which()? else {
2421                    unreachable!();
2422                };
2423                let return_::Results(payload) = ret?.which()? else {
2424                    unreachable!()
2425                };
2426                let mut content = payload?.get_content();
2427                content.imbue_mut(cap_table);
2428                Ok(content)
2429            }
2430            Some(ResultsVariant::LocallyRedirected(ref mut message, ref mut cap_table)) => {
2431                let mut result: any_pointer::Builder = message.get_root()?;
2432                result.imbue_mut(cap_table);
2433                Ok(result)
2434            }
2435        }
2436    }
2437
2438    fn set_pipeline(&mut self) -> ::capnp::Result<()> {
2439        use ::capnp::traits::ImbueMut;
2440        let root = self.get()?;
2441        let size = root.target_size()?;
2442        let mut message2 = capnp::message::Builder::new(
2443            capnp::message::HeapAllocator::new().first_segment_words(size.word_count as u32 + 1),
2444        );
2445        let mut root2: capnp::any_pointer::Builder = message2.init_root();
2446        let mut cap_table2 = vec![];
2447        root2.imbue_mut(&mut cap_table2);
2448        root2.set_as(root.into_reader())?;
2449        let hook =
2450            Box::new(local::ResultsDone::new(message2, cap_table2)) as Box<dyn ResultsDoneHook>;
2451        let Some(ref mut inner) = self.inner else {
2452            unreachable!();
2453        };
2454        let Some(sender) = inner.pipeline_sender.take() else {
2455            return Err(Error::failed("set_pipeline() called twice".into()));
2456        };
2457        sender.complete(Box::new(local::Pipeline::new(hook)));
2458        Ok(())
2459    }
2460
2461    fn tail_call(self: Box<Self>, _request: Box<dyn RequestHook>) -> Promise<(), Error> {
2462        unimplemented!()
2463    }
2464
2465    fn direct_tail_call(
2466        mut self: Box<Self>,
2467        request: Box<dyn RequestHook>,
2468    ) -> (Promise<(), Error>, Box<dyn PipelineHook>) {
2469        if let (Some(inner), Some(fulfiller)) =
2470            (self.inner.take(), self.results_done_fulfiller.take())
2471        {
2472            let state = inner.connection_state.clone();
2473            if request.get_brand() == state.get_brand() && !inner.redirect_results {
2474                // The tail call is headed towards the peer that called us in the first place, so we can
2475                // optimize out the return trip.
2476                if let Some((question_id, promise, pipeline)) = request.tail_send() {
2477                    let mut message = state.new_outgoing_message(100).expect("no connection?"); // size hint?
2478
2479                    {
2480                        let root: message::Builder = message.get_body().unwrap().init_as();
2481                        let mut ret = root.init_return();
2482                        ret.set_answer_id(inner.answer_id);
2483                        ret.set_release_param_caps(false);
2484                        ret.set_take_from_other_question(question_id);
2485                    }
2486                    let _ = message.send();
2487
2488                    // TODO cleanupanswertable
2489
2490                    let _ = fulfiller.send(inner); // ??
2491                    return (promise, pipeline);
2492                }
2493                unimplemented!()
2494            } else {
2495                unimplemented!()
2496            }
2497        } else {
2498            unreachable!();
2499        }
2500    }
2501
2502    fn allow_cancellation(&self) {
2503        unimplemented!()
2504    }
2505}
2506
2507enum ResultsDoneVariant {
2508    Rpc(
2509        Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>,
2510        Vec<Option<Box<dyn ClientHook>>>,
2511    ),
2512    LocallyRedirected(
2513        ::capnp::message::Builder<::capnp::message::HeapAllocator>,
2514        Vec<Option<Box<dyn ClientHook>>>,
2515    ),
2516}
2517
2518struct ResultsDone {
2519    inner: Rc<ResultsDoneVariant>,
2520}
2521
2522impl ResultsDone {
2523    fn from_results_inner<VatId>(
2524        results_inner: Result<ResultsInner<VatId>, Error>,
2525        call_status: Result<(), Error>,
2526        pipeline_sender: queued::PipelineInnerSender,
2527    ) -> Result<Box<dyn ResultsDoneHook>, Error>
2528    where
2529        VatId: 'static,
2530    {
2531        match results_inner {
2532            Err(e) => {
2533                pipeline_sender.complete(Box::new(crate::broken::Pipeline::new(e.clone())));
2534                Err(e)
2535            }
2536            Ok(mut results_inner) => {
2537                results_inner.ensure_initialized();
2538                let ResultsInner {
2539                    connection_state,
2540                    variant,
2541                    answer_id,
2542                    finish_received,
2543                    ..
2544                } = results_inner;
2545                match variant {
2546                    None => unreachable!(),
2547                    Some(ResultsVariant::Rpc(mut message, cap_table)) => {
2548                        match (finish_received.get(), call_status) {
2549                            (true, _) => {
2550                                let hook = Box::new(Self::rpc(Rc::new(message.take()), cap_table))
2551                                    as Box<dyn ResultsDoneHook>;
2552                                pipeline_sender
2553                                    .complete(Box::new(local::Pipeline::new(hook.clone())));
2554
2555                                // Send a Canceled return.
2556                                if let Ok(connection) =
2557                                    connection_state.connection.borrow_mut().as_mut()
2558                                {
2559                                    let mut message = connection.new_outgoing_message(10);
2560                                    {
2561                                        let root: message::Builder =
2562                                            message.get_body()?.get_as()?;
2563                                        let mut ret = root.init_return();
2564                                        ret.set_answer_id(answer_id);
2565                                        ret.set_release_param_caps(false);
2566                                        ret.set_canceled(());
2567                                    }
2568                                    let _ = message.send();
2569                                }
2570
2571                                connection_state.answer_has_sent_return(answer_id, Vec::new());
2572                                Ok(hook)
2573                            }
2574                            (false, Ok(())) => {
2575                                let exports = {
2576                                    let root: message::Builder = message.get_body()?.get_as()?;
2577                                    let message::Return(Ok(mut ret)) = root.which()? else {
2578                                        unreachable!()
2579                                    };
2580                                    if cap_table.is_empty() {
2581                                        ret.set_no_finish_needed(true);
2582                                        finish_received.set(true);
2583                                    }
2584                                    let crate::rpc_capnp::return_::Results(Ok(payload)) =
2585                                        ret.which()?
2586                                    else {
2587                                        unreachable!()
2588                                    };
2589                                    ConnectionState::write_descriptors(
2590                                        &connection_state,
2591                                        &cap_table,
2592                                        payload,
2593                                    )
2594                                };
2595
2596                                let (_promise, m) = message.send();
2597                                connection_state.answer_has_sent_return(answer_id, exports);
2598                                let hook =
2599                                    Box::new(Self::rpc(m, cap_table)) as Box<dyn ResultsDoneHook>;
2600                                pipeline_sender
2601                                    .complete(Box::new(local::Pipeline::new(hook.clone())));
2602                                Ok(hook)
2603                            }
2604                            (false, Err(e)) => {
2605                                // Send an error return.
2606                                if let Ok(connection) =
2607                                    connection_state.connection.borrow_mut().as_mut()
2608                                {
2609                                    let mut message = connection.new_outgoing_message(50); // XXX size hint
2610                                    {
2611                                        let root: message::Builder =
2612                                            message.get_body()?.get_as()?;
2613                                        let mut ret = root.init_return();
2614                                        ret.set_answer_id(answer_id);
2615                                        ret.set_release_param_caps(false);
2616                                        let mut exc = ret.init_exception();
2617                                        from_error(&e, exc.reborrow());
2618                                    }
2619                                    let _ = message.send();
2620                                }
2621                                connection_state.answer_has_sent_return(answer_id, Vec::new());
2622
2623                                pipeline_sender
2624                                    .complete(Box::new(crate::broken::Pipeline::new(e.clone())));
2625
2626                                Err(e)
2627                            }
2628                        }
2629                    }
2630                    Some(ResultsVariant::LocallyRedirected(results_done, cap_table)) => {
2631                        let hook = Box::new(Self::redirected(results_done, cap_table))
2632                            as Box<dyn ResultsDoneHook>;
2633                        pipeline_sender
2634                            .complete(Box::new(crate::local::Pipeline::new(hook.clone())));
2635                        Ok(hook)
2636                    }
2637                }
2638            }
2639        }
2640    }
2641
2642    fn rpc(
2643        message: Rc<::capnp::message::Builder<::capnp::message::HeapAllocator>>,
2644        cap_table: Vec<Option<Box<dyn ClientHook>>>,
2645    ) -> Self {
2646        Self {
2647            inner: Rc::new(ResultsDoneVariant::Rpc(message, cap_table)),
2648        }
2649    }
2650
2651    fn redirected(
2652        message: ::capnp::message::Builder<::capnp::message::HeapAllocator>,
2653        cap_table: Vec<Option<Box<dyn ClientHook>>>,
2654    ) -> Self {
2655        Self {
2656            inner: Rc::new(ResultsDoneVariant::LocallyRedirected(message, cap_table)),
2657        }
2658    }
2659}
2660
2661impl ResultsDoneHook for ResultsDone {
2662    fn add_ref(&self) -> Box<dyn ResultsDoneHook> {
2663        Box::new(Self {
2664            inner: self.inner.clone(),
2665        })
2666    }
2667    fn get(&self) -> ::capnp::Result<any_pointer::Reader<'_>> {
2668        use ::capnp::traits::Imbue;
2669        match *self.inner {
2670            ResultsDoneVariant::Rpc(ref message, ref cap_table) => {
2671                let root: message::Reader = message.get_root_as_reader()?;
2672                let message::Return(ret) = root.which()? else {
2673                    unreachable!();
2674                };
2675                let crate::rpc_capnp::return_::Results(payload) = ret?.which()? else {
2676                    unreachable!();
2677                };
2678                let mut content = payload?.get_content();
2679                content.imbue(cap_table);
2680                Ok(content)
2681            }
2682            ResultsDoneVariant::LocallyRedirected(ref message, ref cap_table) => {
2683                let mut result: any_pointer::Reader = message.get_root_as_reader()?;
2684                result.imbue(cap_table);
2685                Ok(result)
2686            }
2687        }
2688    }
2689}
2690
2691enum ClientVariant<VatId>
2692where
2693    VatId: 'static,
2694{
2695    Import(Rc<RefCell<ImportClient<VatId>>>),
2696    Pipeline(Rc<RefCell<PipelineClient<VatId>>>),
2697    Promise(Rc<RefCell<PromiseClient<VatId>>>),
2698}
2699
2700struct Client<VatId>
2701where
2702    VatId: 'static,
2703{
2704    connection_state: Rc<ConnectionState<VatId>>,
2705    variant: ClientVariant<VatId>,
2706    flow_controller: Rc<RefCell<Option<Box<dyn crate::FlowController>>>>,
2707}
2708
2709enum WeakClientVariant<VatId>
2710where
2711    VatId: 'static,
2712{
2713    Import(Weak<RefCell<ImportClient<VatId>>>),
2714    Pipeline(Weak<RefCell<PipelineClient<VatId>>>),
2715    Promise(Weak<RefCell<PromiseClient<VatId>>>),
2716}
2717
2718struct WeakClient<VatId>
2719where
2720    VatId: 'static,
2721{
2722    connection_state: Weak<ConnectionState<VatId>>,
2723    variant: WeakClientVariant<VatId>,
2724    flow_controller: Weak<RefCell<Option<Box<dyn crate::FlowController>>>>,
2725}
2726
2727impl<VatId> WeakClient<VatId>
2728where
2729    VatId: 'static,
2730{
2731    fn upgrade(&self) -> Option<Client<VatId>> {
2732        let variant = match &self.variant {
2733            WeakClientVariant::Import(ic) => ClientVariant::Import(ic.upgrade()?),
2734            WeakClientVariant::Pipeline(pc) => ClientVariant::Pipeline(pc.upgrade()?),
2735            WeakClientVariant::Promise(pc) => ClientVariant::Promise(pc.upgrade()?),
2736        };
2737        let connection_state = self.connection_state.upgrade()?;
2738        let flow_controller = self.flow_controller.upgrade()?;
2739        Some(Client {
2740            connection_state,
2741            variant,
2742            flow_controller,
2743        })
2744    }
2745}
2746
2747struct ImportClient<VatId>
2748where
2749    VatId: 'static,
2750{
2751    connection_state: Rc<ConnectionState<VatId>>,
2752    import_id: ImportId,
2753
2754    /// Number of times we've received this import from the peer.
2755    remote_ref_count: u32,
2756}
2757
2758impl<VatId> Drop for ImportClient<VatId> {
2759    fn drop(&mut self) {
2760        let connection_state = self.connection_state.clone();
2761
2762        assert!(connection_state
2763            .client_downcast_map
2764            .borrow_mut()
2765            .remove(&((self) as *const _ as usize))
2766            .is_some());
2767
2768        // Remove the corresponding entry of the imports table.
2769        // Note: the C++ implementation checks here pointer equality between self and
2770        // the entry in the imports table, but as far as I can tell the check should
2771        // always pass because of how we construct ImportClient in import().
2772        connection_state
2773            .imports
2774            .borrow_mut()
2775            .slots
2776            .remove(&self.import_id);
2777
2778        // Send a message releasing our remote references.
2779        let mut tmp = connection_state.connection.borrow_mut();
2780        if let (true, Ok(c)) = (self.remote_ref_count > 0, tmp.as_mut()) {
2781            let mut message = c.new_outgoing_message(10);
2782            {
2783                let root: message::Builder = message.get_body().unwrap().init_as();
2784                let mut release = root.init_release();
2785                release.set_id(self.import_id);
2786                release.set_reference_count(self.remote_ref_count);
2787            }
2788            let _ = message.send();
2789        }
2790    }
2791}
2792
2793impl<VatId> ImportClient<VatId>
2794where
2795    VatId: 'static,
2796{
2797    fn new(
2798        connection_state: &Rc<ConnectionState<VatId>>,
2799        import_id: ImportId,
2800    ) -> Rc<RefCell<Self>> {
2801        Rc::new(RefCell::new(Self {
2802            connection_state: connection_state.clone(),
2803            import_id,
2804            remote_ref_count: 0,
2805        }))
2806    }
2807
2808    fn add_remote_ref(&mut self) {
2809        self.remote_ref_count += 1;
2810    }
2811}
2812
2813impl<VatId> From<Rc<RefCell<ImportClient<VatId>>>> for Client<VatId> {
2814    fn from(client: Rc<RefCell<ImportClient<VatId>>>) -> Self {
2815        let connection_state = client.borrow().connection_state.clone();
2816        Self::new(&connection_state, ClientVariant::Import(client))
2817    }
2818}
2819
2820/// A `ClientHook` representing a pipelined promise.  Always wrapped in `PromiseClient`.
2821struct PipelineClient<VatId>
2822where
2823    VatId: 'static,
2824{
2825    connection_state: Rc<ConnectionState<VatId>>,
2826    question_ref: Rc<RefCell<QuestionRef<VatId>>>,
2827    ops: Vec<PipelineOp>,
2828}
2829
2830impl<VatId> PipelineClient<VatId>
2831where
2832    VatId: 'static,
2833{
2834    fn new(
2835        connection_state: &Rc<ConnectionState<VatId>>,
2836        question_ref: Rc<RefCell<QuestionRef<VatId>>>,
2837        ops: Vec<PipelineOp>,
2838    ) -> Rc<RefCell<Self>> {
2839        Rc::new(RefCell::new(Self {
2840            connection_state: connection_state.clone(),
2841            question_ref,
2842            ops,
2843        }))
2844    }
2845}
2846
2847impl<VatId> From<Rc<RefCell<PipelineClient<VatId>>>> for Client<VatId> {
2848    fn from(client: Rc<RefCell<PipelineClient<VatId>>>) -> Self {
2849        let connection_state = client.borrow().connection_state.clone();
2850        Self::new(&connection_state, ClientVariant::Pipeline(client))
2851    }
2852}
2853
2854impl<VatId> Drop for PipelineClient<VatId> {
2855    fn drop(&mut self) {
2856        assert!(self
2857            .connection_state
2858            .client_downcast_map
2859            .borrow_mut()
2860            .remove(&((self) as *const _ as usize))
2861            .is_some());
2862    }
2863}
2864
2865/// A `ClientHook` that initially wraps one client and then, later on, redirects
2866/// to some other client.
2867struct PromiseClient<VatId>
2868where
2869    VatId: 'static,
2870{
2871    connection_state: Rc<ConnectionState<VatId>>,
2872    is_resolved: bool,
2873    cap: Box<dyn ClientHook>,
2874    import_id: Option<ImportId>,
2875    received_call: bool,
2876    resolution_waiters: crate::sender_queue::SenderQueue<(), Box<dyn ClientHook>>,
2877}
2878
2879impl<VatId> PromiseClient<VatId> {
2880    fn new(
2881        connection_state: &Rc<ConnectionState<VatId>>,
2882        initial: Box<dyn ClientHook>,
2883        import_id: Option<ImportId>,
2884    ) -> Rc<RefCell<Self>> {
2885        Rc::new(RefCell::new(Self {
2886            connection_state: connection_state.clone(),
2887            is_resolved: false,
2888            cap: initial,
2889            import_id,
2890            received_call: false,
2891            resolution_waiters: crate::sender_queue::SenderQueue::new(),
2892        }))
2893    }
2894
2895    fn resolve(&mut self, replacement: Result<Box<dyn ClientHook>, Error>) {
2896        let (mut replacement, is_error) = match replacement {
2897            Ok(v) => (v, false),
2898            Err(e) => (broken::new_cap(e), true),
2899        };
2900        let connection_state = self.connection_state.clone();
2901        let is_connected = connection_state.connection.borrow().is_ok();
2902        let replacement_brand = replacement.get_brand();
2903        if replacement_brand != connection_state.get_brand()
2904            && self.received_call
2905            && !is_error
2906            && is_connected
2907        {
2908            // The new capability is hosted locally, not on the remote machine.  And, we had made calls
2909            // to the promise.  We need to make sure those calls echo back to us before we allow new
2910            // calls to go directly to the local capability, so we need to set a local embargo and send
2911            // a `Disembargo` to echo through the peer.
2912            let (fulfiller, promise) = oneshot::channel::<Result<(), Error>>();
2913            let promise = promise
2914                .map_err(crate::canceled_to_error)
2915                .and_then(future::ready);
2916            let embargo = Embargo::new(fulfiller);
2917            let embargo_id = connection_state.embargoes.borrow_mut().push(embargo);
2918
2919            let mut message = connection_state
2920                .new_outgoing_message(50)
2921                .expect("no connection?"); // XXX size hint
2922            {
2923                let root: message::Builder = message.get_body().unwrap().init_as();
2924                let mut disembargo = root.init_disembargo();
2925                disembargo
2926                    .reborrow()
2927                    .init_context()
2928                    .set_sender_loopback(embargo_id);
2929                let target = disembargo.init_target();
2930
2931                let redirect = connection_state.write_target(&*self.cap, target);
2932                if redirect.is_some() {
2933                    panic!("Original promise target should always be from this RPC connection.")
2934                }
2935            }
2936
2937            // Make a promise which resolves to `replacement` as soon as the `Disembargo` comes back.
2938            let embargo_promise = promise.map_ok(move |()| replacement);
2939
2940            let mut queued_client = queued::Client::new(None);
2941            let weak_queued = Rc::downgrade(&queued_client.inner);
2942
2943            queued_client.drive(embargo_promise.then(move |r| {
2944                if let Some(q) = weak_queued.upgrade() {
2945                    queued::ClientInner::resolve(&q, r);
2946                }
2947                Promise::ok(())
2948            }));
2949
2950            // We need to queue up calls in the meantime, so we'll resolve ourselves to a local promise
2951            // client instead.
2952            replacement = Box::new(queued_client);
2953
2954            let _ = message.send();
2955        }
2956
2957        for ((), waiter) in self.resolution_waiters.drain() {
2958            let _ = waiter.send(replacement.clone());
2959        }
2960
2961        let old_cap = mem::replace(&mut self.cap, replacement);
2962        connection_state.add_task(async move {
2963            drop(old_cap);
2964            Ok(())
2965        });
2966
2967        self.is_resolved = true;
2968    }
2969}
2970
2971impl<VatId> Drop for PromiseClient<VatId> {
2972    fn drop(&mut self) {
2973        let self_ptr = (self) as *const _ as usize;
2974
2975        if let Some(id) = self.import_id {
2976            // This object is representing an import promise.  That means the import table may still
2977            // contain a pointer back to it.  Remove that pointer.  Note that we have to verify that
2978            // the import still exists and the pointer still points back to this object because this
2979            // object may actually outlive the import.
2980            let slots = &mut self.connection_state.imports.borrow_mut().slots;
2981            if let Some(import) = slots.get_mut(&id) {
2982                if let Some(c) = &import.app_client {
2983                    if let Some(cs) = c.upgrade() {
2984                        if cs.get_ptr() == self_ptr {
2985                            import.app_client = None;
2986                        }
2987                    }
2988                }
2989            }
2990        }
2991
2992        assert!(self
2993            .connection_state
2994            .client_downcast_map
2995            .borrow_mut()
2996            .remove(&self_ptr)
2997            .is_some());
2998    }
2999}
3000
3001impl<VatId> From<Rc<RefCell<PromiseClient<VatId>>>> for Client<VatId> {
3002    fn from(client: Rc<RefCell<PromiseClient<VatId>>>) -> Self {
3003        let connection_state = client.borrow().connection_state.clone();
3004        Self::new(&connection_state, ClientVariant::Promise(client))
3005    }
3006}
3007
3008impl<VatId> Client<VatId> {
3009    fn new(connection_state: &Rc<ConnectionState<VatId>>, variant: ClientVariant<VatId>) -> Self {
3010        let client = Self {
3011            connection_state: connection_state.clone(),
3012            variant,
3013            flow_controller: Rc::new(RefCell::new(None)),
3014        };
3015        let weak = client.downgrade();
3016
3017        // XXX arguably, this should go in each of the variant's constructors.
3018        connection_state
3019            .client_downcast_map
3020            .borrow_mut()
3021            .insert(client.get_ptr(), weak);
3022        client
3023    }
3024    fn downgrade(&self) -> WeakClient<VatId> {
3025        let variant = match &self.variant {
3026            ClientVariant::Import(import_client) => {
3027                WeakClientVariant::Import(Rc::downgrade(import_client))
3028            }
3029            ClientVariant::Pipeline(pipeline_client) => {
3030                WeakClientVariant::Pipeline(Rc::downgrade(pipeline_client))
3031            }
3032            ClientVariant::Promise(promise_client) => {
3033                WeakClientVariant::Promise(Rc::downgrade(promise_client))
3034            }
3035        };
3036        WeakClient {
3037            connection_state: Rc::downgrade(&self.connection_state),
3038            variant,
3039            flow_controller: Rc::downgrade(&self.flow_controller),
3040        }
3041    }
3042
3043    fn from_ptr(ptr: usize, connection_state: &ConnectionState<VatId>) -> Option<Self> {
3044        match connection_state.client_downcast_map.borrow().get(&ptr) {
3045            Some(c) => c.upgrade(),
3046            None => None,
3047        }
3048    }
3049
3050    fn write_target(
3051        &self,
3052        mut target: crate::rpc_capnp::message_target::Builder,
3053    ) -> Option<Box<dyn ClientHook>> {
3054        match &self.variant {
3055            ClientVariant::Import(import_client) => {
3056                target.set_imported_cap(import_client.borrow().import_id);
3057                None
3058            }
3059            ClientVariant::Pipeline(pipeline_client) => {
3060                let mut builder = target.init_promised_answer();
3061                let question_ref = &pipeline_client.borrow().question_ref;
3062                builder.set_question_id(question_ref.borrow().id);
3063                let mut transform =
3064                    builder.init_transform(pipeline_client.borrow().ops.len() as u32);
3065                for idx in 0..pipeline_client.borrow().ops.len() {
3066                    if let ::capnp::private::capability::PipelineOp::GetPointerField(ordinal) =
3067                        pipeline_client.borrow().ops[idx]
3068                    {
3069                        transform
3070                            .reborrow()
3071                            .get(idx as u32)
3072                            .set_get_pointer_field(ordinal);
3073                    }
3074                }
3075                None
3076            }
3077            ClientVariant::Promise(promise_client) => {
3078                promise_client.borrow_mut().received_call = true;
3079                self.connection_state
3080                    .write_target(&*promise_client.borrow().cap, target)
3081            }
3082        }
3083    }
3084
3085    fn write_descriptor(&self, mut descriptor: cap_descriptor::Builder) -> Option<u32> {
3086        match &self.variant {
3087            ClientVariant::Import(import_client) => {
3088                descriptor.set_receiver_hosted(import_client.borrow().import_id);
3089                None
3090            }
3091            ClientVariant::Pipeline(pipeline_client) => {
3092                let mut promised_answer = descriptor.init_receiver_answer();
3093                let question_ref = &pipeline_client.borrow().question_ref;
3094                promised_answer.set_question_id(question_ref.borrow().id);
3095                let mut transform =
3096                    promised_answer.init_transform(pipeline_client.borrow().ops.len() as u32);
3097                for idx in 0..pipeline_client.borrow().ops.len() {
3098                    if let ::capnp::private::capability::PipelineOp::GetPointerField(ordinal) =
3099                        pipeline_client.borrow().ops[idx]
3100                    {
3101                        transform
3102                            .reborrow()
3103                            .get(idx as u32)
3104                            .set_get_pointer_field(ordinal);
3105                    }
3106                }
3107
3108                None
3109            }
3110            ClientVariant::Promise(promise_client) => {
3111                promise_client.borrow_mut().received_call = true;
3112
3113                ConnectionState::write_descriptor(
3114                    &self.connection_state.clone(),
3115                    promise_client.borrow().cap.clone(),
3116                    descriptor,
3117                )
3118                .unwrap()
3119            }
3120        }
3121    }
3122}
3123
3124impl<VatId> Clone for Client<VatId> {
3125    fn clone(&self) -> Self {
3126        let variant = match &self.variant {
3127            ClientVariant::Import(import_client) => ClientVariant::Import(import_client.clone()),
3128            ClientVariant::Pipeline(pipeline_client) => {
3129                ClientVariant::Pipeline(pipeline_client.clone())
3130            }
3131            ClientVariant::Promise(promise_client) => {
3132                ClientVariant::Promise(promise_client.clone())
3133            }
3134        };
3135        Self {
3136            connection_state: self.connection_state.clone(),
3137            variant,
3138            flow_controller: self.flow_controller.clone(),
3139        }
3140    }
3141}
3142
3143impl<VatId> ClientHook for Client<VatId> {
3144    fn add_ref(&self) -> Box<dyn ClientHook> {
3145        Box::new(self.clone())
3146    }
3147    fn new_call(
3148        &self,
3149        interface_id: u64,
3150        method_id: u16,
3151        size_hint: Option<::capnp::MessageSize>,
3152    ) -> ::capnp::capability::Request<any_pointer::Owned, any_pointer::Owned> {
3153        let request: Box<dyn RequestHook> =
3154            match Request::new(self.connection_state.clone(), size_hint, self.clone()) {
3155                Ok(mut request) => {
3156                    {
3157                        let mut call_builder = request.init_call();
3158                        call_builder.set_interface_id(interface_id);
3159                        call_builder.set_method_id(method_id);
3160                    }
3161                    Box::new(request)
3162                }
3163                Err(e) => Box::new(broken::Request::new(e, None)),
3164            };
3165
3166        ::capnp::capability::Request::new(request)
3167    }
3168
3169    fn call(
3170        &self,
3171        interface_id: u64,
3172        method_id: u16,
3173        params: Box<dyn ParamsHook>,
3174        mut results: Box<dyn ResultsHook>,
3175    ) -> Promise<(), Error> {
3176        // Implement call() by copying params and results messages.
3177
3178        let maybe_request = params.get().and_then(|p| {
3179            let mut request = p
3180                .target_size()
3181                .map(|s| self.new_call(interface_id, method_id, Some(s)))?;
3182            request.get().set_as(p)?;
3183            Ok(request)
3184        });
3185
3186        match maybe_request {
3187            Err(e) => Promise::err(e),
3188            Ok(request) => {
3189                let ::capnp::capability::RemotePromise { promise, .. } = request.send();
3190
3191                Promise::from_future(async move {
3192                    let response = promise.await?;
3193                    results.get()?.set_as(response.get()?)?;
3194                    Ok(())
3195                })
3196            }
3197        }
3198        // TODO implement this in terms of direct tail call.
3199        // We can and should propagate cancellation.
3200        // (TODO ?)
3201        // context -> allowCancellation();
3202
3203        //results.direct_tail_call(request.hook)
3204    }
3205
3206    fn get_ptr(&self) -> usize {
3207        match &self.variant {
3208            ClientVariant::Import(import_client) => (&*import_client.borrow()) as *const _ as usize,
3209            ClientVariant::Pipeline(pipeline_client) => {
3210                (&*pipeline_client.borrow()) as *const _ as usize
3211            }
3212            ClientVariant::Promise(promise_client) => {
3213                (&*promise_client.borrow()) as *const _ as usize
3214            }
3215        }
3216    }
3217
3218    fn get_brand(&self) -> usize {
3219        self.connection_state.get_brand()
3220    }
3221
3222    fn get_resolved(&self) -> Option<Box<dyn ClientHook>> {
3223        match &self.variant {
3224            ClientVariant::Import(_import_client) => None,
3225            ClientVariant::Pipeline(_pipeline_client) => None,
3226            ClientVariant::Promise(promise_client) => {
3227                if promise_client.borrow().is_resolved {
3228                    Some(promise_client.borrow().cap.clone())
3229                } else {
3230                    None
3231                }
3232            }
3233        }
3234    }
3235
3236    fn when_more_resolved(&self) -> Option<Promise<Box<dyn ClientHook>, Error>> {
3237        match &self.variant {
3238            ClientVariant::Import(_import_client) => None,
3239            ClientVariant::Pipeline(_pipeline_client) => None,
3240            ClientVariant::Promise(promise_client) => {
3241                Some(promise_client.borrow_mut().resolution_waiters.push(()))
3242            }
3243        }
3244    }
3245
3246    fn when_resolved(&self) -> Promise<(), Error> {
3247        default_when_resolved_impl(self)
3248    }
3249}
3250
3251pub(crate) fn default_when_resolved_impl<C>(client: &C) -> Promise<(), Error>
3252where
3253    C: ClientHook,
3254{
3255    match client.when_more_resolved() {
3256        Some(promise) => {
3257            Promise::from_future(promise.and_then(|resolution| resolution.when_resolved()))
3258        }
3259        None => Promise::ok(()),
3260    }
3261}
3262
3263// ===================================
3264
3265struct SingleCapPipeline {
3266    cap: Box<dyn ClientHook>,
3267}
3268
3269impl SingleCapPipeline {
3270    fn new(cap: Box<dyn ClientHook>) -> Self {
3271        Self { cap }
3272    }
3273}
3274
3275impl PipelineHook for SingleCapPipeline {
3276    fn add_ref(&self) -> Box<dyn PipelineHook> {
3277        Box::new(Self {
3278            cap: self.cap.clone(),
3279        })
3280    }
3281    fn get_pipelined_cap(&self, ops: &[PipelineOp]) -> Box<dyn ClientHook> {
3282        if ops.is_empty() {
3283            self.cap.add_ref()
3284        } else {
3285            broken::new_cap(Error::failed("Invalid pipeline transform.".to_string()))
3286        }
3287    }
3288}