Restructure protobuf message parsing to avoid panics

src/blob.rs

@@ -10,6 +10,7 @@ use proto::fileformat;
 use std::fs::File;
 use std::io::{BufReader, ErrorKind, Read};
 use std::path::Path;
+use util::{parse_message_from_bytes, parse_message_from_reader};
 
 #[cfg(feature = "system-libz")]
 use flate2::read::ZlibDecoder;
@@ -179,7 +180,7 @@ impl<R: Read> Iterator for BlobReader<R> {
             },
         };
 
-        let header: fileformat::BlobHeader = match protobuf::parse_from_reader(&mut self.reader.by_ref().take(size)) {
+        let header: fileformat::BlobHeader = match parse_message_from_reader(&mut self.reader.by_ref().take(size)) {
             Ok(header) => header,
             Err(e) => {
                 self.last_blob_ok = false;
@@ -187,7 +188,7 @@ impl<R: Read> Iterator for BlobReader<R> {
             },
         };
 
-        let blob: fileformat::Blob = match protobuf::parse_from_reader(&mut self.reader.by_ref().take(header.get_datasize() as u64)) {
+        let blob: fileformat::Blob = match parse_message_from_reader(&mut self.reader.by_ref().take(header.get_datasize() as u64)) {
             Ok(blob) => blob,
             Err(e) => {
                 self.last_blob_ok = false;
@@ -203,10 +204,10 @@ impl<R: Read> Iterator for BlobReader<R> {
 pub(crate) fn decode_blob<T>(blob: &fileformat::Blob) -> Result<T>
     where T: protobuf::Message + protobuf::MessageStatic {
     if blob.has_raw() {
-        protobuf::parse_from_bytes(blob.get_raw()).chain_err(|| "Could not parse raw data")
+        parse_message_from_bytes(blob.get_raw()).chain_err(|| "Could not parse raw data")
     } else if blob.has_zlib_data() {
         let mut decoder = ZlibDecoder::new(blob.get_zlib_data());
-        protobuf::parse_from_reader(&mut decoder).chain_err(|| "Could not parse zlib data")
+        parse_message_from_reader(&mut decoder).chain_err(|| "Could not parse zlib data")
     } else {
         bail!("Blob is missing fields 'raw' and 'zlib_data")
     }
@@ -216,10 +217,10 @@ pub(crate) fn decode_blob<T>(blob: &fileformat::Blob) -> Result<T>
 pub(crate) fn decode_blob<T>(blob: &fileformat::Blob) -> Result<T>
     where T: protobuf::Message + protobuf::MessageStatic {
     if blob.has_raw() {
-        protobuf::parse_from_bytes(blob.get_raw()).chain_err(|| "Could not parse raw data")
+        parse_message_from_bytes(blob.get_raw()).chain_err(|| "Could not parse raw data")
     } else if blob.has_zlib_data() {
         let mut decoder = DeflateDecoder::from_zlib(blob.get_zlib_data());
-        protobuf::parse_from_reader(&mut decoder).chain_err(|| "Could not parse zlib data")
+        parse_message_from_reader(&mut decoder).chain_err(|| "Could not parse zlib data")
     } else {
         bail!("Blob is missing fields 'raw' and 'zlib_data")
     }

src/lib.rs

@@ -101,3 +101,4 @@ pub mod block;
 pub mod dense;
 pub mod elements;
 pub mod mmap_blob;
+mod util;

src/mmap_blob.rs

@@ -5,13 +5,14 @@ extern crate byteorder;
 extern crate memmap;
 
 use blob::{BlobDecode, BlobType, decode_blob};
+use block::{HeaderBlock, PrimitiveBlock};
 use byteorder::ByteOrder;
 use errors::*;
-use block::{HeaderBlock, PrimitiveBlock};
 use proto::{fileformat, osmformat};
 use self::fileformat::BlobHeader;
 use std::fs::File;
 use std::path::Path;
+use util::parse_message_from_bytes;
 
 
 /// A read-only memory map.
@@ -84,7 +85,7 @@ impl<'a> MmapBlob<'a> {
     /// Decodes the Blob and tries to obtain the inner content (usually a `HeaderBlock` or a
     /// `PrimitiveBlock`). This operation might involve an expensive decompression step.
     pub fn decode(&'a self) -> Result<BlobDecode<'a>> {
-        let blob: fileformat::Blob = protobuf::parse_from_bytes(self.data)
+        let blob: fileformat::Blob = parse_message_from_bytes(self.data)
             .chain_err(|| "failed to parse Blob")?;
         match self.header.get_field_type() {
             "OSMHeader" => {
@@ -169,7 +170,7 @@ impl<'a> Iterator for MmapBlobReader<'a> {
             return Some(Err(Error::from_kind(ErrorKind::Io(io_error))));
         }
 
-        let header: BlobHeader = match protobuf::parse_from_bytes(&slice[4..(4 + header_size)]) {
+        let header: BlobHeader = match parse_message_from_bytes(&slice[4..(4 + header_size)]) {
             Ok(x) => x,
             Err(e) => {
                 self.last_blob_ok = false;

src/util.rs

@@ -0,0 +1,34 @@
+use errors::*;
+use std::io::Read;
+
+
+pub(crate) fn parse_message_from_bytes<M>(bytes: &[u8]) -> Result<M>
+    where M: ::protobuf::Message + ::protobuf::MessageStatic
+{
+    let mut stream = ::protobuf::CodedInputStream::from_bytes(bytes);
+    let mut message: M = ::protobuf::MessageStatic::new();
+    message.merge_from(&mut stream)?;
+
+    if message.is_initialized() {
+        Ok(message)
+    } else {
+        Err(::protobuf::ProtobufError::message_not_initialized("").into())
+    }
+}
+
+pub(crate) fn parse_message_from_reader<R, M>(reader: &mut R) -> Result<M>
+    where R: Read,
+          M: ::protobuf::Message + ::protobuf::MessageStatic,
+{
+    let mut stream = ::protobuf::CodedInputStream::new(reader);
+    let mut message: M = ::protobuf::MessageStatic::new();
+    message.merge_from(&mut stream)?;
+
+    stream.check_eof()?;
+
+    if message.is_initialized() {
+        Ok(message)
+    } else {
+        Err(::protobuf::ProtobufError::message_not_initialized("").into())
+    }
+}

tests/read.rs

@@ -119,3 +119,20 @@ fn mmap_read() {
         panic!("Unexpected blob type");
     }
 }
+
+#[test]
+fn decode_blob() {
+    let reader = BlobReader::from_path(TEST_FILE_PATH).unwrap();
+    let blobs = reader.collect::<Result<Vec<_>>>().unwrap();
+
+    assert_eq!(blobs.len(), 2);
+    assert_eq!(blobs[0].get_type(), BlobType::OsmHeader);
+    assert_eq!(blobs[1].get_type(), BlobType::OsmData);
+
+    // decoding to the wrong blob type should not panic, but produce an Err.
+    assert!(blobs[0].to_primitiveblock().is_err());
+    assert!(blobs[1].to_headerblock().is_err());
+
+    assert!(blobs[0].to_headerblock().is_ok());
+    assert!(blobs[1].to_primitiveblock().is_ok());
+}