Selaa lähdekoodia

Restructure protobuf message parsing to avoid panics

Johannes Hofmann 8 vuotta sitten
vanhempi
commit
33b42843f1
5 muutettua tiedostoa jossa 63 lisäystä ja 9 poistoa
  1. 7
    6
      src/blob.rs
  2. 1
    0
      src/lib.rs
  3. 4
    3
      src/mmap_blob.rs
  4. 34
    0
      src/util.rs
  5. 17
    0
      tests/read.rs

+ 7
- 6
src/blob.rs Näytä tiedosto

10
 use std::fs::File;
10
 use std::fs::File;
11
 use std::io::{BufReader, ErrorKind, Read};
11
 use std::io::{BufReader, ErrorKind, Read};
12
 use std::path::Path;
12
 use std::path::Path;
13
+use util::{parse_message_from_bytes, parse_message_from_reader};
13
 
14
 
14
 #[cfg(feature = "system-libz")]
15
 #[cfg(feature = "system-libz")]
15
 use flate2::read::ZlibDecoder;
16
 use flate2::read::ZlibDecoder;
179
             },
180
             },
180
         };
181
         };
181
 
182
 
182
-        let header: fileformat::BlobHeader = match protobuf::parse_from_reader(&mut self.reader.by_ref().take(size)) {
183
+        let header: fileformat::BlobHeader = match parse_message_from_reader(&mut self.reader.by_ref().take(size)) {
183
             Ok(header) => header,
184
             Ok(header) => header,
184
             Err(e) => {
185
             Err(e) => {
185
                 self.last_blob_ok = false;
186
                 self.last_blob_ok = false;
187
             },
188
             },
188
         };
189
         };
189
 
190
 
190
-        let blob: fileformat::Blob = match protobuf::parse_from_reader(&mut self.reader.by_ref().take(header.get_datasize() as u64)) {
191
+        let blob: fileformat::Blob = match parse_message_from_reader(&mut self.reader.by_ref().take(header.get_datasize() as u64)) {
191
             Ok(blob) => blob,
192
             Ok(blob) => blob,
192
             Err(e) => {
193
             Err(e) => {
193
                 self.last_blob_ok = false;
194
                 self.last_blob_ok = false;
203
 pub(crate) fn decode_blob<T>(blob: &fileformat::Blob) -> Result<T>
204
 pub(crate) fn decode_blob<T>(blob: &fileformat::Blob) -> Result<T>
204
     where T: protobuf::Message + protobuf::MessageStatic {
205
     where T: protobuf::Message + protobuf::MessageStatic {
205
     if blob.has_raw() {
206
     if blob.has_raw() {
206
-        protobuf::parse_from_bytes(blob.get_raw()).chain_err(|| "Could not parse raw data")
207
+        parse_message_from_bytes(blob.get_raw()).chain_err(|| "Could not parse raw data")
207
     } else if blob.has_zlib_data() {
208
     } else if blob.has_zlib_data() {
208
         let mut decoder = ZlibDecoder::new(blob.get_zlib_data());
209
         let mut decoder = ZlibDecoder::new(blob.get_zlib_data());
209
-        protobuf::parse_from_reader(&mut decoder).chain_err(|| "Could not parse zlib data")
210
+        parse_message_from_reader(&mut decoder).chain_err(|| "Could not parse zlib data")
210
     } else {
211
     } else {
211
         bail!("Blob is missing fields 'raw' and 'zlib_data")
212
         bail!("Blob is missing fields 'raw' and 'zlib_data")
212
     }
213
     }
216
 pub(crate) fn decode_blob<T>(blob: &fileformat::Blob) -> Result<T>
217
 pub(crate) fn decode_blob<T>(blob: &fileformat::Blob) -> Result<T>
217
     where T: protobuf::Message + protobuf::MessageStatic {
218
     where T: protobuf::Message + protobuf::MessageStatic {
218
     if blob.has_raw() {
219
     if blob.has_raw() {
219
-        protobuf::parse_from_bytes(blob.get_raw()).chain_err(|| "Could not parse raw data")
220
+        parse_message_from_bytes(blob.get_raw()).chain_err(|| "Could not parse raw data")
220
     } else if blob.has_zlib_data() {
221
     } else if blob.has_zlib_data() {
221
         let mut decoder = DeflateDecoder::from_zlib(blob.get_zlib_data());
222
         let mut decoder = DeflateDecoder::from_zlib(blob.get_zlib_data());
222
-        protobuf::parse_from_reader(&mut decoder).chain_err(|| "Could not parse zlib data")
223
+        parse_message_from_reader(&mut decoder).chain_err(|| "Could not parse zlib data")
223
     } else {
224
     } else {
224
         bail!("Blob is missing fields 'raw' and 'zlib_data")
225
         bail!("Blob is missing fields 'raw' and 'zlib_data")
225
     }
226
     }

+ 1
- 0
src/lib.rs Näytä tiedosto

101
 pub mod dense;
101
 pub mod dense;
102
 pub mod elements;
102
 pub mod elements;
103
 pub mod mmap_blob;
103
 pub mod mmap_blob;
104
+mod util;

+ 4
- 3
src/mmap_blob.rs Näytä tiedosto

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

+ 34
- 0
src/util.rs Näytä tiedosto

1
+use errors::*;
2
+use std::io::Read;
3
+
4
+
5
+pub(crate) fn parse_message_from_bytes<M>(bytes: &[u8]) -> Result<M>
6
+    where M: ::protobuf::Message + ::protobuf::MessageStatic
7
+{
8
+    let mut stream = ::protobuf::CodedInputStream::from_bytes(bytes);
9
+    let mut message: M = ::protobuf::MessageStatic::new();
10
+    message.merge_from(&mut stream)?;
11
+
12
+    if message.is_initialized() {
13
+        Ok(message)
14
+    } else {
15
+        Err(::protobuf::ProtobufError::message_not_initialized("").into())
16
+    }
17
+}
18
+
19
+pub(crate) fn parse_message_from_reader<R, M>(reader: &mut R) -> Result<M>
20
+    where R: Read,
21
+          M: ::protobuf::Message + ::protobuf::MessageStatic,
22
+{
23
+    let mut stream = ::protobuf::CodedInputStream::new(reader);
24
+    let mut message: M = ::protobuf::MessageStatic::new();
25
+    message.merge_from(&mut stream)?;
26
+
27
+    stream.check_eof()?;
28
+
29
+    if message.is_initialized() {
30
+        Ok(message)
31
+    } else {
32
+        Err(::protobuf::ProtobufError::message_not_initialized("").into())
33
+    }
34
+}

+ 17
- 0
tests/read.rs Näytä tiedosto

119
         panic!("Unexpected blob type");
119
         panic!("Unexpected blob type");
120
     }
120
     }
121
 }
121
 }
122
+
123
+#[test]
124
+fn decode_blob() {
125
+    let reader = BlobReader::from_path(TEST_FILE_PATH).unwrap();
126
+    let blobs = reader.collect::<Result<Vec<_>>>().unwrap();
127
+
128
+    assert_eq!(blobs.len(), 2);
129
+    assert_eq!(blobs[0].get_type(), BlobType::OsmHeader);
130
+    assert_eq!(blobs[1].get_type(), BlobType::OsmData);
131
+
132
+    // decoding to the wrong blob type should not panic, but produce an Err.
133
+    assert!(blobs[0].to_primitiveblock().is_err());
134
+    assert!(blobs[1].to_headerblock().is_err());
135
+
136
+    assert!(blobs[0].to_headerblock().is_ok());
137
+    assert!(blobs[1].to_primitiveblock().is_ok());
138
+}