Middleware for metrics

examples/Cargo.toml

@@ -18,6 +18,10 @@ tokio = { version = "1", features = ["full"] }
 name = "http"
 path = "http.rs"
 
+[[example]]
+name = "middleware"
+path = "middleware.rs"
+
 [[example]]
 name = "ws"
 path = "ws.rs"

examples/middleware.rs

@@ -0,0 +1,94 @@
+// 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 jsonrpsee::{
+	types::traits::Client,
+	utils::server::middleware,
+	ws_client::WsClientBuilder,
+	ws_server::{RpcModule, WsServerBuilder},
+};
+use std::net::SocketAddr;
+use std::sync::atomic;
+
+#[derive(Default)]
+struct ManInTheMiddle {
+	when: atomic::AtomicU64,
+}
+
+impl Clone for ManInTheMiddle {
+	fn clone(&self) -> Self {
+		ManInTheMiddle { when: atomic::AtomicU64::new(self.when.load(atomic::Ordering::SeqCst)) }
+	}
+}
+
+impl middleware::Middleware for ManInTheMiddle {
+	type Instant = u64;
+	fn on_request(&self) -> Self::Instant {
+		self.when.fetch_add(1, atomic::Ordering::SeqCst)
+	}
+
+	fn on_call(&self, name: &str) {
+		println!("They called '{}'", name);
+	}
+
+	fn on_result(&self, name: &str, succeess: bool, started_at: Self::Instant) {
+		println!("call={}, worked? {}, when? {}", name, succeess, started_at);
+	}
+
+	fn on_response(&self, started_at: Self::Instant) {
+		println!("Response started_at={}", started_at);
+	}
+}
+
+#[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 addr = run_server().await?;
+	let url = format!("ws://{}", addr);
+
+	let client = WsClientBuilder::default().build(&url).await?;
+	let response: String = client.request("say_hello", None).await?;
+	tracing::info!("response: {:?}", response);
+	// TODO: This prints `They called 'blabla'` but nothing more. I expected the `on_response` callback to be called too?
+	let _response: Result<String, _> = client.request("blabla", None).await;
+	let _ = client.request::<String>("say_hello", None).await?;
+
+	Ok(())
+}
+
+async fn run_server() -> anyhow::Result<SocketAddr> {
+	let m = ManInTheMiddle::default();
+	let server = WsServerBuilder::with_middleware(m).build("127.0.0.1:0").await?;
+	let mut module = RpcModule::new(());
+	module.register_method("say_hello", |_, _| Ok("lo"))?;
+	let addr = server.local_addr()?;
+	server.start(module)?;
+	Ok(addr)
+}

http-server/src/server.rs

@@ -39,9 +39,9 @@ use jsonrpsee_types::{
 };
 use jsonrpsee_utils::http_helpers::read_body;
 use jsonrpsee_utils::server::{
-	helpers::{collect_batch_response, prepare_error, send_error},
+	helpers::{collect_batch_response, prepare_error, MethodSink},
 	resource_limiting::Resources,
-	rpc_module::Methods,
+	rpc_module::{MethodResult, Methods},
 };
 
 use serde_json::value::RawValue;
