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