1+ use std:: ffi:: OsString ;
12use std:: fs:: File ;
23use std:: io:: BufReader ;
34use std:: path:: { Path , PathBuf } ;
45use std:: {
56 net:: { IpAddr , SocketAddr , ToSocketAddrs } ,
6- os:: unix:: prelude:: AsRawFd ,
7+ os:: unix:: prelude:: { AsRawFd , OsStringExt } ,
78 time:: Duration ,
89} ;
910
1011use anyhow:: { anyhow, Context } ;
1112use clap:: { Parser , Subcommand } ;
1213use futures:: { future, SinkExt , StreamExt } ;
14+ use regex:: bytes:: Regex ;
1315use propolis_client:: handmade:: {
1416 api:: {
1517 DiskRequest , InstanceEnsureRequest , InstanceMigrateInitiateRequest ,
@@ -90,6 +92,29 @@ enum Command {
9092 /// Defaults to the most recent 16 KiB of console output (-16384).
9193 #[ clap( long, short) ]
9294 byte_offset : Option < i64 > ,
95+
96+ /// If this sequence of bytes is typed, the client will exit.
97+ /// Defaults to "^]^C" (Ctrl+], Ctrl+C). Note that the string passed
98+ /// for this argument is used verbatim without any parsing; in most
99+ /// shells, if you wish to include a special character (such as Enter
100+ /// or a Ctrl+letter combo), you can insert the character by preceding
101+ /// it with Ctrl+V at the command line.
102+ #[ clap( long, short, default_value = "\x1d \x03 " ) ]
103+ escape_string : OsString ,
104+
105+ /// The number of bytes from the beginning of the escape string to pass
106+ /// to the VM before beginning to buffer inputs until a mismatch.
107+ /// Defaults to 0, such that input matching the escape string does not
108+ /// get sent to the VM at all until a non-matching character is typed.
109+ /// To mimic the escape sequence for exiting SSH (Enter, tilde, dot),
110+ /// you may pass `-e '^M~.' --escape-prefix-length=1` such that normal
111+ /// Enter presses are sent to the VM immediately.
112+ #[ clap( long, default_value = "0" ) ]
113+ escape_prefix_length : usize ,
114+
115+ /// Disable escape string altogether (to exit, use pkill or similar).
116+ #[ clap( long, short = 'E' ) ]
117+ no_escape : bool ,
93118 } ,
94119
95120 /// Migrate instance to new propolis-server
@@ -221,60 +246,81 @@ async fn put_instance(
221246async fn stdin_to_websockets_task (
222247 mut stdinrx : tokio:: sync:: mpsc:: Receiver < Vec < u8 > > ,
223248 wstx : tokio:: sync:: mpsc:: Sender < Vec < u8 > > ,
249+ escape_vector : Option < Vec < u8 > > ,
250+ escape_prefix_length : usize ,
224251) {
225- // next_raw must live outside loop, because Ctrl-A should work across
226- // multiple inbuf reads.
227- let mut next_raw = false ;
252+ if let Some ( esc_sequence) = & escape_vector {
253+ // esc_pos must live outside loop, because escape string should work
254+ // across multiple inbuf reads.
255+ let mut esc_pos = 0 ;
228256
229- loop {
230- let inbuf = if let Some ( inbuf) = stdinrx. recv ( ) . await {
231- inbuf
232- } else {
233- continue ;
234- } ;
257+ // matches partial increments of "\x1b[14;30R"
258+ let ansi_curs_pat = Regex :: new ( "^\x1b (\\ [([0-9]{1,2}(;([0-9]{1,2}R?)?)?)?)?$" ) . unwrap ( ) ;
259+ let mut ansi_curs_check = Vec :: new ( ) ;
235260
236- // Put bytes from inbuf to outbuf, but don't send Ctrl-A unless
237- // next_raw is true.
238- let mut outbuf = Vec :: with_capacity ( inbuf. len ( ) ) ;
239-
240- let mut exit = false ;
241- for c in inbuf {
242- match c {
243- // Ctrl-A means send next one raw
244- b'\x01' => {
245- if next_raw {
246- // Ctrl-A Ctrl-A should be sent as Ctrl-A
247- outbuf. push ( c) ;
248- next_raw = false ;
261+ loop {
262+ let inbuf = if let Some ( inbuf) = stdinrx. recv ( ) . await {
263+ inbuf
264+ } else {
265+ continue ;
266+ } ;
267+
268+ // Put bytes from inbuf to outbuf, but don't send characters in the
269+ // escape string sequence unless we bail.
270+ let mut outbuf = Vec :: with_capacity ( inbuf. len ( ) ) ;
271+
272+ let mut exit = false ;
273+ for c in inbuf {
274+ // ignore ANSI escape sequence for the cursor position
275+ // response sent by xterm-alikes in response to shells
276+ // requesting one after receiving a newline.
277+ if esc_pos > 0 && esc_pos <= escape_prefix_length
278+ && b"\r \n " . contains ( & esc_sequence[ esc_pos - 1 ] )
279+ {
280+ ansi_curs_check. push ( c) ;
281+ if ansi_curs_pat. is_match ( & ansi_curs_check) {
282+ if c == b'R' { // end of the sequence
283+ ansi_curs_check. clear ( ) ;
284+ }
285+ continue ;
249286 } else {
250- next_raw = true ;
287+ ansi_curs_check . clear ( ) ;
251288 }
252289 }
253- b'\x03' => {
254- if !next_raw {
255- // Exit on non-raw Ctrl-C
290+
291+ if c == esc_sequence[ esc_pos] {
292+ esc_pos += 1 ;
293+ if esc_pos == esc_sequence. len ( ) {
294+ // Exit on completed escape string
256295 exit = true ;
257296 break ;
258- } else {
259- // Otherwise send Ctrl-C
297+ } else if esc_pos <= escape_prefix_length {
298+ // let through incomplete prefix up to the given limit
260299 outbuf. push ( c) ;
261- next_raw = false ;
262300 }
263- }
264- _ => {
301+ } else {
302+ // they bailed from the sequence,
303+ // feed everything that matched so far through
304+ if esc_pos != 0 {
305+ outbuf. extend ( & esc_sequence[ escape_prefix_length..esc_pos] )
306+ }
307+ esc_pos = 0 ;
265308 outbuf. push ( c) ;
266- next_raw = false ;
267309 }
268310 }
269- }
270311
271- // Send what we have, even if there's a Ctrl-C at the end .
272- if !outbuf. is_empty ( ) {
273- wstx. send ( outbuf) . await . unwrap ( ) ;
274- }
312+ // Send what we have, even if we're about to exit .
313+ if !outbuf. is_empty ( ) {
314+ wstx. send ( outbuf) . await . unwrap ( ) ;
315+ }
275316
276- if exit {
277- break ;
317+ if exit {
318+ break ;
319+ }
320+ }
321+ } else {
322+ while let Some ( buf) = stdinrx. recv ( ) . await {
323+ wstx. send ( buf) . await . unwrap ( ) ;
278324 }
279325 }
280326}
@@ -286,7 +332,10 @@ async fn test_stdin_to_websockets_task() {
286332 let ( stdintx, stdinrx) = tokio:: sync:: mpsc:: channel ( 16 ) ;
287333 let ( wstx, mut wsrx) = tokio:: sync:: mpsc:: channel ( 16 ) ;
288334
289- tokio:: spawn ( async move { stdin_to_websockets_task ( stdinrx, wstx) . await } ) ;
335+ let escape_vector = Some ( vec ! [ 0x1d , 0x03 ] ) ;
336+ tokio:: spawn ( async move {
337+ stdin_to_websockets_task ( stdinrx, wstx, escape_vector, 0 ) . await
338+ } ) ;
290339
291340 // send characters, receive characters
292341 stdintx
@@ -296,33 +345,22 @@ async fn test_stdin_to_websockets_task() {
296345 let actual = wsrx. recv ( ) . await . unwrap ( ) ;
297346 assert_eq ! ( String :: from_utf8( actual) . unwrap( ) , "test post please ignore" ) ;
298347
299- // don't send ctrl-a
300- stdintx. send ( "\x01 " . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
348+ // don't send a started escape sequence
349+ stdintx. send ( "\x1d " . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
301350 assert_eq ! ( wsrx. try_recv( ) , Err ( TryRecvError :: Empty ) ) ;
302351
303- // the "t" here is sent "raw" because of last ctrl-a but that doesn't change anything
352+ // since we didn't enter the \x03, the previous \x1d shows up here
304353 stdintx. send ( "test" . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
305354 let actual = wsrx. recv ( ) . await . unwrap ( ) ;
306- assert_eq ! ( String :: from_utf8( actual) . unwrap( ) , "test" ) ;
307-
308- // ctrl-a ctrl-c = only ctrl-c sent
309- stdintx. send ( "\x01 \x03 " . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
310- let actual = wsrx. recv ( ) . await . unwrap ( ) ;
311- assert_eq ! ( String :: from_utf8( actual) . unwrap( ) , "\x03 " ) ;
355+ assert_eq ! ( String :: from_utf8( actual) . unwrap( ) , "\x1d test" ) ;
312356
313- // same as above, across two messages
314- stdintx. send ( "\x01 " . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
357+ // \x03 gets sent if not preceded by \x1d
315358 stdintx. send ( "\x03 " . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
316- assert_eq ! ( wsrx. try_recv( ) , Err ( TryRecvError :: Empty ) ) ;
317359 let actual = wsrx. recv ( ) . await . unwrap ( ) ;
318360 assert_eq ! ( String :: from_utf8( actual) . unwrap( ) , "\x03 " ) ;
319361
320- // ctrl-a ctrl-a = only ctrl-a sent
321- stdintx. send ( "\x01 \x01 " . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
322- let actual = wsrx. recv ( ) . await . unwrap ( ) ;
323- assert_eq ! ( String :: from_utf8( actual) . unwrap( ) , "\x01 " ) ;
324-
325- // ctrl-c on its own means exit
362+ // \x1d followed by \x03 means exit, even if they're separate messages
363+ stdintx. send ( "\x1d " . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
326364 stdintx. send ( "\x03 " . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
327365 assert_eq ! ( wsrx. try_recv( ) , Err ( TryRecvError :: Empty ) ) ;
328366
@@ -333,6 +371,8 @@ async fn test_stdin_to_websockets_task() {
333371async fn serial (
334372 addr : SocketAddr ,
335373 byte_offset : Option < i64 > ,
374+ escape_vector : Option < Vec < u8 > > ,
375+ escape_prefix_length : usize ,
336376) -> anyhow:: Result < ( ) > {
337377 let client = propolis_client:: Client :: new ( & format ! ( "http://{}" , addr) ) ;
338378 let mut req = client. instance_serial ( ) ;
@@ -375,7 +415,13 @@ async fn serial(
375415 }
376416 } ) ;
377417
378- tokio:: spawn ( async move { stdin_to_websockets_task ( stdinrx, wstx) . await } ) ;
418+ let escape_len = escape_vector. as_ref ( ) . map ( |x| x. len ( ) ) . unwrap_or ( 0 ) ;
419+ if escape_prefix_length > escape_len {
420+ anyhow:: bail!( "prefix length {} is greater than length of escape string ({})" , escape_prefix_length, escape_len) ;
421+ }
422+ tokio:: spawn ( async move {
423+ stdin_to_websockets_task ( stdinrx, wstx, escape_vector, escape_prefix_length) . await
424+ } ) ;
379425
380426 loop {
381427 tokio:: select! {
@@ -569,7 +615,14 @@ async fn main() -> anyhow::Result<()> {
569615 }
570616 Command :: Get => get_instance ( & client) . await ?,
571617 Command :: State { state } => put_instance ( & client, state) . await ?,
572- Command :: Serial { byte_offset } => serial ( addr, byte_offset) . await ?,
618+ Command :: Serial { byte_offset, escape_string, escape_prefix_length, no_escape } => {
619+ let escape_vector = if no_escape || escape_string. is_empty ( ) {
620+ None
621+ } else {
622+ Some ( escape_string. into_vec ( ) )
623+ } ;
624+ serial ( addr, byte_offset, escape_vector, escape_prefix_length) . await ?
625+ }
573626 Command :: Migrate { dst_server, dst_port, dst_uuid } => {
574627 let dst_addr = SocketAddr :: new ( dst_server, dst_port) ;
575628 let dst_client = Client :: new ( dst_addr, log. clone ( ) ) ;
0 commit comments