Improved RPC handling. WIP

beacon_node/eth2-libp2p/src/rpc/protocol.rs

@@ -1,10 +1,13 @@
 use super::methods::*;
+use super::request_response::{rpc_request_response, RPCRequestResponse};
+use futures::future::Future;
 use libp2p::core::{upgrade, InboundUpgrade, OutboundUpgrade, UpgradeInfo};
 use ssz::{Decode, Encode};
-use std::hash::Hasher;
 use std::io;
-use std::iter;
+use std::time::Duration;
 use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::prelude::future::MapErr;
+use tokio::util::FutureExt;
 
 /// The maximum bytes that can be sent across the RPC.
 const MAX_RPC_SIZE: usize = 4_194_304; // 4M
@@ -33,42 +36,6 @@ impl UpgradeInfo for RPCProtocol {
     }
 }
 
-/* Inbound upgrade */
-
-// The inbound protocol reads the request, decodes it and returns the stream to the protocol
-// handler to respond to once ready.
-
-type FnDecodeRPCEvent<TSocket> = fn(
-    upgrade::Negotiated<TSocket>,
-    Vec<u8>,
-    (),
-) -> Result<(upgrade::Negotiated<TSocket>, RPCEvent), RPCError>;
-
-impl<TSocket> InboundUpgrade<TSocket> for RPCProtocol
-where
-    TSocket: AsyncRead + AsyncWrite,
-{
-    type Output = (upgrade::Negotiated<TSocket>, RPCEvent);
-    type Error = RPCError;
-    type Future = upgrade::ReadRespond<upgrade::Negotiated<TSocket>, (), FnDecodeRPCEvent<TSocket>>;
-
-    fn upgrade_inbound(
-        self,
-        socket: upgrade::Negotiated<TSocket>,
-        protocol: Self::Info,
-    ) -> Self::Future {
-        upgrade::read_respond(socket, MAX_RPC_SIZE, (), |socket, packet, ()| {
-            Ok((socket, decode_request(packet, protocol)?))
-        })
-        .timeout(Duration::from_secs(RESPONSE_TIMEOUT))
-    }
-}
-
-/* Outbound request */
-
-// Combines all the RPC requests into a single enum to implement `UpgradeInfo` and
-// `OutboundUpgrade`
-
 /// The raw protocol id sent over the wire.
 type RawProtocolId = Vec<u8>;
 
