diff --git a/src/api/icons.rs b/src/api/icons.rs index cf2f5e7..ab69213 100644 --- a/src/api/icons.rs +++ b/src/api/icons.rs @@ -213,7 +213,7 @@ fn get_icon_url(domain: &str) -> Result<(Vec, String), Error> { let mut cookie_str = String::new(); let resp = get_page(&ssldomain).or_else(|_| get_page(&httpdomain)); - if let Ok(content) = resp { + if let Ok(mut content) = resp { // Extract the URL from the respose in case redirects occured (like @ gitlab.com) let url = content.url().clone(); @@ -233,7 +233,11 @@ fn get_icon_url(domain: &str) -> Result<(Vec, String), Error> { // Add the default favicon.ico to the list with the domain the content responded from. iconlist.push(Icon::new(35, url.join("/favicon.ico").unwrap().into_string())); - let soup = Soup::from_reader(content)?; + // 512KB should be more than enough for the HTML, though as we only really need + // the HTML header, it could potentially be reduced even further + let limited_reader = crate::util::LimitedReader::new(&mut content, 512 * 1024); + + let soup = Soup::from_reader(limited_reader)?; // Search for and filter let favicons = soup .tag("link") diff --git a/src/util.rs b/src/util.rs index 4057918..aa3feb5 100644 --- a/src/util.rs +++ b/src/util.rs @@ -216,6 +216,33 @@ pub fn delete_file(path: &str) -> IOResult<()> { res } +pub struct LimitedReader<'a> { + reader: &'a mut dyn std::io::Read, + limit: usize, // In bytes + count: usize, +} +impl<'a> LimitedReader<'a> { + pub fn new(reader: &'a mut dyn std::io::Read, limit: usize) -> LimitedReader<'a> { + LimitedReader { + reader, + limit, + count: 0, + } + } +} + +impl<'a> std::io::Read for LimitedReader<'a> { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + self.count += buf.len(); + + if self.count > self.limit { + Ok(0) // End of the read + } else { + self.reader.read(buf) + } + } +} + const UNITS: [&str; 6] = ["bytes", "KB", "MB", "GB", "TB", "PB"]; pub fn get_display_size(size: i32) -> String {