From a69614a89aab501a94433c1b565a6a22ab8e1fcb Mon Sep 17 00:00:00 2001
From: Titouan Rigoudy <titouan.rigoudy@gmail.com>
Date: Wed, 27 Dec 2017 17:08:24 -0500
Subject: [PATCH] Return an error when decoding unknown peer message code.

---
 src/proto/peer/message.rs | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/proto/peer/message.rs b/src/proto/peer/message.rs
index 140358c..2af7a4c 100644
--- a/src/proto/peer/message.rs
+++ b/src/proto/peer/message.rs
@@ -47,12 +47,14 @@ impl ProtoDecode for Message {
             CODE_PIERCE_FIREWALL => {
                 let val = decoder.decode_u32()?;
                 Message::PierceFirewall(val)
-            }
+            },
             CODE_PEER_INIT => {
                 let peer_init = PeerInit::decode(decoder)?;
                 Message::PeerInit(peer_init)
-            }
-            code => Message::Unknown(code),
+            },
+            _ => {
+                return Err(DecodeError::UnknownCodeError(code));
+            },
         };
         Ok(message)
     }
@@ -64,11 +66,11 @@ impl ProtoEncode for Message {
             Message::PierceFirewall(token) => {
                 encoder.encode_u32(CODE_PIERCE_FIREWALL)?;
                 encoder.encode_u32(token)?;
-            }
+            },
             Message::PeerInit(ref request) => {
                 encoder.encode_u32(CODE_PEER_INIT)?;
                 request.encode(encoder)?;
-            }
+            },
             Message::Unknown(_) => unreachable!(),
         }
         Ok(())
@@ -166,6 +168,16 @@ mod tests {
         assert_eq!(output, input);
     }
 
+    #[test]
+    #[should_panic]
+    fn invalid_code() {
+        let mut bytes = BytesMut::new();
+        ProtoEncoder::new(&mut bytes).encode_u32(1337).unwrap();
+
+        let mut cursor = io::Cursor::new(bytes);
+        Message::decode(&mut ProtoDecoder::new(&mut cursor)).unwrap();
+    }
+
     #[test]
     fn roundtrip_pierce_firewall() {
         roundtrip(Message::PierceFirewall(1337))