#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
use futures::Future;
use super::*;
mod compress {
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use crate::keys::ssh_key::rand_core::OsRng;
use keys::PrivateKeyWithHashAlg;
use log::debug;
use ssh_key::PrivateKey;
use super::server::{Server as _, Session};
use super::*;
use crate::server::Msg;
#[tokio::test]
async fn compress_local_test() {
let _ = env_logger::try_init();
let client_key = PrivateKey::random(&mut OsRng, ssh_key::Algorithm::Ed25519).unwrap();
let mut config = server::Config::default();
config.preferred = Preferred::COMPRESSED;
config.inactivity_timeout = None; config.auth_rejection_time = std::time::Duration::from_secs(3);
config
.keys
.push(PrivateKey::random(&mut OsRng, ssh_key::Algorithm::Ed25519).unwrap());
let config = Arc::new(config);
let mut sh = Server {
clients: Arc::new(Mutex::new(HashMap::new())),
id: 0,
};
let socket = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = socket.local_addr().unwrap();
tokio::spawn(async move {
let (socket, _) = socket.accept().await.unwrap();
let server = sh.new_client(socket.peer_addr().ok());
server::run_stream(config, socket, server).await.unwrap();
});
let mut config = client::Config::default();
config.preferred = Preferred::COMPRESSED;
let config = Arc::new(config);
let mut session = client::connect(config, addr, Client {}).await.unwrap();
let authenticated = session
.authenticate_publickey(
std::env::var("USER").unwrap_or("user".to_owned()),
PrivateKeyWithHashAlg::new(
Arc::new(client_key),
session.best_supported_rsa_hash().await.unwrap().flatten(),
),
)
.await
.unwrap()
.success();
assert!(authenticated);
let mut channel = session.channel_open_session().await.unwrap();
let data = &b"Hello, world!"[..];
channel.data(data).await.unwrap();
let msg = channel.wait().await.unwrap();
match msg {
ChannelMsg::Data { data: msg_data } => {
assert_eq!(*data, *msg_data)
}
msg => panic!("Unexpected message {msg:?}"),
}
}
#[derive(Clone)]
struct Server {
clients: Arc<Mutex<HashMap<(usize, ChannelId), super::server::Handle>>>,
id: usize,
}
impl server::Server for Server {
type Handler = Self;
fn new_client(&mut self, _: Option<std::net::SocketAddr>) -> Self {
let s = self.clone();
self.id += 1;
s
}
}
impl server::Handler for Server {
type Error = super::Error;
async fn channel_open_session(
&mut self,
channel: Channel<Msg>,
session: &mut Session,
) -> Result<bool, Self::Error> {
{
let mut clients = self.clients.lock().unwrap();
clients.insert((self.id, channel.id()), session.handle());
}
Ok(true)
}
async fn auth_publickey(
&mut self,
_: &str,
_: &crate::keys::ssh_key::PublicKey,
) -> Result<server::Auth, Self::Error> {
debug!("auth_publickey");
Ok(server::Auth::Accept)
}
async fn data(
&mut self,
channel: ChannelId,
data: &[u8],
session: &mut Session,
) -> Result<(), Self::Error> {
debug!("server data = {:?}", std::str::from_utf8(data));
session.data(channel, CryptoVec::from_slice(data))?;
Ok(())
}
}
struct Client {}
impl client::Handler for Client {
type Error = super::Error;
async fn check_server_key(
&mut self,
_server_public_key: &crate::keys::ssh_key::PublicKey,
) -> Result<bool, Self::Error> {
Ok(true)
}
}
}
mod channels {
use elliptic_curve::rand_core::OsRng;
use keys::PrivateKeyWithHashAlg;
use server::Session;
use ssh_key::PrivateKey;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use super::*;
use crate::CryptoVec;
async fn test_session<RC, RS, CH, SH, F1, F2>(
client_handler: CH,
server_handler: SH,
run_client: RC,
run_server: RS,
) where
RC: FnOnce(crate::client::Handle<CH>) -> F1 + Send + Sync + 'static,
RS: FnOnce(crate::server::Handle) -> F2 + Send + Sync + 'static,
F1: Future<Output = crate::client::Handle<CH>> + Send + Sync + 'static,
F2: Future<Output = crate::server::Handle> + Send + Sync + 'static,
CH: crate::client::Handler + Send + Sync + 'static,
SH: crate::server::Handler + Send + Sync + 'static,
{
use std::sync::Arc;
use crate::*;
let _ = env_logger::try_init();
let client_key = PrivateKey::random(&mut OsRng, ssh_key::Algorithm::Ed25519).unwrap();
let mut config = server::Config::default();
config.inactivity_timeout = None;
config.auth_rejection_time = std::time::Duration::from_secs(3);
config
.keys
.push(PrivateKey::random(&mut OsRng, ssh_key::Algorithm::Ed25519).unwrap());
let config = Arc::new(config);
let socket = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = socket.local_addr().unwrap();
let server_join = tokio::spawn(async move {
let (socket, _) = socket.accept().await.unwrap();
server::run_stream(config, socket, server_handler)
.await
.map_err(|_| ())
.unwrap()
});
let client_join = tokio::spawn(async move {
let config = Arc::new(client::Config::default());
let mut session = client::connect(config, addr, client_handler)
.await
.map_err(|_| ())
.unwrap();
let authenticated = session
.authenticate_publickey(
std::env::var("USER").unwrap_or("user".to_owned()),
PrivateKeyWithHashAlg::new(Arc::new(client_key), None),
)
.await
.unwrap();
assert!(authenticated.success());
session
});
let (server_session, client_session) = tokio::join!(server_join, client_join);
let client_handle = tokio::spawn(run_client(client_session.unwrap()));
let server_handle = tokio::spawn(run_server(server_session.unwrap().handle()));
let (server_session, client_session) = tokio::join!(server_handle, client_handle);
assert!(server_session.is_ok());
assert!(client_session.is_ok());
drop(client_session);
drop(server_session);
}
#[tokio::test]
async fn test_server_channels() {
#[derive(Debug)]
struct Client {}
impl client::Handler for Client {
type Error = crate::Error;
async fn check_server_key(
&mut self,
_server_public_key: &crate::keys::ssh_key::PublicKey,
) -> Result<bool, Self::Error> {
Ok(true)
}
async fn data(
&mut self,
channel: ChannelId,
data: &[u8],
session: &mut client::Session,
) -> Result<(), Self::Error> {
assert_eq!(data, &b"hello world!"[..]);
session.data(channel, CryptoVec::from_slice(&b"hey there!"[..]))?;
Ok(())
}
}
struct ServerHandle {
did_auth: Option<tokio::sync::oneshot::Sender<()>>,
}
impl ServerHandle {
fn get_auth_waiter(&mut self) -> tokio::sync::oneshot::Receiver<()> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.did_auth = Some(tx);
rx
}
}
impl server::Handler for ServerHandle {
type Error = crate::Error;
async fn auth_publickey(
&mut self,
_: &str,
_: &crate::keys::ssh_key::PublicKey,
) -> Result<server::Auth, Self::Error> {
Ok(server::Auth::Accept)
}
async fn auth_succeeded(&mut self, _session: &mut Session) -> Result<(), Self::Error> {
if let Some(a) = self.did_auth.take() {
a.send(()).unwrap();
}
Ok(())
}
}
let mut sh = ServerHandle { did_auth: None };
let a = sh.get_auth_waiter();
test_session(
Client {},
sh,
|c| async move { c },
|s| async move {
a.await.unwrap();
let mut ch = s.channel_open_session().await.unwrap();
ch.data(&b"hello world!"[..]).await.unwrap();
let msg = ch.wait().await.unwrap();
if let ChannelMsg::Data { data } = msg {
assert_eq!(data.as_ref(), &b"hey there!"[..]);
} else {
panic!("Unexpected message {msg:?}");
}
s
},
)
.await;
}
#[tokio::test]
async fn test_channel_streams() {
#[derive(Debug)]
struct Client {}
impl client::Handler for Client {
type Error = crate::Error;
async fn check_server_key(
&mut self,
_server_public_key: &crate::keys::ssh_key::PublicKey,
) -> Result<bool, Self::Error> {
Ok(true)
}
}
struct ServerHandle {
channel: Option<tokio::sync::oneshot::Sender<Channel<server::Msg>>>,
}
impl ServerHandle {
fn get_channel_waiter(
&mut self,
) -> tokio::sync::oneshot::Receiver<Channel<server::Msg>> {
let (tx, rx) = tokio::sync::oneshot::channel::<Channel<server::Msg>>();
self.channel = Some(tx);
rx
}
}
impl server::Handler for ServerHandle {
type Error = crate::Error;
async fn auth_publickey(
&mut self,
_: &str,
_: &crate::keys::ssh_key::PublicKey,
) -> Result<server::Auth, Self::Error> {
Ok(server::Auth::Accept)
}
async fn channel_open_session(
&mut self,
channel: Channel<server::Msg>,
_session: &mut server::Session,
) -> Result<bool, Self::Error> {
if let Some(a) = self.channel.take() {
println!("channel open session {a:?}");
a.send(channel).unwrap();
}
Ok(true)
}
}
let mut sh = ServerHandle { channel: None };
let scw = sh.get_channel_waiter();
test_session(
Client {},
sh,
|client| async move {
let ch = client.channel_open_session().await.unwrap();
let mut stream = ch.into_stream();
stream.write_all(&b"request"[..]).await.unwrap();
let mut buf = Vec::new();
stream.read_buf(&mut buf).await.unwrap();
assert_eq!(&buf, &b"response"[..]);
stream.write_all(&b"reply"[..]).await.unwrap();
client
},
|server| async move {
let channel = scw.await.unwrap();
let mut stream = channel.into_stream();
let mut buf = Vec::new();
stream.read_buf(&mut buf).await.unwrap();
assert_eq!(&buf, &b"request"[..]);
stream.write_all(&b"response"[..]).await.unwrap();
buf.clear();
stream.read_buf(&mut buf).await.unwrap();
assert_eq!(&buf, &b"reply"[..]);
server
},
)
.await;
}
#[tokio::test]
async fn test_channel_objects() {
#[derive(Debug)]
struct Client {}
impl client::Handler for Client {
type Error = crate::Error;
async fn check_server_key(
&mut self,
_server_public_key: &crate::keys::ssh_key::PublicKey,
) -> Result<bool, Self::Error> {
Ok(true)
}
}
struct ServerHandle {}
impl ServerHandle {}
impl server::Handler for ServerHandle {
type Error = crate::Error;
async fn auth_publickey(
&mut self,
_: &str,
_: &crate::keys::ssh_key::PublicKey,
) -> Result<server::Auth, Self::Error> {
Ok(server::Auth::Accept)
}
async fn channel_open_session(
&mut self,
mut channel: Channel<server::Msg>,
_session: &mut Session,
) -> Result<bool, Self::Error> {
tokio::spawn(async move {
while let Some(msg) = channel.wait().await {
match msg {
ChannelMsg::Data { data } => {
channel.data(&data[..]).await.unwrap();
channel.close().await.unwrap();
break;
}
_ => {}
}
}
});
Ok(true)
}
}
let sh = ServerHandle {};
test_session(
Client {},
sh,
|c| async move {
let mut ch = c.channel_open_session().await.unwrap();
ch.data(&b"hello world!"[..]).await.unwrap();
let msg = ch.wait().await.unwrap();
if let ChannelMsg::Data { data } = msg {
assert_eq!(data.as_ref(), &b"hello world!"[..]);
} else {
panic!("Unexpected message {msg:?}");
}
assert!(ch.wait().await.is_none());
c
},
|s| async move { s },
)
.await;
}
#[tokio::test]
async fn test_channel_window_size() {
#[derive(Debug)]
struct Client {}
impl client::Handler for Client {
type Error = crate::Error;
async fn check_server_key(
&mut self,
_server_public_key: &crate::keys::ssh_key::PublicKey,
) -> Result<bool, Self::Error> {
Ok(true)
}
}
struct ServerHandle {
channel: Option<tokio::sync::oneshot::Sender<Channel<server::Msg>>>,
}
impl ServerHandle {
fn get_channel_waiter(
&mut self,
) -> tokio::sync::oneshot::Receiver<Channel<server::Msg>> {
let (tx, rx) = tokio::sync::oneshot::channel::<Channel<server::Msg>>();
self.channel = Some(tx);
rx
}
}
impl server::Handler for ServerHandle {
type Error = crate::Error;
async fn auth_publickey(
&mut self,
_: &str,
_: &crate::keys::ssh_key::PublicKey,
) -> Result<server::Auth, Self::Error> {
Ok(server::Auth::Accept)
}
async fn channel_open_session(
&mut self,
channel: Channel<server::Msg>,
_session: &mut server::Session,
) -> Result<bool, Self::Error> {
if let Some(a) = self.channel.take() {
println!("channel open session {a:?}");
a.send(channel).unwrap();
}
Ok(true)
}
}
let mut sh = ServerHandle { channel: None };
let scw = sh.get_channel_waiter();
test_session(
Client {},
sh,
|client| async move {
let ch = client.channel_open_session().await.unwrap();
let mut writer_1 = ch.make_writer();
let jh_1 = tokio::spawn(async move {
let buf = [1u8; 1024 * 64];
assert!(writer_1.write_all(&buf).await.is_ok());
});
let mut writer_2 = ch.make_writer();
let jh_2 = tokio::spawn(async move {
let buf = [2u8; 1024 * 64];
assert!(writer_2.write_all(&buf).await.is_ok());
});
assert!(tokio::try_join!(jh_1, jh_2).is_ok());
client
},
|server| async move {
let mut channel = scw.await.unwrap();
let mut total_data = 2 * 1024 * 64;
while let Some(msg) = channel.wait().await {
match msg {
ChannelMsg::Data { data } => {
total_data -= data.len();
if total_data == 0 {
break;
}
}
_ => panic!("Unexpected message {msg:?}"),
}
}
server
},
)
.await;
}
}
mod server_kex_junk {
use std::sync::Arc;
use tokio::io::AsyncWriteExt;
use super::server::Server as _;
use super::*;
#[tokio::test]
async fn server_kex_junk_test() {
let _ = env_logger::try_init();
let config = server::Config::default();
let config = Arc::new(config);
let mut sh = Server {};
let socket = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = socket.local_addr().unwrap();
tokio::spawn(async move {
let mut client_stream = tokio::net::TcpStream::connect(addr).await.unwrap();
client_stream
.write_all(b"SSH-2.0-Client_1.0\r\n")
.await
.unwrap();
client_stream.write_all(&[0, 0, 0, 2, 0, 99]).await.unwrap();
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
});
let (socket, _) = socket.accept().await.unwrap();
let server = sh.new_client(socket.peer_addr().ok());
let rs = server::run_stream(config, socket, server).await.unwrap();
assert!(rs.await.is_err());
}
#[derive(Clone)]
struct Server {}
impl server::Server for Server {
type Handler = Self;
fn new_client(&mut self, _: Option<std::net::SocketAddr>) -> Self {
self.clone()
}
}
impl server::Handler for Server {
type Error = super::Error;
}
}