server: Register raw method with connection ID

core/src/server/rpc_module.rs

@@ -60,6 +60,10 @@ pub type SyncMethod = Arc<dyn Send + Sync + Fn(Id, Params, MaxResponseSize) -> M
 /// Similar to [`SyncMethod`], but represents an asynchronous handler.
 pub type AsyncMethod<'a> =
 	Arc<dyn Send + Sync + Fn(Id<'a>, Params<'a>, ConnectionId, MaxResponseSize) -> BoxFuture<'a, MethodResponse>>;
+/// Similar to [`SyncMethod`], but represents an asynchronous handler with connection details.
+pub type AsyncMethodWithDetails<'a> =
+	Arc<dyn Send + Sync + Fn(Id<'a>, Params<'a>, ConnectionDetails, MaxResponseSize) -> BoxFuture<'a, MethodResponse>>;
+/// Similar to [`SyncMethod`], but represents a raw handler that has access to the connection Id.
 /// Method callback for subscriptions.
 pub type SubscriptionMethod<'a> =
 	Arc<dyn Send + Sync + Fn(Id, Params, MethodSink, SubscriptionState) -> BoxFuture<'a, MethodResponse>>;
@@ -79,6 +83,45 @@ pub type MaxResponseSize = usize;
 ///   - a [`mpsc::UnboundedReceiver<String>`] to receive future subscription results
 pub type RawRpcResponse = (String, mpsc::Receiver<String>);
 
