Harden Windows terminal byte handling

src/bin/dosh-client.rs

@@ -10845,6 +10845,8 @@ fn windows_vt_input_mode(mode: u32) -> u32 {
 struct WindowsConsoleModeGuard {
     input: Option<(windows_sys::Win32::Foundation::HANDLE, u32)>,
     output: Option<(windows_sys::Win32::Foundation::HANDLE, u32)>,
+    input_code_page: Option<u32>,
+    output_code_page: Option<u32>,
 }
 
 #[cfg(windows)]
@@ -10853,51 +10855,91 @@ impl WindowsConsoleModeGuard {
         unsafe {
             use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
             use windows_sys::Win32::System::Console::{
-                GetConsoleMode, GetStdHandle, STD_INPUT_HANDLE, STD_OUTPUT_HANDLE, SetConsoleMode,
+                GetConsoleCP, GetConsoleMode, GetConsoleOutputCP, GetStdHandle, STD_INPUT_HANDLE,
+                STD_OUTPUT_HANDLE, SetConsoleCP, SetConsoleMode, SetConsoleOutputCP,
+            };
+
+            const CP_UTF8: u32 = 65001;
+
+            let mut guard = Self {
+                input: None,
+                output: None,
+                input_code_page: None,
+                output_code_page: None,
             };
 
             let input_handle = GetStdHandle(STD_INPUT_HANDLE);
-            let input = if input_handle.is_null() || input_handle == INVALID_HANDLE_VALUE {
+            if !input_handle.is_null() && input_handle != INVALID_HANDLE_VALUE {
-                None
-            } else {
                 let mut original = 0u32;
-                if GetConsoleMode(input_handle, &mut original) == 0 {
+                if GetConsoleMode(input_handle, &mut original) != 0 {
-                    None
-                } else {
                     let desired = windows_vt_input_mode(original);
                     if desired != original && SetConsoleMode(input_handle, desired) == 0 {
                         return Err(std::io::Error::last_os_error())
                             .context("enable Windows virtual-terminal input mode");
                     }
-                    Some((input_handle, original))
+                    guard.input = Some((input_handle, original));
+
+                    let original_code_page = GetConsoleCP();
+                    if original_code_page == 0 {
+                        let err = std::io::Error::last_os_error();
+                        guard.restore();
+                        return Err(err).context("read Windows console input code page");
+                    }
+                    guard.input_code_page = Some(original_code_page);
+                    if original_code_page != CP_UTF8 && SetConsoleCP(CP_UTF8) == 0 {
+                        let err = std::io::Error::last_os_error();
+                        guard.restore();
+                        return Err(err).context("set Windows console input to UTF-8");
+                    }
+                }
             }
-            };
 
             let output_handle = GetStdHandle(STD_OUTPUT_HANDLE);
-            let output = if output_handle.is_null() || output_handle == INVALID_HANDLE_VALUE {
+            if !output_handle.is_null() && output_handle != INVALID_HANDLE_VALUE {
-                None
-            } else {
                 let mut original = 0u32;
-                if GetConsoleMode(output_handle, &mut original) == 0 {
+                if GetConsoleMode(output_handle, &mut original) != 0 {
-                    None
-                } else {
                     let desired = windows_vt_output_mode(original);
                     if desired != original && SetConsoleMode(output_handle, desired) == 0 {
                         let err = std::io::Error::last_os_error();
-                        if let Some((handle, mode)) = input {
+                        guard.restore();
-                            let _ = SetConsoleMode(handle, mode);
-                        }
                         return Err(err).context("enable Windows virtual-terminal output mode");
                     }
-                    Some((output_handle, original))
+                    guard.output = Some((output_handle, original));
-                }
-            };
 
-            Ok(Self { input, output })
+                    let original_code_page = GetConsoleOutputCP();
+                    if original_code_page == 0 {
+                        let err = std::io::Error::last_os_error();
+                        guard.restore();
+                        return Err(err).context("read Windows console output code page");
+                    }
+                    guard.output_code_page = Some(original_code_page);
+                    if original_code_page != CP_UTF8 && SetConsoleOutputCP(CP_UTF8) == 0 {
+                        let err = std::io::Error::last_os_error();
+                        guard.restore();
+                        return Err(err).context("set Windows console output to UTF-8");
+                    }
+                }
+            }
+
+            Ok(guard)
         }
     }
 
     fn restore(&self) {
+        if let Some(code_page) = self.output_code_page {
+            unsafe {
+                use windows_sys::Win32::System::Console::SetConsoleOutputCP;
+
+                let _ = SetConsoleOutputCP(code_page);
+            }
+        }
+        if let Some(code_page) = self.input_code_page {
+            unsafe {
+                use windows_sys::Win32::System::Console::SetConsoleCP;
+
+                let _ = SetConsoleCP(code_page);
+            }
+        }
         if let Some((handle, mode)) = self.input {
             unsafe {
                 use windows_sys::Win32::System::Console::SetConsoleMode;

tests/client_terminal_runtime.rs

@@ -0,0 +1,434 @@
+use dosh::config::ClientConfig;
+use dosh::protocol::{self, CLIENT_TO_SERVER, Frame, Input, PacketKind, Resize, SERVER_TO_CLIENT};
+use portable_pty::{CommandBuilder, NativePtySystem, PtySize, PtySystem};
+use serde::Serialize;
+use std::fs;
+use std::io::{Read, Write};
+use std::net::UdpSocket;
+use std::path::Path;
+use std::sync::{Arc, Mutex, mpsc};
+use std::thread;
+use std::time::{Duration, Instant};
+
+const SESSION: &str = "terminal-parity";
+const CLIENT_ID: [u8; 16] = [0x45; 16];
+const SESSION_KEY: [u8; 32] = [0x91; 32];
+const INPUT_BYTES: &[u8] = concat!(
+    "\x1b[A",
+    "\x1b[B",
+    "\x1b[200~",
+    "paste-λ-界",
+    "\x1b[201~",
+    "\x1b[<0;12;7M",
+    "\x1b[<0;12;7m"
+)
+.as_bytes();
+const OUTPUT_CHUNKS: &[&[u8]] = &[
+    b"\x1b[3;4H\x1b[38;2;12;200;90mDOSH_GRAPH_\xe2",
+    b"\xa3\xbf\xe2\xa3",
+    b"\xb7\xe2\xa3\xa4_\xce",
+    b"\xbb_\xe7",
+    b"\x95\x8c\x1b[0m",
+];
+
+#[derive(Serialize)]
+struct CachedCredentialWire {
+    server: String,
+    session: String,
+    mode: String,
+    udp_host: String,
+    udp_port: u16,
+    client_id: [u8; 16],
+    session_key: [u8; 32],
+    session_key_id: [u8; 16],
+    attach_ticket: Vec<u8>,
+    attach_ticket_psk: [u8; 32],
+    last_rendered_seq: u64,
+}
+
+#[derive(Debug)]
+enum ServerObservation {
+    Input(Vec<u8>),
+    Resize(u16, u16),
+}
+
+#[test]
+fn native_client_terminal_round_trip_is_platform_complete() {
+    let dir = tempfile::tempdir().unwrap();
+    let home = dir.path().join("home");
+    let cache = dir.path().join("credentials");
+    fs::create_dir_all(home.join(".config/dosh")).unwrap();
+    fs::create_dir_all(&cache).unwrap();
+
+    let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
+    socket
+        .set_read_timeout(Some(Duration::from_secs(8)))
+        .unwrap();
+    let port = socket.local_addr().unwrap().port();
+    write_client_fixture(&home, &cache, port);
+
+    let (observation_tx, observation_rx) = mpsc::channel();
+    let server = thread::spawn(move || run_fake_terminal_server(socket, observation_tx));
+
+    let pty = NativePtySystem::default();
+    let pair = pty
+        .openpty(PtySize {
+            rows: 24,
+            cols: 80,
+            pixel_width: 0,
+            pixel_height: 0,
+        })
+        .unwrap();
+    let mut reader = pair.master.try_clone_reader().unwrap();
+    let mut writer = pair.master.take_writer().unwrap();
+    let mut command = client_command(dir.path(), port);
+    command.env("HOME", home.to_string_lossy().to_string());
+    command.env("USERPROFILE", home.to_string_lossy().to_string());
+    command.env("APPDATA", home.to_string_lossy().to_string());
+    command.env("LOCALAPPDATA", home.to_string_lossy().to_string());
+    command.env("TERM", "xterm-256color");
+    command.env("COLORTERM", "truecolor");
+    let mut child = pair.slave.spawn_command(command).unwrap();
+    drop(pair.slave);
+
+    let output = Arc::new(Mutex::new(Vec::new()));
+    let reader_output = Arc::clone(&output);
+    let reader_thread = thread::spawn(move || {
+        let mut buf = [0u8; 4096];
+        loop {
+            match reader.read(&mut buf) {
+                Ok(0) | Err(_) => break,
+                Ok(n) => reader_output.lock().unwrap().extend_from_slice(&buf[..n]),
+            }
+        }
+    });
+
+    wait_for_output(&output, b"DOSH_PARITY_READY", Duration::from_secs(5));
+    wait_for_output(
+        &output,
+        "DOSH_GRAPH_⣿⣷⣤_λ_界".as_bytes(),
+        Duration::from_secs(5),
+    );
+
+    pair.master
+        .resize(PtySize {
+            rows: 31,
+            cols: 100,
+            pixel_width: 0,
+            pixel_height: 0,
+        })
+        .unwrap();
+    wait_for_resize(&observation_rx, (100, 31), Duration::from_secs(3));
+
+    writer.write_all(INPUT_BYTES).unwrap();
+    writer.flush().unwrap();
+    wait_for_input(&observation_rx, INPUT_BYTES, Duration::from_secs(3));
+
+    wait_for_output(&output, b"DOSH_PARITY_DONE", Duration::from_secs(5));
+    let status = child.wait().unwrap();
+    drop(writer);
+    drop(pair.master);
+    reader_thread.join().unwrap();
+    server.join().unwrap();
+
+    assert!(status.success(), "Dosh client exited with {status:?}");
+    let output = output.lock().unwrap();
+    assert!(
+        contains(&output, b"\x1b[38;2;12;200;90m"),
+        "true-color control sequence was not preserved: {output:?}"
+    );
+    assert!(
+        contains(&output, b"\x1b[?1049h"),
+        "alternate-screen mode was not preserved: {output:?}"
+    );
+    assert!(
+        contains(&output, b"\x1b[?1006h"),
+        "SGR mouse mode was not preserved: {output:?}"
+    );
+    assert!(
+        contains(&output, b"\x1b[?25h"),
+        "terminal cleanup did not restore the cursor: {output:?}"
+    );
+    #[cfg(windows)]
+    {
+        let code_page = fs::read(dir.path().join("restored-code-page.txt")).unwrap();
+        assert!(
+            contains(&code_page, b"437"),
+            "Dosh did not restore the original Windows console code page: {code_page:?}"
+        );
+    }
+}
+
+fn write_client_fixture(home: &Path, cache: &Path, port: u16) {
+    let config = ClientConfig {
+        server: "local".to_string(),
+        dosh_host: Some("127.0.0.1".to_string()),
+        dosh_port: port,
+        cache_attach_tickets: false,
+        credential_cache: cache.to_string_lossy().to_string(),
+        auth_preference: "ssh".to_string(),
+        predict: false,
+        predict_mode: "off".to_string(),
+        disconnect_status: false,
+        ..ClientConfig::default()
+    };
+    fs::write(
+        home.join(".config/dosh/client.toml"),
+        toml::to_string(&config).unwrap(),
+    )
+    .unwrap();
+
+    let credential = CachedCredentialWire {
+        server: "local".to_string(),
+        session: SESSION.to_string(),
+        mode: "read-write".to_string(),
+        udp_host: "127.0.0.1".to_string(),
+        udp_port: port,
+        client_id: CLIENT_ID,
+        session_key: SESSION_KEY,
+        session_key_id: protocol::session_key_id(&SESSION_KEY),
+        attach_ticket: Vec::new(),
+        attach_ticket_psk: [0x33; 32],
+        last_rendered_seq: 9,
+    };
+    fs::write(
+        cache.join("local_terminal_parity_read_write.bin"),
+        bincode::serialize(&credential).unwrap(),
+    )
+    .unwrap();
+}
+
+fn client_command(_dir: &Path, port: u16) -> CommandBuilder {
+    let client = env!("CARGO_BIN_EXE_dosh-client");
+    let args = [
+        "--local-auth",
+        "--predict-never",
+        "--session",
+        SESSION,
+        "--dosh-host",
+        "127.0.0.1",
+        "--dosh-port",
+        &port.to_string(),
+        "local",
+    ];
+
+    #[cfg(windows)]
+    {
+        let script = _dir.join("run-client.cmd");
+        let quoted_args = args
+            .iter()
+            .map(|arg| format!("\"{}\"", arg.replace('"', "\"\"")))
+            .collect::<Vec<_>>()
+            .join(" ");
+        fs::write(
+            &script,
+            format!(
+                "@echo off\r\nchcp 437 >nul\r\n\"{}\" {quoted_args}\r\nchcp > \"{}\"\r\n",
+                client.replace('"', "\"\""),
+                _dir.join("restored-code-page.txt").display()
+            ),
+        )
+        .unwrap();
+        let mut command = CommandBuilder::new("cmd.exe");
+        command.args(["/d", "/c", &script.to_string_lossy()]);
+        command
+    }
+
+    #[cfg(not(windows))]
+    {
+        let mut command = CommandBuilder::new(client);
+        command.args(args);
+        command
+    }
+}
+
+fn run_fake_terminal_server(socket: UdpSocket, observations: mpsc::Sender<ServerObservation>) {
+    let mut peer = None;
+    let mut next_output_chunk = 0usize;
+    let mut received_input = Vec::new();
+    let mut buf = [0u8; 65535];
+    loop {
+        let (n, source) = socket.recv_from(&mut buf).unwrap();
+        let packet = protocol::decode(&buf[..n]).unwrap();
+        match packet.header.kind {
+            PacketKind::ResumeRequest => {
+                let plain =
+                    protocol::decrypt_body(&packet, &SESSION_KEY, CLIENT_TO_SERVER).unwrap();
+                let request: protocol::ResumeRequest = protocol::from_body(&plain).unwrap();
+                assert_eq!(request.session, SESSION);
+                peer = Some(source);
+                send_frame(
+                    &socket,
+                    source,
+                    PacketKind::ResumeOk,
+                    1,
+                    10,
+                    concat!(
+                        "\x1b[?1049h",
+                        "\x1b[?1003h",
+                        "\x1b[?1006h",
+                        "\x1b[?2004h",
+                        "\x1b[2J\x1b[H",
+                        "DOSH_PARITY_READY"
+                    )
+                    .as_bytes(),
+                    true,
+                    false,
+                );
+            }
+            PacketKind::Ack if next_output_chunk < OUTPUT_CHUNKS.len() => {
+                let source = peer.unwrap_or(source);
+                send_frame(
+                    &socket,
+                    source,
+                    PacketKind::Frame,
+                    2 + next_output_chunk as u64,
+                    11 + next_output_chunk as u64,
+                    OUTPUT_CHUNKS[next_output_chunk],
+                    false,
+                    false,
+                );
+                next_output_chunk += 1;
+            }
+            PacketKind::Input => {
+                let plain =
+                    protocol::decrypt_body(&packet, &SESSION_KEY, CLIENT_TO_SERVER).unwrap();
+                let input: Input = protocol::from_body(&plain).unwrap();
+                received_input.extend_from_slice(&input.bytes);
+                observations
+                    .send(ServerObservation::Input(input.bytes))
+                    .unwrap();
+                if contains(&received_input, INPUT_BYTES) {
+                    send_frame(
+                        &socket,
+                        peer.unwrap_or(source),
+                        PacketKind::Frame,
+                        20,
+                        16,
+                        b"\x1b[?1003l\x1b[?1006l\x1b[?1049lDOSH_PARITY_DONE",
+                        false,
+                        true,
+                    );
+                }
+            }
+            PacketKind::Resize => {
+                let plain =
+                    protocol::decrypt_body(&packet, &SESSION_KEY, CLIENT_TO_SERVER).unwrap();
+                let resize: Resize = protocol::from_body(&plain).unwrap();
+                observations
+                    .send(ServerObservation::Resize(resize.cols, resize.rows))
+                    .unwrap();
+            }
+            PacketKind::Ping => {
+                let pong = protocol::encode_encrypted(
+                    PacketKind::Pong,
+                    CLIENT_ID,
+                    4,
+                    11,
+                    &SESSION_KEY,
+                    SERVER_TO_CLIENT,
+                    b"",
+                )
+                .unwrap();
+                socket.send_to(&pong, source).unwrap();
+            }
+            PacketKind::Detach => break,
+            _ => {}
+        }
+        if received_input
+            .windows(INPUT_BYTES.len())
+            .any(|bytes| bytes == INPUT_BYTES)
+        {
+            break;
+        }
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn send_frame(
+    socket: &UdpSocket,
+    peer: std::net::SocketAddr,
+    kind: PacketKind,
+    packet_seq: u64,
+    output_seq: u64,
+    bytes: &[u8],
+    snapshot: bool,
+    closed: bool,
+) {
+    let frame = Frame {
+        session: SESSION.to_string(),
+        output_seq,
+        bytes: bytes.to_vec(),
+        snapshot,
+        closed,
+    };
+    let body = protocol::to_body(&frame).unwrap();
+    let packet = protocol::encode_encrypted(
+        kind,
+        CLIENT_ID,
+        packet_seq,
+        output_seq.saturating_sub(1),
+        &SESSION_KEY,
+        SERVER_TO_CLIENT,
+        &body,
+    )
+    .unwrap();
+    socket.send_to(&packet, peer).unwrap();
+}
+
+fn wait_for_output(output: &Arc<Mutex<Vec<u8>>>, needle: &[u8], timeout: Duration) {
+    let deadline = Instant::now() + timeout;
+    while Instant::now() < deadline {
+        if contains(&output.lock().unwrap(), needle) {
+            return;
+        }
+        thread::sleep(Duration::from_millis(10));
+    }
+    panic!(
+        "terminal output did not contain {:?}: {:?}",
+        String::from_utf8_lossy(needle),
+        output.lock().unwrap()
+    );
+}
+
+fn wait_for_resize(
+    observations: &mpsc::Receiver<ServerObservation>,
+    expected: (u16, u16),
+    timeout: Duration,
+) {
+    let deadline = Instant::now() + timeout;
+    while let Some(remaining) = deadline.checked_duration_since(Instant::now()) {
+        match observations.recv_timeout(remaining) {
+            Ok(ServerObservation::Resize(cols, rows)) if (cols, rows) == expected => return,
+            Ok(_) => {}
+            Err(err) => panic!("terminal resize {expected:?} was not delivered: {err}"),
+        }
+    }
+    panic!("terminal resize {expected:?} was not delivered");
+}
+
+fn wait_for_input(
+    observations: &mpsc::Receiver<ServerObservation>,
+    expected: &[u8],
+    timeout: Duration,
+) {
+    let deadline = Instant::now() + timeout;
+    let mut input = Vec::new();
+    while let Some(remaining) = deadline.checked_duration_since(Instant::now()) {
+        match observations.recv_timeout(remaining) {
+            Ok(ServerObservation::Input(bytes)) => {
+                input.extend_from_slice(&bytes);
+                if contains(&input, expected) {
+                    return;
+                }
+            }
+            Ok(_) => {}
+            Err(err) => panic!("terminal input was not delivered: {err}; got {input:?}"),
+        }
+    }
+    panic!("terminal input was not delivered; got {input:?}");
+}
+
+fn contains(haystack: &[u8], needle: &[u8]) -> bool {
+    !needle.is_empty() && haystack.windows(needle.len()).any(|bytes| bytes == needle)
+}