vw_small

commit 25454697130a2bcbb76f7f29cdf8d1d382de96c7
parent f09996a21d8dec277e994edaa79fcca1a9890d27
Author: Daniel García <dani-garcia@users.noreply.github.com>
Date:   Thu, 19 Dec 2019 00:37:16 +0100

Fix crash when page URL points to huge file

Diffstat:
M
src/api/icons.rs
 | 
8
++++++--
M
src/util.rs
 | 
27
+++++++++++++++++++++++++++

2 files changed, 33 insertions(+), 2 deletions(-)
diff --git a/src/api/icons.rs b/src/api/icons.rs
@@ -213,7 +213,7 @@ fn get_icon_url(domain: &str) -> Result<(Vec<Icon>, 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<Icon>, 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
@@ -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<usize> {
+        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 {