migrate cluelessh-faked to cluelessh-tokio

This commit is contained in:
nora 2024-08-25 15:01:22 +02:00
parent 01d6a861f1
commit b6d0675976
11 changed files with 513 additions and 137 deletions

View file

@ -32,7 +32,7 @@ enum ChannelState {
our_window_size: u32,
/// For validation only.
our_max_packet_size: u32,
update_message: ChannelOpen,
update_message: ChannelKind,
},
Open(Channel),
}
@ -71,7 +71,7 @@ pub struct ChannelUpdate {
pub enum ChannelUpdateKind {
Success,
Failure,
Open(ChannelOpen),
Open(ChannelKind),
OpenFailed { code: u32, message: String },
Request(ChannelRequest),
Data { data: Vec<u8> },
@ -80,7 +80,7 @@ pub enum ChannelUpdateKind {
Closed,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ChannelOpen {
pub enum ChannelKind {
Session,
}
#[derive(Debug)]
@ -173,7 +173,7 @@ impl ChannelsState {
debug!(%channel_type, %sender_channel, "Receving channel open");
let update_message = match channel_type {
"session" => ChannelOpen::Session,
"session" => ChannelKind::Session,
_ => {
self.packets_to_send
.push_back(Packet::new_msg_channel_open_failure(
@ -512,7 +512,7 @@ impl ChannelsState {
}
/// Create a new channel
pub fn create_channel(&mut self, kind: ChannelOpen) -> ChannelNumber {
pub fn create_channel(&mut self, kind: ChannelKind) -> ChannelNumber {
let our_number = self.next_channel_id;
self.next_channel_id = ChannelNumber(
self.next_channel_id
@ -521,7 +521,7 @@ impl ChannelsState {
.expect("created too many channels"),
);
assert_eq!(kind, ChannelOpen::Session, "TODO");
assert_eq!(kind, ChannelKind::Session, "TODO");
let our_window_size = 2097152; // same as OpenSSH
let our_max_packet_size = 32768; // same as OpenSSH

View file

@ -44,8 +44,9 @@ impl ServerConnection {
self.transport.send_plaintext_packet(to_send);
}
if auth.is_authenticated() {
self.state =
ServerConnectionState::Open(cluelessh_connection::ChannelsState::new(true));
self.state = ServerConnectionState::Open(
cluelessh_connection::ChannelsState::new(true),
);
}
}
ServerConnectionState::Open(con) => {
@ -94,6 +95,20 @@ impl ServerConnection {
}
}
}
pub fn channels(&mut self) -> Option<&mut cluelessh_connection::ChannelsState> {
match &mut self.state {
ServerConnectionState::Open(channels) => Some(channels),
_ => None,
}
}
pub fn auth(&mut self) -> Option<&mut auth::BadAuth> {
match &mut self.state {
ServerConnectionState::Auth(auth) => Some(auth),
_ => None,
}
}
}
pub struct ClientConnection {
@ -108,7 +123,10 @@ enum ClientConnectionState {
}
impl ClientConnection {
pub fn new(transport: cluelessh_transport::client::ClientConnection, auth: auth::ClientAuth) -> Self {
pub fn new(
transport: cluelessh_transport::client::ClientConnection,
auth: auth::ClientAuth,
) -> Self {
Self {
transport,
state: ClientConnectionState::Setup(Some(auth)),
@ -139,8 +157,9 @@ impl ClientConnection {
self.transport.send_plaintext_packet(to_send);
}
if auth.is_authenticated() {
self.state =
ClientConnectionState::Open(cluelessh_connection::ChannelsState::new(false));
self.state = ClientConnectionState::Open(
cluelessh_connection::ChannelsState::new(false),
);
}
}
ClientConnectionState::Open(con) => {
@ -227,6 +246,18 @@ pub mod auth {
is_authenticated: bool,
}
pub enum ServerRequest {
VerifyPassword {
user: String,
password: String,
},
VerifyPubkey {
session_identifier: [u8; 32],
user: String,
pubkey: Vec<u8>,
},
}
impl BadAuth {
pub fn new() -> Self {
Self {
@ -320,6 +351,10 @@ pub mod auth {
self.is_authenticated
}
pub fn server_requests(&mut self) -> impl Iterator<Item = ServerRequest> + '_ {
[].into_iter()
}
fn queue_packet(&mut self, packet: Packet) {
self.packets_to_send.push_back(packet);
}

View file

@ -1,13 +1,15 @@
use cluelessh_connection::{ChannelNumber, ChannelOpen, ChannelOperation, ChannelOperationKind};
use cluelessh_connection::{ChannelKind, ChannelNumber, ChannelOperation, ChannelOperationKind};
use std::{collections::HashMap, pin::Pin, sync::Arc};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use cluelessh_protocol::{ChannelUpdateKind, SshStatus};
use eyre::{bail, ContextCompat, OptionExt, Result, WrapErr};
use futures::future::BoxFuture;
use cluelessh_protocol::{ChannelUpdateKind, SshStatus};
use tokio::io::{AsyncRead, AsyncWrite};
use tracing::{debug, info, warn};
use crate::Channel;
pub struct ClientConnection<S> {
stream: Pin<Box<S>>,
buf: [u8; 1024],
@ -55,11 +57,6 @@ pub struct PendingChannel {
ready_recv: tokio::sync::oneshot::Receiver<Result<(), String>>,
channel: Channel,
}
pub struct Channel {
number: ChannelNumber,
updates_recv: tokio::sync::mpsc::Receiver<ChannelUpdateKind>,
ops_send: tokio::sync::mpsc::Sender<ChannelOperation>,
}
impl<S: AsyncRead + AsyncWrite> ClientConnection<S> {
pub async fn connect(stream: S, auth: ClientAuth) -> Result<Self> {
@ -75,7 +72,9 @@ impl<S: AsyncRead + AsyncWrite> ClientConnection<S> {
channel_ops_recv,
channels: HashMap::new(),
proto: cluelessh_protocol::ClientConnection::new(
cluelessh_transport::client::ClientConnection::new(cluelessh_protocol::ThreadRngRand),
cluelessh_transport::client::ClientConnection::new(
cluelessh_protocol::ThreadRngRand,
),
cluelessh_protocol::auth::ClientAuth::new(auth.username.as_bytes().to_vec()),
),
auth,
@ -245,14 +244,14 @@ impl<S: AsyncRead + AsyncWrite> ClientConnection<S> {
Ok(())
}
pub fn open_channel(&mut self, kind: ChannelOpen) -> PendingChannel {
pub fn open_channel(&mut self, kind: ChannelKind) -> PendingChannel {
let Some(channels) = self.proto.channels() else {
panic!("connection not ready yet")
};
let (updates_send, updates_recv) = tokio::sync::mpsc::channel(10);
let (ready_send, ready_recv) = tokio::sync::oneshot::channel();
let number = channels.create_channel(kind);
let number = channels.create_channel(kind.clone());
self.channels.insert(
number,
@ -268,6 +267,7 @@ impl<S: AsyncRead + AsyncWrite> ClientConnection<S> {
number,
updates_recv,
ops_send: self.channel_ops_send.clone(),
kind,
},
}
}
@ -290,11 +290,4 @@ impl Channel {
.await
.map_err(Into::into)
}
pub async fn next_update(&mut self) -> Result<ChannelUpdateKind> {
self.updates_recv
.recv()
.await
.ok_or_eyre("channel has been closed")
}
}

View file

@ -1 +1,33 @@
pub mod client;
pub mod server;
use cluelessh_connection::{ChannelKind, ChannelNumber, ChannelOperation, ChannelOperationKind};
use cluelessh_protocol::ChannelUpdateKind;
use eyre::{OptionExt, Result};
pub struct Channel {
number: ChannelNumber,
updates_recv: tokio::sync::mpsc::Receiver<ChannelUpdateKind>,
ops_send: tokio::sync::mpsc::Sender<ChannelOperation>,
kind: ChannelKind,
}
impl Channel {
pub async fn send(&mut self, op: ChannelOperationKind) -> Result<()> {
self.ops_send
.send(self.number.construct_op(op))
.await
.map_err(Into::into)
}
pub async fn next_update(&mut self) -> Result<ChannelUpdateKind> {
self.updates_recv
.recv()
.await
.ok_or_eyre("channel has been closed")
}
pub fn kind(&self) -> &ChannelKind {
&self.kind
}
}

View file

@ -0,0 +1,327 @@
use cluelessh_connection::{ChannelKind, ChannelNumber, ChannelOperation};
use std::{
collections::{HashMap, VecDeque},
net::SocketAddr,
pin::Pin,
};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
};
use cluelessh_protocol::{ChannelUpdateKind, SshStatus};
use eyre::{eyre, ContextCompat, Result, WrapErr};
use tokio::io::{AsyncRead, AsyncWrite};
use tracing::info;
use crate::Channel;
pub struct ServerListener {
listener: TcpListener,
// todo ratelimits etc
}
pub struct ServerConnection<S> {
stream: Pin<Box<S>>,
peer_addr: SocketAddr,
buf: [u8; 1024],
proto: cluelessh_protocol::ServerConnection,
operations_send: tokio::sync::mpsc::Sender<Operation>,
operations_recv: tokio::sync::mpsc::Receiver<Operation>,
/// Cloned and passed on to channels.
channel_ops_send: tokio::sync::mpsc::Sender<ChannelOperation>,
channel_ops_recv: tokio::sync::mpsc::Receiver<ChannelOperation>,
channels: HashMap<ChannelNumber, ChannelState>,
/// New channels opened by the peer.
new_channels: VecDeque<Channel>,
}
enum ChannelState {
Pending {
ready_send: tokio::sync::oneshot::Sender<Result<(), String>>,
updates_send: tokio::sync::mpsc::Sender<ChannelUpdateKind>,
},
Ready(tokio::sync::mpsc::Sender<ChannelUpdateKind>),
}
enum Operation {
VerifyPassword {
user: String,
password: String,
},
VerifyPubkey {
session_identifier: [u8; 32],
user: String,
pubkey: Vec<u8>,
},
}
pub struct SignatureResult {
pub key_alg_name: &'static str,
pub public_key: Vec<u8>,
pub signature: Vec<u8>,
}
pub struct PendingChannel {
ready_recv: tokio::sync::oneshot::Receiver<Result<(), String>>,
channel: Channel,
}
pub enum Error {
SshStatus(SshStatus),
ServerError(eyre::Report),
}
impl From<eyre::Report> for Error {
fn from(value: eyre::Report) -> Self {
Self::ServerError(value)
}
}
impl ServerListener {
pub fn new(listener: TcpListener) -> Self {
Self { listener }
}
pub async fn accept(&mut self) -> Result<ServerConnection<TcpStream>> {
let (conn, peer_addr) = self.listener.accept().await?;
Ok(ServerConnection::new(conn, peer_addr))
}
}
impl<S: AsyncRead + AsyncWrite> ServerConnection<S> {
pub fn new(stream: S, peer_addr: SocketAddr) -> Self {
let (operations_send, operations_recv) = tokio::sync::mpsc::channel(15);
let (channel_ops_send, channel_ops_recv) = tokio::sync::mpsc::channel(15);
Self {
stream: Box::pin(stream),
peer_addr,
buf: [0; 1024],
operations_send,
operations_recv,
channel_ops_send,
channel_ops_recv,
channels: HashMap::new(),
proto: cluelessh_protocol::ServerConnection::new(
cluelessh_transport::server::ServerConnection::new(
cluelessh_protocol::ThreadRngRand,
),
),
new_channels: VecDeque::new(),
}
}
pub fn peer_addr(&self) -> SocketAddr {
self.peer_addr
}
/// Executes one loop iteration of the main loop.
// IMPORTANT: no operations on this struct should ever block the main loop, except this one.
pub async fn progress(&mut self) -> Result<(), Error> {
if let Some(auth) = self.proto.auth() {
for req in auth.server_requests() {
match req {
cluelessh_protocol::auth::ServerRequest::VerifyPassword { user, password } => {
let send = self.operations_send.clone();
tokio::spawn(async move {
let _ = send
.send(Operation::VerifyPassword { user, password })
.await;
});
}
cluelessh_protocol::auth::ServerRequest::VerifyPubkey {
session_identifier,
pubkey,
user,
} => {
let send = self.operations_send.clone();
tokio::spawn(async move {
let _ = send
.send(Operation::VerifyPubkey {
session_identifier,
user,
pubkey,
})
.await;
});
}
}
}
}
if let Some(channels) = self.proto.channels() {
while let Some(update) = channels.next_channel_update() {
match &update.kind {
ChannelUpdateKind::Open(channel_kind) => {
let channel = self.channels.get_mut(&update.number);
match channel {
// We opened.
Some(ChannelState::Pending { updates_send, .. }) => {
let updates_send = updates_send.clone();
let old = self
.channels
.insert(update.number, ChannelState::Ready(updates_send));
match old.unwrap() {
ChannelState::Pending { ready_send, .. } => {
let _ = ready_send.send(Ok(()));
}
_ => unreachable!(),
}
}
Some(ChannelState::Ready(_)) => {
return Err(Error::ServerError(eyre!(
"attemping to open channel twice: {}",
update.number
)))
}
// They opened.
None => {
let (updates_send, updates_recv) = tokio::sync::mpsc::channel(10);
let number = update.number;
self.channels
.insert(number, ChannelState::Ready(updates_send));
let channel = Channel {
number,
updates_recv,
ops_send: self.channel_ops_send.clone(),
kind: channel_kind.clone(),
};
self.new_channels.push_back(channel);
}
}
}
ChannelUpdateKind::OpenFailed { message, .. } => {
let channel = self
.channels
.get_mut(&update.number)
.wrap_err("unknown channel")?;
match channel {
ChannelState::Pending { .. } => {
let old = self.channels.remove(&update.number);
match old.unwrap() {
ChannelState::Pending { ready_send, .. } => {
let _ = ready_send.send(Err(message.clone()));
}
_ => unreachable!(),
}
}
ChannelState::Ready(_) => {
return Err(Error::ServerError(eyre!(
"attemping to open channel twice: {}",
update.number
)))
}
}
}
_ => {
let channel = self
.channels
.get_mut(&update.number)
.wrap_err("unknown channel")?;
match channel {
ChannelState::Pending { .. } => {
return Err(Error::ServerError(eyre!("channel not ready yet")))
}
ChannelState::Ready(updates_send) => {
let _ = updates_send.send(update.kind).await;
}
}
}
}
}
}
// Make sure that we send all queued messages before going into the select, waiting for things to happen.
self.send_off_data().await?;
tokio::select! {
read = self.stream.read(&mut self.buf) => {
let read = read.wrap_err("reading from connection")?;
if read == 0 {
info!("Did not read any bytes from TCP stream, EOF");
return Ok(());
}
if let Err(err) = self.proto.recv_bytes(&self.buf[..read]) {
return Err(Error::SshStatus(err));
}
}
channel_op = self.channel_ops_recv.recv() => {
let channels = self.proto.channels().expect("connection not ready");
if let Some(channel_op) = channel_op {
channels.do_operation(channel_op);
}
}
op = self.operations_recv.recv() => {
match op {
Some(Operation::VerifyPubkey { .. }) => todo!(),
Some(Operation::VerifyPassword { .. }) => todo!(),
None => {}
}
self.send_off_data().await?;
}
}
Ok(())
}
async fn send_off_data(&mut self) -> Result<()> {
self.proto.progress();
while let Some(msg) = self.proto.next_msg_to_send() {
self.stream
.write_all(&msg.to_bytes())
.await
.wrap_err("writing response")?;
}
Ok(())
}
pub fn open_channel(&mut self, kind: ChannelKind) -> PendingChannel {
let Some(channels) = self.proto.channels() else {
panic!("connection not ready yet")
};
let (updates_send, updates_recv) = tokio::sync::mpsc::channel(10);
let (ready_send, ready_recv) = tokio::sync::oneshot::channel();
let number = channels.create_channel(kind.clone());
self.channels.insert(
number,
ChannelState::Pending {
ready_send,
updates_send,
},
);
PendingChannel {
ready_recv,
channel: Channel {
number,
updates_recv,
ops_send: self.channel_ops_send.clone(),
kind,
},
}
}
pub fn next_new_channel(&mut self) -> Option<Channel> {
self.new_channels.pop_front()
}
}
impl PendingChannel {
pub async fn wait_ready(self) -> Result<Channel, Option<String>> {
match self.ready_recv.await {
Ok(Ok(())) => Ok(self.channel),
Ok(Err(err)) => Err(Some(err)),
Err(_) => Err(None),
}
}
}