@@ -278,38 +278,42 @@ impl Server {
 
 						// NOTE(niklasad1): it's a channel because it's needed for batch requests.
 						let (tx, mut rx) = mpsc::unbounded::<String>();
+						let sink = MethodSink::new_with_limit(tx, max_request_body_size);
 
 						type Notif<'a> = Notification<'a, Option<&'a RawValue>>;
 
 						// Single request or notification
 						if is_single {
 							if let Ok(req) = serde_json::from_slice::<Request>(&body) {
 								// NOTE: we don't need to track connection id on HTTP, so using hardcoded 0 here.
-								if let Some(fut) =
-									methods.execute_with_resources(&tx, req, 0, &resources, max_request_body_size)
+								if let Some((_, MethodResult::Async(fut))) =
+									methods.execute_with_resources(&sink, req, 0, &resources)
 								{
 									fut.await;
 								}
 							} else if let Ok(_req) = serde_json::from_slice::<Notif>(&body) {
 								return Ok::<_, HyperError>(response::ok_response("".into()));
 							} else {
 								let (id, code) = prepare_error(&body);
-								send_error(id, &tx, code.into());
+								sink.send_error(id, code.into());
 							}
 
 						// Batch of requests or notifications
 						} else if let Ok(batch) = serde_json::from_slice::<Vec<Request>>(&body) {
 							if !batch.is_empty() {
 								join_all(batch.into_iter().filter_map(|req| {
-									methods.execute_with_resources(&tx, req, 0, &resources, max_request_body_size)
+									match methods.execute_with_resources(&sink, req, 0, &resources) {
+										Some((_, MethodResult::Async(fut))) => Some(fut),
+										_ => None,
+									}
 								}))
 								.await;
 							} else {
 								// "If the batch rpc call itself fails to be recognized as an valid JSON or as an
 								// Array with at least one value, the response from the Server MUST be a single
 								// Response object." – The Spec.
 								is_single = true;
-								send_error(Id::Null, &tx, ErrorCode::InvalidRequest.into());
+								sink.send_error(Id::Null, ErrorCode::InvalidRequest.into());
 							}
 						} else if let Ok(_batch) = serde_json::from_slice::<Vec<Notif>>(&body) {
 							return Ok::<_, HyperError>(response::ok_response("".into()));
@@ -319,7 +323,7 @@ impl Server {
 							// Response object." – The Spec.
 							is_single = true;
 							let (id, code) = prepare_error(&body);
-							send_error(id, &tx, code.into());
+							sink.send_error(id, code.into());
 						}
 
 						// Closes the receiving half of a channel without dropping it. This prevents any further

jsonrpsee/src/lib.rs

@@ -76,6 +76,10 @@ pub use jsonrpsee_types as types;
 #[cfg(any(feature = "http-server", feature = "ws-server"))]
 pub use jsonrpsee_utils::server::rpc_module::{RpcModule, SubscriptionSink};
 
+/// TODO: (dp) any reason not to export this? narrow the scope to `jsonrpsee_utils::server`?
+#[cfg(any(feature = "http-server", feature = "ws-server"))]
+pub use jsonrpsee_utils as utils;
+
 #[cfg(feature = "http-server")]
 pub use http_server::tracing;
 

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`, `resources`
- --> tests/ui/incorrect/method/method_unexpected_field.rs:6:25
+ --> $DIR/method_unexpected_field.rs:6:25
   |
 6 |     #[method(name = "foo", magic = false)]
   |                            ^^^^^

proc-macros/tests/ui/incorrect/sub/sub_dup_name_override.stderr

@@ -1,5 +1,5 @@
 error: "override" is already defined
- --> tests/ui/incorrect/sub/sub_dup_name_override.rs:9:5
+ --> $DIR/sub_dup_name_override.rs:9:5
   |
 9 |     fn two(&self) -> RpcResult<()>;
   |        ^^^

proc-macros/tests/ui/incorrect/sub/sub_name_override.stderr

@@ -1,5 +1,5 @@
 error: "one" is already defined
- --> tests/ui/incorrect/sub/sub_name_override.rs:7:5
+ --> $DIR/sub_name_override.rs:7:5
   |
 7 |     fn one(&self) -> RpcResult<()>;
   |        ^^^

proc-macros/tests/ui/incorrect/sub/sub_unsupported_field.stderr

@@ -1,5 +1,5 @@
 error: Unknown argument `magic`, expected one of: `aliases`, `item`, `name`, `param_kind`, `unsubscribe_aliases`
- --> tests/ui/incorrect/sub/sub_unsupported_field.rs:6:42
+ --> $DIR/sub_unsupported_field.rs:6:42
   |
 6 |     #[subscription(name = "sub", item = u8, magic = true)]
   |                                             ^^^^^

utils/src/server/helpers.rs

@@ -24,7 +24,6 @@
 // IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
-use crate::server::rpc_module::MethodSink;
 use futures_channel::mpsc;
 use futures_util::stream::StreamExt;
 use jsonrpsee_types::error::{CallError, Error};
@@ -82,67 +81,108 @@ impl<'a> io::Write for &'a mut BoundedWriter {
 	}
 }
 
-/// Helper for sending JSON-RPC responses to the client
-pub fn send_response(id: Id, tx: &MethodSink, result: impl Serialize, max_response_size: u32) {
-	let mut writer = BoundedWriter::new(max_response_size as usize);
+/// Sink that is used to send back the result to the server for a specific method.
+#[derive(Clone, Debug)]
+pub struct MethodSink {
+	/// Channel sender
+	tx: mpsc::UnboundedSender<String>,
+	/// Max response size in bytes for a executed call.
+	max_response_size: u32,
+}
 
-	let json = match serde_json::to_writer(&mut writer, &Response { jsonrpc: TwoPointZero, id: id.clone(), result }) {
-		Ok(_) => {
-			// Safety - serde_json does not emit invalid UTF-8.
-			unsafe { String::from_utf8_unchecked(writer.into_bytes()) }
-		}
-		Err(err) => {
-			tracing::error!("Error serializing response: {:?}", err);
-
-			if err.is_io() {
-				let data = to_json_raw_value(&format!("Exceeded max limit {}", max_response_size)).ok();
-				let err = ErrorObject {
-					code: ErrorCode::ServerError(OVERSIZED_RESPONSE_CODE),
-					message: OVERSIZED_RESPONSE_MSG,
-					data: data.as_deref(),
-				};
-				return send_error(id, tx, err);
-			} else {
-				return send_error(id, tx, ErrorCode::InternalError.into());
+impl MethodSink {
+	/// Create a new `MethodSink` with unlimited response size
+	pub fn new(tx: mpsc::UnboundedSender<String>) -> Self {
+		MethodSink { tx, max_response_size: u32::MAX }
+	}
+
+	/// Create a new `MethodSink` with a limited response size
+	pub fn new_with_limit(tx: mpsc::UnboundedSender<String>, max_response_size: u32) -> Self {
+		MethodSink { tx, max_response_size }
+	}
+
+	/// Send a JSON-RPC response to the client. If the serialization of `result` exceeds `max_response_size`,
+	/// an error will be sent instead.
+	pub fn send_response(&self, id: Id, result: impl Serialize) -> bool {
+		let mut writer = BoundedWriter::new(self.max_response_size as usize);
+
+		let json = match serde_json::to_writer(&mut writer, &Response { jsonrpc: TwoPointZero, id: id.clone(), result })
+		{
+			Ok(_) => {
+				// Safety - serde_json does not emit invalid UTF-8.
+				unsafe { String::from_utf8_unchecked(writer.into_bytes()) }
 			}
-		}
-	};
+			Err(err) => {
+				tracing::error!("Error serializing response: {:?}", err);
+
+				if err.is_io() {
+					let data = to_json_raw_value(&format!("Exceeded max limit {}", self.max_response_size)).ok();
+					let err = ErrorObject {
+						code: ErrorCode::ServerError(OVERSIZED_RESPONSE_CODE),
+						message: OVERSIZED_RESPONSE_MSG,
+						data: data.as_deref(),
+					};
+					return self.send_error(id, err);
+				} else {
+					return self.send_error(id, ErrorCode::InternalError.into());
+				}
+			}
+		};
 
-	if let Err(err) = tx.unbounded_send(json) {
-		tracing::error!("Error sending response to the client: {:?}", err)
+		if let Err(err) = self.tx.unbounded_send(json) {
+			tracing::error!("Error sending response to the client: {:?}", err);
+			false
+		} else {
+			true
+		}
 	}
-}
 
-/// Helper for sending JSON-RPC errors to the client
-pub fn send_error(id: Id, tx: &MethodSink, error: ErrorObject) {
-	let json = match serde_json::to_string(&RpcError { jsonrpc: TwoPointZero, error, id }) {
-		Ok(json) => json,
-		Err(err) => {
-			tracing::error!("Error serializing error message: {:?}", err);
+	/// Send a JSON-RPC error to the client
+	pub fn send_error(&self, id: Id, error: ErrorObject) -> bool {
+		let json = match serde_json::to_string(&RpcError { jsonrpc: TwoPointZero, error, id }) {
+			Ok(json) => json,
+			Err(err) => {
+				tracing::error!("Error serializing error message: {:?}", err);
+
+				return false;
+			}
+		};
 
-			return;
+		if let Err(err) = self.tx.unbounded_send(json) {
+			tracing::error!("Could not send error response to the client: {:?}", err)
 		}
-	};
 
-	if let Err(err) = tx.unbounded_send(json) {
-		tracing::error!("Could not send error response to the client: {:?}", err)
+		false
 	}
-}
 
-/// Helper for sending the general purpose `Error` as a JSON-RPC errors to the client
-pub fn send_call_error(id: Id, tx: &MethodSink, err: Error) {
-	let (code, message, data) = match err {
-		Error::Call(CallError::InvalidParams(e)) => (ErrorCode::InvalidParams, e.to_string(), None),
-		Error::Call(CallError::Failed(e)) => (ErrorCode::ServerError(CALL_EXECUTION_FAILED_CODE), e.to_string(), None),
-		Error::Call(CallError::Custom { code, message, data }) => (code.into(), message, data),
-		// This should normally not happen because the most common use case is to
-		// return `Error::Call` in `register_async_method`.
-		e => (ErrorCode::ServerError(UNKNOWN_ERROR_CODE), e.to_string(), None),
-	};
+	/// Helper for sending the general purpose `Error` as a JSON-RPC errors to the client
+	pub fn send_call_error(&self, id: Id, err: Error) -> bool {
+		let (code, message, data) = match err {
+			Error::Call(CallError::InvalidParams(e)) => (ErrorCode::InvalidParams, e.to_string(), None),
+			Error::Call(CallError::Failed(e)) => {
+				(ErrorCode::ServerError(CALL_EXECUTION_FAILED_CODE), e.to_string(), None)
+			}
+			Error::Call(CallError::Custom { code, message, data }) => (code.into(), message, data),
+			// This should normally not happen because the most common use case is to
+			// return `Error::Call` in `register_async_method`.
+			e => (ErrorCode::ServerError(UNKNOWN_ERROR_CODE), e.to_string(), None),
+		};
+
+		let err = ErrorObject { code, message: &message, data: data.as_deref() };
 
-	let err = ErrorObject { code, message: &message, data: data.as_deref() };
+		self.send_error(id, err)
+	}
 
-	send_error(id, tx, err)
+	/// Send a raw JSON-RPC message to the client, `MethodSink` does not check verify the validity
+	/// of the JSON being sent.
+	pub fn send_raw(&self, raw_json: String) -> Result<(), mpsc::TrySendError<String>> {
+		self.tx.unbounded_send(raw_json)
+	}
+
+	/// Close the channel for any further messages.
+	pub fn close(&self) {
+		self.tx.close_channel();
+	}
 }
 
 /// Figure out if this is a sufficiently complete request that we can extract an [`Id`] out of, or just plain