refactor

src/main.rs

@@ -1,15 +1,16 @@
-use std::net::SocketAddr;
+use std::{collections::HashMap, net::SocketAddr};
 
 use eyre::{Context, Result};
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     net::{TcpListener, TcpStream},
 };
-use tracing::{error, info};
+use tracing::{debug, error, info};
 
 use ssh_protocol::{
+    connection::{ChannelOpen, ChannelOperation, ChannelOperationKind, ChannelRequestKind},
     transport::{self, ThreadRngRand},
-    ServerConnection, SshStatus,
+    ChannelUpdateKind, ServerConnection, SshStatus,
 };
 use tracing_subscriber::EnvFilter;
 
@@ -58,6 +59,8 @@ async fn handle_connection(next: (TcpStream, SocketAddr)) -> Result<()> {
 
     let mut state = ServerConnection::new(transport::ServerConnection::new(ThreadRngRand));
 
+    let mut session_channels = HashMap::new();
+
     loop {
         let mut buf = [0; 1024];
         let read = conn
@@ -83,6 +86,47 @@ async fn handle_connection(next: (TcpStream, SocketAddr)) -> Result<()> {
             }
         }
 
+        while let Some(update) = state.next_channel_update() {
+            match update.kind {
+                ChannelUpdateKind::Open(kind) => match kind {
+                    ChannelOpen::Session => {
+                        session_channels.insert(update.number, 0);
+                    }
+                },
+                ChannelUpdateKind::Request(req) => {
+                    match req.kind {
+                        ChannelRequestKind::PtyReq { .. } => {}
+                        ChannelRequestKind::Shell => {}
+                    };
+                    if req.want_reply {
+                        // TODO: sent the reply.
+                    }
+                }
+                ChannelUpdateKind::Data { data } => {
+                    let is_eof = data.contains(&0x03 /*EOF, Ctrl-C*/);
+
+                    // echo :3
+                    state.do_operation(ChannelOperation {
+                        number: update.number,
+                        kind: ChannelOperationKind::Data(data),
+                    });
+
+                    if is_eof {
+                        debug!(channel = ?update.number, "Received EOF, closing channel");
+
+                        state.do_operation(ChannelOperation {
+                            number: update.number,
+                            kind: ChannelOperationKind::Close,
+                        });
+                    }
+                }
+                ChannelUpdateKind::ExtendedData { .. } | ChannelUpdateKind::Eof => { /* ignore */ }
+                ChannelUpdateKind::Closed => {
+                    session_channels.remove(&update.number);
+                }
+            }
+        }
+
         while let Some(msg) = state.next_msg_to_send() {
             conn.write_all(&msg.to_bytes())
                 .await

ssh-connection/src/lib.rs

@@ -1,4 +1,4 @@
-use std::collections::VecDeque;
+use std::collections::{HashMap, VecDeque};
 use tracing::{debug, warn};
 
 use ssh_transport::client_error;
@@ -7,39 +7,76 @@ use ssh_transport::Result;
 
 pub struct ServerChannelsState {
     packets_to_send: VecDeque<Packet>,
-    channels: Vec<SessionChannel>,
-
     channel_updates: VecDeque<ChannelUpdate>,
+
+    channels: HashMap<u32, Channel>,
+    next_channel_id: u32,
 }
 
-struct SessionChannel {
+struct Channel {
     /// Whether our side has closed this channel.
     we_closed: bool,
+    /// The channel number for the other side.
     peer_channel: u32,
-    has_pty: bool,
-    has_shell: bool,
-    sent_bytes: Vec<u8>,
 }
 
+/// An update from a channel.
+/// The receiver-equivalent of [`ChannelOperation`].
 pub struct ChannelUpdate {
-    pub channel: u32,
+    pub number: u32,
     pub kind: ChannelUpdateKind,
 }
 pub enum ChannelUpdateKind {
-    Create { kind: String, args: Vec<u8> },
+    Open(ChannelOpen),
-    Request { kind: String, args: Vec<u8> },
+    Request(ChannelRequest),
     Data { data: Vec<u8> },
     ExtendedData { code: u32, data: Vec<u8> },
     Eof,
-    ChannelClosed,
+    Closed,
+}
+
+pub enum ChannelOpen {
+    Session,
+}
+
+pub struct ChannelRequest {
+    pub want_reply: bool,
+    pub kind: ChannelRequestKind,
+}
+
+pub enum ChannelRequestKind {
+    PtyReq {
+        term: String,
+        width_chars: u32,
+        height_rows: u32,
+        width_px: u32,
+        height_px: u32,
+        term_modes: Vec<u8>,
+    },
+    Shell,
+}
+
+/// An operation to do on a channel.
+/// The sender-equivalent of [`ChannelUpdate`].
+pub struct ChannelOperation {
+    pub number: u32,
+    pub kind: ChannelOperationKind,
+}
+
+pub enum ChannelOperationKind {
+    Success,
+    Failure,
+    Data(Vec<u8>),
+    Close,
 }
 
 impl ServerChannelsState {
     pub fn new() -> Self {
         ServerChannelsState {
             packets_to_send: VecDeque::new(),
-            channels: Vec::new(),
+            channels: HashMap::new(),
             channel_updates: VecDeque::new(),
+            next_channel_id: 0,
         }
     }
 
@@ -64,28 +101,8 @@ impl ServerChannelsState {
 
                 debug!(?channel_type, ?sender_channel, "Opening channel");
 
-                match channel_type {
+                let update_message = match channel_type {
-                    "session" => {
+                    "session" => ChannelOpen::Session,
-                        let our_number = self.channels.len() as u32;
-
-                        self.packets_to_send
-                            .push_back(Packet::new_msg_channel_open_confirmation(
-                                our_number,
-                                sender_channel,
-                                initial_window_size,
-                                max_packet_size,
-                            ));
-
-                        self.channels.push(SessionChannel {
-                            we_closed: false,
-                            peer_channel: sender_channel,
-                            has_pty: false,
-                            has_shell: false,
-                            sent_bytes: Vec::new(),
-                        });
-
-                        debug!(?channel_type, ?our_number, "Successfully opened channel");
-                    }
                     _ => {
                         self.packets_to_send
                             .push_back(Packet::new_msg_channel_open_failure(
@@ -94,30 +111,50 @@ impl ServerChannelsState {
                                 b"unknown channel type",
                                 b"",
                             ));
+                        return Ok(());
                     }
-                }
+                };
+
+                let our_number = self.next_channel_id;
+                self.next_channel_id = self.next_channel_id.checked_add(1).ok_or_else(|| {
+                    client_error!("created too many channels, overflowed the counter")
+                })?;
+
+                self.packets_to_send
+                    .push_back(Packet::new_msg_channel_open_confirmation(
+                        our_number,
+                        sender_channel,
+                        initial_window_size,
+                        max_packet_size,
+                    ));
+
+                self.channels.insert(
+                    our_number,
+                    Channel {
+                        we_closed: false,
+                        peer_channel: sender_channel,
+                    },
+                );
+
+                self.channel_updates.push_back(ChannelUpdate {
+                    number: our_number,
+                    kind: ChannelUpdateKind::Open(update_message),
+                });
+
+                debug!(?channel_type, ?our_number, "Successfully opened channel");
             }
             Packet::SSH_MSG_CHANNEL_DATA => {
                 let our_channel = packet.u32()?;
                 let data = packet.string()?;
 
-                let channel = self.channel(our_channel)?;
+                let _ = self.channel(our_channel)?;
-                channel.recv_bytes(data);
 
-                let peer = channel.peer_channel;
+                self.channel_updates.push_back(ChannelUpdate {
-                // echo :3
+                    number: our_channel,
-                self.packets_to_send
+                    kind: ChannelUpdateKind::Data {
-                    .push_back(Packet::new_msg_channel_data(peer, data));
+                        data: data.to_owned(),
-
+                    },
-                if data.contains(&0x03 /*EOF, Ctrl-C*/) {
+                });
-                    debug!(?our_channel, "Received EOF, closing channel");
-                    // <https://datatracker.ietf.org/doc/html/rfc4254#section-5.3>
-                    self.packets_to_send
-                        .push_back(Packet::new_msg_channel_close(peer));
-
-                    let channel = self.channel(our_channel)?;
-                    channel.we_closed = true;
-                }
             }
             Packet::SSH_MSG_CHANNEL_CLOSE => {
                 // <https://datatracker.ietf.org/doc/html/rfc4254#section-5.3>
@@ -128,7 +165,12 @@ impl ServerChannelsState {
                     self.packets_to_send.push_back(close);
                 }
 
-                self.channels.remove(our_channel as usize);
+                self.channels.remove(&our_channel);
+
+                self.channel_updates.push_back(ChannelUpdate {
+                    number: our_channel,
+                    kind: ChannelUpdateKind::Closed,
+                });
 
                 debug!("Channel has been closed");
             }
@@ -142,14 +184,14 @@ impl ServerChannelsState {
                 let channel = self.channel(our_channel)?;
                 let peer_channel = channel.peer_channel;
 
-                match request_type {
+                let channel_request_kind = match request_type {
                     "pty-req" => {
                         let term = packet.utf8_string()?;
                         let width_chars = packet.u32()?;
                         let height_rows = packet.u32()?;
-                        let _width_px = packet.u32()?;
+                        let width_px = packet.u32()?;
-                        let _height_px = packet.u32()?;
+                        let height_px = packet.u32()?;
-                        let _term_modes = packet.string()?;
+                        let term_modes = packet.string()?;
 
                         debug!(
                             ?our_channel,
@@ -159,37 +201,41 @@ impl ServerChannelsState {
                             "Trying to open a terminal"
                         );
 
-                        // Faithfully allocate the PTY.
+                        ChannelRequestKind::PtyReq {
-                        channel.has_pty = true;
+                            term: term.to_owned(),
-
+                            width_chars,
-                        if want_reply {
+                            height_rows,
-                            self.send_channel_success(peer_channel);
+                            width_px,
+                            height_px,
+                            term_modes: term_modes.to_owned(),
                         }
                     }
                     "shell" => {
-                        if !channel.has_pty {
+                        let _ = self.channel(our_channel)?;
-                            self.send_channel_failure(peer_channel);
-                        }
-
-                        // Sure! (reborrow)
-                        let channel = self.channel(our_channel)?;
-                        channel.has_shell = true;
 
                         debug!(?our_channel, "Opening shell");
 
-                        if want_reply {
+                        ChannelRequestKind::Shell
-                            self.send_channel_success(peer_channel);
-                        }
                     }
                     "signal" => {
                         debug!(?our_channel, "Received signal");
                         // Ignore signals, something we can do.
+                        return Ok(());
                     }
                     _ => {
                         warn!(?request_type, ?our_channel, "Unknown channel request");
                         self.send_channel_failure(peer_channel);
+                        return Ok(());
                     }
-                }
+                };
+
+                self.channel_updates.push_back(ChannelUpdate {
+                    number: our_channel,
+                    kind: ChannelUpdateKind::Request(ChannelRequest {
+                        want_reply,
+                        kind: channel_request_kind,
+                    }),
+                })
             }
             _ => {
                 todo!("{packet_type}");
@@ -203,8 +249,31 @@ impl ServerChannelsState {
         self.packets_to_send.drain(..)
     }
 
-    pub(crate) fn channel_updates(&mut self) -> impl Iterator<Item = ChannelUpdate> + '_ {
+    pub fn next_channel_update(&mut self) -> Option<ChannelUpdate> {
-        self.channel_updates.drain(..)
+        self.channel_updates.pop_front()
+    }
+
+    pub fn do_operation(&mut self, op: ChannelOperation) {
+        let peer = self
+            .channel(op.number)
+            .expect("passed channel ID that does not exist")
+            .peer_channel;
+        match op.kind {
+            ChannelOperationKind::Success => self.send_channel_success(peer),
+            ChannelOperationKind::Failure => self.send_channel_failure(peer),
+            ChannelOperationKind::Data(data) => {
+                self.packets_to_send
+                    .push_back(Packet::new_msg_channel_data(peer, &data));
+            }
+            ChannelOperationKind::Close => {
+                // <https://datatracker.ietf.org/doc/html/rfc4254#section-5.3>
+                self.packets_to_send
+                    .push_back(Packet::new_msg_channel_close(peer));
+
+                let channel = self.channel(op.number).unwrap();
+                channel.we_closed = true;
+            }
+        }
     }
 
     fn send_channel_success(&mut self, recipient_channel: u32) {
@@ -217,15 +286,9 @@ impl ServerChannelsState {
             .push_back(Packet::new_msg_channel_failure(recipient_channel));
     }
 
-    fn channel(&mut self, number: u32) -> Result<&mut SessionChannel> {
+    fn channel(&mut self, number: u32) -> Result<&mut Channel> {
         self.channels
-            .get_mut(number as usize)
+            .get_mut(&number)
             .ok_or_else(|| client_error!("unknown channel: {number}"))
     }
 }
-
-impl SessionChannel {
-    fn recv_bytes(&mut self, bytes: &[u8]) {
-        self.sent_bytes.extend_from_slice(bytes);
-    }
-}

ssh-protocol/src/lib.rs

@@ -1,3 +1,6 @@
+pub use ssh_connection as connection;
+use ssh_connection::ChannelOperation;
+pub use ssh_connection::{ChannelUpdate, ChannelUpdateKind};
 pub use ssh_transport as transport;
 pub use ssh_transport::{Result, SshStatus};
 
@@ -30,7 +33,8 @@ impl ServerConnection {
                         self.transport.send_plaintext_packet(to_send);
                     }
                     if auth.is_authenticated() {
-                        self.state = ServerConnectionState::Open(ssh_connection::ServerChannelsState::new());
+                        self.state =
+                            ServerConnectionState::Open(ssh_connection::ServerChannelsState::new());
                     }
                 }
                 ServerConnectionState::Open(con) => {
@@ -48,6 +52,25 @@ impl ServerConnection {
     pub fn next_msg_to_send(&mut self) -> Option<ssh_transport::Msg> {
         self.transport.next_msg_to_send()
     }
+
+    pub fn next_channel_update(&mut self) -> Option<ssh_connection::ChannelUpdate> {
+        match &mut self.state {
+            ServerConnectionState::Auth(_) => None,
+            ServerConnectionState::Open(con) => con.next_channel_update(),
+        }
+    }
+
+    pub fn do_operation(&mut self, op: ChannelOperation) {
+        match &mut self.state {
+            ServerConnectionState::Auth(_) => panic!("tried to get connection during auth"),
+            ServerConnectionState::Open(con) => {
+                con.do_operation(op);
+                for to_send in con.packets_to_send() {
+                    self.transport.send_plaintext_packet(to_send);
+                }
+            }
+        }
+    }
 }
 
 /// <https://datatracker.ietf.org/doc/html/rfc4252>