Skip to main content

bitwarden_ipc/
ipc_client_ext.rs

1use bitwarden_threading::cancellation_token::CancellationToken;
2use serde::{Serialize, de::DeserializeOwned};
3
4use crate::{
5    RpcHandler,
6    broadcast::{BroadcastGroup, EndpointSendResult},
7    endpoint::Endpoint,
8    error::{BroadcastError, RequestError, SubscribeError},
9    ipc_client::IpcClientTypedSubscription,
10    ipc_client_trait::IpcClient,
11    message::{OutgoingMessage, PayloadTypeName, TypedOutgoingMessage},
12    rpc::{
13        error::RpcError,
14        request::RpcRequest,
15        request_message::{RPC_REQUEST_PAYLOAD_TYPE_NAME, RpcRequestMessage},
16        response_message::IncomingRpcResponseMessage,
17    },
18    serde_utils,
19};
20
21/// Extension trait providing generic convenience methods on any [`IpcClient`].
22///
23/// This trait is automatically implemented for all types that implement [`IpcClient`],
24/// including `dyn IpcClient`. It provides typed subscriptions, handler registration,
25/// and RPC request functionality with full static type safety.
26pub trait IpcClientExt: IpcClient {
27    /// Register a new RPC handler for processing incoming RPC requests.
28    /// The handler will be executed by the IPC client when an RPC request is received and
29    /// the response will be sent back over IPC.
30    fn register_rpc_handler<H>(&self, handler: H) -> impl std::future::Future<Output = ()> + Send
31    where
32        H: RpcHandler + Send + Sync + 'static,
33    {
34        async move {
35            self.register_rpc_handler_erased(H::Request::NAME, Box::new(handler))
36                .await;
37        }
38    }
39
40    /// Send a message with a payload of any serializable type to the specified destination.
41    fn send_typed<Payload>(
42        &self,
43        payload: Payload,
44        destination: Endpoint,
45    ) -> impl std::future::Future<Output = Result<(), RequestError>> + Send
46    where
47        Payload: Serialize + PayloadTypeName + Send,
48    {
49        async move {
50            let message = TypedOutgoingMessage {
51                payload,
52                destination,
53            }
54            .try_into()
55            .map_err(|e: serde_utils::DeserializeError| {
56                RequestError::Rpc(RpcError::RequestSerialization(e.to_string()))
57            })?;
58
59            self.send(message).await.map_err(RequestError::from)
60        }
61    }
62
63    /// Send a message with a payload of any serializable type to every member of `group`. See
64    /// [`IpcClient::broadcast`] for how the group is resolved and how failures are reported.
65    fn broadcast_typed<Payload>(
66        &self,
67        payload: Payload,
68        group: BroadcastGroup,
69    ) -> impl std::future::Future<Output = Result<Vec<EndpointSendResult>, BroadcastError>> + Send
70    where
71        Payload: Serialize + PayloadTypeName + Send,
72    {
73        async move {
74            let payload =
75                serde_utils::to_vec(&payload).map_err(|e: serde_utils::SerializeError| {
76                    BroadcastError::Serialization(e.to_string())
77                })?;
78
79            self.broadcast(group, payload, Some(Payload::PAYLOAD_TYPE_NAME.to_owned()))
80                .await
81        }
82    }
83
84    /// Create a subscription to receive messages that can be deserialized into the provided
85    /// payload type.
86    fn subscribe_typed<Payload>(
87        &self,
88    ) -> impl std::future::Future<
89        Output = Result<IpcClientTypedSubscription<Payload>, SubscribeError>,
90    > + Send
91    where
92        Payload: DeserializeOwned + PayloadTypeName,
93    {
94        async move {
95            Ok(IpcClientTypedSubscription::new(
96                self.subscribe(Some(Payload::PAYLOAD_TYPE_NAME.to_owned()))
97                    .await?,
98            ))
99        }
100    }
101
102    /// Send a request to the specified destination and wait for a response.
103    /// The destination must have a registered RPC handler for the request type, otherwise
104    /// an error will be returned by the remote endpoint.
105    ///
106    /// The response is accepted only from `destination`, as identified by the source the
107    /// transport reports for each incoming message (the SDK trusts the transport, and any relay
108    /// in between, to label sources correctly). A response from any other source is ignored, so a
109    /// caller that needs a bound on the wait should pass a `cancellation_token`.
110    fn request<Request>(
111        &self,
112        request: Request,
113        destination: Endpoint,
114        cancellation_token: Option<CancellationToken>,
115    ) -> impl std::future::Future<Output = Result<Request::Response, RequestError>> + Send
116    where
117        Request: RpcRequest + Send,
118        Request::Response: Send,
119    {
120        async move {
121            let request_payload = RpcRequestMessage::new(request);
122
123            // Each request gets its own response topic, subscribed to before sending, and the
124            // handler publishes its reply there. The subscription receives every message on that
125            // topic; the loop below accepts only the one from the destination. A deserialization
126            // failure of that message is therefore a malformed response to this request.
127            let mut response_subscription = self
128                .subscribe(Some(request_payload.response_topic.clone()))
129                .await?;
130
131            // Requests are dispatched by a single fixed topic that the receiver matches on to route
132            // them to the handler registry.
133            let payload = serde_utils::to_vec(&request_payload)
134                .map_err(|e| RequestError::Rpc(RpcError::RequestSerialization(e.to_string())))?;
135            let message = OutgoingMessage {
136                payload,
137                destination: destination.clone(),
138                topic: Some(RPC_REQUEST_PAYLOAD_TYPE_NAME.to_owned()),
139            };
140
141            self.send(message).await.map_err(RequestError::from)?;
142
143            // Only the endpoint the request was sent to may answer it. A message on the response
144            // topic from any other source is skipped without failing the request, so another peer
145            // cannot forge the response or end the wait by posting on the response topic.
146            let received = loop {
147                let received = response_subscription
148                    .receive(cancellation_token.clone())
149                    .await
150                    .map_err(|e| RequestError::Receive(e.into()))?;
151
152                if received.source.to_endpoint() == destination {
153                    break received;
154                }
155
156                tracing::warn!(
157                    source = ?received.source,
158                    ?destination,
159                    "Ignoring RPC response from a source other than the request destination"
160                );
161            };
162
163            let response: IncomingRpcResponseMessage<Request::Response> =
164                serde_utils::from_slice(&received.payload).map_err(|e| {
165                    RequestError::Rpc(RpcError::ResponseDeserialization(e.to_string()))
166                })?;
167
168            Ok(response.result?)
169        }
170    }
171}
172
173/// Blanket implementation: every [`IpcClient`] gets the extension methods for free.
174impl<T: IpcClient + ?Sized> IpcClientExt for T {}