Make a fine distinction between enabled and connected
This is part of remembering which ports should be enabled and which not; track ports even if they've gone away or between connection and disconnection, &c.
This commit is contained in:
parent
c6acf52b32
commit
a3727812a2
1 changed files with 153 additions and 91 deletions
244
src/ui.rs
244
src/ui.rs
|
|
@ -15,7 +15,7 @@ use crossterm::{
|
|||
use log::{error, info, Level, Metadata, Record};
|
||||
use open;
|
||||
use std::collections::vec_deque::VecDeque;
|
||||
use std::collections::HashMap;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::io::{stdout, Write};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
|
|
@ -57,15 +57,79 @@ impl log::Log for Logger {
|
|||
fn flush(&self) {}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Listener {
|
||||
enabled: bool,
|
||||
stop: Option<oneshot::Sender<()>>,
|
||||
desc: Option<PortDesc>,
|
||||
}
|
||||
|
||||
impl Listener {
|
||||
pub fn from_desc(desc: PortDesc) -> Listener {
|
||||
Listener {
|
||||
enabled: false,
|
||||
stop: None,
|
||||
desc: Some(desc),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn enabled(&self) -> bool {
|
||||
self.enabled
|
||||
}
|
||||
|
||||
pub fn set_enabled(&mut self, socks_port: Option<u16>, enabled: bool) {
|
||||
if enabled {
|
||||
self.enabled = true;
|
||||
self.start(socks_port);
|
||||
} else {
|
||||
self.enabled = false;
|
||||
self.stop = None;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn connect(&mut self, socks_port: Option<u16>, desc: PortDesc) {
|
||||
self.desc = Some(desc);
|
||||
self.start(socks_port);
|
||||
}
|
||||
|
||||
pub fn disconnect(&mut self) {
|
||||
self.desc = None;
|
||||
self.stop = None;
|
||||
}
|
||||
|
||||
pub fn start(&mut self, socks_port: Option<u16>) {
|
||||
if self.enabled {
|
||||
if let (Some(desc), Some(socks_port), None) =
|
||||
(&self.desc, socks_port, &self.stop)
|
||||
{
|
||||
let (l, stop) = oneshot::channel();
|
||||
let port = desc.port;
|
||||
tokio::spawn(async move {
|
||||
let result = tokio::select! {
|
||||
r = client_listen(port, socks_port) => r,
|
||||
_ = stop => Ok(()),
|
||||
};
|
||||
if let Err(e) = result {
|
||||
error!("Error listening on port {port}: {e:?}");
|
||||
} else {
|
||||
info!("Stopped listening on port {port}");
|
||||
}
|
||||
});
|
||||
|
||||
self.stop = Some(l);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct UI {
|
||||
events: mpsc::Receiver<UIEvent>,
|
||||
listeners: HashMap<u16, oneshot::Sender<()>>,
|
||||
socks_port: u16,
|
||||
ports: HashMap<u16, Listener>,
|
||||
socks_port: Option<u16>,
|
||||
running: bool,
|
||||
show_logs: bool,
|
||||
selection: usize,
|
||||
ports: Option<Vec<PortDesc>>,
|
||||
lines: VecDeque<String>,
|
||||
alternate_screen: bool,
|
||||
raw_mode: bool,
|
||||
|
|
@ -75,12 +139,11 @@ impl UI {
|
|||
pub fn new(events: mpsc::Receiver<UIEvent>) -> UI {
|
||||
UI {
|
||||
events,
|
||||
listeners: HashMap::new(),
|
||||
socks_port: 0,
|
||||
ports: HashMap::new(),
|
||||
socks_port: None,
|
||||
running: true,
|
||||
show_logs: false,
|
||||
selection: 0,
|
||||
ports: None,
|
||||
lines: VecDeque::with_capacity(1024),
|
||||
alternate_screen: false,
|
||||
raw_mode: false,
|
||||
|
|
@ -160,30 +223,31 @@ impl UI {
|
|||
)
|
||||
.negative()
|
||||
);
|
||||
if let Some(ports) = &self.ports {
|
||||
let max_ports: usize = if self.show_logs {
|
||||
(rows / 2) - 1
|
||||
|
||||
let max_ports: usize = if self.show_logs {
|
||||
(rows / 2) - 1
|
||||
} else {
|
||||
rows - 2
|
||||
}
|
||||
.into();
|
||||
|
||||
let ports = self.get_ui_ports();
|
||||
for (index, port) in ports.into_iter().take(max_ports).enumerate() {
|
||||
let listener = self.ports.get(&port).unwrap();
|
||||
|
||||
let caret = if index == self.selection {
|
||||
"\u{2B46}"
|
||||
} 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 == self.selection {
|
||||
"\u{2B46}"
|
||||
} else {
|
||||
" "
|
||||
},
|
||||
if self.listeners.contains_key(&port.port) {
|
||||
"+"
|
||||
} else {
|
||||
" "
|
||||
},
|
||||
port.port,
|
||||
port.desc
|
||||
);
|
||||
}
|
||||
" "
|
||||
};
|
||||
|
||||
let state = if listener.enabled { "+" } else { " " };
|
||||
let desc: &str = match &listener.desc {
|
||||
Some(port_desc) => &port_desc.desc,
|
||||
None => "",
|
||||
};
|
||||
|
||||
print!("{caret} {state:>enabled_width$} {port:port_width$} {desc:description_width$}\r\n");
|
||||
}
|
||||
|
||||
// Log
|
||||
|
|
@ -215,40 +279,29 @@ impl UI {
|
|||
}
|
||||
|
||||
fn connected(&self) -> bool {
|
||||
self.socks_port != 0
|
||||
match self.socks_port {
|
||||
Some(_) => true,
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn get_selected_port(&self) -> Option<&PortDesc> {
|
||||
if let Some(p) = &self.ports {
|
||||
if self.selection < p.len() {
|
||||
return Some(&p[self.selection]);
|
||||
}
|
||||
fn get_ui_ports(&self) -> Vec<u16> {
|
||||
let mut ports: Vec<u16> = self.ports.keys().copied().collect();
|
||||
ports.sort();
|
||||
ports
|
||||
}
|
||||
|
||||
fn get_selected_port(&self) -> Option<u16> {
|
||||
if self.selection < self.ports.len() {
|
||||
Some(self.get_ui_ports()[self.selection])
|
||||
} else {
|
||||
None
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn enable_disable_port(&mut self, port: u16) {
|
||||
if self.socks_port != 0 {
|
||||
if let Some(_) = self.listeners.remove(&port) {
|
||||
return; // We disabled the listener.
|
||||
}
|
||||
|
||||
// We need to enable the listener.
|
||||
let (l, stop) = oneshot::channel();
|
||||
self.listeners.insert(port, l);
|
||||
|
||||
let socks_port = self.socks_port;
|
||||
tokio::spawn(async move {
|
||||
let result = tokio::select! {
|
||||
r = client_listen(port, socks_port) => r,
|
||||
_ = stop => Ok(()),
|
||||
};
|
||||
if let Err(e) = result {
|
||||
error!("Error listening on port {port}: {e:?}");
|
||||
} else {
|
||||
info!("Stopped listening on port {port}");
|
||||
}
|
||||
});
|
||||
if let Some(listener) = self.ports.get_mut(&port) {
|
||||
listener.set_enabled(self.socks_port, !listener.enabled());
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -313,7 +366,7 @@ impl UI {
|
|||
}
|
||||
KeyEvent { code: KeyCode::Char('e'), .. } => {
|
||||
if let Some(p) = self.get_selected_port() {
|
||||
self.enable_disable_port(p.port);
|
||||
self.enable_disable_port(p);
|
||||
}
|
||||
}
|
||||
KeyEvent { code: KeyCode::Up, .. }
|
||||
|
|
@ -324,15 +377,13 @@ impl UI {
|
|||
}
|
||||
KeyEvent { code: KeyCode::Down, .. }
|
||||
| KeyEvent { code: KeyCode::Char('k'), .. } => {
|
||||
if let Some(p) = &self.ports {
|
||||
if self.selection != p.len() - 1 {
|
||||
self.selection += 1;
|
||||
}
|
||||
if self.selection != self.ports.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));
|
||||
_ = open::that(format!("http://127.0.0.1:{}/", p));
|
||||
}
|
||||
}
|
||||
_ => (),
|
||||
|
|
@ -346,40 +397,51 @@ impl UI {
|
|||
fn handle_internal_event(&mut self, event: Option<UIEvent>) {
|
||||
match event {
|
||||
Some(UIEvent::Disconnected) => {
|
||||
self.socks_port = 0;
|
||||
self.listeners = HashMap::new(); // Bye.
|
||||
self.socks_port = None;
|
||||
for port in self.ports.values_mut() {
|
||||
port.disconnect();
|
||||
}
|
||||
}
|
||||
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() && p.len() > 0 {
|
||||
self.selection = p.len() - 1;
|
||||
info!("Socks port {sp}");
|
||||
self.socks_port = Some(sp);
|
||||
for port in self.ports.values_mut() {
|
||||
port.start(self.socks_port);
|
||||
}
|
||||
}
|
||||
Some(UIEvent::Ports(p)) => {
|
||||
let mut leftover_ports: HashSet<u16> =
|
||||
HashSet::from_iter(self.ports.keys().copied());
|
||||
|
||||
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);
|
||||
}
|
||||
for port_desc in p.into_iter() {
|
||||
leftover_ports.remove(&port_desc.port);
|
||||
if let Some(listener) = self.ports.get_mut(&port_desc.port)
|
||||
{
|
||||
listener.connect(self.socks_port, port_desc);
|
||||
} else {
|
||||
self.ports.insert(
|
||||
port_desc.port,
|
||||
Listener::from_desc(port_desc),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// 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);
|
||||
for port in leftover_ports {
|
||||
let mut enabled = false;
|
||||
if let Some(listener) = self.ports.get_mut(&port) {
|
||||
enabled = listener.enabled;
|
||||
listener.disconnect();
|
||||
}
|
||||
|
||||
if !enabled {
|
||||
// Just... forget it? Or leave it around forever?
|
||||
self.ports.remove(&port);
|
||||
}
|
||||
}
|
||||
|
||||
if self.selection >= self.ports.len() && self.ports.len() > 0 {
|
||||
self.selection = self.ports.len() - 1;
|
||||
}
|
||||
}
|
||||
Some(UIEvent::ServerLine(line)) => {
|
||||
while self.lines.len() >= 1024 {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue