bitwarden_ipc/
ipc_client_ext.rs1use 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
21pub trait IpcClientExt: IpcClient {
27 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 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 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 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 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 let mut response_subscription = self
128 .subscribe(Some(request_payload.response_topic.clone()))
129 .await?;
130
131 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 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
173impl<T: IpcClient + ?Sized> IpcClientExt for T {}