jsonrpsee_server/middleware/
rpc.rs

1// Copyright 2019-2021 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any
4// person obtaining a copy of this software and associated
5// documentation files (the "Software"), to deal in the
6// Software without restriction, including without
7// limitation the rights to use, copy, modify, merge,
8// publish, distribute, sublicense, and/or sell copies of
9// the Software, and to permit persons to whom the Software
10// is furnished to do so, subject to the following
11// conditions:
12//
13// The above copyright notice and this permission notice
14// shall be included in all copies or substantial portions
15// of the Software.
16//
17// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
18// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
19// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
20// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
21// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
22// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
23// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
24// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
25// DEALINGS IN THE SOFTWARE.
26
27//! JSON-RPC service middleware.
28
29pub use jsonrpsee_core::{middleware::*, server::MethodResponse};
30
31use std::{panic::AssertUnwindSafe, sync::Arc};
32
33use crate::ConnectionId;
34use futures_util::FutureExt;
35use jsonrpsee_core::{
36    server::{
37        BatchResponseBuilder, BoundedSubscriptions, MethodCallback, MethodSink, Methods,
38        SubscriptionState,
39    },
40    traits::IdProvider,
41};
42use jsonrpsee_types::{
43    error::{reject_too_many_subscriptions, ErrorCode},
44    ErrorObject, Id,
45};
46
47/// JSON-RPC service middleware.
48#[derive(Clone, Debug)]
49pub struct RpcService {
50    conn_id: ConnectionId,
51    methods: Methods,
52    max_response_body_size: usize,
53    /// `None` if the transport can't send subscription notifications.
54    subscriptions: Option<Subscriptions>,
55}
56
57/// State required to serve subscriptions on a connection.
58#[derive(Clone, Debug)]
59struct Subscriptions {
60    bounded: BoundedSubscriptions,
61    sink: MethodSink,
62    id_provider: Arc<dyn IdProvider>,
63}
64
65impl RpcService {
66    /// Creates a service for a connection that supports subscriptions.
67    ///
68    /// Subscription notifications are sent to `sink`, and at most `bounded_subscriptions` of them
69    /// can be active at once.
70    pub fn new(
71        methods: Methods,
72        max_response_body_size: usize,
73        conn_id: ConnectionId,
74        bounded_subscriptions: BoundedSubscriptions,
75        sink: MethodSink,
76        id_provider: Arc<dyn IdProvider>,
77    ) -> Self {
78        Self {
79            methods,
80            max_response_body_size,
81            conn_id,
82            subscriptions: Some(Subscriptions {
83                bounded: bounded_subscriptions,
84                sink,
85                id_provider,
86            }),
87        }
88    }
89
90    /// Creates a service that answers subscription calls with an internal error.
91    pub(crate) fn without_subscriptions(
92        methods: Methods,
93        max_response_body_size: usize,
94        conn_id: ConnectionId,
95    ) -> Self {
96        Self { methods, max_response_body_size, conn_id, subscriptions: None }
97    }
98}
99
100impl RpcServiceT for RpcService {
101    fn call<'a>(&self, req: Request<'a>) -> impl Future<Output = MethodResponse> + Send + 'a {
102        let conn_id = self.conn_id;
103        let max_response_body_size = self.max_response_body_size;
104
105        let Request { id, method, params, .. } = req;
106        let params = jsonrpsee_types::Params::new(
107            params.as_ref().map(|p| serde_json::value::RawValue::get(p)),
108        );
109
110        let Some(method) = self.methods.method(&method) else {
111            return ResponseFuture::ready(MethodResponse::error(
112                id,
113                ErrorObject::from(ErrorCode::MethodNotFound),
114            ));
115        };
116
117        match method {
118            MethodCallback::Async(callback) => ResponseFuture::future((callback)(
119                id.into_owned(),
120                params.into_owned(),
121                conn_id,
122                max_response_body_size,
123            )),
124            MethodCallback::Sync(callback) => {
125                ResponseFuture::ready((callback)(id, params, max_response_body_size))
126            }
127            MethodCallback::Subscription(callback) => {
128                let Some(subscriptions) = &self.subscriptions else {
129                    tracing::debug!("Subscriptions not supported");
130                    return ResponseFuture::ready(MethodResponse::error(
131                        id,
132                        ErrorObject::from(ErrorCode::InternalError),
133                    ));
134                };
135
136                let Some(subscription_permit) = subscriptions.bounded.acquire() else {
137                    let max = subscriptions.bounded.max();
138                    return ResponseFuture::ready(MethodResponse::error(
139                        id,
140                        reject_too_many_subscriptions(max),
141                    ));
142                };
143
144                let conn_state = SubscriptionState {
145                    conn_id,
146                    id_provider: &*subscriptions.id_provider,
147                    subscription_permit,
148                };
149                ResponseFuture::future((callback)(
150                    id,
151                    params,
152                    subscriptions.sink.clone(),
153                    conn_state,
154                ))
155            }
156            MethodCallback::Unsubscription(callback) => {
157                // Don't adhere to any resource or subscription limits; always let unsubscribing
158                // happen!
159                if self.subscriptions.is_none() {
160                    tracing::debug!("Subscriptions not supported");
161                    return ResponseFuture::ready(MethodResponse::error(
162                        id,
163                        ErrorObject::from(ErrorCode::InternalError),
164                    ));
165                }
166
167                ResponseFuture::ready(callback(id, params, conn_id, max_response_body_size))
168            }
169        }
170    }
171
172    fn batch<'a>(&self, batch: Batch<'a>) -> impl Future<Output = MethodResponse> + Send + 'a {
173        let mut batch_rp = BatchResponseBuilder::new_with_limit(self.max_response_body_size);
174        let service = self.clone();
175        async move {
176            // Notifications and subscription calls get no entry in the batch response.
177            let mut omitted_entry = false;
178
179            // Entries run one after another, so a batch occupies at most one call at a time.
180            for batch_entry in batch.into_iter() {
181                // Entries that complete without ever returning `Pending` (sync methods, unknown
182                // methods, invalid entries) would otherwise run the whole batch in one poll.
183                // This yields only once the task's budget is used up, so a long batch can't hold
184                // a worker or ignore cancellation.
185                tokio::task::coop::consume_budget().await;
186
187                match batch_entry {
188                    Ok(BatchEntry::Call(req)) => {
189                        let rp = catch_call_panic(req.id(), || service.call(req)).await;
190                        // Accepting or rejecting a subscription already sent its response through
191                        // the subscription sink.
192                        if rp.is_subscription() {
193                            omitted_entry = true;
194                            continue;
195                        }
196                        if let Err(err) = batch_rp.append(rp) {
197                            return err;
198                        }
199                    }
200                    Ok(BatchEntry::Notification(n)) => {
201                        omitted_entry = true;
202                        service.notification(n).await;
203                    }
204                    Err(err) => {
205                        let (err, id) = err.into_parts();
206                        let rp = MethodResponse::error(id, err);
207                        if let Err(err) = batch_rp.append(rp) {
208                            return err;
209                        }
210                    }
211                }
212            }
213
214            // If every entry was omitted, there is nothing to send back.
215            if batch_rp.is_empty() && omitted_entry {
216                MethodResponse::notification()
217            }
218            // An empty batch is regarded as an invalid request here.
219            else {
220                batch_rp.finish()
221            }
222        }
223    }
224
225    fn notification<'a>(
226        &self,
227        _n: Notification<'a>,
228    ) -> impl Future<Output = MethodResponse> + Send + 'a {
229        // Notifications are not answered.
230        std::future::ready(MethodResponse::notification())
231    }
232}
233
234/// Answers a panicking call with an internal error, otherwise the panic would only surface as a
235/// failed connection or call task and the client would never receive a response for `id`.
236///
237/// `call` builds the call's future inside the guard, because sync methods run while
238/// [`RpcServiceT::call`] builds it, not when it is polled.
239pub(crate) async fn catch_call_panic<F>(id: Id<'_>, call: impl FnOnce() -> F) -> MethodResponse
240where
241    F: Future<Output = MethodResponse>,
242{
243    AssertUnwindSafe(async move { call().await })
244        .catch_unwind()
245        .await
246        .unwrap_or_else(|_| MethodResponse::error(id, ErrorObject::from(ErrorCode::InternalError)))
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252    use crate::RpcModule;
253    use jsonrpsee_types::Request;
254    use std::{
255        pin::pin,
256        sync::atomic::{AtomicUsize, Ordering},
257        task::Poll,
258    };
259
260    #[tokio::test]
261    async fn batch_of_ready_entries_yields() {
262        const ENTRIES: usize = 4096;
263
264        let calls = Arc::new(AtomicUsize::new(0));
265        let mut module = RpcModule::new(calls.clone());
266        module.register_method("tick", |_, calls| calls.fetch_add(1, Ordering::SeqCst)).unwrap();
267        let service =
268            RpcService::without_subscriptions(module.into(), usize::MAX, ConnectionId::default());
269
270        let batch = (0..ENTRIES as u64)
271            .map(|id| Ok(BatchEntry::Call(Request::borrowed("tick", None, Id::Number(id)))))
272            .collect::<Vec<_>>();
273        let mut response = pin!(service.batch(Batch::from(batch)));
274
275        // Every entry is ready, so only the task budget can make the batch yield.
276        let first_poll = std::future::poll_fn(|cx| Poll::Ready(response.as_mut().poll(cx))).await;
277        assert!(first_poll.is_pending());
278        assert!(calls.load(Ordering::SeqCst) < ENTRIES);
279
280        let expected = (0..ENTRIES)
281            .map(|i| format!(r#"{{"jsonrpc":"2.0","id":{i},"result":{i}}}"#))
282            .collect::<Vec<_>>()
283            .join(",");
284        assert_eq!(response.await.into_json(), format!("[{expected}]"));
285        assert_eq!(calls.load(Ordering::SeqCst), ENTRIES);
286    }
287}
jsonrpsee_server/middleware/rpc.rs

jsonrpsee_server/middleware/
rpc.rs
