1
+ use std:: ffi:: OsString ;
1
2
use std:: fs:: File ;
2
3
use std:: io:: BufReader ;
3
4
use std:: path:: { Path , PathBuf } ;
4
5
use std:: {
5
6
net:: { IpAddr , SocketAddr , ToSocketAddrs } ,
6
- os:: unix:: prelude:: AsRawFd ,
7
+ os:: unix:: prelude:: { AsRawFd , OsStringExt } ,
7
8
time:: Duration ,
8
9
} ;
9
10
10
11
use anyhow:: { anyhow, Context } ;
11
12
use clap:: { Parser , Subcommand } ;
12
13
use futures:: { future, SinkExt , StreamExt } ;
14
+ use regex:: bytes:: Regex ;
13
15
use propolis_client:: handmade:: {
14
16
api:: {
15
17
DiskRequest , InstanceEnsureRequest , InstanceMigrateInitiateRequest ,
@@ -90,6 +92,29 @@ enum Command {
90
92
/// Defaults to the most recent 16 KiB of console output (-16384).
91
93
#[ clap( long, short) ]
92
94
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 ,
93
118
} ,
94
119
95
120
/// Migrate instance to new propolis-server
@@ -221,60 +246,81 @@ async fn put_instance(
221
246
async fn stdin_to_websockets_task (
222
247
mut stdinrx : tokio:: sync:: mpsc:: Receiver < Vec < u8 > > ,
223
248
wstx : tokio:: sync:: mpsc:: Sender < Vec < u8 > > ,
249
+ escape_vector : Option < Vec < u8 > > ,
250
+ escape_prefix_length : usize ,
224
251
) {
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 ;
228
256
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 ( ) ;
235
260
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 ;
249
286
} else {
250
- next_raw = true ;
287
+ ansi_curs_check . clear ( ) ;
251
288
}
252
289
}
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
256
295
exit = true ;
257
296
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
260
299
outbuf. push ( c) ;
261
- next_raw = false ;
262
300
}
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 ;
265
308
outbuf. push ( c) ;
266
- next_raw = false ;
267
309
}
268
310
}
269
- }
270
311
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
+ }
275
316
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 ( ) ;
278
324
}
279
325
}
280
326
}
@@ -286,7 +332,10 @@ async fn test_stdin_to_websockets_task() {
286
332
let ( stdintx, stdinrx) = tokio:: sync:: mpsc:: channel ( 16 ) ;
287
333
let ( wstx, mut wsrx) = tokio:: sync:: mpsc:: channel ( 16 ) ;
288
334
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
+ } ) ;
290
339
291
340
// send characters, receive characters
292
341
stdintx
@@ -296,33 +345,22 @@ async fn test_stdin_to_websockets_task() {
296
345
let actual = wsrx. recv ( ) . await . unwrap ( ) ;
297
346
assert_eq ! ( String :: from_utf8( actual) . unwrap( ) , "test post please ignore" ) ;
298
347
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 ( ) ;
301
350
assert_eq ! ( wsrx. try_recv( ) , Err ( TryRecvError :: Empty ) ) ;
302
351
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
304
353
stdintx. send ( "test" . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
305
354
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" ) ;
312
356
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
315
358
stdintx. send ( "\x03 " . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
316
- assert_eq ! ( wsrx. try_recv( ) , Err ( TryRecvError :: Empty ) ) ;
317
359
let actual = wsrx. recv ( ) . await . unwrap ( ) ;
318
360
assert_eq ! ( String :: from_utf8( actual) . unwrap( ) , "\x03 " ) ;
319
361
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 ( ) ;
326
364
stdintx. send ( "\x03 " . chars ( ) . map ( |c| c as u8 ) . collect ( ) ) . await . unwrap ( ) ;
327
365
assert_eq ! ( wsrx. try_recv( ) , Err ( TryRecvError :: Empty ) ) ;
328
366
@@ -333,6 +371,8 @@ async fn test_stdin_to_websockets_task() {
333
371
async fn serial (
334
372
addr : SocketAddr ,
335
373
byte_offset : Option < i64 > ,
374
+ escape_vector : Option < Vec < u8 > > ,
375
+ escape_prefix_length : usize ,
336
376
) -> anyhow:: Result < ( ) > {
337
377
let client = propolis_client:: Client :: new ( & format ! ( "http://{}" , addr) ) ;
338
378
let mut req = client. instance_serial ( ) ;
@@ -375,7 +415,13 @@ async fn serial(
375
415
}
376
416
} ) ;
377
417
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
+ } ) ;
379
425
380
426
loop {
381
427
tokio:: select! {
@@ -569,7 +615,14 @@ async fn main() -> anyhow::Result<()> {
569
615
}
570
616
Command :: Get => get_instance ( & client) . await ?,
571
617
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
+ }
573
626
Command :: Migrate { dst_server, dst_port, dst_uuid } => {
574
627
let dst_addr = SocketAddr :: new ( dst_server, dst_port) ;
575
628
let dst_client = Client :: new ( dst_addr, log. clone ( ) ) ;
0 commit comments