diff --git a/src/png.rs b/src/png.rs index 455d3d44..276c5183 100644 --- a/src/png.rs +++ b/src/png.rs @@ -11,7 +11,7 @@ use itertools::Itertools; use reduction::*; use std::collections::{HashMap, HashSet}; use std::fs::File; -use std::io::Read; +use std::io::{Read, Seek, SeekFrom}; use std::iter::Iterator; use std::path::Path; @@ -195,8 +195,19 @@ impl PngData { Ok(f) => f, Err(_) => return Err(PngError::new("Failed to open file for reading")), }; - let mut byte_data: Vec = Vec::new(); + // Check file for PNG header + let mut header = [0; 8]; + if file.read_exact(&mut header).is_err() { + return Err(PngError::new("Not a PNG file: too small")); + } + if !file_header_is_valid(&header) { + return Err(PngError::new("Invalid PNG header detected")); + } + if file.seek(SeekFrom::Start(0)).is_err() { + return Err(PngError::new("Failed to read from file")); + } // Read raw png data into memory + let mut byte_data: Vec = Vec::new(); match file.read_to_end(&mut byte_data) { Ok(_) => (), Err(_) => return Err(PngError::new("Failed to read from file")),