diff --git a/p2p/Cargo.toml b/p2p/Cargo.toml index 80af62f02..f0062bd01 100644 --- a/p2p/Cargo.toml +++ b/p2p/Cargo.toml @@ -10,6 +10,8 @@ log = "^0.3" rand = "^0.3" mioco = "^0.8" time = "^0.1" +enum_primitive = "^0.1.0" +num = "^0.1.36" grin_core = { path = "../core" } diff --git a/p2p/src/lib.rs b/p2p/src/lib.rs index ed1dd38d5..924685207 100644 --- a/p2p/src/lib.rs +++ b/p2p/src/lib.rs @@ -24,6 +24,8 @@ #[macro_use] extern crate bitflags; #[macro_use] +extern crate enum_primitive; +#[macro_use] extern crate grin_core as core; #[macro_use] extern crate log; @@ -31,6 +33,7 @@ extern crate log; extern crate mioco; extern crate rand; extern crate time; +extern crate num; mod types; mod msg; diff --git a/p2p/src/msg.rs b/p2p/src/msg.rs index 8a283be95..c8ad81e56 100644 --- a/p2p/src/msg.rs +++ b/p2p/src/msg.rs @@ -16,6 +16,8 @@ use std::net::*; +use num::FromPrimitive; + use core::ser::{self, Writeable, Readable, Writer, Reader, Error}; use types::*; @@ -34,6 +36,7 @@ pub enum ErrCodes { } /// Types of messages +enum_from_primitive! { #[derive(Clone, Copy)] pub enum Type { Error, @@ -43,8 +46,7 @@ pub enum Type { Pong, GetPeerAddrs, PeerAddrs, - /// Never used over the network but to detect unrecognized types. - MaxMsgType, +} } /// Header of any protocol message, used to identify incoming messages. @@ -62,7 +64,7 @@ impl MsgHeader { } pub fn acceptable(&self) -> bool { - (self.msg_type as u8) < (Type::MaxMsgType as u8) + Type::from_u8(self.msg_type as u8).is_some() } } @@ -81,23 +83,10 @@ impl Readable for MsgHeader { try!(reader.expect_u8(MAGIC[0])); try!(reader.expect_u8(MAGIC[1])); let t = try!(reader.read_u8()); - if t >= (Type::MaxMsgType as u8) { - return Err(ser::Error::CorruptedData); + match Type::from_u8(t) { + Some(ty) => Ok(MsgHeader {magic: MAGIC, msg_type: ty}), + None => Err(ser::Error::CorruptedData) } - Ok(MsgHeader { - magic: MAGIC, - msg_type: match t { - // TODO this is rather ugly, think of a better way - 0 => Type::Error, - 1 => Type::Hand, - 2 => Type::Shake, - 3 => Type::Ping, - 4 => Type::Pong, - 5 => Type::GetPeerAddrs, - 6 => Type::PeerAddrs, - _ => panic!(), - }, - }) } }