lsp_max/jsonrpc/
router.rs

1//! Lightweight JSON-RPC router service.
2
3use rustc_hash::FxHashMap;
4use std::convert::Infallible;
5use std::fmt::{self, Debug, Formatter};
6use std::future::Future;
7use std::marker::PhantomData;
8use std::sync::Arc;
9use std::task::{Context, Poll};
10
11use futures::future::{self, BoxFuture, FutureExt};
12use serde::{de::DeserializeOwned, Serialize};
13use serde_json::Value;
14use tower::{util::BoxService, Layer, Service};
15
16use crate::jsonrpc::ErrorCode;
17
18use super::{Error, Id, Request, Response};
19
20/// Catch-all handler type for unrecognised JSON-RPC methods.
21type FallbackFn<E> =
22    Arc<dyn Fn(Request) -> BoxFuture<'static, Result<Option<Response>, E>> + Send + Sync>;
23
24/// A modular JSON-RPC 2.0 request router service.
25pub struct Router<S, E = Infallible> {
26    server: Arc<S>,
27    methods: FxHashMap<&'static str, BoxService<Request, Option<Response>, E>>,
28    /// Optional catch-all invoked when no registered method matches the incoming request.
29    /// When absent, unrecognised requests receive a `MethodNotFound` error response.
30    fallback: Option<FallbackFn<E>>,
31}
32
33impl<S: Send + Sync + 'static, E> Router<S, E> {
34    /// Creates a new `Router` with the given shared state.
35    pub fn new(server: S) -> Self {
36        Router {
37            server: Arc::new(server),
38            methods: FxHashMap::default(),
39            fallback: None,
40        }
41    }
42
43    /// Registers a catch-all handler invoked for any method not in the routing table.
44    ///
45    /// The compositor uses this to forward unknown LSP requests to the primary child server
46    /// rather than returning `MethodNotFound` to the editor client.
47    pub fn set_fallback<F>(&mut self, f: F) -> &mut Self
48    where
49        F: Fn(Request) -> BoxFuture<'static, Result<Option<Response>, E>> + Send + Sync + 'static,
50    {
51        self.fallback = Some(Arc::new(f));
52        self
53    }
54
55    /// Returns a reference to the inner server.
56    pub fn inner(&self) -> &S {
57        self.server.as_ref()
58    }
59
60    /// Registers a new RPC method which constructs a response with the given `callback`.
61    ///
62    /// The `layer` argument can be used to inject middleware into the method handler, if desired.
63    pub fn method<P, R, F, L>(&mut self, name: &'static str, callback: F, layer: L) -> &mut Self
64    where
65        P: FromParams,
66        R: IntoResponse,
67        F: for<'a> Method<&'a S, P, R> + Clone + Send + Sync + 'static,
68        L: Layer<MethodHandler<P, R, E>>,
69        L::Service: Service<Request, Response = Option<Response>, Error = E> + Send + 'static,
70        <L::Service as Service<Request>>::Future: Send + 'static,
71    {
72        let server = &self.server;
73        self.methods.entry(name).or_insert_with(|| {
74            let server = server.clone();
75            let handler = MethodHandler::new(move |params| {
76                let callback = callback.clone();
77                let server = server.clone();
78                async move { callback.invoke(&*server, params).await }
79            });
80
81            BoxService::new(layer.layer(handler))
82        });
83
84        self
85    }
86}
87
88impl<S: Debug, E> Debug for Router<S, E> {
89    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
90        f.debug_struct("Router")
91            .field("server", &self.server)
92            .field("methods", &self.methods.keys())
93            .finish_non_exhaustive()
94    }
95}
96
97impl<S, E: Send + 'static> Service<Request> for Router<S, E> {
98    type Response = Option<Response>;
99    type Error = E;
100    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
101
102    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
103        Poll::Ready(Ok(()))
104    }
105
106    fn call(&mut self, req: Request) -> Self::Future {
107        if let Some(handler) = self.methods.get_mut(req.method()) {
108            handler.call(req)
109        } else if let Some(fallback) = &self.fallback {
110            fallback(req)
111        } else {
112            let (method, id, _) = req.into_parts();
113            future::ok(id.map(|id| {
114                let mut error = Error::method_not_found();
115                error.data = Some(Value::from(method));
116                Response::from_error(id, error)
117            }))
118            .boxed()
119        }
120    }
121}
122
123/// Opaque JSON-RPC method handler.
124pub struct MethodHandler<P, R, E> {
125    f: Arc<dyn Fn(P) -> BoxFuture<'static, R> + Send + Sync>,
126    _marker: PhantomData<E>,
127}
128
129impl<P, R, E> Clone for MethodHandler<P, R, E> {
130    fn clone(&self) -> Self {
131        MethodHandler {
132            f: self.f.clone(),
133            _marker: PhantomData,
134        }
135    }
136}
137
138impl<P: FromParams, R: IntoResponse, E> MethodHandler<P, R, E> {
139    fn new<F, Fut>(handler: F) -> Self
140    where
141        F: Fn(P) -> Fut + Send + Sync + 'static,
142        Fut: Future<Output = R> + Send + 'static,
143    {
144        MethodHandler {
145            f: Arc::new(move |p| handler(p).boxed()),
146            _marker: PhantomData,
147        }
148    }
149}
150
151impl<P, R, E> Service<Request> for MethodHandler<P, R, E>
152where
153    P: FromParams,
154    R: IntoResponse,
155    E: Send + 'static,
156{
157    type Response = Option<Response>;
158    type Error = E;
159    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
160
161    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
162        Poll::Ready(Ok(()))
163    }
164
165    fn call(&mut self, req: Request) -> Self::Future {
166        let (_, id, params) = req.into_parts();
167
168        match id {
169            Some(_) if R::is_notification() => return future::ok(().into_response(id)).boxed(),
170            None if !R::is_notification() => return future::ok(None).boxed(),
171            _ => {}
172        }
173
174        let params = match P::from_params(params) {
175            Ok(params) => params,
176            Err(err) => return future::ok(id.map(|id| Response::from_error(id, err))).boxed(),
177        };
178
179        (self.f)(params)
180            .map(move |r| Ok(r.into_response(id)))
181            .boxed()
182    }
183}
184
185/// A trait implemented by all valid JSON-RPC method handlers.
186///
187/// This trait abstracts over the following classes of functions and/or closures:
188///
189/// Signature                                            | Description
190/// -----------------------------------------------------|---------------------------------
191/// `async fn f(&self) -> jsonrpc::Result<R>`            | Request without parameters
192/// `async fn f(&self, params: P) -> jsonrpc::Result<R>` | Request with required parameters
193/// `async fn f(&self)`                                  | Notification without parameters
194/// `async fn f(&self, params: P)`                       | Notification with parameters
195pub trait Method<S, P, R>: private::Sealed {
196    /// The future response value.
197    type Future: Future<Output = R> + Send;
198
199    /// Invokes the method with the given `server` receiver and parameters.
200    fn invoke(&self, server: S, params: P) -> Self::Future;
201}
202
203/// Support parameter-less JSON-RPC methods.
204impl<F, S, R, Fut> Method<S, (), R> for F
205where
206    F: Fn(S) -> Fut,
207    Fut: Future<Output = R> + Send,
208{
209    type Future = Fut;
210
211    #[inline]
212    fn invoke(&self, server: S, _: ()) -> Self::Future {
213        self(server)
214    }
215}
216
217/// Support JSON-RPC methods with `params`.
218impl<F, S, P, R, Fut> Method<S, (P,), R> for F
219where
220    F: Fn(S, P) -> Fut,
221    P: DeserializeOwned,
222    Fut: Future<Output = R> + Send,
223{
224    type Future = Fut;
225
226    #[inline]
227    fn invoke(&self, server: S, params: (P,)) -> Self::Future {
228        self(server, params.0)
229    }
230}
231
232/// A trait implemented by all JSON-RPC method parameters.
233pub trait FromParams: private::Sealed + Send + Sized + 'static {
234    /// Attempts to deserialize `Self` from the `params` value extracted from [`Request`].
235    fn from_params(params: Option<Value>) -> super::Result<Self>;
236}
237
238/// Deserialize non-existent JSON-RPC parameters.
239impl FromParams for () {
240    fn from_params(_params: Option<Value>) -> super::Result<Self> {
241        Ok(())
242    }
243}
244
245/// Deserialize required JSON-RPC parameters.
246impl<P: DeserializeOwned + Send + 'static> FromParams for (P,) {
247    fn from_params(params: Option<Value>) -> super::Result<Self> {
248        let p = params.unwrap_or(Value::Null);
249        serde_json::from_value(p)
250            .map(|p| (p,))
251            .map_err(|e| Error::invalid_params(e.to_string()))
252    }
253}
254
255/// A trait implemented by all JSON-RPC response types.
256pub trait IntoResponse: private::Sealed + Send + 'static {
257    /// Attempts to construct a [`Response`] using `Self` and a corresponding [`Id`].
258    fn into_response(self, id: Option<Id>) -> Option<Response>;
259
260    /// Returns `true` if this is a notification response type.
261    fn is_notification() -> bool;
262}
263
264/// Support JSON-RPC notification methods.
265impl IntoResponse for () {
266    fn into_response(self, id: Option<Id>) -> Option<Response> {
267        id.map(|id| Response::from_error(id, Error::invalid_request()))
268    }
269
270    #[inline]
271    fn is_notification() -> bool {
272        true
273    }
274}
275
276/// Support JSON-RPC request methods.
277impl<R: Serialize + Send + 'static> IntoResponse for Result<R, Error> {
278    fn into_response(self, id: Option<Id>) -> Option<Response> {
279        debug_assert!(id.is_some(), "Requests always contain an `id` field");
280        if let Some(id) = id {
281            let result = self.and_then(|r| {
282                serde_json::to_value(r).map_err(|e| Error {
283                    code: ErrorCode::InternalError,
284                    message: e.to_string().into(),
285                    data: None,
286                })
287            });
288            Some(Response::from_parts(id, result))
289        } else {
290            None
291        }
292    }
293
294    #[inline]
295    fn is_notification() -> bool {
296        false
297    }
298}
299
300mod private {
301    pub trait Sealed {}
302    impl<T> Sealed for T {}
303}
304
305#[cfg(test)]
306mod tests {
307    use serde::{Deserialize, Serialize};
308    use serde_json::json;
309    use tower::layer::layer_fn;
310    use tower::ServiceExt;
311
312    use super::*;
313
314    #[derive(Deserialize, Serialize)]
315    struct Params {
316        foo: i32,
317        bar: String,
318    }
319
320    struct Mock;
321
322    impl Mock {
323        async fn request(&self) -> Result<Value, Error> {
324            Ok(Value::Null)
325        }
326
327        async fn request_params(&self, params: Params) -> Result<Params, Error> {
328            Ok(params)
329        }
330
331        async fn notification(&self) {}
332
333        async fn notification_params(&self, _params: Params) {}
334    }
335
336    #[tokio::test(flavor = "current_thread")]
337    async fn routes_requests() {
338        let mut router: Router<Mock> = Router::new(Mock);
339        router
340            .method("first", Mock::request, layer_fn(|s| s))
341            .method("second", Mock::request_params, layer_fn(|s| s));
342
343        let request = Request::build("first").id(0).finish();
344        let response = router.ready().await.unwrap().call(request).await;
345        assert_eq!(response, Ok(Some(Response::from_ok(0.into(), Value::Null))));
346
347        let params = json!({"foo": -123i32, "bar": "hello world"});
348        let with_params = Request::build("second")
349            .params(params.clone())
350            .id(1)
351            .finish();
352        let response = router.ready().await.unwrap().call(with_params).await;
353        assert_eq!(response, Ok(Some(Response::from_ok(1.into(), params))));
354    }
355
356    #[tokio::test(flavor = "current_thread")]
357    async fn routes_notifications() {
358        let mut router: Router<Mock> = Router::new(Mock);
359        router
360            .method("first", Mock::notification, layer_fn(|s| s))
361            .method("second", Mock::notification_params, layer_fn(|s| s));
362
363        let request = Request::build("first").finish();
364        let response = router.ready().await.unwrap().call(request).await;
365        assert_eq!(response, Ok(None));
366
367        let params = json!({"foo": -123i32, "bar": "hello world"});
368        let with_params = Request::build("second").params(params).finish();
369        let response = router.ready().await.unwrap().call(with_params).await;
370        assert_eq!(response, Ok(None));
371    }
372
373    #[tokio::test(flavor = "current_thread")]
374    async fn rejects_request_with_invalid_params() {
375        let mut router: Router<Mock> = Router::new(Mock);
376        router.method("request", Mock::request_params, layer_fn(|s| s));
377
378        let invalid_params = Request::build("request")
379            .params(json!("wrong"))
380            .id(0)
381            .finish();
382
383        let response = router.ready().await.unwrap().call(invalid_params).await;
384        assert_eq!(
385            response,
386            Ok(Some(Response::from_error(
387                0.into(),
388                Error::invalid_params("invalid type: string \"wrong\", expected struct Params"),
389            )))
390        );
391    }
392
393    #[tokio::test(flavor = "current_thread")]
394    async fn ignores_notification_with_invalid_params() {
395        let mut router: Router<Mock> = Router::new(Mock);
396        router.method("notification", Mock::request_params, layer_fn(|s| s));
397
398        let invalid_params = Request::build("notification")
399            .params(json!("wrong"))
400            .finish();
401
402        let response = router.ready().await.unwrap().call(invalid_params).await;
403        assert_eq!(response, Ok(None));
404    }
405
406    #[tokio::test(flavor = "current_thread")]
407    async fn handles_incorrect_request_types() {
408        let mut router: Router<Mock> = Router::new(Mock);
409        router
410            .method("first", Mock::request, layer_fn(|s| s))
411            .method("second", Mock::notification, layer_fn(|s| s));
412
413        let request = Request::build("first").finish();
414        let response = router.ready().await.unwrap().call(request).await;
415        assert_eq!(response, Ok(None));
416
417        let request = Request::build("second").id(0).finish();
418        let response = router.ready().await.unwrap().call(request).await;
419        assert_eq!(
420            response,
421            Ok(Some(Response::from_error(
422                0.into(),
423                Error::invalid_request(),
424            )))
425        );
426    }
427
428    #[tokio::test(flavor = "current_thread")]
429    async fn responds_to_nonexistent_request() {
430        let mut router: Router<Mock> = Router::new(Mock);
431
432        let request = Request::build("foo").id(0).finish();
433        let response = router.ready().await.unwrap().call(request).await;
434        let mut error = Error::method_not_found();
435        error.data = Some("foo".into());
436        assert_eq!(response, Ok(Some(Response::from_error(0.into(), error))));
437    }
438
439    #[tokio::test(flavor = "current_thread")]
440    async fn ignores_nonexistent_notification() {
441        let mut router: Router<Mock> = Router::new(Mock);
442
443        let request = Request::build("foo").finish();
444        let response = router.ready().await.unwrap().call(request).await;
445        assert_eq!(response, Ok(None));
446    }
447}
lsp_max/jsonrpc/router.rs

lsp_max/jsonrpc/
router.rs
