Merge pull request #496 from stevenroose/exhaustive-enums

Make Inventory and NetworkMessage enums exhaustive

src/consensus/encode.rs

@@ -80,12 +80,6 @@ pub enum Error {
     ParseFailed(&'static str),
     /// Unsupported Segwit flag
     UnsupportedSegwitFlag(u8),
-    /// Unrecognized network command with its length
-    UnrecognizedNetworkCommand(String, usize),
-    /// Invalid Inventory type
-    UnknownInventoryType(u32),
-    /// The network command is longer than the maximum allowed (12 chars)
-    NetworkCommandTooLong(String),
 }
 
 impl fmt::Display for Error {
@@ -104,10 +98,6 @@ impl fmt::Display for Error {
             Error::ParseFailed(ref e) => write!(f, "parse failed: {}", e),
             Error::UnsupportedSegwitFlag(ref swflag) => write!(f,
                 "unsupported segwit version: {}", swflag),
-            Error::UnrecognizedNetworkCommand(ref nwcmd, _) => write!(f,
-                "unrecognized network command: {}", nwcmd),
-            Error::UnknownInventoryType(ref tp) => write!(f, "Unknown Inventory type: {}", tp),
-            Error::NetworkCommandTooLong(ref cmd) => write!(f, "Network Command too long: {}", cmd),
         }
     }
 }
@@ -123,10 +113,7 @@ impl error::Error for Error {
             | Error::NonMinimalVarInt
             | Error::UnknownNetworkMagic(..)
             | Error::ParseFailed(..)
-            | Error::UnsupportedSegwitFlag(..)
+            | Error::UnsupportedSegwitFlag(..) => None,
-            | Error::UnrecognizedNetworkCommand(..)
-            | Error::UnknownInventoryType(..)
-            | Error::NetworkCommandTooLong(..) => None,
         }
     }
 }

src/network/message.rs

@@ -19,7 +19,7 @@
 //! also defines (de)serialization routines for many primitives.
 //!
 
-use std::{io, iter, mem, fmt};
+use std::{fmt, io, iter, mem, str};
 use std::borrow::Cow;
 use std::io::Cursor;
 
@@ -42,24 +42,31 @@ pub const MAX_INV_SIZE: usize = 50_000;
 #[derive(PartialEq, Eq, Clone, Debug)]
 pub struct CommandString(Cow<'static, str>);
 
+impl CommandString {
+    /// Convert from various string types into a [CommandString].
+    ///
+    /// Supported types are:
+    /// - `&'static str`
+    /// - `String`
+    ///
+    /// Returns an empty error if and only if the string is
+    /// larger than 12 characters in length.
+    pub fn try_from<S: Into<Cow<'static, str>>>(s: S) -> Result<CommandString, ()> {
+        let cow = s.into();
+        if cow.as_ref().len() > 12 {
+            Err(())
+        } else {
+            Ok(CommandString(cow))
+        }
+    }
+}
+
 impl fmt::Display for CommandString {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.write_str(self.0.as_ref())
     }
 }
 
