Limit line buffer length

Discards long lines to prevent out-of-memory events.

src/controller.rs

@@ -5,7 +5,7 @@ use crate::config::{Config, VisibleLines};
 #[cfg(feature = "git")]
 use crate::diff::{get_git_diff, LineChanges};
 use crate::error::*;
-use crate::input::{Input, InputReader, OpenedInput};
+use crate::input::{Input, InputReader, OpenedInput, ReaderError};
 #[cfg(feature = "lessopen")]
 use crate::lessopen::LessOpenPreprocessor;
 #[cfg(feature = "git")]
@@ -249,13 +249,30 @@ impl<'b> Controller<'b> {
 
         let style_snip = self.config.style_components.snip();
 
-        while reader.read_line(&mut line_buffer)? {
+        loop {
+            let mut soft_limit_hit = false;
+            let read_result = reader.read_line(&mut line_buffer);
+            match read_result {
+                Ok(res) => {
+                    if !res {
+                        break;
+                    }
+                }
+                Err(err) => match err {
+                    ReaderError::IoError(io_err) => return Err(io_err.into()),
+                    ReaderError::SoftLimitHit => soft_limit_hit = true,
+                    ReaderError::HardLimitHit => return Err(Error::LineTooLong(line_number)),
+                },
+            };
+
             match line_ranges.check(line_number) {
                 RangeCheckResult::BeforeOrBetweenRanges => {
-                    // Call the printer in case we need to call the syntax highlighter
-                    // for this line. However, set `out_of_range` to `true`.
-                    printer.print_line(true, writer, line_number, &line_buffer)?;
-                    mid_range = false;
+                    if !soft_limit_hit {
+                        // Call the printer in case we need to call the syntax highlighter
+                        // for this line. However, set `out_of_range` to `true`.
+                        printer.print_line(true, writer, line_number, &line_buffer)?;
+                        mid_range = false;
+                    }
                 }
 
                 RangeCheckResult::InRange => {
@@ -268,8 +285,11 @@ impl<'b> Controller<'b> {
                             printer.print_snip(writer)?;
                         }
                     }
-
-                    printer.print_line(false, writer, line_number, &line_buffer)?;
+                    if soft_limit_hit {
+                        printer.print_replaced_line(writer, line_number, "<line too long>")?;
+                    } else {
+                        printer.print_line(false, writer, line_number, &line_buffer)?;
+                    }
                 }
                 RangeCheckResult::AfterLastRange => {
                     break;

src/error.rs

@@ -28,6 +28,8 @@ pub enum Error {
     InvalidPagerValueBat,
     #[error("{0}")]
     Msg(String),
+    #[error("Line {0} is too long")]
+    LineTooLong(usize),
     #[cfg(feature = "lessopen")]
     #[error(transparent)]
     VarError(#[from] ::std::env::VarError),

src/input.rs

@@ -6,6 +6,7 @@ use std::path::{Path, PathBuf};
 
 use clircle::{Clircle, Identifier};
 use content_inspector::{self, ContentType};
+use once_cell::unsync::Lazy;
 
 use crate::error::*;
 
@@ -250,34 +251,50 @@ impl<'a> Input<'a> {
 }
 
 pub(crate) struct InputReader<'a> {
-    inner: Box<dyn BufRead + 'a>,
+    inner: LimitBuf<'a>,
     pub(crate) first_line: Vec<u8>,
     pub(crate) content_type: Option<ContentType>,
 }
 
 impl<'a> InputReader<'a> {
-    pub(crate) fn new<R: BufRead + 'a>(mut reader: R) -> InputReader<'a> {
-        let mut first_line = vec![];
-        reader.read_until(b'\n', &mut first_line).ok();
+    pub(crate) fn new<R: Read + 'a>(reader: R) -> InputReader<'a> {
+        let mut input_reader = InputReader {
+            inner: LimitBuf::new(reader, 4096, 1024 * 64, 1024 * 256),
+            first_line: vec![],
+            content_type: None,
+        };
 
-        let content_type = if first_line.is_empty() {
+        input_reader.read_first_line().ok();
+
+        let content_type = if input_reader.first_line.is_empty() {
             None
         } else {
-            Some(content_inspector::inspect(&first_line[..]))
+            Some(content_inspector::inspect(&input_reader.first_line[..]))
         };
 
         if content_type == Some(ContentType::UTF_16LE) {
-            reader.read_until(0x00, &mut first_line).ok();
+            input_reader
+                .inner
+                .read_until(0x00, &mut input_reader.first_line)
+                .ok();
         }
 
-        InputReader {
-            inner: Box::new(reader),
-            first_line,
-            content_type,
-        }
+        input_reader.content_type = content_type;
+        input_reader
     }
 
-    pub(crate) fn read_line(&mut self, buf: &mut Vec<u8>) -> io::Result<bool> {
+    fn read_first_line(&mut self) -> std::result::Result<bool, ReaderError> {
+        let mut first_line = vec![];
+        let res = self.read_line(&mut first_line);
+        self.first_line = first_line;
+
+        res
+    }
+
+    pub(crate) fn read_line(
+        &mut self,
+        buf: &mut Vec<u8>,
+    ) -> std::result::Result<bool, ReaderError> {
         if !self.first_line.is_empty() {
             buf.append(&mut self.first_line);
             return Ok(true);
@@ -293,6 +310,95 @@ impl<'a> InputReader<'a> {
     }
 }
 
+struct LimitBuf<'a> {
+    reader: Box<dyn Read + 'a>,
+    inner: Box<[u8]>,
+    start: usize,
+    len: usize,
+    soft_limit: usize,
+    hard_limit: usize,
+}
+
+impl<'a> LimitBuf<'a> {
+    pub fn new<R: Read + 'a>(
+        reader: R,
+        buf_size: usize,
+        soft_limit: usize,
+        hard_limit: usize,
+    ) -> Self {
+        Self {
+            reader: Box::new(reader),
+            inner: vec![0u8; buf_size].into_boxed_slice(),
+            start: 0,
+            len: 0,
+            soft_limit,
+            hard_limit,
+        }
+    }
+
+    pub fn read_until(
+        &mut self,
+        byte: u8,
+        buf: &mut Vec<u8>,
+    ) -> std::result::Result<usize, ReaderError> {
+        let mut end_byte_reached = false;
+        let mut total_bytes = 0;
+
+        let mut soft_limit_hit = false;
+        let capacity = self.inner.len();
+        let mut drop_buf = Lazy::new(|| Vec::with_capacity(capacity));
+
+        while !end_byte_reached {
+            if self.len == 0 {
+                let bytes = self.reader.read(&mut self.inner)?;
+                self.len += bytes;
+                self.start = 0;
+
+                if bytes == 0 {
+                    break;
+                }
+            }
+
+            let bytes = (&self.inner[self.start..self.start + self.len])
+                .read_until(byte, if soft_limit_hit { &mut drop_buf } else { buf })?;
+            end_byte_reached = self.inner[self.start + bytes - 1] == byte;
+
+            if soft_limit_hit {
+                drop_buf.clear();
+            }
+
+            self.len -= bytes;
+            self.start += bytes;
+            total_bytes += bytes;
+
+            if total_bytes > self.hard_limit {
+                return Err(ReaderError::HardLimitHit);
+            } else if total_bytes > self.soft_limit {
+                soft_limit_hit = true;
+            }
+        }
+
+        if soft_limit_hit {
+            Err(ReaderError::SoftLimitHit)
+        } else {
+            Ok(total_bytes)
+        }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) enum ReaderError {
+    SoftLimitHit,
+    HardLimitHit,
+    IoError(io::Error),
+}
+
+impl From<io::Error> for ReaderError {
+    fn from(value: io::Error) -> Self {
+        Self::IoError(value)
+    }
+}
+
 #[test]
 fn basic() {
     let content = b"#!/bin/bash\necho hello";

src/printer.rs

@@ -21,6 +21,7 @@ use unicode_width::UnicodeWidthChar;
 
 use crate::assets::{HighlightingAssets, SyntaxReferenceInSet};
 use crate::config::Config;
+use crate::decorations;
 #[cfg(feature = "git")]
 use crate::decorations::LineChangesDecoration;
 use crate::decorations::{Decoration, GridBorderDecoration, LineNumberDecoration};
@@ -97,6 +98,13 @@ pub(crate) trait Printer {
         line_number: usize,
         line_buffer: &[u8],
     ) -> Result<()>;
+
+    fn print_replaced_line(
+        &mut self,
+        handle: &mut OutputHandle,
+        line_number: usize,
+        replace_text: &str,
+    ) -> Result<()>;
 }
 
 pub struct SimplePrinter<'a> {
@@ -179,6 +187,15 @@ impl<'a> Printer for SimplePrinter<'a> {
         }
         Ok(())
     }
+
+    fn print_replaced_line(
+        &mut self,
+        _handle: &mut OutputHandle,
+        _line_number: usize,
+        _replace_text: &str,
+    ) -> Result<()> {
+        Ok(())
+    }
 }
 
 struct HighlighterFromSet<'a> {
@@ -819,6 +836,54 @@ impl<'a> Printer for InteractivePrinter<'a> {
 
         Ok(())
     }
+
+    fn print_replaced_line(
+        &mut self,
+        handle: &mut OutputHandle,
+        line_number: usize,
+        replace_text: &str,
+    ) -> Result<()> {
+        if let Some(ContentType::BINARY) | None = self.content_type {
+            return Ok(());
+        }
+
+        if self.panel_width > 0 {
+            let mut width = 0;
+            if self.config.style_components.numbers() {
+                let line_numbers = decorations::LineNumberDecoration::new(&self.colors);
+
+                width += line_numbers.width();
+                write!(
+                    handle,
+                    "{} ",
+                    line_numbers.generate(line_number, false, self).text
+                )?;
+            }
+
+            if self.config.style_components.grid() {
+                write!(handle, "{}", self.colors.error_indicator.paint("!"))?;
+            }
+
+            if width < self.panel_width {
+                write!(
+                    handle,
+                    "{}",
+                    " ".repeat(self.panel_width.saturating_sub(width).saturating_sub(2))
+                )?;
+            }
+
+            if self.config.style_components.grid() {
+                let grid = decorations::GridBorderDecoration::new(&self.colors);
+                write!(handle, "{} ", grid.generate(line_number, false, self).text)?;
+            }
+        }
+        writeln!(
+            handle,
+            "{}",
+            self.colors.error_indicator.paint(replace_text)
+        )?;
+        Ok(())
+    }
 }
 
 const DEFAULT_GUTTER_COLOR: u8 = 238;
@@ -832,6 +897,7 @@ pub struct Colors {
     pub git_removed: Style,
     pub git_modified: Style,
     pub line_number: Style,
+    pub error_indicator: Style,
 }
 
 impl Colors {
@@ -861,6 +927,7 @@ impl Colors {
             git_removed: Red.normal(),
             git_modified: Yellow.normal(),
             line_number: gutter_style,
+            error_indicator: Red.normal(),
         }
     }
 }