server: Register raw method with connection ID

core/src/server/rpc_module.rs

@@ -60,6 +60,8 @@ 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 a raw handler that has access to the connection Id.
+pub type RawMethod = Arc<dyn Send + Sync + Fn(Id, Params, ConnectionId, MaxResponseSize) -> MethodResponse>;
 /// Method callback for subscriptions.
 pub type SubscriptionMethod<'a> =
 	Arc<dyn Send + Sync + Fn(Id, Params, MethodSink, SubscriptionState) -> BoxFuture<'a, MethodResponse>>;
@@ -131,6 +133,8 @@ pub enum MethodCallback {
 	Sync(SyncMethod),
 	/// Asynchronous method handler.
 	Async(AsyncMethod<'static>),
+	/// Raw method handler.
+	Raw(RawMethod),
 	/// Subscription method handler.
 	Subscription(SubscriptionMethod<'static>),
 	/// Unsubscription method handler.
@@ -185,6 +189,7 @@ impl Debug for MethodCallback {
 		match self {
 			Self::Async(_) => write!(f, "Async"),
 			Self::Sync(_) => write!(f, "Sync"),
+			Self::Raw(_) => write!(f, "Raw"),
 			Self::Subscription(_) => write!(f, "Subscription"),
 			Self::Unsubscription(_) => write!(f, "Unsubscription"),
 		}
@@ -355,6 +360,7 @@ 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::Raw(cb)) => (cb)(id.into_owned(), params.into_owned(), 0, usize::MAX),
 			Some(MethodCallback::Subscription(cb)) => {
 				let conn_state =
 					SubscriptionState { conn_id: 0, id_provider: &RandomIntegerIdProvider, subscription_permit };
@@ -598,6 +604,37 @@ 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_raw_method("say_hello", |_params, _connection_id, _ctx| "lo").unwrap();
+	/// ```
+	pub fn register_raw_method<R, F>(
+		&mut self,
+		method_name: &'static str,
+		callback: F,
+	) -> Result<&mut MethodCallback, RegisterMethodError>
+	where
+		Context: Send + Sync + 'static,
+		R: IntoResponse + 'static,
+		F: Fn(Params, ConnectionId, Arc<Context>) -> R + Clone + Send + Sync + 'static,
+	{
+		let ctx = self.ctx.clone();
+		self.methods.verify_and_insert(
+			method_name,
+			MethodCallback::Raw(Arc::new(move |id, params, connection_id, max_response_size| {
+				let ctx = ctx.clone();
+				let rp = callback(params, connection_id, ctx).into_response();
+				MethodResponse::response(id, rp, max_response_size)
+			})),
+		)
+	}
+
 	/// 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::ConnectionId;
+
+#[rpc(server, client)]
+pub trait Rpc {
+	/// Raw method with connection ID.
+	#[method(name = "connectionIdMethod", raw_method)]
+	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 {
+	fn raw_method(
+		&self,
+		connection_id: ConnectionId,
+		_first_param: usize,
+		_second_param: u16,
+	) -> Result<usize, ErrorObjectOwned> {
+		// Return the connection ID from which this method was called.
+		Ok(connection_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

@@ -62,7 +62,15 @@ impl RpcDescription {
 	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! { ConnectionId });
+				// 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
@@ -133,23 +141,32 @@ impl RpcDescription {
 
 				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_raw_method(#rpc_method_name, |params, connection_id, context| {
 							#parsing
-							#into_response::into_response(context.as_ref().#rust_method_name(#params_seq).await)
+							#into_response::into_response(context.#rust_method_name(connection_id, #params_seq))
 						})
 					})
 				} 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::TraitItemMethod,
 	pub aliases: Vec<String>,
+	pub raw_method: bool,
 }
 
 impl RpcMethod {
 	pub fn from_item(attr: Attribute, mut method: syn::TraitItemMethod) -> 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_some() && method_data.raw_method {
+						return Err(syn::Error::new_spanned(
+							method,
+							"Methods must be synchronous when used with `raw_method`; use `fn` instead of `async fn`",
+						));
+					}
+
 					methods.push(method_data);
 				}
 				if let Some(attr) = find_attr(&method.attrs, "subscription") {

proc-macros/tests/ui/correct/server_with_raw_methods.rs

@@ -0,0 +1,15 @@
+//! Example of using proc macro to generate working client.
+
+use jsonrpsee::proc_macros::rpc;
+use jsonrpsee::types::ErrorObjectOwned;
+
+#[rpc(server)]
+pub trait Rpc {
+	#[method(name = "foo")]
+	async fn async_method(&self, param_a: u8, param_b: String) -> Result<u16, ErrorObjectOwned>;
+
+	#[method(name = "bar", raw_method)]
+	fn sync_method(&self) -> Result<u16, ErrorObjectOwned>;
+}
+
+fn main() {}

proc-macros/tests/ui/incorrect/method/method_async_raw_incompatible.rs

@@ -0,0 +1,9 @@
+use jsonrpsee::proc_macros::rpc;
+
+#[rpc(server)]
+pub trait AsyncMethodCannotBeRaw {
+	#[method(name = "a", raw_method)]
+	async fn a(&self, param: Vec<u8>);
+}
+
+fn main() {}

proc-macros/tests/ui/incorrect/method/method_async_raw_incompatible.stderr

@@ -0,0 +1,6 @@
+error: Methods must be synchronous when used with `raw_method`; use `fn` instead of `async fn`
+ --> tests/ui/incorrect/method/method_async_raw_incompatible.rs:5:2
+  |
+5 | /     #[method(name = "a", raw_method)]
+6 | |     async fn a(&self, param: Vec<u8>);
+  | |______________________________________^

proc-macros/tests/ui/incorrect/method/method_blocking_raw_incompatible.rs

@@ -0,0 +1,9 @@
+use jsonrpsee::proc_macros::rpc;
+
+#[rpc(server)]
+pub trait BlockingMethodCannotBeRaw {
+	#[method(name = "a", blocking, raw_method)]
+	fn a(&self, param: Vec<u8>);
+}
+
+fn main() {}

proc-macros/tests/ui/incorrect/method/method_blocking_raw_incompatible.stderr

@@ -0,0 +1,6 @@
+error: Methods cannot be blocking when used with `raw_method`; remove `blocking` attribute or `raw_method` attribute
+ --> tests/ui/incorrect/method/method_blocking_raw_incompatible.rs:5:2
+  |
+5 | /     #[method(name = "a", blocking, raw_method)]
+6 | |     fn a(&self, param: Vec<u8>);
+  | |________________________________^

proc-macros/tests/ui/incorrect/method/method_unexpected_field.stderr

@@ -1,5 +1,5 @@
-error: Unknown argument `magic`, expected one of: `aliases`, `blocking`, `name`, `param_kind`
- --> $DIR/method_unexpected_field.rs:6:25
+error: Unknown argument `magic`, expected one of: `aliases`, `blocking`, `name`, `param_kind`, `raw_method`
+ --> tests/ui/incorrect/method/method_unexpected_field.rs:6:25
   |
 6 |     #[method(name = "foo", magic = false)]
   |                            ^^^^^