-
-
Notifications
You must be signed in to change notification settings - Fork 117
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
2 changed files
with
422 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,212 @@ | ||
use async_trait::async_trait; | ||
use ratatui::{ | ||
backend::CrosstermBackend, | ||
layout::Rect, | ||
style::{Color, Style}, | ||
widgets::{Block, Borders, Clear, Paragraph}, | ||
Terminal, | ||
}; | ||
use russh::{server::*, Channel, ChannelId}; | ||
use russh_keys::key::PublicKey; | ||
use std::collections::HashMap; | ||
use std::sync::Arc; | ||
use tokio::sync::Mutex; | ||
|
||
type SshTerminal = Terminal<CrosstermBackend<TerminalHandle>>; | ||
|
||
struct App { | ||
pub counter: usize, | ||
} | ||
|
||
impl App { | ||
pub fn new() -> App { | ||
Self { counter: 0 } | ||
} | ||
} | ||
|
||
#[derive(Clone)] | ||
struct TerminalHandle { | ||
handle: Handle, | ||
// The sink collects the data which is finally flushed to the handle. | ||
sink: Vec<u8>, | ||
channel_id: ChannelId, | ||
} | ||
|
||
// The crossterm backend writes to the terminal handle. | ||
impl std::io::Write for TerminalHandle { | ||
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> { | ||
self.sink.extend_from_slice(buf); | ||
Ok(buf.len()) | ||
} | ||
|
||
fn flush(&mut self) -> std::io::Result<()> { | ||
let handle = self.handle.clone(); | ||
let channel_id = self.channel_id.clone(); | ||
let data = self.sink.clone().into(); | ||
futures::executor::block_on(async move { | ||
let result = handle.data(channel_id, data).await; | ||
if result.is_err() { | ||
eprintln!("Failed to send data: {:?}", result); | ||
} | ||
}); | ||
|
||
self.sink.clear(); | ||
Ok(()) | ||
} | ||
} | ||
|
||
#[derive(Clone)] | ||
struct AppServer { | ||
clients: Arc<Mutex<HashMap<usize, (SshTerminal, App)>>>, | ||
id: usize, | ||
} | ||
|
||
impl AppServer { | ||
pub fn new() -> Self { | ||
Self { | ||
clients: Arc::new(Mutex::new(HashMap::new())), | ||
id: 0, | ||
} | ||
} | ||
|
||
pub async fn run(&mut self) -> Result<(), anyhow::Error> { | ||
let clients = self.clients.clone(); | ||
tokio::spawn(async move { | ||
loop { | ||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; | ||
|
||
for (_, (terminal, app)) in clients.lock().await.iter_mut() { | ||
app.counter += 1; | ||
|
||
terminal | ||
.draw(|f| { | ||
let size = f.size(); | ||
f.render_widget(Clear, size); | ||
let style = match app.counter % 3 { | ||
0 => Style::default().fg(Color::Red), | ||
1 => Style::default().fg(Color::Green), | ||
_ => Style::default().fg(Color::Blue), | ||
}; | ||
let paragraph = Paragraph::new(format!("Counter: {}", app.counter)) | ||
.alignment(ratatui::layout::Alignment::Center) | ||
.style(style); | ||
let block = Block::default() | ||
.title("Press 'c' to reset the counter!") | ||
.borders(Borders::ALL); | ||
f.render_widget(paragraph.block(block), size); | ||
}) | ||
.unwrap(); | ||
} | ||
} | ||
}); | ||
|
||
let config = Config { | ||
inactivity_timeout: Some(std::time::Duration::from_secs(3600)), | ||
auth_rejection_time: std::time::Duration::from_secs(3), | ||
auth_rejection_time_initial: Some(std::time::Duration::from_secs(0)), | ||
keys: vec![russh_keys::key::KeyPair::generate_ed25519().unwrap()], | ||
..Default::default() | ||
}; | ||
|
||
self.run_on_address(Arc::new(config), ("0.0.0.0", 2222)) | ||
.await?; | ||
Ok(()) | ||
} | ||
} | ||
|
||
impl Server for AppServer { | ||
type Handler = Self; | ||
fn new_client(&mut self, _: Option<std::net::SocketAddr>) -> Self { | ||
let s = self.clone(); | ||
self.id += 1; | ||
s | ||
} | ||
} | ||
|
||
#[async_trait] | ||
impl Handler for AppServer { | ||
type Error = anyhow::Error; | ||
|
||
async fn channel_open_session( | ||
&mut self, | ||
channel: Channel<Msg>, | ||
session: &mut Session, | ||
) -> Result<bool, Self::Error> { | ||
{ | ||
let mut clients = self.clients.lock().await; | ||
let terminal_handle = TerminalHandle { | ||
handle: session.handle(), | ||
sink: Vec::new(), | ||
channel_id: channel.id(), | ||
}; | ||
|
||
let backend = CrosstermBackend::new(terminal_handle.clone()); | ||
let terminal = Terminal::new(backend)?; | ||
let app = App::new(); | ||
|
||
clients.insert(self.id, (terminal, app)); | ||
} | ||
|
||
Ok(true) | ||
} | ||
|
||
async fn auth_publickey(&mut self, _: &str, _: &PublicKey) -> Result<Auth, Self::Error> { | ||
Ok(Auth::Accept) | ||
} | ||
|
||
async fn data( | ||
&mut self, | ||
channel: ChannelId, | ||
data: &[u8], | ||
session: &mut Session, | ||
) -> Result<(), Self::Error> { | ||
match data { | ||
// Pressing 'q' closes the connection. | ||
b"q" => { | ||
self.clients.lock().await.remove(&self.id); | ||
session.close(channel); | ||
} | ||
// Pressing 'c' resets the counter for the app. | ||
// Only the client with the id sees the counter reset. | ||
b"c" => { | ||
let mut clients = self.clients.lock().await; | ||
let (_, app) = clients.get_mut(&self.id).unwrap(); | ||
app.counter = 0; | ||
} | ||
_ => {} | ||
} | ||
|
||
Ok(()) | ||
} | ||
|
||
/// The client's window size has changed. | ||
async fn window_change_request( | ||
&mut self, | ||
_: ChannelId, | ||
col_width: u32, | ||
row_height: u32, | ||
_: u32, | ||
_: u32, | ||
_: &mut Session, | ||
) -> Result<(), Self::Error> { | ||
{ | ||
let mut clients = self.clients.lock().await; | ||
let (terminal, _) = clients.get_mut(&self.id).unwrap(); | ||
let rect = Rect { | ||
x: 0, | ||
y: 0, | ||
width: col_width as u16, | ||
height: row_height as u16, | ||
}; | ||
terminal.resize(rect)?; | ||
} | ||
|
||
Ok(()) | ||
} | ||
} | ||
|
||
#[tokio::main] | ||
async fn main() { | ||
let mut server = AppServer::new(); | ||
server.run().await.expect("Failed running server"); | ||
} |
Oops, something went wrong.