@@ -110,25 +110,44 @@ impl UbootShell {
110110 Ok ( interrupt_line)
111111 }
112112
113- async fn clear_shell ( & mut self ) -> Result < ( ) > {
113+ async fn read_until_idle ( & mut self ) -> Result < Vec < u8 > > {
114+ let mut bytes = Vec :: new ( ) ;
114115 loop {
115116 match self
116117 . read_byte_with_timeout ( Duration :: from_millis ( 300 ) )
117118 . await
118119 {
119- Ok ( _ ) => { }
120- Err ( err) if err. kind ( ) == ErrorKind :: TimedOut => return Ok ( ( ) ) ,
120+ Ok ( byte ) => bytes . push ( byte ) ,
121+ Err ( err) if err. kind ( ) == ErrorKind :: TimedOut => return Ok ( bytes ) ,
121122 Err ( err) => return Err ( err) ,
122123 }
123124 }
124125 }
125126
127+ async fn clear_shell ( & mut self ) -> Result < ( ) > {
128+ self . read_until_idle ( ) . await . map ( |_| ( ) )
129+ }
130+
126131 async fn wait_for_shell ( & mut self ) -> Result < ( ) > {
127132 let mut line = self . wait_for_interrupt ( ) . await ?;
128133 debug ! ( "got {}" , String :: from_utf8_lossy( & line) ) ;
129134 line. resize ( line. len ( ) . saturating_sub ( INT . len ( ) ) , 0 ) ;
130- self . perfix = String :: from_utf8_lossy ( & line) . to_string ( ) ;
131- self . clear_shell ( ) . await ?;
135+ if line. is_empty ( ) {
136+ let prompt = self . read_until_idle ( ) . await ?;
137+ let prompt = String :: from_utf8_lossy ( & prompt) ;
138+ self . perfix = prompt
139+ . lines ( )
140+ . rev ( )
141+ . find ( |line| !line. trim ( ) . is_empty ( ) )
142+ . unwrap_or_default ( )
143+ . to_string ( ) ;
144+ if self . perfix . is_empty ( ) {
145+ return Err ( Error :: new ( ErrorKind :: InvalidData , "U-Boot prompt is empty" ) ) ;
146+ }
147+ } else {
148+ self . perfix = String :: from_utf8_lossy ( & line) . to_string ( ) ;
149+ self . clear_shell ( ) . await ?;
150+ }
132151 Ok ( ( ) )
133152 }
134153
@@ -432,6 +451,103 @@ mod tests {
432451 sync:: { Arc , Mutex } ,
433452 } ;
434453
454+ #[ derive( Clone ) ]
455+ struct InterruptTx {
456+ response : Arc < Mutex < VecDeque < u8 > > > ,
457+ interrupt_output : & ' static [ u8 ] ,
458+ }
459+
460+ impl AsyncWrite for InterruptTx {
461+ fn poll_write (
462+ self : Pin < & mut Self > ,
463+ _cx : & mut Context < ' _ > ,
464+ buf : & [ u8 ] ,
465+ ) -> Poll < Result < usize > > {
466+ if buf == [ CTRL_C ] {
467+ self . response . lock ( ) . unwrap ( ) . extend ( self . interrupt_output ) ;
468+ }
469+ Poll :: Ready ( Ok ( buf. len ( ) ) )
470+ }
471+
472+ fn poll_flush ( self : Pin < & mut Self > , _cx : & mut Context < ' _ > ) -> Poll < Result < ( ) > > {
473+ Poll :: Ready ( Ok ( ( ) ) )
474+ }
475+
476+ fn poll_close ( self : Pin < & mut Self > , _cx : & mut Context < ' _ > ) -> Poll < Result < ( ) > > {
477+ Poll :: Ready ( Ok ( ( ) ) )
478+ }
479+ }
480+
481+ struct InterruptRx {
482+ response : Arc < Mutex < VecDeque < u8 > > > ,
483+ }
484+
485+ impl AsyncRead for InterruptRx {
486+ fn poll_read (
487+ self : Pin < & mut Self > ,
488+ _cx : & mut Context < ' _ > ,
489+ buf : & mut [ u8 ] ,
490+ ) -> Poll < Result < usize > > {
491+ let mut response = self . response . lock ( ) . unwrap ( ) ;
492+ if response. is_empty ( ) {
493+ return Poll :: Pending ;
494+ }
495+ let count = buf. len ( ) . min ( response. len ( ) ) ;
496+ for slot in & mut buf[ ..count] {
497+ * slot = response. pop_front ( ) . unwrap ( ) ;
498+ }
499+ Poll :: Ready ( Ok ( count) )
500+ }
501+ }
502+
503+ #[ tokio:: test]
504+ async fn detects_prompt_printed_before_interrupt_marker ( ) -> Result < ( ) > {
505+ let response = Arc :: new ( Mutex :: new ( VecDeque :: new ( ) ) ) ;
506+ let shell = UbootShell :: new (
507+ InterruptTx {
508+ response : response. clone ( ) ,
509+ interrupt_output : b"=> <INTERRUPT>\n => " ,
510+ } ,
511+ InterruptRx { response } ,
512+ )
513+ . await ?;
514+
515+ assert_eq ! ( shell. perfix, "=> " ) ;
516+ Ok ( ( ) )
517+ }
518+
519+ #[ tokio:: test]
520+ async fn detects_prompt_printed_after_interrupt_line ( ) -> Result < ( ) > {
521+ let response = Arc :: new ( Mutex :: new ( VecDeque :: new ( ) ) ) ;
522+ let shell = UbootShell :: new (
523+ InterruptTx {
524+ response : response. clone ( ) ,
525+ interrupt_output : b"<INTERRUPT>\n => " ,
526+ } ,
527+ InterruptRx { response } ,
528+ )
529+ . await ?;
530+
531+ assert_eq ! ( shell. perfix, "=> " ) ;
532+ Ok ( ( ) )
533+ }
534+
535+ #[ tokio:: test]
536+ async fn detects_prompt_after_repeated_interrupt_output ( ) -> Result < ( ) > {
537+ let response = Arc :: new ( Mutex :: new ( VecDeque :: new ( ) ) ) ;
538+ let shell = UbootShell :: new (
539+ InterruptTx {
540+ response : response. clone ( ) ,
541+ interrupt_output : b"<INTERRUPT>\n => <INTERRUPT>\n => " ,
542+ } ,
543+ InterruptRx { response } ,
544+ )
545+ . await ?;
546+
547+ assert_eq ! ( shell. perfix, "=> " ) ;
548+ Ok ( ( ) )
549+ }
550+
435551 #[ derive( Default ) ]
436552 struct LoadyScript {
437553 reads : VecDeque < u8 > ,
0 commit comments