Просмотр исходного кода

rpc: Cleanup client and server and refactor into smaller functions.

* Implement a read timeout to avoid a DoS path
* Use dynamic buffer allocation with a max size limit
* Make the JSON-RPC a line-based protocol
parazyd 3 лет назад
Родитель
Сommit
827ec53f63
7 измененных файлов с 428 добавлено и 172 удалено
  1. 11 0
      Cargo.toml
  2. 68 0
      example/rpc/client.rs
  3. 101 0
      example/rpc/server.rs
  4. 31 0
      src/error.rs
  5. 129 132
      src/rpc/client.rs
  6. 2 1
      src/rpc/clock_sync.rs
  7. 86 39
      src/rpc/server.rs

+ 11 - 0
Cargo.toml

@@ -208,6 +208,7 @@ net = [
     "p2p-transport-tcp",
     "p2p-transport-tcp",
     "p2p-transport-tor",
     "p2p-transport-tor",
     "p2p-transport-nym",
     "p2p-transport-nym",
+    "p2p-transport-unix",
 ]
 ]
 
 
 rpc = [
 rpc = [
@@ -295,6 +296,16 @@ name = "zk-inclusion-proof"
 path = "example/zk-inclusion-proof.rs"
 path = "example/zk-inclusion-proof.rs"
 required-features = ["zk"]
 required-features = ["zk"]
 
 
+[[example]]
+name = "rpc-server"
+path = "example/rpc/server.rs"
+required-features = ["async-runtime", "rpc"]
+
+[[example]]
+name = "rpc-client"
+path = "example/rpc/client.rs"
+required-features = ["async-runtime", "rpc"]
+
 [patch.crates-io]
 [patch.crates-io]
 halo2_proofs = {git="https://github.com/parazyd/halo2", branch="v4"}
 halo2_proofs = {git="https://github.com/parazyd/halo2", branch="v4"}
 halo2_gadgets = {git="https://github.com/parazyd/halo2", branch="v4"}
 halo2_gadgets = {git="https://github.com/parazyd/halo2", branch="v4"}

+ 68 - 0
example/rpc/client.rs

@@ -0,0 +1,68 @@
+/* This file is part of DarkFi (https://dark.fi)
+ *
+ * Copyright (C) 2020-2023 Dyne.org foundation
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+use async_std::sync::Arc;
+use serde_json::json;
+use smol::Executor;
+use url::Url;
+
+use darkfi::{
+    rpc::{client::RpcClient, jsonrpc::JsonRequest},
+    Result,
+};
+
+async fn realmain(ex: Arc<Executor<'_>>) -> Result<()> {
+    let endpoint = Url::parse("tcp://127.0.0.1:55422").unwrap();
+
+    let client = RpcClient::new(endpoint, Some(ex)).await?;
+
+    let req = JsonRequest::new("ping", json!([]));
+    let rep = client.request(req).await?;
+
+    println!("{:#?}", rep);
+
+    let req = JsonRequest::new("kill", json!([]));
+    let rep = client.request(req).await?;
+
+    println!("{:#?}", rep);
+
+    Ok(())
+}
+
+fn main() -> Result<()> {
+    simplelog::TermLogger::init(
+        simplelog::LevelFilter::Debug,
+        simplelog::ConfigBuilder::new().build(),
+        simplelog::TerminalMode::Mixed,
+        simplelog::ColorChoice::Auto,
+    )?;
+
+    let n_threads = std::thread::available_parallelism().unwrap().get();
+    let ex = Arc::new(Executor::new());
+    let (signal, shutdown) = smol::channel::unbounded::<()>();
+    let (_, result) = easy_parallel::Parallel::new()
+        .each(0..n_threads, |_| smol::future::block_on(ex.run(shutdown.recv())))
+        .finish(|| {
+            smol::future::block_on(async {
+                realmain(ex.clone()).await?;
+                drop(signal);
+                Ok::<(), darkfi::Error>(())
+            })
+        });
+
+    result
+}

+ 101 - 0
example/rpc/server.rs

@@ -0,0 +1,101 @@
+/* This file is part of DarkFi (https://dark.fi)
+ *
+ * Copyright (C) 2020-2023 Dyne.org foundation
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+use async_std::sync::Arc;
+use async_trait::async_trait;
+use serde_json::{json, Value};
+use smol::{
+    channel::{Receiver, Sender},
+    Executor,
+};
+use url::Url;
+
+use darkfi::{
+    rpc::{
+        jsonrpc::{ErrorCode, JsonError, JsonRequest, JsonResponse, JsonResult},
+        server::{listen_and_serve, RequestHandler},
+    },
+    Result,
+};
+
+struct RpcSrv {
+    stop_sub: (Sender<()>, Receiver<()>),
+}
+
+impl RpcSrv {
+    async fn pong(&self, id: Value, _params: &[Value]) -> JsonResult {
+        JsonResponse::new(json!("pong"), id).into()
+    }
+
+    async fn kill(&self, id: Value, _params: &[Value]) -> JsonResult {
+        self.stop_sub.0.send(()).await.unwrap();
+        JsonResponse::new(json!("Bye"), id).into()
+    }
+}
+
+#[async_trait]
+impl RequestHandler for RpcSrv {
+    async fn handle_request(&self, req: JsonRequest) -> JsonResult {
+        if !req.params.is_array() {
+            return JsonError::new(ErrorCode::InvalidParams, None, req.id).into()
+        }
+
+        let params = req.params.as_array().unwrap();
+
+        match req.method.as_str() {
+            Some("ping") => return self.pong(req.id, params).await,
+            Some("kill") => return self.kill(req.id, params).await,
+            Some(_) | None => return JsonError::new(ErrorCode::MethodNotFound, None, req.id).into(),
+        }
+    }
+}
+
+async fn realmain(ex: Arc<Executor<'_>>) -> Result<()> {
+    let rpcsrv = Arc::new(RpcSrv { stop_sub: smol::channel::unbounded::<()>() });
+    let rpc_listen = Url::parse("tcp://127.0.0.1:55422").unwrap();
+
+    let _ex = ex.clone();
+    ex.spawn(listen_and_serve(rpc_listen, rpcsrv.clone(), _ex)).detach();
+
+    rpcsrv.stop_sub.1.recv().await?;
+
+    Ok(())
+}
+
+fn main() -> Result<()> {
+    simplelog::TermLogger::init(
+        simplelog::LevelFilter::Debug,
+        simplelog::ConfigBuilder::new().build(),
+        simplelog::TerminalMode::Mixed,
+        simplelog::ColorChoice::Auto,
+    )?;
+
+    let n_threads = std::thread::available_parallelism().unwrap().get();
+    let ex = Arc::new(Executor::new());
+    let (signal, shutdown) = smol::channel::unbounded::<()>();
+    let (_, result) = easy_parallel::Parallel::new()
+        .each(0..n_threads, |_| smol::future::block_on(ex.run(shutdown.recv())))
+        .finish(|| {
+            smol::future::block_on(async {
+                realmain(ex.clone()).await?;
+                drop(signal);
+                Ok::<(), darkfi::Error>(())
+            })
+        });
+
+    result
+}

+ 31 - 0
src/error.rs

@@ -227,6 +227,10 @@ pub enum Error {
     #[error("JSON-RPC error: {0}")]
     #[error("JSON-RPC error: {0}")]
     JsonRpcError(String),
     JsonRpcError(String),
 
 
+    #[cfg(feature = "rpc")]
+    #[error(transparent)]
+    RpcServerError(RpcError),
+
     #[error("Unexpected JSON-RPC data received: {0}")]
     #[error("Unexpected JSON-RPC data received: {0}")]
     UnexpectedJsonRpc(String),
     UnexpectedJsonRpc(String),
 
 
@@ -553,6 +557,33 @@ pub enum ClientFailed {
     VerifyError(String),
     VerifyError(String),
 }
 }
 
 
+#[cfg(feature = "rpc")]
+#[derive(Clone, Debug, thiserror::Error)]
+pub enum RpcError {
+    #[error("Connection closed: {0}")]
+    ConnectionClosed(String),
+
+    #[error("Invalid JSON: {0}")]
+    InvalidJson(String),
+
+    #[error("IO Error: {0}")]
+    IoError(std::io::ErrorKind),
+}
+
+#[cfg(feature = "rpc")]
+impl From<std::io::Error> for RpcError {
+    fn from(err: std::io::Error) -> Self {
+        Self::IoError(err.kind())
+    }
+}
+
+#[cfg(feature = "rpc")]
+impl From<RpcError> for Error {
+    fn from(err: RpcError) -> Self {
+        Self::RpcServerError(err)
+    }
+}
+
 impl From<Error> for ClientFailed {
 impl From<Error> for ClientFailed {
     fn from(err: Error) -> Self {
     fn from(err: Error) -> Self {
         Self::InternalError(err.to_string())
         Self::InternalError(err.to_string())

+ 129 - 132
src/rpc/client.rs

@@ -19,84 +19,135 @@
 //! JSON-RPC client-side implementation.
 //! JSON-RPC client-side implementation.
 use std::time::Duration;
 use std::time::Duration;
 
 
-use async_std::io::timeout;
-use futures::{select, AsyncReadExt, AsyncWriteExt, FutureExt};
+use async_std::{
+    io::{timeout, ReadExt, WriteExt},
+    sync::Arc,
+};
+use futures::{select, FutureExt};
 use log::{debug, error};
 use log::{debug, error};
 use serde_json::{json, Value};
 use serde_json::{json, Value};
 use url::Url;
 use url::Url;
 
 
 use super::jsonrpc::{ErrorCode, JsonError, JsonRequest, JsonResult};
 use super::jsonrpc::{ErrorCode, JsonError, JsonRequest, JsonResult};
 use crate::{
 use crate::{
+    error::RpcError,
     net::transport::{Dialer, PtStream},
     net::transport::{Dialer, PtStream},
-    system::SubscriberPtr,
     Error, Result,
     Error, Result,
 };
 };
 
 
+const INIT_BUF_SIZE: usize = 1024; // 1K
+const MAX_BUF_SIZE: usize = 1024 * 8192; // 1M
+const READ_TIMEOUT: Duration = Duration::from_secs(60);
+
 /// JSON-RPC client implementation using asynchronous channels.
 /// JSON-RPC client implementation using asynchronous channels.
 pub struct RpcClient {
 pub struct RpcClient {
     send: smol::channel::Sender<(Value, bool)>,
     send: smol::channel::Sender<(Value, bool)>,
     recv: smol::channel::Receiver<JsonResult>,
     recv: smol::channel::Receiver<JsonResult>,
     stop_signal: smol::channel::Sender<()>,
     stop_signal: smol::channel::Sender<()>,
-    url: Url,
+    endpoint: Url,
 }
 }
 
 
 impl RpcClient {
 impl RpcClient {
-    /// Instantiate a new JSON-RPC client that will connect to the given URL.
-    pub async fn new(url: Url) -> Result<Self> {
-        let (send, recv, stop_signal) = Self::open_channels(&url).await?;
-        Ok(Self { send, recv, stop_signal, url })
+    /// Instantiate a new JSON-RPC client that will connect to the given endpoint
+    pub async fn new(endpoint: Url, executor: Option<Arc<smol::Executor<'_>>>) -> Result<Self> {
+        let (send, recv, stop_signal) = Self::open_channels(&endpoint, executor.clone()).await?;
+        Ok(Self { send, recv, stop_signal, endpoint })
     }
     }
 
 
-    /// Close the channels of an instantiated [`RpcClient`].
+    /// Instantiate channels for a new [`RpcClient`]
+    async fn open_channels(
+        endpoint: &Url,
+        executor: Option<Arc<smol::Executor<'_>>>,
+    ) -> Result<(
+        smol::channel::Sender<(Value, bool)>,
+        smol::channel::Receiver<JsonResult>,
+        smol::channel::Sender<()>,
+    )> {
+        let (data_send, data_recv) = smol::channel::unbounded();
+        let (result_send, result_recv) = smol::channel::unbounded();
+        let (stop_send, stop_recv) = smol::channel::unbounded();
+
+        let dialer = Dialer::new(endpoint.clone()).await?;
+        let stream = dialer.dial(None).await?;
+
+        if let Some(ex) = executor {
+            ex.spawn(Self::reqrep_loop(stream, result_send, data_recv, stop_recv)).detach();
+        } else {
+            smol::spawn(Self::reqrep_loop(stream, result_send, data_recv, stop_recv)).detach();
+        }
+
+        Ok((data_send, result_recv, stop_send))
+    }
+
+    /// Close the channels of an instantiated [`RpcClient`]
     pub async fn close(&self) -> Result<()> {
     pub async fn close(&self) -> Result<()> {
         self.stop_signal.send(()).await?;
         self.stop_signal.send(()).await?;
         Ok(())
         Ok(())
     }
     }
 
 
-    /// Listen instantiated client for notifications.
-    /// NOTE: Subscriber listeners must perform response handling.
-    pub async fn subscribe(
-        &self,
-        req: JsonRequest,
-        subscriber: SubscriberPtr<JsonResult>,
-    ) -> Result<()> {
-        // Perform initial request.
-        debug!(target: "rpc::client", "--> {}", serde_json::to_string(&req)?);
-        // If the connection is closed, the sender will get an error for sending to a closed channel.
-        if let Err(e) = self.send.send((json!(req), false)).await {
-            error!(target: "rpc::client", "JSON-RPC client unable to send to {} (channels closed): {}", self.url, e);
-            return Err(Error::NetworkOperationFailed)
+    /// Internal read function called from `reqrep_loop` that reads from the
+    /// active stream.
+    async fn read_from_stream(stream: &mut Box<dyn PtStream>, buf: &mut Vec<u8>) -> Result<usize> {
+        debug!(target: "rpc::client", "Reading from stream...");
+        let mut total_read = 0;
+
+        while total_read < MAX_BUF_SIZE {
+            buf.resize(total_read + INIT_BUF_SIZE, 0);
+
+            match timeout(READ_TIMEOUT, stream.read(&mut buf[total_read..])).await {
+                Ok(0) if total_read == 0 => {
+                    return Err(RpcError::ConnectionClosed("Connection closed".to_string()).into())
+                }
+                Ok(0) => break, // Finished reading
+                Ok(n) => {
+                    total_read += n;
+                    if buf[total_read - 1] == b'\n' {
+                        break
+                    }
+                }
+                Err(e) => return Err(RpcError::IoError(e.kind()).into()),
+            }
         }
         }
 
 
+        // Truncate buffer to actual data size
+        buf.truncate(total_read);
+        debug!(target: "rpc::client", "Finished reading {} bytes", total_read);
+        Ok(total_read)
+    }
+
+    /// Internal function that loops on a given stream and multiplexes the data.
+    async fn reqrep_loop(
+        mut stream: Box<dyn PtStream>,
+        result_send: smol::channel::Sender<JsonResult>,
+        data_recv: smol::channel::Receiver<(Value, bool)>,
+        stop_recv: smol::channel::Receiver<()>,
+    ) -> Result<()> {
+        let mut buf = Vec::with_capacity(INIT_BUF_SIZE);
+
         loop {
         loop {
-            // If the connection is closed, the receiver will get an error for waiting on a closed channel.
-            let notification = self.recv.recv().await;
-            if notification.is_err() {
-                error!(target: "rpc::client", "JSON-RPC client unable to recv from {} (channels closed)", self.url);
-                break
-            }
+            buf.clear();
 
 
-            // Notify subscribed channels
-            let notification = notification?;
-            debug!(target: "rpc::client", "<-- {}", serde_json::to_string(&notification)?);
+            select! {
+                tuple = data_recv.recv().fuse() => {
+                    let (data, _with_timeout) = tuple?;
+                    let data_bytes = serde_json::to_vec(&data)?;
+                    stream.write_all(&data_bytes).await?;
+                    stream.write_all(&[b'\n']).await?;
 
 
-            subscriber.notify(notification.clone()).await;
+                    let _ = Self::read_from_stream(&mut stream, &mut buf).await?;
 
 
-            // Stop listenning on error
-            match notification {
-                JsonResult::Notification(_) => {}
-                _ => break,
-            }
+                    let r: JsonResult = serde_json::from_slice(&buf).map_err(
+                        |e| RpcError::InvalidJson(e.to_string())
+                    )?;
+
+                    result_send.send(r).await?;
+                }
 
 
-            // Triggering next consume
-            if let Err(e) = self.send.send((json!(req), false)).await {
-                error!(target: "rpc::client", "JSON-RPC client unable to send to {} (channels closed): {}", self.url, e);
-                break
+                _ = stop_recv.recv().fuse() => break
             }
             }
         }
         }
 
 
-        subscriber.notify(JsonError::new(ErrorCode::InternalError, None, req.id).into()).await;
-        Err(Error::NetworkOperationFailed)
+        Ok(())
     }
     }
 
 
     /// Send a given JSON-RPC request over the instantiated client.
     /// Send a given JSON-RPC request over the instantiated client.
@@ -105,125 +156,71 @@ impl RpcClient {
 
 
         debug!(target: "rpc::client", "--> {}", serde_json::to_string(&value)?);
         debug!(target: "rpc::client", "--> {}", serde_json::to_string(&value)?);
 
 
-        // If the connection is closed, the sender will get an error for
-        // sending to a closed channel.
+        // If the connection is closed, the sender will get an error
+        // for sending to a closed channel.
         if let Err(e) = self.send.send((json!(value), true)).await {
         if let Err(e) = self.send.send((json!(value), true)).await {
-            error!(target: "rpc::client", "JSON-RPC client unable to send to {} (channels closed): {}", self.url, e);
+            error!(
+                target: "rpc::client", "[RPC] Client unable to send to {}: {}",
+                self.endpoint, e
+            );
             return Err(Error::NetworkOperationFailed)
             return Err(Error::NetworkOperationFailed)
         }
         }
 
 
-        // If the connection is closed, the receiver will get an error for
-        // waiting on a closed channel.
+        // If the connection is closed, the receiver will get an error
+        // for waiting on a closed channel.
         let reply = self.recv.recv().await;
         let reply = self.recv.recv().await;
-        if reply.is_err() {
-            error!(target: "rpc::client", "JSON-RPC client unable to recv from {} (channels closed)", self.url);
+        if let Err(e) = reply {
+            error!(
+                target: "rpc::client", "[RPC] Client unable to recv from {}: {}",
+                self.endpoint, e
+            );
             return Err(Error::NetworkOperationFailed)
             return Err(Error::NetworkOperationFailed)
         }
         }
 
 
-        match reply? {
+        match reply.unwrap() {
             JsonResult::Response(r) => {
             JsonResult::Response(r) => {
+                debug!(target: "rpc::client", "<-- {}", serde_json::to_string(&r)?);
+
                 // Check if the IDs match
                 // Check if the IDs match
-                let resp_id = r.id.as_u64();
-                if resp_id.is_none() {
-                    let e = JsonError::new(ErrorCode::InvalidId, None, r.id);
-                    return Err(Error::JsonRpcError(e.error.message.to_string()))
-                }
+                match r.id.as_u64() {
+                    Some(id) => {
+                        if id != req_id {
+                            let e = JsonError::new(ErrorCode::InvalidId, None, r.id);
+                            return Err(Error::JsonRpcError(e.error.message.to_string()))
+                        }
+                    }
 
 
-                if resp_id.unwrap() != req_id {
-                    let e = JsonError::new(ErrorCode::InvalidId, None, r.id);
-                    return Err(Error::JsonRpcError(e.error.message.to_string()))
+                    None => {
+                        let e = JsonError::new(ErrorCode::InvalidId, None, r.id);
+                        return Err(Error::JsonRpcError(e.error.message.to_string()))
+                    }
                 }
                 }
 
 
-                debug!(target: "rpc::client", "<-- {}", serde_json::to_string(&r)?);
                 Ok(r.result)
                 Ok(r.result)
             }
             }
+
             JsonResult::Error(e) => {
             JsonResult::Error(e) => {
                 debug!(target: "rpc::client", "<-- {}", serde_json::to_string(&e)?);
                 debug!(target: "rpc::client", "<-- {}", serde_json::to_string(&e)?);
                 Err(Error::JsonRpcError(e.error.message.to_string()))
                 Err(Error::JsonRpcError(e.error.message.to_string()))
             }
             }
+
             JsonResult::Notification(n) => {
             JsonResult::Notification(n) => {
                 debug!(target: "rpc::client", "<-- {}", serde_json::to_string(&n)?);
                 debug!(target: "rpc::client", "<-- {}", serde_json::to_string(&n)?);
                 Err(Error::JsonRpcError("Unexpected reply".to_string()))
                 Err(Error::JsonRpcError("Unexpected reply".to_string()))
             }
             }
-            JsonResult::Subscriber(_) => Err(Error::JsonRpcError("Unexpected reply".to_string())),
+
+            JsonResult::Subscriber(_) => {
+                // When?
+                Err(Error::JsonRpcError("Unexpected reply".to_string()))
+            }
         }
         }
     }
     }
 
 
     /// Oneshot send a given JSON-RPC request over the instantiated client
     /// Oneshot send a given JSON-RPC request over the instantiated client
-    /// and close the channels on reply.
+    /// and immediately close the channels upon receiving a reply.
     pub async fn oneshot_request(&self, value: JsonRequest) -> Result<Value> {
     pub async fn oneshot_request(&self, value: JsonRequest) -> Result<Value> {
         let rep = self.request(value).await?;
         let rep = self.request(value).await?;
         self.stop_signal.send(()).await?;
         self.stop_signal.send(()).await?;
         Ok(rep)
         Ok(rep)
     }
     }
-
-    /// Instantiate channels for a new [`RpcClient`].
-    async fn open_channels(
-        uri: &Url,
-    ) -> Result<(
-        smol::channel::Sender<(Value, bool)>,
-        smol::channel::Receiver<JsonResult>,
-        smol::channel::Sender<()>,
-    )> {
-        let (data_send, data_recv) = smol::channel::unbounded();
-        let (result_send, result_recv) = smol::channel::unbounded();
-        let (stop_send, stop_recv) = smol::channel::unbounded();
-
-        let dialer = Dialer::new(uri.clone()).await?;
-        // TODO: Revisit the timeout here
-        let stream = dialer.dial(None).await?;
-        smol::spawn(Self::reqrep_loop(stream, result_send, data_recv, stop_recv)).detach();
-
-        Ok((data_send, result_recv, stop_send))
-    }
-
-    /// Internal function that loops on a given stream and multiplexes the data.
-    async fn reqrep_loop(
-        mut stream: Box<dyn PtStream>,
-        result_send: smol::channel::Sender<JsonResult>,
-        data_recv: smol::channel::Receiver<(Value, bool)>,
-        stop_recv: smol::channel::Receiver<()>,
-    ) -> Result<()> {
-        // If timeout is enabled and we don't get a reply within 30 seconds, we'll fail.
-        let read_timeout = Duration::from_secs(30);
-
-        loop {
-            // FIXME: Nasty size. 8M
-            let mut buf = vec![0; 1024 * 8192];
-
-            select! {
-                tuple = data_recv.recv().fuse() => {
-                    let (data, with_timeout) = tuple?;
-                    let data_bytes = serde_json::to_vec(&data)?;
-                    stream.write_all(&data_bytes).await?;
-                    // Since we are using async read and write,
-                    // the other side might not have finished writing
-                    // to the stream. To mitigate this, we perform a read
-                    // and check if data can be converted to a JsonResult.
-                    // If data is incomplete, this will fail, therefore,
-                    // we re-execute read and write after previous read in the buffer,
-                    // and repeat until the data in buffer can be converted.
-                    let mut n = 0;
-                    loop {
-                        n += if with_timeout {
-                            timeout(read_timeout, async { stream.read(&mut buf[n..]).await }).await?
-                        } else {
-                            stream.read(&mut buf[n..]).await?
-                        };
-                        match serde_json::from_slice(&buf[0..n]) {
-                            Ok(reply) => {
-                                result_send.send(reply).await?;
-                                break
-                            },
-                            Err(e) => debug!(target: "rpc::client", "JSON-RPC client retrying failed convertion with error: {}", e),
-                        }
-                    }
-                }
-
-                _ = stop_recv.recv().fuse() => break
-            }
-        }
-
-        Ok(())
-    }
 }
 }

+ 2 - 1
src/rpc/clock_sync.rs

@@ -35,6 +35,7 @@ const EPOCH: u32 = 2208988800; // 1900
 
 
 /// JSON-RPC request to a network peer (randomly selected), to
 /// JSON-RPC request to a network peer (randomly selected), to
 /// retrieve their current system clock.
 /// retrieve their current system clock.
+// TODO: This needs executor passed for rpc client
 async fn peer_request(peers: &[Url]) -> Result<Option<Timestamp>> {
 async fn peer_request(peers: &[Url]) -> Result<Option<Timestamp>> {
     // Select peer, None if vector is empty.
     // Select peer, None if vector is empty.
     let peer = peers.choose(&mut OsRng);
     let peer = peers.choose(&mut OsRng);
@@ -42,7 +43,7 @@ async fn peer_request(peers: &[Url]) -> Result<Option<Timestamp>> {
         None => Ok(None),
         None => Ok(None),
         Some(p) => {
         Some(p) => {
             // Create RPC client
             // Create RPC client
-            let rpc_client = RpcClient::new(p.clone()).await?;
+            let rpc_client = RpcClient::new(p.clone(), None).await?;
 
 
             // Execute request
             // Execute request
             let req = JsonRequest::new("clock", json!([]));
             let req = JsonRequest::new("clock", json!([]));

+ 86 - 39
src/rpc/server.rs

@@ -17,14 +17,17 @@
  */
  */
 
 
 //! JSON-RPC server-side implementation.
 //! JSON-RPC server-side implementation.
-use async_std::sync::Arc;
+use std::time::Duration;
+
+use async_std::{io::timeout, sync::Arc};
 use async_trait::async_trait;
 use async_trait::async_trait;
 use futures::{AsyncReadExt, AsyncWriteExt};
 use futures::{AsyncReadExt, AsyncWriteExt};
-use log::{debug, error, info, warn};
+use log::{debug, error, info};
 use url::Url;
 use url::Url;
 
 
 use super::jsonrpc::{JsonRequest, JsonResult};
 use super::jsonrpc::{JsonRequest, JsonResult};
 use crate::{
 use crate::{
+    error::RpcError,
     net::transport::{Listener, PtListener, PtStream},
     net::transport::{Listener, PtListener, PtStream},
     Result,
     Result,
 };
 };
@@ -37,43 +40,59 @@ pub trait RequestHandler: Sync + Send {
     async fn handle_request(&self, req: JsonRequest) -> JsonResult;
     async fn handle_request(&self, req: JsonRequest) -> JsonResult;
 }
 }
 
 
+const INIT_BUF_SIZE: usize = 1024; // 1K
+const MAX_BUF_SIZE: usize = 1024 * 8192; // 8M
+const READ_TIMEOUT: Duration = Duration::from_secs(30);
+
+/// Internal read function called from `accept` that reads from the active stream.
+async fn read_from_stream(stream: &mut Box<dyn PtStream>, buf: &mut Vec<u8>) -> Result<usize> {
+    debug!(target: "rpc::server", "Reading from stream...");
+    let mut total_read = 0;
+
+    while total_read < MAX_BUF_SIZE {
+        buf.resize(total_read + INIT_BUF_SIZE, 0);
+
+        match timeout(READ_TIMEOUT, stream.read(&mut buf[total_read..])).await {
+            Ok(0) if total_read == 0 => {
+                return Err(RpcError::ConnectionClosed("Connection closed".to_string()).into())
+            }
+            Ok(0) => break, // Finished reading
+            Ok(n) => {
+                total_read += n;
+                if buf[total_read - 1] == b'\n' {
+                    break
+                }
+            }
+            Err(e) => return Err(RpcError::IoError(e.kind()).into()),
+        }
+    }
+
+    // Truncate buffer to actual data size
+    buf.truncate(total_read);
+    debug!(target: "rpc::server", "Finished reading {} bytes", total_read);
+    Ok(total_read)
+}
+
 /// Internal accept function that runs inside a loop for accepting incoming
 /// Internal accept function that runs inside a loop for accepting incoming
 /// JSON-RPC requests and passing them to the [`RequestHandler`].
 /// JSON-RPC requests and passing them to the [`RequestHandler`].
 async fn accept(
 async fn accept(
     mut stream: Box<dyn PtStream>,
     mut stream: Box<dyn PtStream>,
-    peer_addr: Url,
+    addr: Url,
     rh: Arc<impl RequestHandler + 'static>,
     rh: Arc<impl RequestHandler + 'static>,
 ) -> Result<()> {
 ) -> Result<()> {
+    let mut buf = Vec::with_capacity(INIT_BUF_SIZE);
     loop {
     loop {
-        // FIXME: Nasty size. 8M
-        let mut buf = vec![0; 1024 * 8192];
+        buf.clear();
 
 
-        let n = match stream.read(&mut buf).await {
-            Ok(0) => {
-                debug!(target: "rpc::server", "Closed connection for {}", peer_addr);
-                break
-            }
-            Ok(n) => n,
-            Err(e) => {
-                error!(target: "rpc::server", "JSON-RPC server failed reading from {} socket: {}", peer_addr, e);
-                debug!(target: "rpc::server", "Closed connection for {}", peer_addr);
-                break
-            }
-        };
+        let _ = read_from_stream(&mut stream, &mut buf).await?;
 
 
-        let r: JsonRequest = match serde_json::from_slice(&buf[0..n]) {
-            Ok(r) => {
-                debug!(target: "rpc::server", "{} --> {}", peer_addr, String::from_utf8_lossy(&buf));
-                r
-            }
-            Err(e) => {
-                warn!(target: "rpc::server", "JSON-RPC server received invalid JSON from {}: {}", peer_addr, e);
-                debug!(target: "rpc::server", "Closed connection for {}", peer_addr);
-                break
-            }
-        };
+        let r: JsonRequest =
+            serde_json::from_slice(&buf).map_err(|e| RpcError::InvalidJson(e.to_string()))?;
+
+        debug!(target: "rpc::server", "{} --> {}", addr, String::from_utf8_lossy(&buf));
 
 
         let reply = rh.handle_request(r).await;
         let reply = rh.handle_request(r).await;
+
         match reply {
         match reply {
             JsonResult::Subscriber(sub) => {
             JsonResult::Subscriber(sub) => {
                 // Subscribe to the inner method subscriber
                 // Subscribe to the inner method subscriber
@@ -84,30 +103,58 @@ async fn accept(
 
 
                     // Push notification
                     // Push notification
                     let j = serde_json::to_string(&notification).unwrap();
                     let j = serde_json::to_string(&notification).unwrap();
-                    debug!(target: "rpc::server", "{} <-- {}", peer_addr, j);
+                    debug!(target: "rpc::server", "{} <-- {}", addr, j);
 
 
                     if let Err(e) = stream.write_all(j.as_bytes()).await {
                     if let Err(e) = stream.write_all(j.as_bytes()).await {
-                        error!(target: "rpc::server", "JSON-RPC server failed writing to {} socket: {}", peer_addr, e);
-                        debug!(target: "rpc::server", "Closed connection for {}", peer_addr);
+                        error!(target: "rpc::server", "[RPC] Server failed writing to {} socket: {}", addr, e);
+                        debug!(target: "rpc::server", "Closed connection for {}", addr);
+                        break
+                    }
+
+                    if let Err(e) = stream.write_all(&[b'\n']).await {
+                        error!(target: "rpc::server", "[RPC] Server failed writing to {} socket: {}", addr, e);
+                        debug!(target: "rpc::server", "Closed connection for {}", addr);
                         break
                         break
                     }
                     }
                 }
                 }
                 subscription.unsubscribe().await;
                 subscription.unsubscribe().await;
             }
             }
             _ => {
             _ => {
-                let j = serde_json::to_string(&reply).unwrap();
-                debug!(target: "rpc::server", "{} <-- {}", peer_addr, j);
+                let j = serde_json::to_string(&reply)
+                    .map_err(|e| RpcError::InvalidJson(e.to_string()))?;
+
+                debug!(target: "rpc::server", "{} <-- {}", addr, j);
 
 
                 if let Err(e) = stream.write_all(j.as_bytes()).await {
                 if let Err(e) = stream.write_all(j.as_bytes()).await {
-                    error!(target: "rpc::server", "JSON-RPC server failed writing to {} socket: {}", peer_addr, e);
-                    debug!(target: "rpc::server", "Closed connection for {}", peer_addr);
-                    break
+                    error!(
+                        target: "rpc::server", "[RPC] Server failed writing to {} socket: {}",
+                        addr, e
+                    );
+                    return close_conn(
+                        &addr,
+                        RpcError::ConnectionClosed("Socket write error".to_string()),
+                    )
+                }
+
+                if let Err(e) = stream.write_all(&[b'\n']).await {
+                    error!(
+                        target: "rpc::server", "[RPC] Server failed writing to {} socket: {}",
+                        addr, e
+                    );
+                    return close_conn(
+                        &addr,
+                        RpcError::ConnectionClosed("Socket write error".to_string()),
+                    )
                 }
                 }
             }
             }
         }
         }
     }
     }
+}
 
 
-    Ok(())
+/// Helper function for connection closing
+fn close_conn(peer_addr: &Url, reason: RpcError) -> Result<()> {
+    debug!(target: "rpc::server", "Closed connection for {}", peer_addr);
+    Err(reason.into())
 }
 }
 
 
 /// Wrapper function around [`accept()`] to take the incoming connection and
 /// Wrapper function around [`accept()`] to take the incoming connection and
@@ -118,12 +165,12 @@ async fn run_accept_loop(
     ex: Arc<smol::Executor<'_>>,
     ex: Arc<smol::Executor<'_>>,
 ) -> Result<()> {
 ) -> Result<()> {
     while let Ok((stream, peer_addr)) = listener.next().await {
     while let Ok((stream, peer_addr)) = listener.next().await {
-        info!(target: "rpc::server", "JSON-RPC server accepted connection from {}", peer_addr);
+        info!(target: "rpc::server", "[RPC] Server accepted connection from {}", peer_addr);
         // Detaching requests handling
         // Detaching requests handling
         let _rh = rh.clone();
         let _rh = rh.clone();
         ex.spawn(async move {
         ex.spawn(async move {
             if let Err(e) = accept(stream, peer_addr.clone(), _rh).await {
             if let Err(e) = accept(stream, peer_addr.clone(), _rh).await {
-                error!(target: "rpc::server", "JSON-RPC server error on handling request of {}: {}", peer_addr, e);
+                error!(target: "rpc::server", "[RPC] Server error on handling request of {}: {}", peer_addr, e);
             }
             }
         }).detach();
         }).detach();
     }
     }