From 5f2a07f1871d1ab8d339fbab825df0625f4f6b4f Mon Sep 17 00:00:00 2001 From: einfachIrgendwer0815 <85333734+einfachIrgendwer0815@users.noreply.github.com> Date: Sun, 10 Mar 2024 13:23:55 +0100 Subject: [PATCH] Limit line buffer length Discards long lines to prevent out-of-memory events. --- src/controller.rs | 36 ++++++++++--- src/error.rs | 2 + src/input.rs | 132 +++++++++++++++++++++++++++++++++++++++++----- src/printer.rs | 67 +++++++++++++++++++++++ 4 files changed, 216 insertions(+), 21 deletions(-) diff --git a/src/controller.rs b/src/controller.rs index ffc5dd5b..2be3124d 100644 --- a/src/controller.rs +++ b/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, "")?; + } else { + printer.print_line(false, writer, line_number, &line_buffer)?; + } } RangeCheckResult::AfterLastRange => { break; diff --git a/src/error.rs b/src/error.rs index 007737b0..3f42ea5f 100644 --- a/src/error.rs +++ b/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), diff --git a/src/input.rs b/src/input.rs index 0ebaa4ce..cadbaa40 100644 --- a/src/input.rs +++ b/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, + inner: LimitBuf<'a>, pub(crate) first_line: Vec, pub(crate) content_type: Option, } impl<'a> InputReader<'a> { - pub(crate) fn new(mut reader: R) -> InputReader<'a> { - let mut first_line = vec![]; - reader.read_until(b'\n', &mut first_line).ok(); + pub(crate) fn new(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) -> io::Result { + fn read_first_line(&mut self) -> std::result::Result { + 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, + ) -> std::result::Result { 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, + inner: Box<[u8]>, + start: usize, + len: usize, + soft_limit: usize, + hard_limit: usize, +} + +impl<'a> LimitBuf<'a> { + pub fn new( + 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, + ) -> std::result::Result { + 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 for ReaderError { + fn from(value: io::Error) -> Self { + Self::IoError(value) + } +} + #[test] fn basic() { let content = b"#!/bin/bash\necho hello"; diff --git a/src/printer.rs b/src/printer.rs index e9bea3fd..4e17b546 100644 --- a/src/printer.rs +++ b/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}; @@ -99,6 +100,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> { @@ -181,6 +189,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> { @@ -851,6 +868,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; @@ -864,6 +929,7 @@ pub struct Colors { pub git_removed: Style, pub git_modified: Style, pub line_number: Style, + pub error_indicator: Style, } impl Colors { @@ -893,6 +959,7 @@ impl Colors { git_removed: Red.normal(), git_modified: Yellow.normal(), line_number: gutter_style, + error_indicator: Red.normal(), } } }