DeCarabas/fwd
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Try to make the UI better when unconnected

Browse source
This commit is contained in:
John Doty 2022-10-16 08:55:30 -07:00
parent c412e81b26
commit 6cc17e7fca
3 changed files with 321 additions and 187 deletions
Download patch file Download diff file Expand all files Collapse all files

2
rustfmt.toml Normal file
View file

@ -0,0 +1,2 @@
max_width = 80
struct_lit_width = 40

66
src/lib.rs
View file

@ -4,8 +4,8 @@ use log::{debug, error, info, warn};
use message::{Message, MessageReader, MessageWriter};
use std::net::{Ipv4Addr, SocketAddrV4};
use tokio::io::{
AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader,
BufWriter,
AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite,
AsyncWriteExt, BufReader, BufWriter,
};
use tokio::net::{TcpListener, TcpStream};
use tokio::process;
@ -119,7 +119,9 @@ async fn server_main<Reader: AsyncRead + Unpin, Writer: AsyncWrite + Unpin>(
}
}
async fn client_sync<Read: AsyncRead + Unpin>(reader: &mut Read) -> Result<(), tokio::io::Error> {
async fn client_sync<Read: AsyncRead + Unpin>(
reader: &mut Read,
) -> Result<(), tokio::io::Error> {
info!("Waiting for synchronization marker...");
// Run these two loops in parallel; the copy of stdin should stop when
@ -149,7 +151,11 @@ async fn client_sync<Read: AsyncRead + Unpin>(reader: &mut Read) -> Result<(), t
/// This contains a very simplified implementation of a SOCKS5 connector,
/// enough to work with the SSH I have. I would have liked it to be SOCKS4,
/// which is a much simpler protocol, but somehow it didn't work.
async fn client_handle_connection(socks_port: u16, port: u16, socket: TcpStream) -> Result<()> {
async fn client_handle_connection(
socks_port: u16,
port: u16,
socket: TcpStream,
) -> Result<()> {
debug!("Handling connection!");
let dest_addr = SocketAddrV4::new(Ipv4Addr::LOCALHOST, socks_port);
@ -257,14 +263,18 @@ async fn client_handle_connection(socks_port: u16, port: u16, socket: TcpStream)
async fn client_listen(port: u16, socks_port: u16) -> Result<()> {
loop {
let listener = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)).await?;
let listener =
TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port))
.await?;
loop {
// The second item contains the IP and port of the new
// connection, but we don't care.
let (socket, _) = listener.accept().await?;
tokio::spawn(async move {
if let Err(e) = client_handle_connection(socks_port, port, socket).await {
if let Err(e) =
client_handle_connection(socks_port, port, socket).await
{
error!("Error handling connection: {:?}", e);
} else {
debug!("Done???");
@ -402,7 +412,9 @@ pub async fn run_server() {
}
}
async fn spawn_ssh(server: &str) -> Result<(tokio::process::Child, u16), std::io::Error> {
async fn spawn_ssh(
server: &str,
) -> Result<(tokio::process::Child, u16), std::io::Error> {
let socks_port = {
let listener = TcpListener::bind("127.0.0.1:0").await?;
listener.local_addr()?.port()
@ -423,11 +435,30 @@ async fn spawn_ssh(server: &str) -> Result<(tokio::process::Child, u16), std::io
Ok((child, socks_port))
}
#[cfg(target_family = "windows")]
fn is_sigint(status: std::process::ExitStatus) -> bool {
match status.code() {
Some(255) => true,
_ => false,
}
}
#[cfg(target_family = "unix")]
fn is_sigint(status: std::process::ExitStatus) -> bool {
use std::os::unix::process::ExitStatusExt;
match status.signal() {
Some(2) => true,
Some(_) => false,
None => false,
}
}
async fn client_connect_loop(remote: &str, events: mpsc::Sender<ui::UIEvent>) {
loop {
_ = events.send(ui::UIEvent::Disconnected).await;
let (mut child, socks_port) = spawn_ssh(remote).await.expect("failed to spawn");
let (mut child, socks_port) =
spawn_ssh(remote).await.expect("failed to spawn");
let mut stderr = BufReader::new(
child
@ -450,6 +481,20 @@ async fn client_connect_loop(remote: &str, events: mpsc::Sender<ui::UIEvent>) {
if let Err(e) = client_sync(&mut reader).await {
error!("Error synchronizing: {:?}", e);
match child.wait().await {
Ok(status) => {
if is_sigint(status) {
return;
} else {
match status.code() {
Some(127) => eprintln!("Cannot find `fwd` remotely, make sure it is installed"),
_ => (),
};
}
}
Err(_) => (),
};
tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
continue;
}
@ -462,7 +507,10 @@ async fn client_connect_loop(remote: &str, events: mpsc::Sender<ui::UIEvent>) {
client_pipe_stderr(&mut stderr, sec).await;
});
if let Err(e) = client_main(socks_port, &mut reader, &mut writer, events.clone()).await {
if let Err(e) =
client_main(socks_port, &mut reader, &mut writer, events.clone())
.await
{
error!("Server disconnected with error: {:?}", e);
} else {
warn!("Disconnected from server, reconnecting...");

440
src/ui.rs
View file

@ -2,13 +2,14 @@ use crate::client_listen;
use crate::message::PortDesc;
use anyhow::Result;
use crossterm::{
cursor::MoveTo,
event::{Event, EventStream, KeyCode, KeyEvent},
cursor::{MoveTo, RestorePosition, SavePosition},
event::{Event, EventStream, KeyCode, KeyEvent, KeyModifiers},
execute, queue,
style::{Color, PrintStyledContent, Stylize},
terminal::{
disable_raw_mode, enable_raw_mode, size, Clear, ClearType, DisableLineWrap, EnableLineWrap,
EnterAlternateScreen, LeaveAlternateScreen,
disable_raw_mode, enable_raw_mode, size, Clear, ClearType,
DisableLineWrap, EnableLineWrap, EnterAlternateScreen,
LeaveAlternateScreen,
},
};
use log::{error, info, Level, Metadata, Record};
@ -56,10 +57,17 @@ impl log::Log for Logger {
fn flush(&self) {}
}
#[derive(Debug)]
pub struct UI {
events: mpsc::Receiver<UIEvent>,
listeners: HashMap<u16, oneshot::Sender<()>>,
socks_port: u16,
running: bool,
show_logs: bool,
selection: usize,
ports: Option<Vec<PortDesc>>,
lines: VecDeque<String>,
alternate_screen: bool,
}
impl UI {
@ -68,162 +76,107 @@ impl UI {
events,
listeners: HashMap::new(),
socks_port: 0,
running: true,
show_logs: false,
selection: 0,
ports: None,
lines: VecDeque::with_capacity(1024),
alternate_screen: false,
}
}
pub async fn run(&mut self) -> Result<()> {
enable_raw_mode()?;
let result = self.run_core().await;
execute!(stdout(), EnableLineWrap, LeaveAlternateScreen)?;
disable_raw_mode()?;
_ = self.leave_alternate_screen();
result
}
async fn run_core(&mut self) -> Result<()> {
let mut stdout = stdout();
execute!(stdout, EnterAlternateScreen, DisableLineWrap)?;
let mut console_events = EventStream::new();
let mut connected = false;
let mut selection = 0;
let mut show_logs = false;
let mut lines: VecDeque<String> = VecDeque::with_capacity(1024);
let mut ports: Option<Vec<PortDesc>> = None;
loop {
tokio::select! {
ev = console_events.next() => {
match ev {
Some(Ok(Event::Key(ev))) => {
match ev {
KeyEvent {code:KeyCode::Esc, ..}
| KeyEvent {code:KeyCode::Char('q'), ..} => { break; },
KeyEvent {code:KeyCode::Char('l'), ..} => {
show_logs = !show_logs;
}
KeyEvent {code:KeyCode::Char('e'), ..} => {
if let Some(ports) = &ports {
if selection < ports.len() {
let p = ports[selection].port;
self.enable_disable_port(p);
}
}
}
KeyEvent { code:KeyCode::Up, ..}
| KeyEvent { code:KeyCode::Char('j'), ..} => {
if selection > 0 {
selection -= 1;
}
}
KeyEvent { code:KeyCode::Down, ..}
| KeyEvent { code:KeyCode::Char('k'), ..} => {
if let Some(p) = &ports {
if selection != p.len() - 1 {
selection += 1;
}
}
}
KeyEvent { code:KeyCode::Enter, ..} => {
if let Some(p) = &ports {
if selection < p.len() {
_ = open::that(format!("http://127.0.0.1:{}/", p[selection].port));
}
}
}
_ => ()
}
},
Some(Ok(_)) => (), // Don't care about this event...
Some(Err(_)) => (), // Hmmmmmm.....?
None => (), // ....no events? what?
}
}
ev = self.events.recv() => {
match ev {
Some(UIEvent::Disconnected) => {
self.socks_port = 0;
connected = false;
}
Some(UIEvent::Connected(sp)) => {
self.socks_port = sp;
info!("Socks port {socks_port}", socks_port=self.socks_port);
connected = true;
}
Some(UIEvent::Ports(mut p)) => {
p.sort_by(|a, b| a.port.partial_cmp(&b.port).unwrap());
if selection >= p.len() {
selection = p.len()-1;
}
self.running = true;
while self.running {
self.handle_events(&mut console_events).await;
let mut new_listeners = HashMap::new();
for port in &p {
let port = port.port;
if let Some(l) = self.listeners.remove(&port) {
if !l.is_closed() {
// `l` here is, of course, the channel that we
// use to tell the listener task to stop (see the
// spawn call below). If it isn't closed then
// that means a spawn task is still running so we
// should just let it keep running and re-use the
// existing listener.
new_listeners.insert(port, l);
}
}
}
// This has the side effect of closing any
// listener that didn't get copied over to the
// new listeners table.
self.listeners = new_listeners;
ports = Some(p);
}
Some(UIEvent::ServerLine(line)) => {
while lines.len() >= 1024 {
lines.pop_front();
}
lines.push_back(format!("[SERVER] {line}"));
}
Some(UIEvent::LogLine(_level, line)) => {
while lines.len() >= 1024 {
lines.pop_front();
}
lines.push_back(format!("[CLIENT] {line}"));
}
None => break,
}
}
if self.connected() {
self.render_connected()?;
} else {
self.render_disconnected()?;
}
}
let (columns, rows) = size()?;
let columns: usize = columns.into();
Ok(())
}
queue!(stdout, Clear(ClearType::All), MoveTo(0, 0))?;
if connected {
// List of open ports
// How wide are all the things?
let padding = 1;
let enabled_width = 1; // Just 1 character
let port_width = 5; // 5 characters for 16-bit number
fn render_disconnected(&mut self) -> Result<()> {
self.leave_alternate_screen()?;
let mut stdout = stdout();
let description_width =
columns - (padding + padding + enabled_width + padding + port_width + padding);
let (columns, _) = size()?;
let columns: usize = columns.into();
execute!(
stdout,
SavePosition,
MoveTo(0, 0),
PrintStyledContent(
format!("{:^columns$}", "Not Connected")
.with(Color::Black)
.on(Color::Red)
),
RestorePosition,
)?;
Ok(())
}
fn render_connected(&mut self) -> Result<()> {
self.enter_alternate_screen()?;
let mut stdout = stdout();
let (columns, rows) = size()?;
let columns: usize = columns.into();
queue!(stdout, Clear(ClearType::All), MoveTo(0, 0))?;
// List of open ports
// How wide are all the things?
let padding = 1;
let enabled_width = 1; // Just 1 character
let port_width = 5; // 5 characters for 16-bit number
let description_width = columns
- (padding
+ padding
+ enabled_width
+ padding
+ port_width
+ padding);
print!(
"{}",
format!(
" {enabled:>enabled_width$} {port:>port_width$} {description:<description_width$}\r\n",
enabled = "E",
port = "port",
description = "description"
)
.negative()
);
if let Some(ports) = &self.ports {
let max_ports: usize = if self.show_logs {
(rows / 2) - 1
} else {
rows - 2
}
.into();
for (index, port) in ports.into_iter().take(max_ports).enumerate() {
print!(
"{}",
format!(
" {enabled:>enabled_width$} {port:>port_width$} {description:<description_width$}\r\n",
enabled = "E",
port = "port",
description = "description"
)
.negative()
);
if let Some(ports) = &mut ports {
let max_ports: usize = if show_logs { (rows / 2) - 1 } else { rows - 2 }.into();
for (index, port) in ports.into_iter().take(max_ports).enumerate() {
print!(
"{} {:>enabled_width$} {:port_width$} {:description_width$}\r\n",
if index == selection { "\u{2B46}" } else { " " },
if index == self.selection {
"\u{2B46}"
} else {
" "
},
if self.listeners.contains_key(&port.port) {
"+"
} else {
@ -232,49 +185,50 @@ impl UI {
port.port,
port.desc
);
}
}
} else {
queue!(
stdout,
PrintStyledContent(
format!("{:^columns$}", "Not Connected")
.with(Color::Black)
.on(Color::Red)
)
)?;
}
// Log
if show_logs {
let hr: usize = ((rows / 2) - 2).into();
let start: usize = if lines.len() > hr {
lines.len() - hr
} else {
0
};
queue!(stdout, MoveTo(0, rows / 2))?;
print!("{}", format!("{:columns$}", " Log").negative());
for line in lines.range(start..) {
print!("{}\r\n", line);
}
}
queue!(stdout, MoveTo(0, rows - 1))?;
print!(
"{}",
format!(
"{:columns$}",
" q - quit | l - toggle log | \u{2191}/\u{2193} - select port | e - enable/disable | <enter> - browse"
).negative()
);
stdout.flush()?;
}
// Log
if self.show_logs {
let hr: usize = ((rows / 2) - 2).into();
let start: usize = if self.lines.len() > hr {
self.lines.len() - hr
} else {
0
};
queue!(stdout, MoveTo(0, rows / 2))?;
print!("{}", format!("{:columns$}", " Log").negative());
for line in self.lines.range(start..) {
print!("{}\r\n", line);
}
}
queue!(stdout, MoveTo(0, rows - 1))?;
print!(
"{}",
format!(
"{:columns$}",
" q - quit | l - toggle log | \u{2191}/\u{2193} - select port | e - enable/disable | <enter> - browse"
).negative()
);
stdout.flush()?;
Ok(())
}
fn connected(&self) -> bool {
self.socks_port != 0
}
fn get_selected_port(&self) -> Option<&PortDesc> {
if let Some(p) = &self.ports {
if self.selection < p.len() {
return Some(&p[self.selection]);
}
}
None
}
fn enable_disable_port(&mut self, port: u16) {
if self.socks_port != 0 {
if let Some(_) = self.listeners.remove(&port) {
@ -299,4 +253,134 @@ impl UI {
});
}
}
fn enter_alternate_screen(&mut self) -> Result<()> {
if !self.alternate_screen {
enable_raw_mode()?;
execute!(stdout(), EnterAlternateScreen, DisableLineWrap)?;
self.alternate_screen = true;
}
Ok(())
}
fn leave_alternate_screen(&mut self) -> Result<()> {
if self.alternate_screen {
execute!(stdout(), LeaveAlternateScreen, EnableLineWrap)?;
disable_raw_mode()?;
self.alternate_screen = false;
}
Ok(())
}
async fn handle_events(&mut self, console_events: &mut EventStream) {
tokio::select! {
ev = console_events.next() => self.handle_console_event(ev),
ev = self.events.recv() => self.handle_internal_event(ev),
}
}
fn handle_console_event(
&mut self,
ev: Option<Result<Event, std::io::Error>>,
) {
match ev {
Some(Ok(Event::Key(ev))) => match ev {
KeyEvent { code: KeyCode::Char('c'), .. } => {
if ev.modifiers.intersects(KeyModifiers::CONTROL) {
self.running = false;
}
}
KeyEvent { code: KeyCode::Esc, .. }
| KeyEvent { code: KeyCode::Char('q'), .. } => {
self.running = false;
}
KeyEvent { code: KeyCode::Char('l'), .. } => {
self.show_logs = !self.show_logs;
}
KeyEvent { code: KeyCode::Char('e'), .. } => {
if let Some(p) = self.get_selected_port() {
self.enable_disable_port(p.port);
}
}
KeyEvent { code: KeyCode::Up, .. }
| KeyEvent { code: KeyCode::Char('j'), .. } => {
if self.selection > 0 {
self.selection -= 1;
}
}
KeyEvent { code: KeyCode::Down, .. }
| KeyEvent { code: KeyCode::Char('k'), .. } => {
if let Some(p) = &self.ports {
if self.selection != p.len() - 1 {
self.selection += 1;
}
}
}
KeyEvent { code: KeyCode::Enter, .. } => {
if let Some(p) = self.get_selected_port() {
_ = open::that(format!("http://127.0.0.1:{}/", p.port));
}
}
_ => (),
},
Some(Ok(_)) => (), // Don't care about this event...
Some(Err(_)) => (), // Hmmmmmm.....?
None => (), // ....no events? what?
}
}
fn handle_internal_event(&mut self, event: Option<UIEvent>) {
match event {
Some(UIEvent::Disconnected) => {
self.socks_port = 0;
}
Some(UIEvent::Connected(sp)) => {
self.socks_port = sp;
info!("Socks port {socks_port}", socks_port = self.socks_port);
}
Some(UIEvent::Ports(mut p)) => {
p.sort_by(|a, b| a.port.partial_cmp(&b.port).unwrap());
if self.selection >= p.len() {
self.selection = p.len() - 1;
}
let mut new_listeners = HashMap::new();
for port in &p {
let port = port.port;
if let Some(l) = self.listeners.remove(&port) {
if !l.is_closed() {
// `l` here is, of course, the channel that we
// use to tell the listener task to stop (see the
// spawn call below). If it isn't closed then
// that means a spawn task is still running so we
// should just let it keep running and re-use the
// existing listener.
new_listeners.insert(port, l);
}
}
}
// This has the side effect of closing any
// listener that didn't get copied over to the
// new listeners table.
self.listeners = new_listeners;
self.ports = Some(p);
}
Some(UIEvent::ServerLine(line)) => {
while self.lines.len() >= 1024 {
self.lines.pop_front();
}
self.lines.push_back(format!("[SERVER] {line}"));
}
Some(UIEvent::LogLine(_level, line)) => {
while self.lines.len() >= 1024 {
self.lines.pop_front();
}
self.lines.push_back(format!("[CLIENT] {line}"));
}
None => {
self.running = false;
}
}
}
}