Some Kademlia code cleanup. (#1101)

This commit is contained in:
Roman Borschel
2019-05-08 10:10:58 +02:00
committed by GitHub
parent 8537eb38b9
commit 61b236172b
4 changed files with 151 additions and 241 deletions

View File

@ -18,23 +18,24 @@
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.
//! Provides the `KadRequestMsg` and `KadResponseMsg` enums of all the possible messages
//! transmitted with the Kademlia protocol, and the `KademliaProtocolConfig` connection upgrade.
//! The Kademlia connection protocol upgrade and associated message types.
//!
//! The upgrade's output a `Sink + Stream` of messages.
//!
//! The `Stream` component is used to poll the underlying transport, and the `Sink` component is
//! used to send messages.
//! The connection protocol upgrade is provided by [`KademliaProtocolConfig`], with the
//! request and response types [`KadRequestMsg`] and [`KadResponseMsg`], respectively.
//! The upgrade's output is a `Sink + Stream` of messages. The `Stream` component is used
//! to poll the underlying transport for incoming messages, and the `Sink` component
//! is used to send messages to remote peers.
use bytes::BytesMut;
use crate::protobuf_structs;
use futures::{future, sink, stream, Sink, Stream};
use libp2p_core::{InboundUpgrade, Multiaddr, OutboundUpgrade, PeerId, UpgradeInfo, upgrade::Negotiated};
use codec::UviBytes;
use crate::protobuf_structs::dht as proto;
use futures::{future::{self, FutureResult}, sink, stream, Sink, Stream};
use libp2p_core::{Multiaddr, PeerId};
use libp2p_core::upgrade::{InboundUpgrade, OutboundUpgrade, UpgradeInfo, Negotiated};
use multihash::Multihash;
use protobuf::{self, Message};
use std::convert::TryFrom;
use std::io::{Error as IoError, ErrorKind as IoErrorKind};
use std::iter;
use std::{io, iter};
use tokio_codec::Framed;
use tokio_io::{AsyncRead, AsyncWrite};
use unsigned_varint::codec;
@ -52,10 +53,10 @@ pub enum KadConnectionType {
CannotConnect = 3,
}
impl From<protobuf_structs::dht::Message_ConnectionType> for KadConnectionType {
impl From<proto::Message_ConnectionType> for KadConnectionType {
#[inline]
fn from(raw: protobuf_structs::dht::Message_ConnectionType) -> KadConnectionType {
use crate::protobuf_structs::dht::Message_ConnectionType::{
fn from(raw: proto::Message_ConnectionType) -> KadConnectionType {
use proto::Message_ConnectionType::{
CAN_CONNECT, CANNOT_CONNECT, CONNECTED, NOT_CONNECTED
};
match raw {
@ -67,10 +68,10 @@ impl From<protobuf_structs::dht::Message_ConnectionType> for KadConnectionType {
}
}
impl Into<protobuf_structs::dht::Message_ConnectionType> for KadConnectionType {
impl Into<proto::Message_ConnectionType> for KadConnectionType {
#[inline]
fn into(self) -> protobuf_structs::dht::Message_ConnectionType {
use crate::protobuf_structs::dht::Message_ConnectionType::{
fn into(self) -> proto::Message_ConnectionType {
use proto::Message_ConnectionType::{
CAN_CONNECT, CANNOT_CONNECT, CONNECTED, NOT_CONNECTED
};
match self {
@ -93,19 +94,19 @@ pub struct KadPeer {
pub connection_ty: KadConnectionType,
}
impl KadPeer {
// Builds a `KadPeer` from its raw protobuf equivalent.
// TODO: use TryFrom once stable
fn from_peer(peer: &mut protobuf_structs::dht::Message_Peer) -> Result<KadPeer, IoError> {
// Builds a `KadPeer` from a corresponding protobuf message.
impl TryFrom<&mut proto::Message_Peer> for KadPeer {
type Error = io::Error;
fn try_from(peer: &mut proto::Message_Peer) -> Result<KadPeer, Self::Error> {
// TODO: this is in fact a CID; not sure if this should be handled in `from_bytes` or
// as a special case here
let node_id = PeerId::from_bytes(peer.get_id().to_vec())
.map_err(|_| IoError::new(IoErrorKind::InvalidData, "invalid peer id"))?;
.map_err(|_| invalid_data("invalid peer id"))?;
let mut addrs = Vec::with_capacity(peer.get_addrs().len());
for addr in peer.take_addrs().into_iter() {
let as_ma = Multiaddr::try_from(addr)
.map_err(|err| IoError::new(IoErrorKind::InvalidData, err))?;
let as_ma = Multiaddr::try_from(addr).map_err(invalid_data)?;
addrs.push(as_ma);
}
debug_assert_eq!(addrs.len(), addrs.capacity());
@ -120,9 +121,9 @@ impl KadPeer {
}
}
impl Into<protobuf_structs::dht::Message_Peer> for KadPeer {
fn into(self) -> protobuf_structs::dht::Message_Peer {
let mut out = protobuf_structs::dht::Message_Peer::new();
impl Into<proto::Message_Peer> for KadPeer {
fn into(self) -> proto::Message_Peer {
let mut out = proto::Message_Peer::new();
out.set_id(self.node_id.into_bytes());
for addr in self.multiaddrs {
out.mut_addrs().push(addr.to_vec());
@ -155,23 +156,22 @@ where
C: AsyncRead + AsyncWrite,
{
type Output = KadInStreamSink<Negotiated<C>>;
type Future = future::FutureResult<Self::Output, IoError>;
type Error = IoError;
type Future = FutureResult<Self::Output, io::Error>;
type Error = io::Error;
#[inline]
fn upgrade_inbound(self, incoming: Negotiated<C>, _: Self::Info) -> Self::Future {
let mut codec = codec::UviBytes::default();
let mut codec = UviBytes::default();
codec.set_max_len(4096);
future::ok(
Framed::new(incoming, codec)
.from_err::<IoError>()
.with::<_, fn(_) -> _, _>(|response| -> Result<_, IoError> {
.from_err()
.with::<_, fn(_) -> _, _>(|response| {
let proto_struct = resp_msg_to_proto(response);
proto_struct.write_to_bytes()
.map_err(|err| IoError::new(IoErrorKind::InvalidData, err.to_string()))
proto_struct.write_to_bytes().map_err(invalid_data)
})
.and_then::<fn(_) -> _, _>(|bytes: BytesMut| {
.and_then::<fn(_) -> _, _>(|bytes| {
let request = protobuf::parse_from_bytes(&bytes)?;
proto_to_req_msg(request)
}),
@ -184,25 +184,22 @@ where
C: AsyncRead + AsyncWrite,
{
type Output = KadOutStreamSink<Negotiated<C>>;
type Future = future::FutureResult<Self::Output, IoError>;
type Error = IoError;
type Future = FutureResult<Self::Output, io::Error>;
type Error = io::Error;
#[inline]
fn upgrade_outbound(self, incoming: Negotiated<C>, _: Self::Info) -> Self::Future {
let mut codec = codec::UviBytes::default();
let mut codec = UviBytes::default();
codec.set_max_len(4096);
future::ok(
Framed::new(incoming, codec)
.from_err::<IoError>()
.with::<_, fn(_) -> _, _>(|request| -> Result<_, IoError> {
.from_err()
.with::<_, fn(_) -> _, _>(|request| {
let proto_struct = req_msg_to_proto(request);
match proto_struct.write_to_bytes() {
Ok(msg) => Ok(msg),
Err(err) => Err(IoError::new(IoErrorKind::Other, err.to_string())),
}
proto_struct.write_to_bytes().map_err(invalid_data)
})
.and_then::<fn(_) -> _, _>(|bytes: BytesMut| {
.and_then::<fn(_) -> _, _>(|bytes| {
let response = protobuf::parse_from_bytes(&bytes)?;
proto_to_resp_msg(response)
}),
@ -211,27 +208,20 @@ where
}
/// Sink of responses and stream of requests.
pub type KadInStreamSink<S> = stream::AndThen<
sink::With<
stream::FromErr<Framed<S, codec::UviBytes<Vec<u8>>>, IoError>,
KadResponseMsg,
fn(KadResponseMsg) -> Result<Vec<u8>, IoError>,
Result<Vec<u8>, IoError>,
>,
fn(BytesMut) -> Result<KadRequestMsg, IoError>,
Result<KadRequestMsg, IoError>,
>;
pub type KadInStreamSink<S> = KadStreamSink<S, KadResponseMsg, KadRequestMsg>;
/// Sink of requests and stream of responses.
pub type KadOutStreamSink<S> = stream::AndThen<
pub type KadOutStreamSink<S> = KadStreamSink<S, KadRequestMsg, KadResponseMsg>;
pub type KadStreamSink<S, A, B> = stream::AndThen<
sink::With<
stream::FromErr<Framed<S, codec::UviBytes<Vec<u8>>>, IoError>,
KadRequestMsg,
fn(KadRequestMsg) -> Result<Vec<u8>, IoError>,
Result<Vec<u8>, IoError>,
stream::FromErr<Framed<S, UviBytes<Vec<u8>>>, io::Error>,
A,
fn(A) -> Result<Vec<u8>, io::Error>,
Result<Vec<u8>, io::Error>,
>,
fn(BytesMut) -> Result<KadResponseMsg, IoError>,
Result<KadResponseMsg, IoError>,
fn(BytesMut) -> Result<B, io::Error>,
Result<B, io::Error>,
>;
/// Request that we can send to a peer or that we received from a peer.
@ -284,31 +274,31 @@ pub enum KadResponseMsg {
},
}
// Turns a type-safe Kadmelia message into the corresponding raw protobuf message.
fn req_msg_to_proto(kad_msg: KadRequestMsg) -> protobuf_structs::dht::Message {
/// Converts a `KadRequestMsg` into the corresponding protobuf message for sending.
fn req_msg_to_proto(kad_msg: KadRequestMsg) -> proto::Message {
match kad_msg {
KadRequestMsg::Ping => {
let mut msg = protobuf_structs::dht::Message::new();
msg.set_field_type(protobuf_structs::dht::Message_MessageType::PING);
let mut msg = proto::Message::new();
msg.set_field_type(proto::Message_MessageType::PING);
msg
}
KadRequestMsg::FindNode { key } => {
let mut msg = protobuf_structs::dht::Message::new();
msg.set_field_type(protobuf_structs::dht::Message_MessageType::FIND_NODE);
let mut msg = proto::Message::new();
msg.set_field_type(proto::Message_MessageType::FIND_NODE);
msg.set_key(key.into_bytes());
msg.set_clusterLevelRaw(10);
msg
}
KadRequestMsg::GetProviders { key } => {
let mut msg = protobuf_structs::dht::Message::new();
msg.set_field_type(protobuf_structs::dht::Message_MessageType::GET_PROVIDERS);
let mut msg = proto::Message::new();
msg.set_field_type(proto::Message_MessageType::GET_PROVIDERS);
msg.set_key(key.into_bytes());
msg.set_clusterLevelRaw(10);
msg
}
KadRequestMsg::AddProvider { key, provider_peer } => {
let mut msg = protobuf_structs::dht::Message::new();
msg.set_field_type(protobuf_structs::dht::Message_MessageType::ADD_PROVIDER);
let mut msg = proto::Message::new();
msg.set_field_type(proto::Message_MessageType::ADD_PROVIDER);
msg.set_clusterLevelRaw(10);
msg.set_key(key.into_bytes());
msg.mut_providerPeers().push(provider_peer.into());
@ -317,17 +307,17 @@ fn req_msg_to_proto(kad_msg: KadRequestMsg) -> protobuf_structs::dht::Message {
}
}
// Turns a type-safe Kadmelia message into the corresponding raw protobuf message.
fn resp_msg_to_proto(kad_msg: KadResponseMsg) -> protobuf_structs::dht::Message {
/// Converts a `KadResponseMsg` into the corresponding protobuf message for sending.
fn resp_msg_to_proto(kad_msg: KadResponseMsg) -> proto::Message {
match kad_msg {
KadResponseMsg::Pong => {
let mut msg = protobuf_structs::dht::Message::new();
msg.set_field_type(protobuf_structs::dht::Message_MessageType::PING);
let mut msg = proto::Message::new();
msg.set_field_type(proto::Message_MessageType::PING);
msg
}
KadResponseMsg::FindNode { closer_peers } => {
let mut msg = protobuf_structs::dht::Message::new();
msg.set_field_type(protobuf_structs::dht::Message_MessageType::FIND_NODE);
let mut msg = proto::Message::new();
msg.set_field_type(proto::Message_MessageType::FIND_NODE);
msg.set_clusterLevelRaw(9);
for peer in closer_peers {
msg.mut_closerPeers().push(peer.into());
@ -338,8 +328,8 @@ fn resp_msg_to_proto(kad_msg: KadResponseMsg) -> protobuf_structs::dht::Message
closer_peers,
provider_peers,
} => {
let mut msg = protobuf_structs::dht::Message::new();
msg.set_field_type(protobuf_structs::dht::Message_MessageType::GET_PROVIDERS);
let mut msg = proto::Message::new();
msg.set_field_type(proto::Message_MessageType::GET_PROVIDERS);
msg.set_clusterLevelRaw(9);
for peer in closer_peers {
msg.mut_closerPeers().push(peer.into());
@ -352,96 +342,80 @@ fn resp_msg_to_proto(kad_msg: KadResponseMsg) -> protobuf_structs::dht::Message
}
}
/// Turns a raw Kademlia message into a type-safe message.
fn proto_to_req_msg(mut message: protobuf_structs::dht::Message) -> Result<KadRequestMsg, IoError> {
/// Converts a received protobuf message into a corresponding `KadRequestMsg`.
///
/// Fails if the protobuf message is not a valid and supported Kademlia request message.
fn proto_to_req_msg(mut message: proto::Message) -> Result<KadRequestMsg, io::Error> {
match message.get_field_type() {
protobuf_structs::dht::Message_MessageType::PING => Ok(KadRequestMsg::Ping),
proto::Message_MessageType::PING => Ok(KadRequestMsg::Ping),
protobuf_structs::dht::Message_MessageType::PUT_VALUE => {
Err(IoError::new(
IoErrorKind::InvalidData,
"received a PUT_VALUE message, but this is not supported by rust-libp2p yet",
))
}
proto::Message_MessageType::PUT_VALUE =>
Err(invalid_data("Unsupported: PUT_VALUE message.")),
protobuf_structs::dht::Message_MessageType::GET_VALUE => {
Err(IoError::new(
IoErrorKind::InvalidData,
"received a GET_VALUE message, but this is not supported by rust-libp2p yet",
))
}
proto::Message_MessageType::GET_VALUE =>
Err(invalid_data("Unsupported: GET_VALUE message.")),
protobuf_structs::dht::Message_MessageType::FIND_NODE => {
let key = PeerId::from_bytes(message.take_key()).map_err(|_| {
IoError::new(IoErrorKind::InvalidData, "invalid peer id in FIND_NODE")
})?;
proto::Message_MessageType::FIND_NODE => {
let key = PeerId::from_bytes(message.take_key())
.map_err(|_| invalid_data("Invalid peer id in FIND_NODE"))?;
Ok(KadRequestMsg::FindNode { key })
}
protobuf_structs::dht::Message_MessageType::GET_PROVIDERS => {
let key = Multihash::from_bytes(message.take_key())
.map_err(|err| IoError::new(IoErrorKind::InvalidData, err))?;
proto::Message_MessageType::GET_PROVIDERS => {
let key = Multihash::from_bytes(message.take_key()).map_err(invalid_data)?;
Ok(KadRequestMsg::GetProviders { key })
}
protobuf_structs::dht::Message_MessageType::ADD_PROVIDER => {
proto::Message_MessageType::ADD_PROVIDER => {
// TODO: for now we don't parse the peer properly, so it is possible that we get
// parsing errors for peers even when they are valid; we ignore these
// errors for now, but ultimately we should just error altogether
let provider_peer = message
.mut_providerPeers()
.iter_mut()
.filter_map(|peer| KadPeer::from_peer(peer).ok())
.next();
.find_map(|peer| KadPeer::try_from(peer).ok());
if let Some(provider_peer) = provider_peer {
let key = Multihash::from_bytes(message.take_key())
.map_err(|err| IoError::new(IoErrorKind::InvalidData, err))?;
let key = Multihash::from_bytes(message.take_key()).map_err(invalid_data)?;
Ok(KadRequestMsg::AddProvider { key, provider_peer })
} else {
Err(IoError::new(
IoErrorKind::InvalidData,
"received an ADD_PROVIDER message with no valid peer",
))
Err(invalid_data("ADD_PROVIDER message with no valid peer."))
}
}
}
}
/// Turns a raw Kademlia message into a type-safe message.
fn proto_to_resp_msg(
mut message: protobuf_structs::dht::Message,
) -> Result<KadResponseMsg, IoError> {
/// Converts a received protobuf message into a corresponding `KadResponseMessage`.
///
/// Fails if the protobuf message is not a valid and supported Kademlia response message.
fn proto_to_resp_msg(mut message: proto::Message) -> Result<KadResponseMsg, io::Error> {
match message.get_field_type() {
protobuf_structs::dht::Message_MessageType::PING => Ok(KadResponseMsg::Pong),
proto::Message_MessageType::PING => Ok(KadResponseMsg::Pong),
protobuf_structs::dht::Message_MessageType::GET_VALUE => {
Err(IoError::new(
IoErrorKind::InvalidData,
"received a GET_VALUE message, but this is not supported by rust-libp2p yet",
))
}
proto::Message_MessageType::GET_VALUE =>
Err(invalid_data("Unsupported: GET_VALUE message")),
protobuf_structs::dht::Message_MessageType::FIND_NODE => {
proto::Message_MessageType::FIND_NODE => {
let closer_peers = message
.mut_closerPeers()
.iter_mut()
.filter_map(|peer| KadPeer::from_peer(peer).ok())
.filter_map(|peer| KadPeer::try_from(peer).ok())
.collect::<Vec<_>>();
Ok(KadResponseMsg::FindNode { closer_peers })
}
protobuf_structs::dht::Message_MessageType::GET_PROVIDERS => {
proto::Message_MessageType::GET_PROVIDERS => {
let closer_peers = message
.mut_closerPeers()
.iter_mut()
.filter_map(|peer| KadPeer::from_peer(peer).ok())
.filter_map(|peer| KadPeer::try_from(peer).ok())
.collect::<Vec<_>>();
let provider_peers = message
.mut_providerPeers()
.iter_mut()
.filter_map(|peer| KadPeer::from_peer(peer).ok())
.filter_map(|peer| KadPeer::try_from(peer).ok())
.collect::<Vec<_>>();
Ok(KadResponseMsg::GetProviders {
@ -450,18 +424,22 @@ fn proto_to_resp_msg(
})
}
protobuf_structs::dht::Message_MessageType::PUT_VALUE => Err(IoError::new(
IoErrorKind::InvalidData,
"received an unexpected PUT_VALUE message",
)),
proto::Message_MessageType::PUT_VALUE =>
Err(invalid_data("received an unexpected PUT_VALUE message")),
protobuf_structs::dht::Message_MessageType::ADD_PROVIDER => Err(IoError::new(
IoErrorKind::InvalidData,
"received an unexpected ADD_PROVIDER message",
)),
proto::Message_MessageType::ADD_PROVIDER =>
Err(invalid_data("received an unexpected ADD_PROVIDER message"))
}
}
/// Creates an `io::Error` with `io::ErrorKind::InvalidData`.
fn invalid_data<E>(e: E) -> io::Error
where
E: Into<Box<dyn std::error::Error + Send + Sync>>
{
io::Error::new(io::ErrorKind::InvalidData, e)
}
#[cfg(test)]
mod tests {