Skip to main content

bitwarden_ipc/rpc/exec/
handler_registry.rs

1use 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}