+#[derive(Debug, Clone)]
+#[allow(missing_copy_implementations)]
+/// The connection details exposed to the server methods.
+pub struct ConnectionDetails {
+	id: ConnectionId,
+}
+
+impl ConnectionDetails {
+	/// Construct a new [`ConnectionDetailsBuilder`].
+	pub fn builder() -> ConnectionDetailsBuilder {
+		ConnectionDetailsBuilder { id: 0 }
+	}
+
+	/// Get the connection ID.
+	pub fn id(&self) -> ConnectionId {
+		self.id
+	}
+}
+
+#[derive(Debug, Clone)]
+#[allow(missing_copy_implementations)]
+/// The connection details exposed to the server methods.
+pub struct ConnectionDetailsBuilder {
+	id: ConnectionId,
+}
+
+impl ConnectionDetailsBuilder {
+	/// Set the connection ID.
+	pub fn id(mut self, id: ConnectionId) -> Self {
+		self.id = id;
+		self
+	}
+
+	/// Build the [`ConnectionDetails`].
+	pub fn build(self) -> ConnectionDetails {
+		ConnectionDetails { id: self.id }
+	}
+}
+
 /// The error that can occur when [`Methods::call`] or [`Methods::subscribe`] is invoked.
 #[derive(thiserror::Error, Debug)]
 pub enum MethodsError {
@@ -131,6 +174,8 @@ pub enum MethodCallback {
 	Sync(SyncMethod),
 	/// Asynchronous method handler.
 	Async(AsyncMethod<'static>),
+	/// Asynchronous method handler with details.
+	AsyncWithDetails(AsyncMethodWithDetails<'static>),
 	/// Subscription method handler.
 	Subscription(SubscriptionMethod<'static>),
 	/// Unsubscription method handler.
@@ -184,6 +229,7 @@ impl Debug for MethodCallback {
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 		match self {
 			Self::Async(_) => write!(f, "Async"),
+			Self::AsyncWithDetails(_) => write!(f, "AsyncWithDetails"),
 			Self::Sync(_) => write!(f, "Sync"),
 			Self::Subscription(_) => write!(f, "Subscription"),
 			Self::Unsubscription(_) => write!(f, "Unsubscription"),
@@ -355,6 +401,9 @@ impl Methods {
 			None => MethodResponse::error(req.id, ErrorObject::from(ErrorCode::MethodNotFound)),
 			Some(MethodCallback::Sync(cb)) => (cb)(id, params, usize::MAX),
 			Some(MethodCallback::Async(cb)) => (cb)(id.into_owned(), params.into_owned(), 0, usize::MAX).await,
+			Some(MethodCallback::AsyncWithDetails(cb)) => {
+				(cb)(id.into_owned(), params.into_owned(), ConnectionDetails::builder().build(), usize::MAX).await
+			}
 			Some(MethodCallback::Subscription(cb)) => {
 				let conn_state =
 					SubscriptionState { conn_id: 0, id_provider: &RandomIntegerIdProvider, subscription_permit };
@@ -598,6 +647,42 @@ impl<Context: Send + Sync + 'static> RpcModule<Context> {
 		Ok(callback)
 	}
 
+	/// Register a new raw RPC method, which computes the response with the given callback.
+	///
+	/// ## Examples
+	///
+	/// ```
+	/// use jsonrpsee_core::server::RpcModule;
+	///
+	/// let mut module = RpcModule::new(());
+	/// module.register_async_with_details("say_hello", |_params, _connection_details, _ctx| async { "lo" }).unwrap();
+	/// ```
+	pub fn register_async_with_details<R, Fun, Fut>(
+		&mut self,
+		method_name: &'static str,
+		callback: Fun,
+	) -> Result<&mut MethodCallback, RegisterMethodError>
+	where
+		R: IntoResponse + 'static,
+		Fut: Future<Output = R> + Send,
+		Fun: (Fn(Params<'static>, ConnectionDetails, Arc<Context>) -> Fut) + Clone + Send + Sync + 'static,
+	{
+		let ctx = self.ctx.clone();
+		self.methods.verify_and_insert(
+			method_name,
+			MethodCallback::AsyncWithDetails(Arc::new(move |id, params, connection_details, max_response_size| {
+				let ctx = ctx.clone();
+				let callback = callback.clone();
+
+				let future = async move {
+					let rp = callback(params, connection_details, ctx).await.into_response();
+					MethodResponse::response(id, rp, max_response_size)
+				};
+				future.boxed()
+			})),
+		)
+	}
+
 	/// Register a new publish/subscribe interface using JSON-RPC notifications.
 	///
 	/// It implements the [ethereum pubsub specification](https://geth.ethereum.org/docs/rpc/pubsub)

examples/examples/server_with_raw_method.rs

@@ -0,0 +1,126 @@
+// Copyright 2019-2021 Parity Technologies (UK) Ltd.
+//
+// Permission is hereby granted, free of charge, to any
+// person obtaining a copy of this software and associated
+// documentation files (the "Software"), to deal in the
+// Software without restriction, including without
+// limitation the rights to use, copy, modify, merge,
+// publish, distribute, sublicense, and/or sell copies of
+// the Software, and to permit persons to whom the Software
+// is furnished to do so, subject to the following
+// conditions:
+//
+// The above copyright notice and this permission notice
+// shall be included in all copies or substantial portions
+// of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
+// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
+// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+use std::net::SocketAddr;
+
+use jsonrpsee::core::{async_trait, client::Subscription};
+use jsonrpsee::proc_macros::rpc;
+use jsonrpsee::server::{PendingSubscriptionSink, Server, SubscriptionMessage};
+use jsonrpsee::types::ErrorObjectOwned;
+use jsonrpsee::ws_client::WsClientBuilder;
+use jsonrpsee::ConnectionDetails;
+
+#[rpc(server, client)]
+pub trait Rpc {
+	/// Raw method with connection ID.
+	#[method(name = "connectionIdMethod", raw_method)]
+	async fn raw_method(&self, first_param: usize, second_param: u16) -> Result<usize, ErrorObjectOwned>;
+
+	/// Normal method call example.
+	#[method(name = "normalMethod")]
+	fn normal_method(&self, first_param: usize, second_param: u16) -> Result<usize, ErrorObjectOwned>;
+
+	/// Subscriptions expose the connection ID on the subscription sink.
+	#[subscription(name = "subscribeSync" => "sync", item = usize)]
+	fn sub(&self, first_param: usize);
+}
+
+pub struct RpcServerImpl;
+
+#[async_trait]
+impl RpcServer for RpcServerImpl {
+	async fn raw_method(
+		&self,
+		connection_details: ConnectionDetails,
+		_first_param: usize,
+		_second_param: u16,
+	) -> Result<usize, ErrorObjectOwned> {
+		// Return the connection ID from which this method was called.
+		Ok(connection_details.id())
+	}
+
+	fn normal_method(&self, _first_param: usize, _second_param: u16) -> Result<usize, ErrorObjectOwned> {
+		// The normal method does not have access to the connection ID.
+		Ok(usize::MAX)
+	}
+
+	fn sub(&self, pending: PendingSubscriptionSink, _first_param: usize) {
+		tokio::spawn(async move {
+			// The connection ID can be obtained before or after accepting the subscription
+			let pending_connection_id = pending.connection_id();
+			let sink = pending.accept().await.unwrap();
+			let sink_connection_id = sink.connection_id();
+
+			assert_eq!(pending_connection_id, sink_connection_id);
+
+			let msg = SubscriptionMessage::from_json(&sink_connection_id).unwrap();
+			sink.send(msg).await.unwrap();
+		});
+	}
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+	tracing_subscriber::FmtSubscriber::builder()
+		.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
+		.try_init()
+		.expect("setting default subscriber failed");
+
+	let server_addr = run_server().await?;
+	let url = format!("ws://{}", server_addr);
+
+	let client = WsClientBuilder::default().build(&url).await?;
+	let connection_id_first = client.raw_method(1, 2).await.unwrap();
+
+	// Second call from the same connection ID.
+	assert_eq!(client.raw_method(1, 2).await.unwrap(), connection_id_first);
+
+	// Second client will increment the connection ID.
+	let client_second = WsClientBuilder::default().build(&url).await?;
+	let connection_id_second = client_second.raw_method(1, 2).await.unwrap();
+	assert_ne!(connection_id_first, connection_id_second);
+
+	let mut sub: Subscription<usize> = RpcClient::sub(&client, 0).await.unwrap();
+	assert_eq!(connection_id_first, sub.next().await.transpose().unwrap().unwrap());
+
+	let mut sub: Subscription<usize> = RpcClient::sub(&client_second, 0).await.unwrap();
+	assert_eq!(connection_id_second, sub.next().await.transpose().unwrap().unwrap());
+
+	Ok(())
+}
+
+async fn run_server() -> anyhow::Result<SocketAddr> {
+	let server = Server::builder().build("127.0.0.1:0").await?;
+
+	let addr = server.local_addr()?;
+	let handle = server.start(RpcServerImpl.into_rpc());
+
+	// In this example we don't care about doing shutdown so let's it run forever.
+	// You may use the `ServerHandle` to shut it down or manage it yourself.
+	tokio::spawn(handle.stopped());
+
+	Ok(addr)
+}

proc-macros/src/render_server.rs

@@ -61,7 +61,15 @@
 	fn render_methods(&self) -> Result<TokenStream2, syn::Error> {
 		let methods = self.methods.iter().map(|method| {
 			let docs = &method.docs;
-			let method_sig = &method.signature;
+			let mut method_sig = method.signature.clone();
+
+			if method.raw_method {
+				let context_ty = self.jrps_server_item(quote! { ConnectionDetails });
+				// Add `connection ID` as the second input parameter to the signature.
+				let context: syn::FnArg = syn::parse_quote!(connection_context: #context_ty);
+				method_sig.sig.inputs.insert(1, context);
+			}
+
 			quote! {
 				#docs
 				#method_sig
@@ -132,24 +140,33 @@
 
 				check_name(&rpc_method_name, rust_method_name.span());
 
-				if method.signature.sig.asyncness.is_some() {
+				if method.raw_method {
 					handle_register_result(quote! {
-						rpc.register_async_method(#rpc_method_name, |params, context| async move {
+						rpc.register_async_with_details(#rpc_method_name, |params, connection_details, context| async move {
 							#parsing
-							#into_response::into_response(context.as_ref().#rust_method_name(#params_seq).await)
+							#into_response::into_response(context.as_ref().#rust_method_name(connection_details, #params_seq).await)
 						})
 					})
 				} else {
-					let register_kind =
-						if method.blocking { quote!(register_blocking_method) } else { quote!(register_method) };
+					if method.signature.sig.asyncness.is_some() {
+						handle_register_result(quote! {
+							rpc.register_async_method(#rpc_method_name, |params, context| async move {
+								#parsing
+								#into_response::into_response(context.as_ref().#rust_method_name(#params_seq).await)
+							})
+						})
+					} else {
+						let register_kind =
+							if method.blocking { quote!(register_blocking_method) } else { quote!(register_method) };
 
-					handle_register_result(quote! {
-						rpc.#register_kind(#rpc_method_name, |params, context| {
-							#parsing
-							#into_response::into_response(context.#rust_method_name(#params_seq))
+						handle_register_result(quote! {
+							rpc.#register_kind(#rpc_method_name, |params, context| {
+								#parsing
+								#into_response::into_response(context.#rust_method_name(#params_seq))
+							})
 						})
-					})
+					}
 				}
 			})
 			.collect::<Vec<_>>();
 

proc-macros/src/rpc_macro.rs

@@ -48,17 +48,19 @@ pub struct RpcMethod {
 	pub returns: Option<syn::Type>,
 	pub signature: syn::TraitItemFn,
 	pub aliases: Vec<String>,
+	pub raw_method: bool,
 }
 
 impl RpcMethod {
 	pub fn from_item(attr: Attribute, mut method: syn::TraitItemFn) -> syn::Result<Self> {
-		let [aliases, blocking, name, param_kind] =
-			AttributeMeta::parse(attr)?.retain(["aliases", "blocking", "name", "param_kind"])?;
+		let [aliases, blocking, name, param_kind, raw_method] =
+			AttributeMeta::parse(attr)?.retain(["aliases", "blocking", "name", "param_kind", "raw_method"])?;
 
 		let aliases = parse_aliases(aliases)?;
 		let blocking = optional(blocking, Argument::flag)?.is_some();
 		let name = name?.string()?;
 		let param_kind = parse_param_kind(param_kind)?;
+		let raw_method = optional(raw_method, Argument::flag)?.is_some();
 
 		let sig = method.sig.clone();
 		let docs = extract_doc_comments(&method.attrs);
@@ -98,7 +100,18 @@ impl RpcMethod {
 		// We've analyzed attributes and don't need them anymore.
 		method.attrs.clear();
 
-		Ok(Self { aliases, blocking, name, params, param_kind, returns, signature: method, docs, deprecated })
+		Ok(Self {
+			aliases,
+			blocking,
+			name,
+			params,
+			param_kind,
+			returns,
+			signature: method,
+			docs,
+			deprecated,
+			raw_method,
+		})
 	}
 }
 
@@ -212,7 +225,6 @@ impl RpcDescription {
 		let namespace = optional(namespace, Argument::string)?;
 		let client_bounds = optional(client_bounds, Argument::group)?;
 		let server_bounds = optional(server_bounds, Argument::group)?;
-
 		if !needs_server && !needs_client {
 			return Err(syn::Error::new_spanned(&item.ident, "Either 'server' or 'client' attribute must be applied"));
 		}
@@ -260,6 +272,28 @@ impl RpcDescription {
 					is_method = true;
 
 					let method_data = RpcMethod::from_item(attr.clone(), method.clone())?;
+
+					if method_data.blocking && method_data.raw_method {
+						return Err(syn::Error::new_spanned(
+							method,
+							"Methods cannot be blocking when used with `raw_method`; remove `blocking` attribute or `raw_method` attribute",
+						));
+					}
+
+					if !needs_server && method_data.raw_method {
+						return Err(syn::Error::new_spanned(
+							&item.ident,
+							"Attribute 'raw_method' must be specified with 'server'",
+						));
+					}
+
+					if method.sig.asyncness.is_none() && method_data.raw_method {
+						return Err(syn::Error::new_spanned(
+							method,
+							"Methods must be asynchronous when used with `raw_method`; use `async fn` instead of `fn`",
+						));
+					}
+
 					methods.push(method_data);
 				}
 				if let Some(attr) = find_attr(&method.attrs, "subscription") {