-impl From<&'static str> for CommandString {
-    fn from(f: &'static str) -> Self {
-        CommandString(f.into())
-    }
-}
-
-impl From<String> for CommandString {
-    fn from(f: String) -> Self {
-        CommandString(f.into())
-    }
-}
-
 impl AsRef<str> for CommandString {
     fn as_ref(&self) -> &str {
         self.0.as_ref()
@@ -74,9 +81,7 @@ impl Encodable for CommandString {
     ) -> Result<usize, encode::Error> {
         let mut rawbytes = [0u8; 12];
         let strbytes = self.0.as_bytes();
-        if strbytes.len() > 12 {
+        debug_assert!(strbytes.len() <= 12);
-            return Err(encode::Error::NetworkCommandTooLong(self.0.clone().into_owned()));
-        }
         rawbytes[..strbytes.len()].clone_from_slice(&strbytes[..]);
         rawbytes.consensus_encode(s)
     }
@@ -104,9 +109,9 @@ pub struct RawNetworkMessage {
     pub payload: NetworkMessage
 }
 
-#[derive(Clone, PartialEq, Eq, Debug)]
 /// A Network message payload. Proper documentation is available on at
 /// [Bitcoin Wiki: Protocol Specification](https://en.bitcoin.it/wiki/Protocol_specification)
+#[derive(Clone, PartialEq, Eq, Debug)]
 pub enum NetworkMessage {
     /// `version`
     Version(message_network::VersionMessage),
@@ -168,10 +173,22 @@ pub enum NetworkMessage {
     AddrV2(Vec<AddrV2Message>),
     /// `sendaddrv2`
     SendAddrV2,
+
+    /// Any other message.
+    Unknown {
+        /// The command of this message.
+        command: CommandString,
+        /// The payload of this message.
+        payload: Vec<u8>,
+    }
 }
 
 impl NetworkMessage {
-    /// Return the message command. This is useful for debug outputs.
+    /// Return the message command as a static string reference.
+    ///
+    /// This returns `"unknown"` for [NetworkMessage::Unknown],
+    /// regardless of the actual command in the unknown message.
+    /// Use the [command] method to get the command for unknown messages.
     pub fn cmd(&self) -> &'static str {
         match *self {
             NetworkMessage::Version(_) => "version",
@@ -202,17 +219,25 @@ impl NetworkMessage {
             NetworkMessage::WtxidRelay => "wtxidrelay",
             NetworkMessage::AddrV2(_) => "addrv2",
             NetworkMessage::SendAddrV2 => "sendaddrv2",
+            NetworkMessage::Unknown { .. } => "unknown",
         }
     }
 
     /// Return the CommandString for the message command.
     pub fn command(&self) -> CommandString {
-        self.cmd().into()
+        match *self {
+            NetworkMessage::Unknown { command: ref c, .. } => c.clone(),
+            _ => CommandString::try_from(self.cmd()).expect("cmd returns valid commands")
+        }
     }
 }
 
 impl RawNetworkMessage {
-    /// Return the message command. This is useful for debug outputs.
+    /// Return the message command as a static string reference.
+    ///
+    /// This returns `"unknown"` for [NetworkMessage::Unknown],
+    /// regardless of the actual command in the unknown message.
+    /// Use the [command] method to get the command for unknown messages.
     pub fn cmd(&self) -> &'static str {
         self.payload.cmd()
     }
@@ -276,8 +301,9 @@ impl Encodable for RawNetworkMessage {
             | NetworkMessage::SendHeaders
             | NetworkMessage::MemPool
             | NetworkMessage::GetAddr
-            | NetworkMessage::WtxidRelay => vec![],
+            | NetworkMessage::WtxidRelay
             | NetworkMessage::SendAddrV2 => vec![],
+            NetworkMessage::Unknown { payload: ref data, .. } => serialize(data),
         }).consensus_encode(&mut s)?;
         Ok(len)
     }
@@ -309,12 +335,11 @@ impl Decodable for HeaderDeserializationWrapper {
 impl Decodable for RawNetworkMessage {
     fn consensus_decode<D: io::Read>(mut d: D) -> Result<Self, encode::Error> {
         let magic = Decodable::consensus_decode(&mut d)?;
-        let cmd = CommandString::consensus_decode(&mut d)?.0;
+        let cmd = CommandString::consensus_decode(&mut d)?;
         let raw_payload = CheckedData::consensus_decode(&mut d)?.0;
-        let raw_payload_len = raw_payload.len();
 
         let mut mem_d = Cursor::new(raw_payload);
-        let payload = match &cmd[..] {
+        let payload = match &cmd.0[..] {
             "version" => NetworkMessage::Version(Decodable::consensus_decode(&mut mem_d)?),
             "verack"  => NetworkMessage::Verack,
             "addr"    => NetworkMessage::Addr(Decodable::consensus_decode(&mut mem_d)?),
@@ -345,7 +370,10 @@ impl Decodable for RawNetworkMessage {
             "wtxidrelay" => NetworkMessage::WtxidRelay,
             "addrv2" => NetworkMessage::AddrV2(Decodable::consensus_decode(&mut mem_d)?),
             "sendaddrv2" => NetworkMessage::SendAddrV2,
-            _ => return Err(encode::Error::UnrecognizedNetworkCommand(cmd.into_owned(), 4 + 12 + 4 + 4 + raw_payload_len)), // magic + msg str + payload len + checksum + payload
+            _ => NetworkMessage::Unknown {
+                command: cmd,
+                payload: mem_d.into_inner(),
+            }
         };
         Ok(RawNetworkMessage {
             magic: magic,
@@ -356,11 +384,10 @@ impl Decodable for RawNetworkMessage {
 
 #[cfg(test)]
 mod test {
-    use std::io;
     use std::net::Ipv4Addr;
     use super::{RawNetworkMessage, NetworkMessage, CommandString};
     use network::constants::ServiceFlags;
-    use consensus::encode::{Encodable, deserialize, deserialize_partial, serialize};
+    use consensus::encode::{deserialize, deserialize_partial, serialize};
     use hashes::hex::FromHex;
     use hashes::sha256d::Hash;
     use hashes::Hash as HashTrait;
@@ -407,7 +434,7 @@ mod test {
             NetworkMessage::GetCFCheckpt(GetCFCheckpt{filter_type: 17, stop_hash: hash([25u8; 32]).into()}),
             NetworkMessage::CFCheckpt(CFCheckpt{filter_type: 27, stop_hash: hash([77u8; 32]).into(), filter_headers: vec![hash([3u8; 32]).into(), hash([99u8; 32]).into()]}),
             NetworkMessage::Alert(vec![45,66,3,2,6,8,9,12,3,130]),
-            NetworkMessage::Reject(Reject{message: "Test reject".into(), ccode: RejectReason::Duplicate, reason: "Cause".into(), hash: hash([255u8; 32])}),
+            NetworkMessage::Reject(Reject{message: CommandString::try_from("Test reject").unwrap(), ccode: RejectReason::Duplicate, reason: "Cause".into(), hash: hash([255u8; 32])}),
             NetworkMessage::FeeFilter(1000),
             NetworkMessage::WtxidRelay,
             NetworkMessage::AddrV2(vec![AddrV2Message{ addr: AddrV2::Ipv4(Ipv4Addr::new(127, 0, 0, 1)), port: 0, services: ServiceFlags::NONE, time: 0 }]),
@@ -422,21 +449,20 @@ mod test {
     }
 
     #[test]
-    fn serialize_commandstring_test() {
+    fn commandstring_test() {
+        // Test converting.
+        assert_eq!(CommandString::try_from("AndrewAndrew").unwrap().as_ref(), "AndrewAndrew");
+        assert!(CommandString::try_from("AndrewAndrewA").is_err());
+
+        // Test serializing.
         let cs = CommandString("Andrew".into());
         assert_eq!(serialize(&cs), vec![0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0, 0]);
 
-        // Test oversized one.
+        // Test deserializing
-        let mut encoder = io::Cursor::new(vec![]);
-        assert!(CommandString("AndrewAndrewA".into()).consensus_encode(&mut encoder).is_err());
-    }
-
-    #[test]
-    fn deserialize_commandstring_test() {
         let cs: Result<CommandString, _> = deserialize(&[0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0, 0]);
         assert!(cs.is_ok());
         assert_eq!(cs.as_ref().unwrap().to_string(), "Andrew".to_owned());
-        assert_eq!(cs.unwrap(), "Andrew".into());
+        assert_eq!(cs.unwrap(), CommandString::try_from("Andrew").unwrap());
 
         let short_cs: Result<CommandString, _> = deserialize(&[0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0]);
         assert!(short_cs.is_err());

src/network/message_blockdata.rs

@@ -41,6 +41,13 @@ pub enum Inventory {
     WitnessTransaction(Txid),
     /// Witness Block
     WitnessBlock(BlockHash),
+    /// Unknown inventory type
+    Unknown {
+        /// The inventory item type.
+        inv_type: u32,
+        /// The hash of the inventory item
+        hash: [u8; 32],
+    }
 }
 
 impl Encodable for Inventory {
@@ -62,6 +69,7 @@ impl Encodable for Inventory {
             Inventory::WTx(w) => encode_inv!(5, w),
             Inventory::WitnessTransaction(ref t) => encode_inv!(0x40000001, t),
             Inventory::WitnessBlock(ref b) => encode_inv!(0x40000002, b),
+            Inventory::Unknown { inv_type: t, hash: ref d } => encode_inv!(t, d),
         })
     }
 }
@@ -77,7 +85,10 @@ impl Decodable for Inventory {
             5 => Inventory::WTx(Decodable::consensus_decode(&mut d)?),
             0x40000001 => Inventory::WitnessTransaction(Decodable::consensus_decode(&mut d)?),
             0x40000002 => Inventory::WitnessBlock(Decodable::consensus_decode(&mut d)?),
-            tp => return Err(encode::Error::UnknownInventoryType(tp)),
+            tp => Inventory::Unknown {
+                inv_type: tp,
+                hash: Decodable::consensus_decode(&mut d)?,
+            }
         })
     }
 }

src/network/message_network.rs

@@ -84,8 +84,8 @@ impl_consensus_encoding!(VersionMessage, version, services, timestamp,
                          receiver, sender, nonce,
                          user_agent, start_height, relay);
 
-#[derive(PartialEq, Eq, Clone, Copy, Debug)]
 /// message rejection reason as a code
+#[derive(PartialEq, Eq, Clone, Copy, Debug)]
 pub enum RejectReason {
     /// malformed message
     Malformed = 0x01,

src/network/stream_reader.rs

@@ -68,10 +68,6 @@ impl<R: Read> StreamReader<R> {
                         return Err(encode::Error::Io(io::Error::from(io::ErrorKind::UnexpectedEof)));
                     }
                 },
-                Err(encode::Error::UnrecognizedNetworkCommand(message, len)) => {
-                    self.unparsed.drain(..len);
-                    return Err(encode::Error::UnrecognizedNetworkCommand(message, len))
-                },
                 Err(err) => return Err(err),
                 // We have successfully read from the buffer
                 Ok((message, index)) => {