nicer processing of message types

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" }
 

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;

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<MsgHeader> 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!(),
-			},
-		})
 	}
 }