bitwarden_ipc/rpc/exec/
handler_registry.rs1use erased_serde::Serialize as ErasedSerialize;
2use tokio::sync::RwLock;
3
4use super::handler::{ErasedRpcHandler, RpcHandler, RpcRequestInfo};
5use crate::rpc::{error::RpcError, request::RpcRequest, request_message::RpcRequestPayload};
6
7pub struct RpcHandlerRegistry {
8 handlers: RwLock<std::collections::HashMap<String, Box<dyn ErasedRpcHandler>>>,
9}
10
11impl RpcHandlerRegistry {
12 pub fn new() -> Self {
13 Self {
14 handlers: RwLock::new(std::collections::HashMap::new()),
15 }
16 }
17
18 pub async fn register<H>(&self, handler: H)
19 where
20 H: RpcHandler + ErasedRpcHandler + 'static,
21 {
22 let name = H::Request::NAME.to_owned();
23 self.register_erased(name, Box::new(handler)).await;
24 }
25
26 pub async fn register_erased(&self, name: String, handler: Box<dyn ErasedRpcHandler>) {
27 self.handlers.write().await.insert(name, handler);
28 }
29
30 pub async fn handle(
31 &self,
32 request: &RpcRequestPayload,
33 info: RpcRequestInfo,
34 ) -> Result<Box<dyn ErasedSerialize>, RpcError> {
35 match self.handlers.read().await.get(request.request_type()) {
36 Some(handler) => handler.handle(request, info).await,
37 None => Err(RpcError::NoHandlerFound),
38 }
39 }
40}
41
42#[cfg(test)]
43mod test {
44 use serde::{Deserialize, Serialize, de::DeserializeOwned};
45
46 use super::*;
47 use crate::{
48 endpoint::{HostId, Source},
49 rpc::{request::RpcRequest, request_message::RpcRequestMessage},
50 serde_utils,
51 };
52
53 #[derive(Debug, Clone, Serialize, Deserialize)]
54 struct TestRequest {
55 a: i32,
56 b: i32,
57 }
58
59 #[derive(Debug, Clone, Serialize, Deserialize)]
60 struct TestResponse {
61 result: i32,
62 }
63
64 impl RpcRequest for TestRequest {
65 type Response = TestResponse;
66
67 const NAME: &str = "TestRequest";
68 }
69
70 struct TestHandler;
71
72 impl RpcHandler for TestHandler {
73 type Request = TestRequest;
74
75 async fn handle(&self, request: Self::Request, _info: RpcRequestInfo) -> TestResponse {
76 TestResponse {
77 result: request.a + request.b,
78 }
79 }
80 }
81
82 #[tokio::test]
83 async fn handle_returns_error_when_no_handler_can_be_found() {
84 let registry = RpcHandlerRegistry::new();
85
86 let request = TestRequest { a: 1, b: 2 };
87 let message = RpcRequestMessage {
88 request,
89 request_id: "test_id".to_string(),
90 request_type: "TestRequest".to_string(),
91 response_topic: "RpcResponseMessage:test_id".to_string(),
92 };
93 let serialized_request =
94 RpcRequestPayload::from_slice(serde_utils::to_vec(&message).unwrap()).unwrap();
95
96 let result = registry.handle(&serialized_request, test_info()).await;
97
98 assert!(matches!(result, Err(RpcError::NoHandlerFound)));
99 }
100
101 #[tokio::test]
102 async fn handle_runs_previously_registered_handler() {
103 let registry = RpcHandlerRegistry::new();
104
105 registry.register(TestHandler).await;
106
107 let request = TestRequest { a: 1, b: 2 };
108 let message = RpcRequestMessage {
109 request,
110 request_id: "test_id".to_string(),
111 request_type: "TestRequest".to_string(),
112 response_topic: "RpcResponseMessage:test_id".to_string(),
113 };
114 let serialized_request =
115 RpcRequestPayload::from_slice(serde_utils::to_vec(&message).unwrap()).unwrap();
116
117 let result = registry
118 .handle(&serialized_request, test_info())
119 .await
120 .expect("Failed to handle request");
121 let response: TestResponse = deserialize_erased_object(&result);
122
123 assert_eq!(response.result, 3);
124 }
125
126 fn test_info() -> RpcRequestInfo {
127 RpcRequestInfo {
128 source: Source::BrowserBackground { id: HostId::Own },
129 }
130 }
131
132 fn deserialize_erased_object<T, R>(value: &T) -> R
133 where
134 T: Serialize,
135 R: DeserializeOwned,
136 {
137 let serialized = serde_utils::to_vec(value).expect("Failed to serialize erased serialize");
138
139 serde_utils::from_slice(&serialized).expect("Failed to deserialize erased serialize")
140 }
141}