@@ -86,38 +53,100 @@ pub struct ProtocolId {
 
 /// An RPC protocol ID.
 impl ProtocolId {
-    pub fn new(message_name: String, version: usize, encoding: String) -> Self {
+    pub fn new(message_name: &str, version: usize, encoding: &str) -> Self {
         ProtocolId {
-            message_name,
+            message_name: message_name.into(),
             version,
-            encoding,
+            encoding: encoding.into(),
         }
     }
 
     /// Converts a raw RPC protocol id string into an `RPCProtocolId`
-    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self, RPCError> {
-        let protocol_string = String::from_utf8(bytes.as_vec())
+    pub fn from_bytes(bytes: &[u8]) -> Result<Self, RPCError> {
+        let protocol_string = String::from_utf8(bytes.to_vec())
             .map_err(|_| RPCError::InvalidProtocol("Invalid protocol Id"))?;
-        let protocol_string = protocol_string.as_str().split('/');
+        let protocol_list: Vec<&str> = protocol_string.as_str().split('/').take(5).collect();
+
+        if protocol_list.len() != 5 {
+            return Err(RPCError::InvalidProtocol("Not enough '/'"));
+        }
 
         Ok(ProtocolId {
-            message_name: protocol_string[3],
-            version: protocol_string[4],
-            encoding: protocol_string[5],
+            message_name: protocol_list[3].into(),
+            version: protocol_list[4]
+                .parse()
+                .map_err(|_| RPCError::InvalidProtocol("Invalid version"))?,
+            encoding: protocol_list[5].into(),
         })
     }
 }
 
 impl Into<RawProtocolId> for ProtocolId {
-    fn into(&self) -> [u8] {
-        &format!(
+    fn into(self) -> RawProtocolId {
+        format!(
             "{}/{}/{}/{}",
             PROTOCOL_PREFIX, self.message_name, self.version, self.encoding
         )
         .as_bytes()
+        .to_vec()
     }
 }
 
+/* Inbound upgrade */
+
+// The inbound protocol reads the request, decodes it and returns the stream to the protocol
+// handler to respond to once ready.
+
+type FnDecodeRPCEvent<TSocket> =
+    fn(
+        upgrade::Negotiated<TSocket>,
+        Vec<u8>,
+        &'static [u8], // protocol id
+    ) -> Result<(upgrade::Negotiated<TSocket>, RPCRequest, ProtocolId), RPCError>;
+
+impl<TSocket> InboundUpgrade<TSocket> for RPCProtocol
+where
+    TSocket: AsyncRead + AsyncWrite,
+{
+    type Output = (upgrade::Negotiated<TSocket>, RPCRequest, ProtocolId);
+    type Error = RPCError;
+    type Future = MapErr<
+        tokio_timer::Timeout<
+            upgrade::ReadRespond<
+                upgrade::Negotiated<TSocket>,
+                Self::Info,
+                FnDecodeRPCEvent<TSocket>,
+            >,
+        >,
+        fn(tokio::timer::timeout::Error<RPCError>) -> RPCError,
+    >;
+
+    fn upgrade_inbound(
+        self,
+        socket: upgrade::Negotiated<TSocket>,
+        protocol: &'static [u8],
+    ) -> Self::Future {
+        upgrade::read_respond(socket, MAX_RPC_SIZE, protocol, {
+            |socket, packet, protocol| {
+                let protocol_id = ProtocolId::from_bytes(protocol)?;
+                Ok((
+                    socket,
+                    RPCRequest::decode(packet, protocol_id)?,
+                    protocol_id,
+                ))
+            }
+        }
+            as FnDecodeRPCEvent<TSocket>)
+        .timeout(Duration::from_secs(RESPONSE_TIMEOUT))
+        .map_err(RPCError::from)
+    }
+}
+
+/* Outbound request */
+
+// Combines all the RPC requests into a single enum to implement `UpgradeInfo` and
+// `OutboundUpgrade`
+
 #[derive(Debug, Clone)]
 pub enum RPCRequest {
     Hello(HelloMessage),
@@ -154,18 +183,16 @@ impl RPCRequest {
             RPCRequest::BeaconBlockBodies(_) => {
                 vec![ProtocolId::new("beacon_block_bodies", 1, "ssz").into()]
             }
-            RPCRequest::BeaconBlockState(_) => {
+            RPCRequest::BeaconChainState(_) => {
                 vec![ProtocolId::new("beacon_block_state", 1, "ssz").into()]
             }
         }
     }
 
     /// Encodes the Request object based on the negotiated protocol.
-    pub fn encode(&self, protocol: RawProtocolId) -> Result<Vec<u8>, io::Error> {
-        // Assume select has given a supported protocol.
-        let protocol = ProtocolId::from_bytes(protocol)?;
+    pub fn encode(&self, protocol: ProtocolId) -> Result<Vec<u8>, RPCError> {
         // Match on the encoding and in the future, the version
-        match protocol.encoding {
+        match protocol.encoding.as_str() {
             "ssz" => Ok(self.ssz_encode()),
             _ => {
                 return Err(RPCError::Custom(format!(
@@ -176,7 +203,7 @@ impl RPCRequest {
         }
     }
 
-    fn ssz_encode(&self) {
+    fn ssz_encode(&self) -> Vec<u8> {
         match self {
             RPCRequest::Hello(req) => req.as_ssz_bytes(),
             RPCRequest::Goodbye(req) => req.as_ssz_bytes(),
@@ -186,177 +213,263 @@ impl RPCRequest {
             RPCRequest::BeaconChainState(req) => req.as_ssz_bytes(),
         }
     }
+
+    // This function can be extended to provide further logic for supporting various protocol versions/encoding
+    /// Decodes a request received from our peer.
+    pub fn decode(packet: Vec<u8>, protocol: ProtocolId, response_code: ResponseCode) -> Result<Self, RPCError> {
+
+        match response_code {
+            ResponseCode::
+
+
+
+        match protocol.message_name.as_str() {
+            "hello" => match protocol.version {
+                1 => match protocol.encoding.as_str() {
+                    "ssz" => Ok(RPCRequest::Hello(HelloMessage::from_ssz_bytes(&packet)?)),
+                    _ => Err(RPCError::InvalidProtocol("Unknown HELLO encoding")),
+                },
+                _ => Err(RPCError::InvalidProtocol("Unknown HELLO version")),
+            },
+            "goodbye" => match protocol.version {
+                1 => match protocol.encoding.as_str() {
+                    "ssz" => Ok(RPCRequest::Goodbye(Goodbye::from_ssz_bytes(&packet)?)),
+                    _ => Err(RPCError::InvalidProtocol("Unknown GOODBYE encoding")),
+                },
+                _ => Err(RPCError::InvalidProtocol("Unknown GOODBYE version")),
+            },
+            "beacon_block_roots" => match protocol.version {
+                1 => match protocol.encoding.as_str() {
+                    "ssz" => Ok(RPCRequest::BeaconBlockRoots(
+                        BeaconBlockRootsRequest::from_ssz_bytes(&packet)?,
+                    )),
+                    _ => Err(RPCError::InvalidProtocol(
+                        "Unknown BEACON_BLOCK_ROOTS encoding",
+                    )),
+                },
+                _ => Err(RPCError::InvalidProtocol(
+                    "Unknown BEACON_BLOCK_ROOTS version",
+                )),
+            },
+            "beacon_block_headers" => match protocol.version {
+                1 => match protocol.encoding.as_str() {
+                    "ssz" => Ok(RPCRequest::BeaconBlockHeaders(
+                        BeaconBlockHeadersRequest::from_ssz_bytes(&packet)?,
+                    )),
+                    _ => Err(RPCError::InvalidProtocol(
+                        "Unknown BEACON_BLOCK_HEADERS encoding",
+                    )),
+                },
+                _ => Err(RPCError::InvalidProtocol(
+                    "Unknown BEACON_BLOCK_HEADERS version",
+                )),
+            },
+            "beacon_block_bodies" => match protocol.version {
+                1 => match protocol.encoding.as_str() {
+                    "ssz" => Ok(RPCRequest::BeaconBlockBodies(
+                        BeaconBlockBodiesRequest::from_ssz_bytes(&packet)?,
+                    )),
+                    _ => Err(RPCError::InvalidProtocol(
+                        "Unknown BEACON_BLOCK_BODIES encoding",
+                    )),
+                },
+                _ => Err(RPCError::InvalidProtocol(
+                    "Unknown BEACON_BLOCK_BODIES version",
+                )),
+            },
+            "beacon_chain_state" => match protocol.version {
+                1 => match protocol.encoding.as_str() {
+                    "ssz" => Ok(RPCRequest::BeaconChainState(
+                        BeaconChainStateRequest::from_ssz_bytes(&packet)?,
+                    )),
+                    _ => Err(RPCError::InvalidProtocol(
+                        "Unknown BEACON_CHAIN_STATE encoding",
+                    )),
+                },
+                _ => Err(RPCError::InvalidProtocol(
+                    "Unknown BEACON_CHAIN_STATE version",
+                )),
+            },
+        }
+    }
+}
+
+/* Response Type */
+
+#[derive(Debug, Clone)]
+pub enum RPCResponse {
+    /// A HELLO message.
+    Hello(HelloMessage),
+    /// An empty field returned from sending a GOODBYE request.
+    Goodbye, // empty value - required for protocol handler
+    /// A response to a get BEACON_BLOCK_ROOTS request.
+    BeaconBlockRoots(BeaconBlockRootsResponse),
+    /// A response to a get BEACON_BLOCK_HEADERS request.
+    BeaconBlockHeaders(BeaconBlockHeadersResponse),
+    /// A response to a get BEACON_BLOCK_BODIES request.
+    BeaconBlockBodies(BeaconBlockBodiesResponse),
+    /// A response to a get BEACON_CHAIN_STATE request.
+    BeaconChainState(BeaconChainStateResponse),
+    /// The Error returned from the peer during a request.
+    Error(String),
+}
+
+pub enum ResponseCode {
+    Success = 0,
+    EncodingError = 1,
+    InvalidRequest = 2,
+    ServerError = 3,
+    Unknown,
+}
+
+
+impl From<u64> for ResponseCode {
+    fn from(val: u64) -> ResponseCode {
+        match val {
+            0 => ResponseCode::Success,
+            1 => ResponseCode::EncodingError,
+            2 => ResponseCode::InvalidRequest,
+            3 => ResponseCode::ServerError,
+            _ => ResponseCode::Unknown,
+        }
+    }
+}
+
+impl RPCResponse {
+    /// Decodes a response that was received on the same stream as a request. The response type should
+    /// therefore match the request protocol type.
+    fn decode(packet: Vec<u8>, protocol: ProtocolId) -> Result<Self, RPCError> {
+        match protocol.message_name.as_str() {
+            "hello" => match protocol.version {
+                1 => match protocol.encoding.as_str() {
+                    "ssz" => Ok(RPCResponse::Hello(HelloMessage::from_ssz_bytes(&packet)?)),
+                    _ => Err(RPCError::InvalidProtocol("Unknown HELLO encoding")),
+                },
+                _ => Err(RPCError::InvalidProtocol("Unknown HELLO version")),
+            },
+            "goodbye" => Err(RPCError::Custom(
+                "GOODBYE should not have a response".into(),
+            )),
+            "beacon_block_roots" => match protocol.version {
+                1 => match protocol.encoding.as_str() {
+                    "ssz" => Ok(RPCResponse::BeaconBlockRoots(
+                        BeaconBlockRootsResponse::from_ssz_bytes(&packet)?,
+                    )),
+                    _ => Err(RPCError::InvalidProtocol(
+                        "Unknown BEACON_BLOCK_ROOTS encoding",
+                    )),
+                },
+                _ => Err(RPCError::InvalidProtocol(
+                    "Unknown BEACON_BLOCK_ROOTS version",
+                )),
+            },
+            "beacon_block_headers" => match protocol.version {
+                1 => match protocol.encoding.as_str() {
+                    "ssz" => Ok(RPCResponse::BeaconBlockHeaders(
+                        BeaconBlockHeadersResponse { headers: packet },
+                    )),
+                    _ => Err(RPCError::InvalidProtocol(
+                        "Unknown BEACON_BLOCK_HEADERS encoding",
+                    )),
+                },
+                _ => Err(RPCError::InvalidProtocol(
+                    "Unknown BEACON_BLOCK_HEADERS version",
+                )),
+            },
+            "beacon_block_bodies" => match protocol.version {
+                1 => match protocol.encoding.as_str() {
+                    "ssz" => Ok(RPCResponse::BeaconBlockBodies(BeaconBlockBodiesResponse {
+                        block_bodies: packet,
+                    })),
+                    _ => Err(RPCError::InvalidProtocol(
+                        "Unknown BEACON_BLOCK_BODIES encoding",
+                    )),
+                },
+                _ => Err(RPCError::InvalidProtocol(
+                    "Unknown BEACON_BLOCK_BODIES version",
+                )),
+            },
+            "beacon_chain_state" => match protocol.version {
+                1 => match protocol.encoding.as_str() {
+                    "ssz" => Ok(RPCResponse::BeaconChainState(
+                        BeaconChainStateResponse::from_ssz_bytes(&packet)?,
+                    )),
+                    _ => Err(RPCError::InvalidProtocol(
+                        "Unknown BEACON_CHAIN_STATE encoding",
+                    )),
+                },
+                _ => Err(RPCError::InvalidProtocol(
+                    "Unknown BEACON_CHAIN_STATE version",
+                )),
+            },
+        }
+    }
+
+    /// Encodes the Response object based on the negotiated protocol.
+    pub fn encode(&self, protocol: ProtocolId) -> Result<Vec<u8>, RPCError> {
+        // Match on the encoding and in the future, the version
+        match protocol.encoding.as_str() {
+            "ssz" => Ok(self.ssz_encode()),
+            _ => {
+                return Err(RPCError::Custom(format!(
+                    "Unknown Encoding: {}",
+                    protocol.encoding
+                )))
+            }
+        }
+    }
+
+    fn ssz_encode(&self) -> Vec<u8> {
+        match self {
+            RPCResponse::Hello(res) => res.as_ssz_bytes(),
+            RPCResponse::Goodbye => unreachable!(),
+            RPCResponse::BeaconBlockRoots(res) => res.as_ssz_bytes(),
+            RPCResponse::BeaconBlockHeaders(res) => res.headers, // already raw bytes
+            RPCResponse::BeaconBlockBodies(res) => res.block_bodies, // already raw bytes
+            RPCResponse::BeaconChainState(res) => res.as_ssz_bytes(),
+        }
+    }
 }
 
 /* Outbound upgrades */
 
 impl<TSocket> OutboundUpgrade<TSocket> for RPCRequest
 where
-    TSocket: AsyncWrite,
+    TSocket: AsyncRead + AsyncWrite,
 {
     type Output = RPCResponse;
-    type Error = RPCResponse;
-    type Future = upgrade::RequestResponse<upgrade::Negotiated<TSocket>>;
+    type Error = RPCError;
+    type Future = MapErr<
+        tokio_timer::Timeout<RPCRequestResponse<upgrade::Negotiated<TSocket>, Vec<u8>>>,
+        fn(tokio::timer::timeout::Error<RPCError>) -> RPCError,
+    >;
 
     fn upgrade_outbound(
         self,
         socket: upgrade::Negotiated<TSocket>,
         protocol: Self::Info,
     ) -> Self::Future {
-        let bytes = self.encode(protocol);
-        wait_for_response = if let RPCRequest::Goodbye(_) = self {
-            false
+        let protocol_id = ProtocolId::from_bytes(&protocol)
+            .expect("Protocol ID must be valid for outbound requests");
+
+        let request_bytes = self
+            .encode(protocol_id)
+            .expect("Should be able to encode a supported protocol");
+        // if sending a goodbye, drop the stream and return an empty GOODBYE response
+        let short_circuit_return = if let RPCRequest::Goodbye(_) = self {
+            Some(RPCResponse::Goodbye)
         } else {
-            true
+            None
         };
-        // TODO: Reimplement request_response
-        upgrade::request_response(socket, bytes, MAX_RPC_SIZE, protocol, |packet, protocol| {
-            Ok(decode_response(packet, protocol)?)
-        })
+        rpc_request_response(
+            socket,
+            request_bytes,
+            MAX_RPC_SIZE,
+            short_circuit_return,
+            protocol_id,
+        )
         .timeout(Duration::from_secs(RESPONSE_TIMEOUT))
-    }
-}
-
-/* Decoding for Requests/Responses */
-
-// This function can be extended to provide further logic for supporting various protocol versions/encoding
-fn decode_request(packet: Vec<u8>, protocol: ProtocolId) -> Result<RPCRequest, io::Error> {
-    let protocol_id = ProtocolId::from_bytes(protocol);
-
-    match protocol_id.message_name {
-        "hello" => match protocol_id.version {
-            "1" => match protocol_id.encoding {
-                "ssz" => Ok(RPCRequest::Hello(HelloMessage::from_ssz_bytes(&packet)?)),
-                _ => Err(RPCError::InvalidProtocol("Unknown HELLO encoding")),
-            },
-            _ => Err(RPCError::InvalidProtocol("Unknown HELLO version")),
-        },
-        "goodbye" => match protocol_id.version {
-            "1" => match protocol_id.encoding {
-                "ssz" => Ok(RPCRequest::Goodbye(Goodbye::from_ssz_bytes(&packet)?)),
-                _ => Err(RPCError::InvalidProtocol("Unknown GOODBYE encoding")),
-            },
-            _ => Err(RPCError::InvalidProtocol("Unknown GOODBYE version")),
-        },
-        "beacon_block_roots" => match protocol_id.version {
-            "1" => match protocol_id.encoding {
-                "ssz" => Ok(RPCRequest::BeaconBlockRooots(
-                    BeaconBlockRootsRequest::from_ssz_bytes(&packet)?,
-                )),
-                _ => Err(RPCError::InvalidProtocol(
-                    "Unknown BEACON_BLOCK_ROOTS encoding",
-                )),
-            },
-            _ => Err(RPCError::InvalidProtocol(
-                "Unknown BEACON_BLOCK_ROOTS version",
-            )),
-        },
-        "beacon_block_headers" => match protocol_id.version {
-            "1" => match protocol_id.encoding {
-                "ssz" => Ok(RPCRequest::BeaconBlockHeaders(
-                    BeaconBlockHeadersRequest::from_ssz_bytes(&packet),
-                )),
-                _ => Err(RPCError::InvalidProtocol(
-                    "Unknown BEACON_BLOCK_HEADERS encoding",
-                )),
-            },
-            _ => Err(RPCError::InvalidProtocol(
-                "Unknown BEACON_BLOCK_HEADERS version",
-            )),
-        },
-        "beacon_block_bodies" => match protocol_id.version {
-            "1" => match protocol_id.encoding {
-                "ssz" => Ok(RPCRequest::BeaconBlockBodies(
-                    BeaconBlockBodiesRequest::from_ssz_bytes(&packet)?,
-                )),
-                _ => Err(RPCError::InvalidProtocol(
-                    "Unknown BEACON_BLOCK_BODIES encoding",
-                )),
-            },
-            _ => Err(RPCError::InvalidProtocol(
-                "Unknown BEACON_BLOCK_BODIES version",
-            )),
-        },
-        "beacon_chain_state" => match protocol_id.version {
-            "1" => match protocol_id.encoding {
-                "ssz" => Ok(RPCRequest::BeaconChainState(
-                    BeaconChainStateRequest::from_ssz_bytes(&packet)?,
-                )),
-                _ => Err(RPCError::InvalidProtocol(
-                    "Unknown BEACON_CHAIN_STATE encoding",
-                )),
-            },
-            _ => Err(RPCError::InvalidProtocol(
-                "Unknown BEACON_CHAIN_STATE version",
-            )),
-        },
-    }
-}
-
-/// Decodes a response that was received on the same stream as a request. The response type should
-/// therefore match the request protocol type.
-fn decode_response(packet: Vec<u8>, protocol: RawProtocolId) -> Result<RPCResponse, RPCError> {
-    let protocol_id = ProtocolId::from_bytes(protocol)?;
-
-    match protocol_id.message_name {
-        "hello" => match protocol_id.version {
-            "1" => match protocol_id.encoding {
-                "ssz" => Ok(RPCResponse::Hello(HelloMessage::from_ssz_bytes(&packet)?)),
-                _ => Err(RPCError::InvalidProtocol("Unknown HELLO encoding")),
-            },
-            _ => Err(RPCError::InvalidProtocol("Unknown HELLO version")),
-        },
-        "goodbye" => Err(RPCError::Custom("GOODBYE should not have a response")),
-        "beacon_block_roots" => match protocol_id.version {
-            "1" => match protocol_id.encoding {
-                "ssz" => Ok(RPCResponse::BeaconBlockRoots(
-                    BeaconBlockRootsResponse::from_ssz_bytes(&packet)?,
-                )),
-                _ => Err(RPCError::InvalidProtocol(
-                    "Unknown BEACON_BLOCK_ROOTS encoding",
-                )),
-            },
-            _ => Err(RPCError::InvalidProtocol(
-                "Unknown BEACON_BLOCK_ROOTS version",
-            )),
-        },
-        "beacon_block_headers" => match protocol_id.version {
-            "1" => match protocol_id.encoding {
-                "ssz" => Ok(RPCResponse::BeaconBlockHeaders(
-                    BeaconBlockHeadersResponse { headers: packet },
-                )),
-                _ => Err(RPCError::InvalidProtocol(
-                    "Unknown BEACON_BLOCK_HEADERS encoding",
-                )),
-            },
-            _ => Err(RPCError::InvalidProtocol(
-                "Unknown BEACON_BLOCK_HEADERS version",
-            )),
-        },
-        "beacon_block_bodies" => match protocol_id.version {
-            "1" => match protocol_id.encoding {
-                "ssz" => Ok(RPCResponse::BeaconBlockBodies(BeaconBlockBodiesResponse {
-                    block_bodies: packet,
-                })),
-                _ => Err(RPCError::InvalidProtocol(
-                    "Unknown BEACON_BLOCK_BODIES encoding",
-                )),
-            },
-            _ => Err(RPCError::InvalidProtocol(
-                "Unknown BEACON_BLOCK_BODIES version",
-            )),
-        },
-        "beacon_chain_state" => match protocol_id.version {
-            "1" => match protocol_id.encoding {
-                "ssz" => Ok(BeaconChainStateRequest::from_ssz_bytes(&packet)?),
-                _ => Err(RPCError::InvalidProtocol(
-                    "Unknown BEACON_CHAIN_STATE encoding",
-                )),
-            },
-            _ => Err(RPCError::InvalidProtocol(
-                "Unknown BEACON_CHAIN_STATE version",
-            )),
-        },
+        .map_err(RPCError::from)
     }
 }
 
@@ -367,8 +480,12 @@ pub enum RPCError {
     ReadError(upgrade::ReadOneError),
     /// Error when decoding the raw buffer from ssz.
     SSZDecodeError(ssz::DecodeError),
-    /// Invalid Protocol ID
+    /// Invalid Protocol ID.
     InvalidProtocol(&'static str),
+    /// IO Error.
+    IoError(io::Error),
+    /// Waiting for a request/response timed out, or timer error'd.
+    StreamTimeout,
     /// Custom message.
     Custom(String),
 }
@@ -386,3 +503,46 @@ impl From<ssz::DecodeError> for RPCError {
         RPCError::SSZDecodeError(err)
     }
 }
+
+impl<T> From<tokio::timer::timeout::Error<T>> for RPCError {
+    fn from(err: tokio::timer::timeout::Error<T>) -> Self {
+        if err.is_elapsed() {
+            RPCError::StreamTimeout
+        } else {
+            RPCError::Custom("Stream timer failed".into())
+        }
+    }
+}
+
+impl From<io::Error> for RPCError {
+    fn from(err: io::Error) -> Self {
+        RPCError::IoError(err)
+    }
+}
+
+// Error trait is required for `ProtocolsHandler`
+impl std::fmt::Display for RPCError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match *self {
+            RPCError::ReadError(ref err) => write!(f, "Error while reading from socket: {}", err),
+            RPCError::SSZDecodeError(ref err) => write!(f, "Error while decoding ssz: {:?}", err),
+            RPCError::InvalidProtocol(ref err) => write!(f, "Invalid Protocol: {}", err),
+            RPCError::IoError(ref err) => write!(f, "IO Error: {}", err),
+            RPCError::StreamTimeout => write!(f, "Stream Timeout"),
+            RPCError::Custom(ref err) => write!(f, "{}", err),
+        }
+    }
+}
+
+impl std::error::Error for RPCError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match *self {
+            RPCError::ReadError(ref err) => Some(err),
+            RPCError::SSZDecodeError(ref err) => None,
+            RPCError::InvalidProtocol(ref err) => None,
+            RPCError::IoError(ref err) => Some(err),
+            RPCError::StreamTimeout => None,
+            RPCError::Custom(ref err) => None,
+        }
+    }
+}