whoami_ssh_server.rs

use async_trait::async_trait;
use russh::server::{Msg, Session};
use russh::*;
use russh_keys::*;
use std::net::SocketAddr;
use std::process::Stdio;
use std::sync::Arc;
use tokio::io::AsyncReadExt;
use tokio::io::BufReader;
use tokio::process::Command;

#[tokio::main]
async fn main() {
  // Default logging level does not show debug, info, etc.
  env_logger::builder().filter_level(log::LevelFilter::Info).init();

  let config = russh::server::Config {
    inactivity_timeout: Some(std::time::Duration::from_secs(60 * 60)),
    auth_rejection_time: std::time::Duration::from_secs(5),
    auth_rejection_time_initial: Some(std::time::Duration::from_secs(0)),
    keys: vec![
      // TODO: only do this in dev
      russh_keys::key::KeyPair::Ed25519(ed25519_dalek::SigningKey::from([0; 32])),
    ],
    ..Default::default()
  };
  log::info!("Listening on 0.0.0.0:2222");
  russh::server::run(Arc::new(config), ("0.0.0.0", 2222), Server {})
    .await
    .unwrap();
}

#[derive(Clone)]
struct Server {}

struct ServerHandler {}

impl server::Server for Server {
  type Handler = ServerHandler;
  fn new_client(&mut self, addr: Option<SocketAddr>) -> ServerHandler {
    // TODO(docs): when is addr None?
    log::info!("new client");
    ServerHandler {}
  }
}

#[async_trait]
impl server::Handler for ServerHandler {
  type Error = anyhow::Error;

  async fn channel_open_session(
    self,
    channel: Channel<Msg>,
    session: Session,
  ) -> Result<(Self, bool, Session), Self::Error> {
    log::info!("channel_open_session, channel_id: {}", channel.id());

    // TODO(docs): what effect do these return values have?
    Ok((self, true, session))
  }

  async fn auth_publickey(
    self,
    username: &str,
    public_key: &key::PublicKey,
  ) -> Result<(Self, server::Auth), Self::Error> {
    log::info!("auth_publickey, username: {:?}", username);
    Ok((self, server::Auth::Accept))
  }

  async fn shell_request(self, channel_id: ChannelId, session: Session) -> Result<(Self, Session), Self::Error> {
    log::info!("shell_request, channel_id: {}", channel_id);

    let mut child = Command::new("whoami")
      .stdin(Stdio::piped())
      .stdout(Stdio::piped())
      .stderr(Stdio::piped())
      .spawn()?;
    let stderr = child.stderr.take().unwrap();
    let stdout = child.stdout.take().unwrap();

    // Spawn task for stdout
    let handle = session.handle();
    tokio::spawn(async move {
      let mut stdout = BufReader::new(stdout);
      let mut buffer = vec![0; 1024];
      while let Ok(size) = stdout.read(&mut buffer).await {
        if size == 0 {
          break;
        }
        handle
          .data(channel_id, CryptoVec::from_slice(&buffer[..size]))
          .await
          .unwrap();
      }
    });

    // Spawn task for stderr
    let handle = session.handle();
    tokio::spawn(async move {
      let mut stderr = BufReader::new(stderr);
      let mut buffer = vec![0; 1024];
      while let Ok(size) = stderr.read(&mut buffer).await {
        if size == 0 {
          break;
        }
        handle
          .data(channel_id, CryptoVec::from_slice(&buffer[..size]))
          .await
          .unwrap();
      }
    });

    // Close the channel when child exits
    let handle = session.handle();
    tokio::spawn(async move {
      let status = child.wait().await.unwrap().code().unwrap();
      handle
        .data(channel_id, CryptoVec::from(format!("\r\nExit status: {}\r\n", status)))
        .await
        .unwrap();
      // TODO: the SSH client still exits with a non-zero status code. Why?
      handle.eof(channel_id).await.unwrap();
      handle.close(channel_id).await.unwrap();
    });

    Ok((self, session))
  }
}