11use std:: collections:: BTreeMap ;
22use std:: fs:: File ;
3- use std:: io:: { BufRead , Read , Write } ;
3+ use std:: io:: { BufRead , Read , Seek , SeekFrom , Write } ;
44use std:: path:: { Path , PathBuf } ;
55use std:: process:: ExitCode ;
66
@@ -175,7 +175,11 @@ struct Cli {
175175 ) ]
176176 max_archive_entries : Option < usize > ,
177177
178- /// Wall-clock budget in milliseconds.
178+ /// Maximum source file or stdin size before parsing (separate from physical read budget).
179+ #[ arg( long, value_name = "BYTES" , help_heading = "Resource limits" ) ]
180+ max_input_bytes : Option < u64 > ,
181+
182+ /// Cooperative engine budget in milliseconds; not a hard process deadline.
179183 #[ arg(
180184 short = 'T' ,
181185 long,
@@ -424,7 +428,21 @@ fn run(cli: Cli) -> deckprobe_core::Result<u8> {
424428 None => { }
425429 }
426430
427- let options = probe_options ( & cli) ?;
431+ if cli. max_input_bytes == Some ( 0 ) {
432+ return Err ( DeckProbeError :: InvalidRequest (
433+ "max-input-bytes must be positive" . into ( ) ,
434+ ) ) ;
435+ }
436+ let mut options = probe_options ( & cli) ?;
437+ if cli. input . as_ref ( ) . is_some_and ( |p| p. as_os_str ( ) == "-" )
438+ && let Some ( cap) = cli. max_input_bytes
439+ {
440+ let current = options
441+ . budget
442+ . max_physical_bytes
443+ . unwrap_or ( Budget :: for_level ( options. level ) . max_physical_bytes ) ;
444+ options. budget . max_physical_bytes = Some ( current. min ( cap) ) ;
445+ }
428446 if cli. jsonl {
429447 return run_jsonl ( & cli, & options) ;
430448 }
@@ -447,7 +465,7 @@ fn run(cli: Cli) -> deckprobe_core::Result<u8> {
447465 "--stdin-name is only valid when INPUT is '-'" . to_owned ( ) ,
448466 ) ) ;
449467 }
450- std:: sync:: Arc :: new ( PathSource :: open ( input, None ) ?)
468+ std:: sync:: Arc :: new ( PathSource :: bounded ( input, None , cli . max_input_bytes ) ?)
451469 } ;
452470
453471 let report = probe ( source, options) ?;
@@ -565,7 +583,7 @@ fn run_jsonl(cli: &Cli, options: &ProbeOptions) -> deckprobe_core::Result<u8> {
565583 "JSONL line {line_number} is invalid: {error}"
566584 ) )
567585 } )
568- . and_then ( |input| jsonl_source ( input, options) )
586+ . and_then ( |input| jsonl_source ( input, options, cli . max_input_bytes ) )
569587 . and_then ( |source| probe ( source, options. clone ( ) ) ) ;
570588 match outcome {
571589 Ok ( report) => {
@@ -593,27 +611,46 @@ fn run_jsonl(cli: &Cli, options: &ProbeOptions) -> deckprobe_core::Result<u8> {
593611fn jsonl_source (
594612 input : JsonlInput ,
595613 options : & ProbeOptions ,
614+ max_input_bytes : Option < u64 > ,
596615) -> deckprobe_core:: Result < std:: sync:: Arc < dyn ProbeSource > > {
597616 match input {
598- JsonlInput :: Path ( path) => Ok ( std:: sync:: Arc :: new ( PathSource :: open ( path, None ) ?) ) ,
617+ JsonlInput :: Path ( path) => Ok ( std:: sync:: Arc :: new ( PathSource :: bounded (
618+ path,
619+ None ,
620+ max_input_bytes,
621+ ) ?) ) ,
599622 JsonlInput :: Record {
600623 path : Some ( path) ,
601624 name,
602625 data_base64 : None ,
603- } => Ok ( std:: sync:: Arc :: new ( PathSource :: open ( path, name) ?) ) ,
626+ } => Ok ( std:: sync:: Arc :: new ( PathSource :: bounded (
627+ path,
628+ name,
629+ max_input_bytes,
630+ ) ?) ) ,
604631 JsonlInput :: Record {
605632 path : None ,
606633 name : Some ( name) ,
607634 data_base64 : Some ( data) ,
608635 } => {
636+ let mut budget = Budget :: for_level ( options. level ) ;
637+ options. budget . apply_to ( & mut budget) ;
638+ let limit = max_input_bytes
639+ . unwrap_or ( u64:: MAX )
640+ . min ( budget. max_physical_bytes ) ;
641+ if data. len ( ) as u64 > limit. saturating_add ( 2 ) . saturating_div ( 3 ) . saturating_mul ( 4 ) {
642+ return Err ( DeckProbeError :: BudgetExceeded (
643+ "JSONL input exceeds size limit" . into ( ) ,
644+ ) ) ;
645+ }
609646 let bytes = base64:: engine:: general_purpose:: STANDARD
610647 . decode ( data)
611648 . map_err ( |error| {
612649 DeckProbeError :: InvalidRequest ( format ! ( "invalid JSONL base64 payload: {error}" ) )
613650 } ) ?;
614651 let mut budget = Budget :: for_level ( options. level ) ;
615652 options. budget . apply_to ( & mut budget) ;
616- if bytes. len ( ) as u64 > budget . max_physical_bytes {
653+ if bytes. len ( ) as u64 > limit {
617654 return Err ( DeckProbeError :: BudgetExceeded ( format ! (
618655 "JSONL byte payload size {} exceeds physical read budget {}" ,
619656 bytes. len( ) ,
@@ -636,15 +673,23 @@ struct PathSource {
636673 path : PathBuf ,
637674 display_name : String ,
638675 file_size : u64 ,
676+ max_input_bytes : Option < u64 > ,
677+ modified : Option < std:: time:: SystemTime > ,
639678}
640679
641680impl PathSource {
642- fn open (
681+ fn bounded (
643682 path : impl Into < PathBuf > ,
644683 display_name : Option < String > ,
684+ max_input_bytes : Option < u64 > ,
645685 ) -> deckprobe_core:: Result < Self > {
646686 let path = path. into ( ) ;
647687 let metadata = std:: fs:: metadata ( & path) ?;
688+ if max_input_bytes. is_some_and ( |cap| metadata. len ( ) > cap) {
689+ return Err ( DeckProbeError :: BudgetExceeded (
690+ "source exceeds max-input-bytes" . into ( ) ,
691+ ) ) ;
692+ }
648693 if !metadata. is_file ( ) {
649694 return Err ( DeckProbeError :: InvalidRequest ( format ! (
650695 "input is not a regular file: {}" ,
@@ -661,6 +706,8 @@ impl PathSource {
661706 path,
662707 display_name,
663708 file_size : metadata. len ( ) ,
709+ max_input_bytes,
710+ modified : metadata. modified ( ) . ok ( ) ,
664711 } )
665712 }
666713}
@@ -679,7 +726,46 @@ impl ProbeSource for PathSource {
679726 }
680727
681728 fn open ( & self ) -> deckprobe_core:: Result < BoxedProbeReader > {
682- Ok ( Box :: new ( File :: open ( & self . path ) ?) )
729+ let file = File :: open ( & self . path ) ?;
730+ if self . max_input_bytes . is_some ( ) {
731+ let metadata = file. metadata ( ) ?;
732+ if !metadata. is_file ( )
733+ || metadata. len ( ) != self . file_size
734+ || metadata. modified ( ) . ok ( ) != self . modified
735+ {
736+ return Err ( DeckProbeError :: Io ( std:: io:: Error :: other (
737+ "source changed during probe" ,
738+ ) ) ) ;
739+ }
740+ return Ok ( Box :: new ( CheckedFile {
741+ file,
742+ size : self . file_size ,
743+ modified : self . modified ,
744+ } ) ) ;
745+ }
746+ Ok ( Box :: new ( file) )
747+ }
748+ }
749+
750+ struct CheckedFile {
751+ file : File ,
752+ size : u64 ,
753+ modified : Option < std:: time:: SystemTime > ,
754+ }
755+ impl Read for CheckedFile {
756+ fn read ( & mut self , buf : & mut [ u8 ] ) -> std:: io:: Result < usize > {
757+ let meta = self . file . metadata ( ) ?;
758+ if meta. len ( ) != self . size || meta. modified ( ) . ok ( ) != self . modified {
759+ return Err ( std:: io:: Error :: other ( "source changed during probe" ) ) ;
760+ }
761+ let position = self . file . stream_position ( ) ?;
762+ let len = buf. len ( ) . min ( self . size . saturating_sub ( position) as usize ) ;
763+ self . file . read ( & mut buf[ ..len] )
764+ }
765+ }
766+ impl Seek for CheckedFile {
767+ fn seek ( & mut self , pos : SeekFrom ) -> std:: io:: Result < u64 > {
768+ self . file . seek ( pos)
683769 }
684770}
685771
0 commit comments