fix(ssh): handle invalid key comment encoding

crates/dbx-core/src/db/ssh_tunnel.rs

@@ -1,8 +1,11 @@
 use std::collections::HashMap;
+use std::fs;
 use std::sync::Arc;
 
+use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
+use base64::Engine;
 use russh::client::{self, Config, Handle};
-use russh::keys::{key::PrivateKeyWithHashAlg, load_secret_key};
+use russh::keys::{decode_secret_key, key::PrivateKeyWithHashAlg, PrivateKey};
 use russh::ChannelMsg;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpListener;
@@ -62,7 +65,8 @@ async fn connect_and_authenticate(
         validate_file_path(ssh_key_path, |_| false)?;
 
         let passphrase = if ssh_key_passphrase.is_empty() { None } else { Some(ssh_key_passphrase) };
-        let key_pair = load_secret_key(ssh_key_path, passphrase).map_err(|e| format!("Failed to load SSH key: {e}"))?;
+        let key_pair =
+            load_ssh_private_key(ssh_key_path, passphrase).map_err(|e| format!("Failed to load SSH key: {e}"))?;
         let auth_res = tokio::time::timeout(
             connect_timeout,
             session.authenticate_publickey(
@@ -94,6 +98,141 @@ async fn connect_and_authenticate(
     Ok(session)
 }
 
+fn load_ssh_private_key(path: &str, passphrase: Option<&str>) -> Result<PrivateKey, String> {
+    let secret = fs::read_to_string(path).map_err(|e| e.to_string())?;
+    match decode_secret_key(&secret, passphrase) {
+        Ok(key) => Ok(key),
+        Err(err) if is_ssh_key_character_encoding_error(&err.to_string()) => {
+            let sanitized = sanitize_unencrypted_openssh_comment(&secret)?;
+            decode_secret_key(&sanitized, passphrase).map_err(|retry_err| retry_err.to_string())
+        }
+        Err(err) => Err(err.to_string()),
+    }
+}
+
+fn is_ssh_key_character_encoding_error(error: &str) -> bool {
+    error.contains("SshKey: character encoding invalid")
+}
+
+fn sanitize_unencrypted_openssh_comment(secret: &str) -> Result<String, String> {
+    const BEGIN: &str = "-----BEGIN OPENSSH PRIVATE KEY-----";
+    const END: &str = "-----END OPENSSH PRIVATE KEY-----";
+
+    if !secret.lines().any(|line| line == BEGIN) {
+        return Err("SSH key comment encoding is invalid and the key is not an OpenSSH private key".to_string());
+    }
+
+    let body = secret.lines().filter(|line| !line.starts_with("-----")).collect::<String>();
+    let mut bytes =
+        BASE64_STANDARD.decode(body.as_bytes()).map_err(|e| format!("OpenSSH key base64 decode failed: {e}"))?;
+
+    sanitize_unencrypted_openssh_comment_bytes(&mut bytes)?;
+
+    Ok(format!("{BEGIN}\n{}\n{END}\n", BASE64_STANDARD.encode(bytes)))
+}
+
+fn sanitize_unencrypted_openssh_comment_bytes(bytes: &mut Vec<u8>) -> Result<(), String> {
+    const AUTH_MAGIC: &[u8] = b"openssh-key-v1\0";
+
+    if !bytes.starts_with(AUTH_MAGIC) {
+        return Err("OpenSSH key header is invalid".to_string());
+    }
+
+    let mut pos = AUTH_MAGIC.len();
+    let ciphername = read_ssh_string(bytes, &mut pos)?;
+    if ciphername != b"none" {
+        return Err("SSH key comment encoding is invalid and encrypted OpenSSH keys cannot be sanitized".to_string());
+    }
+
+    let _kdfname = read_ssh_string(bytes, &mut pos)?;
+    let _kdfoptions = read_ssh_string(bytes, &mut pos)?;
+    let key_count = read_u32(bytes, &mut pos)?;
+    if key_count != 1 {
+        return Err("OpenSSH keys with multiple private keys are unsupported".to_string());
+    }
+
+    let _public_key = read_ssh_string(bytes, &mut pos)?;
+    let private_blob_len_pos = pos;
+    let private_blob = read_ssh_string(bytes, &mut pos)?;
+    let patched_private_blob = sanitize_private_blob_comment(private_blob)?;
+    let patched_private_blob_len = (patched_private_blob.len() as u32).to_be_bytes();
+
+    bytes.splice(private_blob_len_pos..pos, patched_private_blob_len.into_iter().chain(patched_private_blob));
+
+    Ok(())
+}
+
+fn sanitize_private_blob_comment(blob: &[u8]) -> Result<Vec<u8>, String> {
+    let unpadded_end = blob
+        .len()
+        .checked_sub(openssh_padding_len(blob)?)
+        .ok_or_else(|| "OpenSSH private key padding is invalid".to_string())?;
+    let comment_len_pos = find_trailing_ssh_string_len_pos(&blob[..unpadded_end])
+        .ok_or_else(|| "OpenSSH private key comment field was not found".to_string())?;
+
+    let mut patched = Vec::with_capacity(blob.len());
+    patched.extend_from_slice(&blob[..comment_len_pos]);
+    patched.extend_from_slice(&0u32.to_be_bytes());
+
+    let padding_len = padding_len_for_block(patched.len(), 8);
+    for value in 1..=padding_len {
+        patched.push(value as u8);
+    }
+
+    Ok(patched)
+}
+
+fn openssh_padding_len(bytes: &[u8]) -> Result<usize, String> {
+    for len in (1..=16).rev() {
+        if bytes.len() >= len
+            && bytes[bytes.len() - len..].iter().enumerate().all(|(index, byte)| *byte == (index + 1) as u8)
+        {
+            return Ok(len);
+        }
+    }
+
+    Err("OpenSSH private key padding is invalid".to_string())
+}
+
+fn find_trailing_ssh_string_len_pos(bytes: &[u8]) -> Option<usize> {
+    (8..bytes.len().saturating_sub(3)).rev().find(|pos| {
+        let Some(len_bytes) = bytes.get(*pos..*pos + 4) else {
+            return false;
+        };
+        let len = u32::from_be_bytes(len_bytes.try_into().expect("slice length checked")) as usize;
+        pos.checked_add(4).and_then(|value| value.checked_add(len)) == Some(bytes.len())
+    })
+}
+
+fn padding_len_for_block(len: usize, block_size: usize) -> usize {
+    let remainder = len % block_size;
+    if remainder == 0 {
+        block_size
+    } else {
+        block_size - remainder
+    }
+}
+
+fn read_ssh_string<'a>(bytes: &'a [u8], pos: &mut usize) -> Result<&'a [u8], String> {
+    let len = read_u32(bytes, pos)? as usize;
+    let end = pos.checked_add(len).ok_or_else(|| "OpenSSH key field length is invalid".to_string())?;
+    if end > bytes.len() {
+        return Err("OpenSSH key field is truncated".to_string());
+    }
+
+    let value = &bytes[*pos..end];
+    *pos = end;
+    Ok(value)
+}
+
+fn read_u32(bytes: &[u8], pos: &mut usize) -> Result<u32, String> {
+    let end = pos.checked_add(4).ok_or_else(|| "OpenSSH key field length is invalid".to_string())?;
+    let value = bytes.get(*pos..end).ok_or_else(|| "OpenSSH key field is truncated".to_string())?;
+    *pos = end;
+
+    Ok(u32::from_be_bytes(value.try_into().map_err(|_| "OpenSSH key field length is invalid".to_string())?))
+}
+
 /// Accept connections on the local listener and forward them through the SSH session.
 /// Returns when the SSH session dies (listener error or session.is_closed()).
 async fn forward_loop(session: &Handle<SshClient>, listener: &TcpListener, remote_host: &str, remote_port: u16) {
@@ -515,9 +654,44 @@ fn plan_chain(
 
 #[cfg(test)]
 mod tests {
-    use super::{effective_hop_timeout, plan_chain, PlannedTunnel, TunnelManager};
+    use super::{
+        effective_hop_timeout, openssh_padding_len, plan_chain, read_ssh_string,
+        sanitize_unencrypted_openssh_comment_bytes, PlannedTunnel, TunnelManager,
+    };
     use crate::models::connection::{default_ssh_connect_timeout_secs, SshTunnelConfig};
 
+    fn push_u32(bytes: &mut Vec<u8>, value: u32) {
+        bytes.extend_from_slice(&value.to_be_bytes());
+    }
+
+    fn push_ssh_string(bytes: &mut Vec<u8>, value: &[u8]) {
+        push_u32(bytes, value.len() as u32);
+        bytes.extend_from_slice(value);
+    }
+
+    fn padded_private_blob(comment: &[u8]) -> Vec<u8> {
+        let mut blob = Vec::new();
+        push_u32(&mut blob, 7);
+        push_u32(&mut blob, 7);
+        blob.extend_from_slice(b"fake-private-key");
+        push_ssh_string(&mut blob, comment);
+        for value in 1..=(8 - (blob.len() % 8)) {
+            blob.push(value as u8);
+        }
+        blob
+    }
+
+    fn openssh_container(private_blob: &[u8]) -> Vec<u8> {
+        let mut bytes = b"openssh-key-v1\0".to_vec();
+        push_ssh_string(&mut bytes, b"none");
+        push_ssh_string(&mut bytes, b"none");
+        push_ssh_string(&mut bytes, b"");
+        push_u32(&mut bytes, 1);
+        push_ssh_string(&mut bytes, b"fake-public-key");
+        push_ssh_string(&mut bytes, private_blob);
+        bytes
+    }
+
     fn hop(id: &str, host: &str, port: u16) -> SshTunnelConfig {
         SshTunnelConfig {
             id: id.to_string(),
@@ -567,6 +741,25 @@ mod tests {
         assert_eq!(effective_hop_timeout(&tunnel), default_ssh_connect_timeout_secs());
     }
 
+    #[test]
+    fn sanitizes_invalid_openssh_private_key_comment() {
+        let mut key = openssh_container(&padded_private_blob(&[0xff, 0xfe, b'a']));
+
+        sanitize_unencrypted_openssh_comment_bytes(&mut key).unwrap();
+
+        let mut pos = b"openssh-key-v1\0".len();
+        assert_eq!(read_ssh_string(&key, &mut pos).unwrap(), b"none");
+        let _kdfname = read_ssh_string(&key, &mut pos).unwrap();
+        let _kdfoptions = read_ssh_string(&key, &mut pos).unwrap();
+        pos += 4;
+        let _public_key = read_ssh_string(&key, &mut pos).unwrap();
+        let private_blob = read_ssh_string(&key, &mut pos).unwrap();
+        let unpadded_end = private_blob.len() - openssh_padding_len(private_blob).unwrap();
+        let comment_len_pos = unpadded_end - 4;
+
+        assert_eq!(&private_blob[comment_len_pos..unpadded_end], &0u32.to_be_bytes());
+    }
+
     #[tokio::test]
     async fn local_port_reuses_existing_chain_entry() {
         let manager = TunnelManager::new();