milksad
/
rust-bitcoin-unsafe-fast
14
0
Fork 0
Code Issues Pull Requests Activity

Merge pull request #496 from stevenroose/exhaustive-enums 

Make Inventory and NetworkMessage enums exhaustive
This commit is contained in:
Andrew Poelstra 2020-12-21 15:17:22 +00:00 committed by GitHub
parent a6264cfca6 767b14f696
commit 65c972b0b0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 77 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
src/consensus/encode.rs
View File

@ -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::UnrecognizedNetworkCommand(..)
| Error::UnknownInventoryType(..)
| Error::NetworkCommandTooLong(..) => None,
| Error::UnsupportedSegwitFlag(..) => None,
}
}
}

100
src/network/message.rs
View File

@ -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 {
return Err(encode::Error::NetworkCommandTooLong(self.0.clone().into_owned()));
}
debug_assert!(strbytes.len() <= 12);
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.
let mut encoder = io::Cursor::new(vec![]);
assert!(CommandString("AndrewAndrewA".into()).consensus_encode(&mut encoder).is_err());
}
#[test]
fn deserialize_commandstring_test() {
// Test deserializing
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());

13
src/network/message_blockdata.rs
View File

@ -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)?,
}
})
}
}

2
src/network/message_network.rs
View File

@ -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,

4
src/network/stream_reader.rs
View File

@ -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)) => {