#[cfg(feature = "alloc")]
use alloc::vec::Vec;
use crate::error::{Error, Result};
use crate::format::Reader;
pub const SSH_MSG_PING: u8 = 192;
pub const SSH_MSG_PONG: u8 = 193;
#[cfg(feature = "alloc")]
pub fn encode_ping(data: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(1 + 4 + data.len());
out.push(SSH_MSG_PING);
out.extend_from_slice(&(data.len() as u32).to_be_bytes());
out.extend_from_slice(data);
out
}
#[cfg(feature = "alloc")]
pub fn encode_pong(data: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(1 + 4 + data.len());
out.push(SSH_MSG_PONG);
out.extend_from_slice(&(data.len() as u32).to_be_bytes());
out.extend_from_slice(data);
out
}
#[cfg(feature = "alloc")]
pub fn pong_for_ping(payload: &[u8]) -> Result<Vec<u8>> {
let mut r = Reader::new(payload);
let msg = r.read_u8()?;
if msg != SSH_MSG_PING {
return Err(Error::Protocol("ping: not a SSH_MSG_PING"));
}
let data = r.read_string()?;
Ok(encode_pong(data))
}
#[cfg(all(test, feature = "alloc"))]
mod tests {
use super::*;
#[test]
fn ping_roundtrips_into_pong_with_echoed_data() {
let ping = encode_ping(b"hello");
assert_eq!(ping[0], SSH_MSG_PING);
let pong = pong_for_ping(&ping).expect("valid ping");
assert_eq!(pong[0], SSH_MSG_PONG);
let mut r = Reader::new(&pong);
assert_eq!(r.read_u8().unwrap(), SSH_MSG_PONG);
assert_eq!(r.read_string().unwrap(), b"hello");
}
#[test]
fn empty_ping_echoes_empty_pong() {
let pong = pong_for_ping(&encode_ping(b"")).unwrap();
assert_eq!(pong, encode_pong(b""));
let mut r = Reader::new(&pong);
assert_eq!(r.read_u8().unwrap(), SSH_MSG_PONG);
assert_eq!(r.read_string().unwrap(), b"");
}
#[test]
fn pong_for_non_ping_is_rejected() {
assert!(pong_for_ping(&encode_pong(b"x")).is_err());
assert!(pong_for_ping(&[]).is_err());
assert!(pong_for_ping(&[SSH_MSG_PING]).is_err());
assert!(pong_for_ping(&[SSH_MSG_PING, 0, 0, 0, 5, b'a']).is_err());
}
}