jsonrpsee_core/server/rpc_module.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
27use std::{
28 collections::hash_map::Entry,
29 fmt::Debug,
30 future::Future,
31 ops::{Deref, DerefMut},
32 sync::{
33 atomic::{AtomicUsize, Ordering},
34 Arc,
35 },
36};
37
38use crate::{
39 error::RegisterMethodError,
40 id_providers::RandomIntegerIdProvider,
41 server::{
42 helpers::MethodSink,
43 subscription::{
44 sub_message_to_json, BoundedSubscriptions, IntoSubscriptionCloseResponse,
45 PendingSubscriptionSink, Subscribers, SubscriptionCloseResponse, SubscriptionKey,
46 SubscriptionPermit, SubscriptionState,
47 },
48 MethodResponse, ResponsePayload, ResponseSizeHint, LOG_TARGET,
49 },
50};
51use futures_util::{future::BoxFuture, FutureExt};
52use jsonrpsee_types::{
53 error::{ErrorCode, ErrorObject},
54 Id, Params, Request, SubscriptionId as RpcSubscriptionId,
55};
56use rustc_hash::FxHashMap;
57use tokio::sync::{mpsc, oneshot};
58
59use super::{sub_err_to_json, IntoResponse};
60
61/// A `MethodCallback` is an RPC endpoint, callable with a standard JSON-RPC request,
62/// implemented as a function pointer to a `Fn` function taking four arguments:
63/// the `id`, `params`, a channel the function uses to communicate the result (or error)
64/// back to `jsonrpsee`, and the connection ID (useful for the websocket transport).
65pub type SyncMethod = Arc<dyn Send + Sync + Fn(Id, Params, MaxResponseSize) -> MethodResponse>;
66/// Similar to [`SyncMethod`], but represents an asynchronous handler.
67pub type AsyncMethod<'a> = Arc<
68 dyn Send
69 + Sync
70 + Fn(Id<'a>, Params<'a>, ConnectionId, MaxResponseSize) -> BoxFuture<'a, MethodResponse>,
71>;
72
73/// Method callback for subscriptions.
74pub type SubscriptionMethod<'a> = Arc<
75 dyn Send
76 + Sync
77 + Fn(Id, Params, MethodSink, SubscriptionState) -> BoxFuture<'a, MethodResponse>,
78>;
79// Method callback to unsubscribe.
80type UnsubscriptionMethod =
81 Arc<dyn Send + Sync + Fn(Id, Params, ConnectionId, MaxResponseSize) -> MethodResponse>;
82
83/// Connection ID.
84#[derive(
85 Debug, Copy, Clone, PartialEq, Eq, Hash, Default, serde::Deserialize, serde::Serialize,
86)]
87pub struct ConnectionId(pub usize);
88
89impl ConnectionId {
90 /// Owner of subscriptions created by in-process calls such as
91 /// [`Methods::raw_json_request`].
92 ///
93 /// Servers count connections upwards from zero and never reach this value, so no transport
94 /// connection can unsubscribe an in-process subscription.
95 pub const IN_PROCESS: Self = Self(usize::MAX);
96
97 /// Returns an id that no other live connection in this process shares.
98 ///
99 /// Together with the subscription id this scopes a subscription, and servers of different
100 /// transports can serve one shared subscription table, so the id has to be unique across every
101 /// server in the process and not only within a single accept loop. Counting per server would
102 /// give the first connection of each transport the same id, and an unsubscribe on one of them
103 /// would then cancel a subscription of another.
104 ///
105 /// A `usize` cannot wrap in practice on 64-bit targets, so ids are never reused and never reach
106 /// [`Self::IN_PROCESS`].
107 pub fn next_unique() -> Self {
108 static NEXT: AtomicUsize = AtomicUsize::new(0);
109
110 Self(NEXT.fetch_add(1, Ordering::Relaxed))
111 }
112}
113
114impl From<usize> for ConnectionId {
115 fn from(id: usize) -> Self {
116 Self(id)
117 }
118}
119
120/// Max response size.
121pub type MaxResponseSize = usize;
122
123/// Callback wrapper that can be either sync or async.
124#[derive(Clone)]
125pub enum MethodCallback {
126 /// Synchronous method handler.
127 Sync(SyncMethod),
128 /// Asynchronous method handler.
129 Async(AsyncMethod<'static>),
130 /// Subscription method handler.
131 Subscription(SubscriptionMethod<'static>),
132 /// Unsubscription method handler.
133 Unsubscription(UnsubscriptionMethod),
134}
135
136impl Debug for MethodCallback {
137 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
138 match self {
139 Self::Async(_) => write!(f, "Async"),
140 Self::Sync(_) => write!(f, "Sync"),
141 Self::Subscription(_) => write!(f, "Subscription"),
142 Self::Unsubscription(_) => write!(f, "Unsubscription"),
143 }
144 }
145}
146
147/// Reference-counted, clone-on-write collection of synchronous and asynchronous methods.
148#[derive(Default, Debug, Clone)]
149pub struct Methods {
150 callbacks: Arc<FxHashMap<&'static str, MethodCallback>>,
151}
152
153impl Methods {
154 /// Creates a new empty [`Methods`].
155 pub fn new() -> Self {
156 Self::default()
157 }
158
159 /// Verifies that the method name is not already taken, and returns an error if it is.
160 pub fn verify_method_name(&mut self, name: &'static str) -> Result<(), RegisterMethodError> {
161 if self.callbacks.contains_key(name) {
162 return Err(RegisterMethodError::AlreadyRegistered(name.into()));
163 }
164
165 Ok(())
166 }
167
168 /// Inserts the method callback for a given name, or returns an error if the name was already
169 /// taken. On success it returns a mut reference to the [`MethodCallback`] just inserted.
170 pub fn verify_and_insert(
171 &mut self,
172 name: &'static str,
173 callback: MethodCallback,
174 ) -> Result<&mut MethodCallback, RegisterMethodError> {
175 match self.mut_callbacks().entry(name) {
176 Entry::Occupied(_) => Err(RegisterMethodError::AlreadyRegistered(name.into())),
177 Entry::Vacant(vacant) => Ok(vacant.insert(callback)),
178 }
179 }
180
181 /// Helper for obtaining a mut ref to the callbacks HashMap.
182 fn mut_callbacks(&mut self) -> &mut FxHashMap<&'static str, MethodCallback> {
183 Arc::make_mut(&mut self.callbacks)
184 }
185
186 /// Merge two [`Methods`]'s by adding all [`MethodCallback`]s from `other` into `self`.
187 /// Fails if any of the methods in `other` is present already.
188 pub fn merge(&mut self, other: impl Into<Methods>) -> Result<(), RegisterMethodError> {
189 let mut other = other.into();
190
191 for name in other.callbacks.keys() {
192 self.verify_method_name(name)?;
193 }
194
195 let callbacks = self.mut_callbacks();
196
197 for (name, callback) in other.mut_callbacks().drain() {
198 callbacks.insert(name, callback);
199 }
200
201 Ok(())
202 }
203
204 /// Returns the method callback.
205 pub fn method(&self, method_name: &str) -> Option<&MethodCallback> {
206 self.callbacks.get(method_name)
207 }
208
209 /// Make a request (JSON-RPC method call or subscription) by using raw JSON.
210 ///
211 /// Returns the raw JSON response to the call and a stream to receive notifications if the call
212 /// was a subscription.
213 ///
214 /// # Examples
215 ///
216 /// ```
217 /// #[tokio::main]
218 /// async fn main() {
219 /// use jsonrpsee::{RpcModule, SubscriptionMessage};
220 /// use jsonrpsee::types::{response::Success, Response};
221 /// use jsonrpsee::core::to_json_raw_value;
222 /// use futures_util::StreamExt;
223 ///
224 /// let mut module = RpcModule::new(());
225 /// module.register_subscription("hi", "hi", "goodbye", |_, pending, _| async {
226 /// let sink = pending.accept().await?;
227 ///
228 /// // see comment above.
229 /// let msg = to_json_raw_value(&"one answer").unwrap();
230 /// sink.send(msg).await?;
231 ///
232 /// Ok(())
233 /// }).unwrap();
234 /// let (resp, mut stream) = module.raw_json_request(r#"{"jsonrpc":"2.0","method":"hi","id":0}"#, 1).await.unwrap();
235 /// // If the response is an error converting it to `Success` will fail.
236 /// let resp: Success<u64> = serde_json::from_str::<Response<u64>>(&resp).unwrap().try_into().unwrap();
237 /// let sub_resp = stream.recv().await.unwrap();
238 /// assert_eq!(
239 /// format!(r#"{{"jsonrpc":"2.0","method":"hi","params":{{"subscription":{},"result":"one answer"}}}}"#, resp.result),
240 /// sub_resp
241 /// );
242 /// }
243 /// ```
244 pub async fn raw_json_request(
245 &self,
246 request: &str,
247 buf_size: usize,
248 ) -> Result<(String, mpsc::Receiver<String>), serde_json::Error> {
249 tracing::trace!("[Methods::raw_json_request] Request: {:?}", request);
250 let req: Request = serde_json::from_str(request)?;
251 let (resp, rx) = self.inner_call(req, buf_size, mock_subscription_permit()).await;
252
253 Ok((resp, rx))
254 }
255
256 /// Execute a callback.
257 async fn inner_call(
258 &self,
259 req: Request<'_>,
260 buf_size: usize,
261 subscription_permit: SubscriptionPermit,
262 ) -> (String, mpsc::Receiver<String>) {
263 let (tx, mut rx) = mpsc::channel(buf_size);
264 let Request { id, method, params, .. } = req;
265 let params = Params::new(params.as_ref().map(|params| params.as_ref().get()));
266 let max_response_size = usize::MAX;
267 let conn_id = ConnectionId::IN_PROCESS;
268
269 let response = match self.method(&method) {
270 None => MethodResponse::error(id, ErrorObject::from(ErrorCode::MethodNotFound)),
271 Some(MethodCallback::Sync(cb)) => (cb)(id, params, max_response_size),
272 Some(MethodCallback::Async(cb)) => {
273 (cb)(id.into_owned(), params.into_owned(), conn_id, max_response_size).await
274 }
275 Some(MethodCallback::Subscription(cb)) => {
276 let conn_state = SubscriptionState {
277 conn_id,
278 id_provider: &RandomIntegerIdProvider,
279 subscription_permit,
280 };
281 let res = (cb)(id, params, MethodSink::new(tx.clone()), conn_state).await;
282
283 // This message is not used because it's used for metrics so we discard in other to
284 // not read once this is used for subscriptions.
285 //
286 // The same information is part of `res` above.
287 let _ =
288 rx.recv().await.expect("Every call must at least produce one response; qed");
289
290 res
291 }
292 Some(MethodCallback::Unsubscription(cb)) => {
293 (cb)(id, params, conn_id, max_response_size)
294 }
295 };
296
297 let rp = response.into_json();
298 tracing::trace!(target: LOG_TARGET, "[Methods::inner_call] Method: {}, response: {}", method, rp);
299
300 (rp, rx)
301 }
302
303 /// Returns an `Iterator` with all the method names registered on this server.
304 pub fn method_names(&self) -> impl Iterator<Item = &'static str> + '_ {
305 self.callbacks.keys().copied()
306 }
307}
308
309impl<Context> Deref for RpcModule<Context> {
310 type Target = Methods;
311
312 fn deref(&self) -> &Methods {
313 &self.methods
314 }
315}
316
317impl<Context> DerefMut for RpcModule<Context> {
318 fn deref_mut(&mut self) -> &mut Methods {
319 &mut self.methods
320 }
321}
322
323/// Sets of JSON-RPC methods can be organized into "module"s that are in turn registered on the
324/// server or, alternatively, merged with other modules to construct a cohesive API. [`RpcModule`]
325/// wraps an additional context argument that can be used to access data during call execution.
326#[derive(Debug, Clone)]
327pub struct RpcModule<Context> {
328 ctx: Arc<Context>,
329 methods: Methods,
330}
331
332impl<Context> RpcModule<Context> {
333 /// Create a new module with a given shared `Context`.
334 pub fn new(ctx: Context) -> Self {
335 Self::from_arc(Arc::new(ctx))
336 }
337
338 /// Create a new module from an already shared `Context`.
339 ///
340 /// This is useful if `Context` needs to be shared outside of an [`RpcModule`].
341 pub fn from_arc(ctx: Arc<Context>) -> Self {
342 Self { ctx, methods: Default::default() }
343 }
344
345 /// Transform a module into an `RpcModule<()>` (unit context).
346 pub fn remove_context(self) -> RpcModule<()> {
347 let mut module = RpcModule::new(());
348 module.methods = self.methods;
349 module
350 }
351}
352
353impl<Context> From<RpcModule<Context>> for Methods {
354 fn from(module: RpcModule<Context>) -> Methods {
355 module.methods
356 }
357}
358
359impl<Context: Send + Sync + 'static> RpcModule<Context> {
360 /// Register a new synchronous RPC method, which computes the response with the given callback.
361 ///
362 /// ## Examples
363 ///
364 /// ```
365 /// use jsonrpsee_core::server::RpcModule;
366 ///
367 /// let mut module = RpcModule::new(());
368 /// module.register_method("say_hello", |_params, _ctx| "lo").unwrap();
369 /// ```
370 pub fn register_method<R, F>(
371 &mut self,
372 method_name: &'static str,
373 callback: F,
374 ) -> Result<&mut MethodCallback, RegisterMethodError>
375 where
376 Context: Send + Sync + 'static,
377 R: IntoResponse + 'static,
378 F: Fn(Params, &Context) -> R + Send + Sync + 'static,
379 {
380 let ctx = self.ctx.clone();
381 let size_hint = ResponseSizeHint::new();
382 self.methods.verify_and_insert(
383 method_name,
384 MethodCallback::Sync(Arc::new(move |id, params, max_response_size| {
385 let rp = callback(params, &*ctx).into_response();
386 MethodResponse::sized_response(id, rp, max_response_size, &size_hint)
387 })),
388 )
389 }
390
391 /// Removes the method if it exists.
392 ///
393 /// Be aware that a subscription consist of two methods, `subscribe` and `unsubscribe` and
394 /// it's the caller responsibility to remove both `subscribe` and `unsubscribe` methods for
395 /// subscriptions.
396 pub fn remove_method(&mut self, method_name: &'static str) -> Option<MethodCallback> {
397 self.methods.mut_callbacks().remove(method_name)
398 }
399
400 /// Register a new asynchronous RPC method, which computes the response with the given callback.
401 ///
402 /// ## Examples
403 ///
404 /// ```
405 /// use jsonrpsee_core::server::RpcModule;
406 ///
407 /// let mut module = RpcModule::new(());
408 /// module.register_async_method("say_hello", |_params, _ctx| async { "lo" }).unwrap();
409 /// ```
410 pub fn register_async_method<R, Fun, Fut>(
411 &mut self,
412 method_name: &'static str,
413 callback: Fun,
414 ) -> Result<&mut MethodCallback, RegisterMethodError>
415 where
416 R: IntoResponse + 'static,
417 Fut: Future<Output = R> + Send,
418 Fun: (Fn(Params<'static>, Arc<Context>) -> Fut) + Clone + Send + Sync + 'static,
419 {
420 let ctx = self.ctx.clone();
421 let size_hint = Arc::new(ResponseSizeHint::new());
422 self.methods.verify_and_insert(
423 method_name,
424 MethodCallback::Async(Arc::new(move |id, params, _, max_response_size| {
425 let ctx = ctx.clone();
426 let callback = callback.clone();
427 let size_hint = size_hint.clone();
428
429 let future = async move {
430 let rp = callback(params, ctx).await.into_response();
431 MethodResponse::sized_response(id, rp, max_response_size, &size_hint)
432 };
433 future.boxed()
434 })),
435 )
436 }
437
438 /// Register a new **blocking** synchronous RPC method, which computes the response with the
439 /// given callback. Unlike the regular [`register_method`](RpcModule::register_method), this
440 /// method runs on tokio's blocking thread pool, so it can block its thread and perform
441 /// expensive computations.
442 pub fn register_blocking_method<R, F>(
443 &mut self,
444 method_name: &'static str,
445 callback: F,
446 ) -> Result<&mut MethodCallback, RegisterMethodError>
447 where
448 Context: Send + Sync + 'static,
449 R: IntoResponse + 'static,
450 F: Fn(Params<'static>, Arc<Context>) -> R + Clone + Send + Sync + 'static,
451 {
452 let ctx = self.ctx.clone();
453 let size_hint = Arc::new(ResponseSizeHint::new());
454 self.methods.verify_and_insert(
455 method_name,
456 MethodCallback::Async(Arc::new(move |id, params, _, max_response_size| {
457 let ctx = ctx.clone();
458 let callback = callback.clone();
459 let size_hint = size_hint.clone();
460 let error_id = id.clone();
461
462 tokio::task::spawn_blocking(move || {
463 let rp = callback(params, ctx).into_response();
464 MethodResponse::sized_response(id, rp, max_response_size, &size_hint)
465 })
466 .map(|result| {
467 result.unwrap_or_else(|err| {
468 tracing::error!(target: LOG_TARGET, "Join error for blocking RPC method: {:?}", err);
469 MethodResponse::error(error_id, ErrorObject::from(ErrorCode::InternalError))
470 })
471 })
472 .boxed()
473 })),
474 )
475 }
476
477 /// Register a new publish/subscribe interface using JSON-RPC notifications.
478 ///
479 /// It implements the [ethereum pubsub specification](https://geth.ethereum.org/docs/rpc/pubsub)
480 /// with an option to choose custom subscription ID generation.
481 ///
482 /// Furthermore, it generates the `unsubscribe implementation` where a `bool` is used as
483 /// the result to indicate whether the subscription was successfully unsubscribed to or not.
484 /// For instance an `unsubscribe call` may fail if a non-existent subscription ID is used in the
485 /// call.
486 ///
487 /// This method ensures that the `subscription_method_name` and `unsubscription_method_name` are
488 /// unique. The `notif_method_name` argument sets the content of the `method` field in the
489 /// JSON document that the server sends back to the client. The uniqueness of this value is
490 /// not machine checked and it's up to the user to ensure it is not used in any other
491 /// [`RpcModule`] used in the server.
492 ///
493 /// # Arguments
494 ///
495 /// * `subscription_method_name` - name of the method to call to initiate a subscription
496 /// * `notif_method_name` - name of method to be used in the subscription payload (technically a
497 /// JSON-RPC notification)
498 /// * `unsubscription_method` - name of the method to call to terminate a subscription
499 /// * `callback` - A callback to invoke on each subscription; it takes three parameters:
500 /// - [`Params`]: JSON-RPC parameters in the subscription call.
501 /// - [`PendingSubscriptionSink`]: A pending subscription waiting to be accepted, in order
502 /// to send out messages on the subscription
503 /// - Context: Any type that can be embedded into the [`RpcModule`].
504 ///
505 /// # Returns
506 ///
507 /// An async block which returns something that implements
508 /// [`crate::server::IntoSubscriptionCloseResponse`] which decides what action to take when
509 /// the subscription ends whether such as to sent out another message on the subscription
510 /// stream before closing down it.
511 ///
512 /// NOTE: The return value is ignored if [`PendingSubscriptionSink`] hasn't been called or is
513 /// unsuccessful, as the subscription is not allowed to send out subscription notifications
514 /// before the actual subscription has been established.
515 ///
516 /// This is implemented for `Result<T, E>` and `()`.
517 ///
518 /// It's recommended to use `Result` if you want to propagate the error as special error
519 /// notification Another option is to implement
520 /// [`crate::server::IntoSubscriptionCloseResponse`] if you want customized behaviour.
521 ///
522 /// The error notification has the following format:
523 ///
524 /// ```json
525 /// {
526 /// "jsonrpc": "2.0",
527 /// "method": "<method>",
528 /// "params": {
529 /// "subscription": "<subscriptionID>",
530 /// "error": <your msg>
531 /// }
532 /// }
533 /// }
534 /// ```
535 ///
536 /// # Examples
537 ///
538 /// ```no_run
539 /// use jsonrpsee_core::server::{RpcModule, SubscriptionMessage, SubscriptionSink};
540 /// use jsonrpsee_types::ErrorObjectOwned;
541 ///
542 /// let mut ctx = RpcModule::new(99_usize);
543 /// ctx.register_subscription("sub", "notif_name", "unsub", |params, pending, ctx| async move {
544 /// let x = match params.one::<usize>() {
545 /// Ok(x) => x,
546 /// Err(e) => {
547 /// pending.reject(ErrorObjectOwned::from(e)).await;
548 /// // If the subscription has not been "accepted" then
549 /// // the return value will be "ignored" as it's not
550 /// // allowed to send out any further notifications on
551 /// // on the subscription.
552 /// return Ok(());
553 /// }
554 /// };
555 ///
556 /// // Mark the subscription is accepted after the params has been parsed successful.
557 /// // This is actually responds the underlying RPC method call and may fail if the
558 /// // connection is closed.
559 /// let sink = pending.accept().await?;
560 /// let sum = x + (*ctx);
561 ///
562 /// // This will send out an error notification if it fails.
563 /// //
564 /// // If you need some other behavior implement or custom format of the error field
565 /// // you need to manually handle that.
566 /// let msg = serde_json::value::to_raw_value(&sum).unwrap();
567 ///
568 /// // This fails only if the connection is closed
569 /// sink.send(msg).await?;
570 ///
571 /// Ok(())
572 /// });
573 /// ```
574 pub fn register_subscription<R, F, Fut>(
575 &mut self,
576 subscribe_method_name: &'static str,
577 notif_method_name: &'static str,
578 unsubscribe_method_name: &'static str,
579 callback: F,
580 ) -> Result<&mut MethodCallback, RegisterMethodError>
581 where
582 Context: Send + Sync + 'static,
583 F: (Fn(Params<'static>, PendingSubscriptionSink, Arc<Context>) -> Fut)
584 + Send
585 + Sync
586 + Clone
587 + 'static,
588 Fut: Future<Output = R> + Send + 'static,
589 R: IntoSubscriptionCloseResponse + Send,
590 {
591 let subscribers =
592 self.verify_and_register_unsubscribe(subscribe_method_name, unsubscribe_method_name)?;
593 let ctx = self.ctx.clone();
594
595 // Subscribe
596 let callback = {
597 self.methods.verify_and_insert(
598 subscribe_method_name,
599 MethodCallback::Subscription(Arc::new(move |id, params, method_sink, conn| {
600 let uniq_sub = SubscriptionKey {
601 conn_id: conn.conn_id,
602 sub_id: conn.id_provider.next_id(),
603 };
604
605 // response to the subscription call.
606 let (tx, rx) = oneshot::channel();
607 let (accepted_tx, accepted_rx) = oneshot::channel();
608
609 let sub_id = uniq_sub.sub_id.clone();
610 let method = notif_method_name;
611
612 let sink = PendingSubscriptionSink {
613 inner: method_sink.clone(),
614 method: notif_method_name,
615 subscribers: subscribers.clone(),
616 uniq_sub,
617 id: id.clone().into_owned(),
618 subscribe: tx,
619 permit: conn.subscription_permit,
620 };
621
622 // The subscription callback is a future from the subscription
623 // definition and not the as same when the subscription call has been
624 // completed.
625 //
626 // This runs until the subscription callback has completed.
627 let sub_fut = callback(params.into_owned(), sink, ctx.clone());
628
629 tokio::spawn(async move {
630 // This will wait for the subscription future to be resolved
631 let response = match futures_util::future::try_join(
632 sub_fut.map(|f| Ok(f)),
633 accepted_rx,
634 )
635 .await
636 {
637 Ok((r, _)) => r.into_response(),
638 // The accept call failed i.e, the subscription was not accepted.
639 Err(_) => return,
640 };
641
642 match response {
643 SubscriptionCloseResponse::Notif(msg) => {
644 let json = sub_message_to_json(msg, &sub_id, method);
645 let _ = method_sink.send(json).await;
646 }
647 SubscriptionCloseResponse::NotifErr(err) => {
648 let json = sub_err_to_json(err, sub_id, method);
649 let _ = method_sink.send(json).await;
650 }
651 SubscriptionCloseResponse::None => (),
652 }
653 });
654
655 let id = id.clone().into_owned();
656
657 Box::pin(async move {
658 match rx.await {
659 Ok(rp) => {
660 // If the subscription was accepted then send a message
661 // to subscription task otherwise rely on the drop impl.
662 if rp.is_success() {
663 let _ = accepted_tx.send(());
664 }
665 rp
666 }
667 Err(_) => MethodResponse::error(id, ErrorCode::InternalError),
668 }
669 })
670 })),
671 )?
672 };
673
674 Ok(callback)
675 }
676
677 /// Similar to [`RpcModule::register_subscription`] but a little lower-level API
678 /// where handling the subscription is managed the user i.e, polling the subscription
679 /// such as spawning a separate task to do so.
680 ///
681 /// This is more efficient as this doesn't require cloning the `params` in the subscription
682 /// and it won't send out a close message. Such things are delegated to the user of this API
683 ///
684 /// # Examples
685 ///
686 /// ```no_run
687 /// use jsonrpsee_core::server::{RpcModule, SubscriptionMessage, SubscriptionSink};
688 /// use jsonrpsee_types::ErrorObjectOwned;
689 ///
690 /// let mut ctx = RpcModule::new(99_usize);
691 /// ctx.register_subscription_raw("sub", "notif_name", "unsub", |params, pending, ctx| {
692 /// // The params are parsed outside the async block below to avoid cloning the bytes.
693 /// let val = match params.one::<usize>() {
694 /// Ok(val) => val,
695 /// Err(e) => {
696 /// // If the subscription has not been "accepted" then
697 /// // the return value will be "ignored" as it's not
698 /// // allowed to send out any further notifications on
699 /// // on the subscription.
700 /// tokio::spawn(pending.reject(ErrorObjectOwned::from(e)));
701 /// return;
702 /// }
703 /// };
704 ///
705 /// tokio::spawn(async move {
706 /// // Mark the subscription is accepted after the params has been parsed successful.
707 /// // This is actually responds the underlying RPC method call and may fail if the
708 /// // connection is closed.
709 /// let sink = pending.accept().await.unwrap();
710 /// let sum = val + (*ctx);
711 ///
712 /// let msg = serde_json::value::to_raw_value(&sum).unwrap();
713 ///
714 /// // This fails only if the connection is closed
715 /// sink.send(msg).await.unwrap();
716 /// });
717 /// });
718 /// ```
719 pub fn register_subscription_raw<R, F>(
720 &mut self,
721 subscribe_method_name: &'static str,
722 notif_method_name: &'static str,
723 unsubscribe_method_name: &'static str,
724 callback: F,
725 ) -> Result<&mut MethodCallback, RegisterMethodError>
726 where
727 Context: Send + Sync + 'static,
728 F: (Fn(Params, PendingSubscriptionSink, Arc<Context>) -> R) + Send + Sync + Clone + 'static,
729 R: IntoSubscriptionCloseResponse,
730 {
731 let subscribers =
732 self.verify_and_register_unsubscribe(subscribe_method_name, unsubscribe_method_name)?;
733 let ctx = self.ctx.clone();
734
735 // Subscribe
736 let callback = {
737 self.methods.verify_and_insert(
738 subscribe_method_name,
739 MethodCallback::Subscription(Arc::new(move |id, params, method_sink, conn| {
740 let uniq_sub = SubscriptionKey {
741 conn_id: conn.conn_id,
742 sub_id: conn.id_provider.next_id(),
743 };
744
745 // response to the subscription call.
746 let (tx, rx) = oneshot::channel();
747
748 let sink = PendingSubscriptionSink {
749 inner: method_sink.clone(),
750 method: notif_method_name,
751 subscribers: subscribers.clone(),
752 uniq_sub,
753 id: id.clone().into_owned(),
754 subscribe: tx,
755 permit: conn.subscription_permit,
756 };
757
758 callback(params, sink, ctx.clone());
759
760 let id = id.clone().into_owned();
761
762 Box::pin(async move {
763 match rx.await {
764 Ok(rp) => rp,
765 Err(_) => MethodResponse::error(id, ErrorCode::InternalError),
766 }
767 })
768 })),
769 )?
770 };
771
772 Ok(callback)
773 }
774
775 /// Helper to verify the subscription can be created
776 /// and register the unsubscribe handler.
777 fn verify_and_register_unsubscribe(
778 &mut self,
779 subscribe_method_name: &'static str,
780 unsubscribe_method_name: &'static str,
781 ) -> Result<Subscribers, RegisterMethodError> {
782 if subscribe_method_name == unsubscribe_method_name {
783 return Err(RegisterMethodError::SubscriptionNameConflict(subscribe_method_name.into()));
784 }
785
786 self.methods.verify_method_name(subscribe_method_name)?;
787 self.methods.verify_method_name(unsubscribe_method_name)?;
788
789 let subscribers = Subscribers::default();
790
791 // Unsubscribe
792 {
793 let subscribers = subscribers.clone();
794 self.methods.mut_callbacks().insert(
795 unsubscribe_method_name,
796 MethodCallback::Unsubscription(Arc::new(move |id, params, conn_id, max_response_size| {
797 let sub_id = match params.one::<RpcSubscriptionId>() {
798 Ok(sub_id) => sub_id,
799 Err(_) => {
800 // The params and request id are client input, so only their size is logged.
801 tracing::debug!(
802 target: LOG_TARGET,
803 "Unsubscribe call `{}` failed: couldn't parse subscription id from {} bytes of params",
804 unsubscribe_method_name,
805 params.len_bytes(),
806 );
807
808 return MethodResponse::response(id, ResponsePayload::success(false), max_response_size);
809 }
810 };
811
812 let key = SubscriptionKey { conn_id, sub_id: sub_id.into_owned() };
813 let mut subs = subscribers.lock();
814 let removed = subs.remove(&key);
815 let result = removed.is_some();
816 // Dropping the entry's receiver marks the sink inactive, which
817 // `SubscriptionSink::drop` checks under the lock, so drop it before unlocking.
818 drop(removed);
819 drop(subs);
820
821 if !result {
822 tracing::debug!(
823 target: LOG_TARGET,
824 "Unsubscribe call `{}` subscription key={:?} not an active subscription",
825 unsubscribe_method_name,
826 key,
827 );
828 }
829
830 MethodResponse::response(id, ResponsePayload::success(result), max_response_size)
831 })),
832 );
833 }
834
835 Ok(subscribers)
836 }
837
838 /// Register an alias for an existing_method. Alias uniqueness is enforced.
839 pub fn register_alias(
840 &mut self,
841 alias: &'static str,
842 existing_method: &'static str,
843 ) -> Result<(), RegisterMethodError> {
844 self.methods.verify_method_name(alias)?;
845
846 let callback = match self.methods.callbacks.get(existing_method) {
847 Some(callback) => callback.clone(),
848 None => return Err(RegisterMethodError::MethodNotFound(existing_method.into())),
849 };
850
851 self.methods.mut_callbacks().insert(alias, callback);
852
853 Ok(())
854 }
855}
856
857fn mock_subscription_permit() -> SubscriptionPermit {
858 BoundedSubscriptions::new(1).acquire().expect("1 permit should exist; qed")
859}