diff --git a/src/connection.rs b/src/connection.rs deleted file mode 100644 index 360ea75..0000000 --- a/src/connection.rs +++ /dev/null @@ -1,323 +0,0 @@ -use crate::message::Message; -use anyhow::Result; -use bytes::{Bytes, BytesMut}; -use std::collections::HashMap; -use std::sync::{Arc, Mutex}; -use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; -use tokio::net::TcpStream; -use tokio::sync::mpsc; -use tokio::sync::oneshot; - -const MAX_PACKET: usize = u16::max_value() as usize; - -/// Read from a socket and convert the reads into Messages to put into the -/// queue until the socket is closed for reading or an error occurs. We read -/// at most 2^16-1 bytes at a time, accepting the overhead of multiple reads -/// to keep one writer from clogging the pipe for everybody else. Each read -/// is converted into a [Message::Data] message that is sent to the `writer` -/// channel. -/// -/// Once we're done reading (either because of a connection error or a clean -/// shutdown) we send a [Message::Close] message on the channel before -/// returning. -/// -/// # Errors -/// If an error occurs reading from `read` we return [Error::IO]. If the -/// message channel is closed before we can send to it then we return -/// [Error::ConnectionReset]. -/// -async fn connection_read( - channel: u64, - read: &mut T, - writer: &mut mpsc::Sender, -) -> Result<(), tokio::io::Error> { - let result = loop { - let mut buffer = BytesMut::with_capacity(MAX_PACKET); - if let Err(e) = read.read_buf(&mut buffer).await { - break Err(e); - } - - if buffer.len() == 0 { - break Ok(()); - } - - if let Err(_) = writer.send(Message::Data(channel, buffer.into())).await { - break Err(tokio::io::Error::from( - tokio::io::ErrorKind::ConnectionReset, - )); - } - - // TODO: Flow control here, wait for the packet to be acknowleged so - // there isn't head-of-line blocking or infinite buffering on the - // remote side. Also buffer re-use! - }; - - // We are effectively closed on this side, send the close to drop the - // corresponding write side on the other end of the pipe. - _ = writer.send(Message::Close(channel)).await; - result -} - -/// Get messages from a queue and write them out to a socket until there are -/// no more messages in the queue or a write fails for some reason. -/// -/// # Errors -/// If a write fails this returns `Error::IO`. -/// -async fn connection_write( - data: &mut mpsc::Receiver, - write: &mut T, -) -> Result<()> { - while let Some(buf) = data.recv().await { - write.write_all(&buf[..]).await?; - } - Ok(()) -} - -/// Handle a connection, from the socket to the multiplexer and from the -/// multiplexer to the socket. Keeps running until both the read and write -/// side are closed. In natural circumstances, we expect the write side to -/// close when the `data` sender is dropped from the connection table (see -/// `ConnectionTable`), and we expect the read side to close when the -/// socket's read half closes (which will cause a `Close` to be sent which -/// should drop the `data` sender on the other side, etc...). -/// -pub async fn process( - channel: u64, - stream: &mut TcpStream, - data: &mut mpsc::Receiver, - writer: &mut mpsc::Sender, -) { - let (mut read_half, mut write_half) = stream.split(); - - let read = connection_read(channel, &mut read_half, writer); - let write = connection_write(data, &mut write_half); - - tokio::pin!(read); - tokio::pin!(write); - - let (mut done_reading, mut done_writing) = (false, false); - while !(done_reading && done_writing) { - tokio::select! { - _ = &mut read, if !done_reading => { done_reading = true; }, - _ = &mut write, if !done_writing => { done_writing = true; }, - } - } -} - -// ---------------------------------------------------------------------------- -// Tables - -/// The connection structure tracks the various channels used to communicate -/// with an "open" connection. -struct Connection { - /// The callback for the connected message, if we haven't already - /// connected across the channel. Realistically, this only ever has a - /// value on the client side, where we wait for the server side to - /// connect and then acknowlege that the connection. - connected: Option>, - - /// The channel where the connection receives [Bytes] to be written to - /// the socket. - data: mpsc::Sender, -} - -struct ConnectionTableState { - next_id: u64, - connections: HashMap, -} - -/// A tracking structure for connections. This structure is thread-safe and -/// so can be used to track new connections from as many concurrent listeners -/// as you would like. -#[derive(Clone)] -pub struct ConnectionTable { - connections: Arc>, -} - -impl ConnectionTable { - /// Create a new, empty connection table. - pub fn new() -> ConnectionTable { - ConnectionTable { - connections: Arc::new(Mutex::new(ConnectionTableState { - next_id: 0, - connections: HashMap::new(), - })), - } - } - - /// Allocate a new connection on the client side. The connection is - /// assigned a new ID, which is returned to the caller. - pub fn alloc( - self: &mut Self, - connected: oneshot::Sender<()>, - data: mpsc::Sender, - ) -> u64 { - let mut tbl = self.connections.lock().unwrap(); - let id = tbl.next_id; - tbl.next_id += 1; - tbl.connections.insert( - id, - Connection { - connected: Some(connected), - data, - }, - ); - id - } - - /// Add a connection to the table on the server side. The client sent us - /// the ID to use, so we don't need to allocate it, and obviously we - /// aren't going to be waiting for the connection to be "connected." - pub fn add(self: &mut Self, id: u64, data: mpsc::Sender) { - let mut tbl = self.connections.lock().unwrap(); - tbl.connections.insert( - id, - Connection { - connected: None, - data, - }, - ); - } - - /// Mark a connection as being "connected", on the client side, where we - /// wait for the server to tell us such things. Note that this gets used - /// for a successful connection; on a failure just call [remove]. - pub fn connected(self: &mut Self, id: u64) { - let connected = { - let mut tbl = self.connections.lock().unwrap(); - if let Some(c) = tbl.connections.get_mut(&id) { - c.connected.take() - } else { - None - } - }; - - if let Some(connected) = connected { - _ = connected.send(()); - } - } - - /// Tell a connection that we have received data. This gets used on both - /// sides of the pipe; if the connection exists and is still active it - /// will send the data out through its socket. - pub async fn receive(self: &Self, id: u64, buf: Bytes) { - let data = { - let tbl = self.connections.lock().unwrap(); - if let Some(connection) = tbl.connections.get(&id) { - Some(connection.data.clone()) - } else { - None - } - }; - - if let Some(data) = data { - _ = data.send(buf).await; - } - } - - /// Remove a connection from the table, effectively closing it. This will - /// close all the pipes that the connection uses to receive data from the - /// other side, performing a cleanup on our "write" side of the socket. - pub fn remove(self: &mut Self, id: u64) { - let mut tbl = self.connections.lock().unwrap(); - tbl.connections.remove(&id); - } -} - -#[cfg(test)] -mod tests { - use super::*; - use tokio::net::TcpListener; - - async fn create_connected_pair() -> (TcpStream, TcpStream) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - - let connect = tokio::spawn(async move { - TcpStream::connect(format!("127.0.0.1:{}", port)) - .await - .unwrap() - }); - - let (server, _) = listener.accept().await.unwrap(); - let client = connect.await.unwrap(); - - (client, server) - } - - #[tokio::test] - async fn test_connected_pair() { - // This is just a sanity test to make sure my socket nonsense is working. - let (mut client, mut server) = create_connected_pair().await; - - let a = tokio::spawn(async move { - let mut d = vec![1, 2, 3]; - client.write_all(&mut d).await.unwrap(); - //eprintln!("Wrote something!"); - }); - - let b = tokio::spawn(async move { - let mut x = BytesMut::with_capacity(3); - server.read_buf(&mut x).await.unwrap(); - //panic!("Read {:?}", x); - }); - - a.await.unwrap(); - b.await.unwrap(); - } - - #[tokio::test] - async fn test_process_connection() { - let (mut client, mut server) = create_connected_pair().await; - - const CHID: u64 = 123; - let (mut msg_writer, mut msg_receiver) = mpsc::channel(32); - let (data_writer, mut data_receiver) = mpsc::channel(32); - - let proc = tokio::spawn(async move { - process(CHID, &mut server, &mut data_receiver, &mut msg_writer).await - }); - - // Any bytes I send through `data_writer` will come into my socket. - let packet = Bytes::from("hello world"); - data_writer.send(packet.clone()).await.unwrap(); - - let mut buffer = BytesMut::with_capacity(packet.len()); - buffer.resize(packet.len(), 0); - client.read_exact(&mut buffer).await.unwrap(); - assert_eq!(packet, buffer); - - // Any bytes I send through client come through on msg_receiver. - client.write_all(&packet[..]).await.unwrap(); - let msg = msg_receiver.recv().await.unwrap(); - assert_eq!(msg, Message::Data(CHID, packet.clone())); - - // When I close the write half of the socket then I get a close - // message. - let (mut read_half, mut write_half) = client.split(); - write_half.shutdown().await.unwrap(); - let msg = msg_receiver.recv().await.unwrap(); - assert_eq!(msg, Message::Close(CHID)); - - // I should still be able to use the read half of the socket. - let packet = Bytes::from("StIlL AlIvE"); - data_writer.send(packet.clone()).await.unwrap(); - - let mut buffer = BytesMut::with_capacity(packet.len()); - buffer.resize(packet.len(), 0); - read_half.read_exact(&mut buffer).await.unwrap(); - assert_eq!(packet, buffer); - - // When I drop the data writer my read half closes. - drop(data_writer); - let mut buffer = BytesMut::with_capacity(1024); - read_half.read_buf(&mut buffer).await.unwrap(); - assert_eq!(buffer.len(), 0); - - drop(read_half); - - // and the processing loop terminates. - proc.await.unwrap(); - } -} diff --git a/src/lib.rs b/src/lib.rs index 6b7f97c..f954731 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,6 @@ use anyhow::{bail, Result}; -use connection::ConnectionTable; use log::LevelFilter; -use log::{error, info, warn}; +use log::{debug, error, info, warn}; use message::{Message, MessageReader, MessageWriter}; use std::net::{Ipv4Addr, SocketAddrV4}; use tokio::io::{ @@ -11,9 +10,7 @@ use tokio::io::{ use tokio::net::{TcpListener, TcpStream}; use tokio::process; use tokio::sync::mpsc; -use tokio::sync::oneshot; -mod connection; mod message; mod refresh; mod ui; @@ -51,29 +48,9 @@ async fn pump_write( // ---------------------------------------------------------------------------- // Server -async fn server_handle_connection( - channel: u64, - port: u16, - writer: mpsc::Sender, - connections: ConnectionTable, -) { - let mut connections = connections; - if let Ok(mut stream) = TcpStream::connect(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)).await { - let (send_data, mut data) = mpsc::channel(32); - connections.add(channel, send_data); - if let Ok(_) = writer.send(Message::Connected(channel)).await { - let mut writer = writer.clone(); - connection::process(channel, &mut stream, &mut data, &mut writer).await; - - info!("< Done server!"); - } - } -} - async fn server_read( reader: &mut MessageReader, writer: mpsc::Sender, - connections: ConnectionTable, ) -> Result<()> { // info!("< Processing packets..."); loop { @@ -82,28 +59,6 @@ async fn server_read( use Message::*; match message { Ping => (), - Connect(channel, port) => { - let (writer, connections) = (writer.clone(), connections.clone()); - tokio::spawn(async move { - server_handle_connection(channel, port, writer, connections).await; - }); - } - Close(channel) => { - let mut connections = connections.clone(); - tokio::spawn(async move { - // Once we get a close the connection becomes unreachable. - // - // NOTE: If all goes well the 'data' channel gets dropped - // here, and we close the write half of the socket. - connections.remove(channel); - }); - } - Data(channel, buf) => { - let connections = connections.clone(); - tokio::spawn(async move { - connections.receive(channel, buf).await; - }); - } Refresh => { let writer = writer.clone(); tokio::spawn(async move { @@ -129,15 +84,13 @@ async fn server_main( reader: &mut MessageReader, writer: &mut MessageWriter, ) -> Result<()> { - let connections = ConnectionTable::new(); - // The first message we send must be an announcement. writer.write(Message::Hello(0, 1, vec![])).await?; // Jump into it... let (msg_sender, mut msg_receiver) = mpsc::channel(32); let writing = pump_write(&mut msg_receiver, writer); - let reading = server_read(reader, msg_sender, connections); + let reading = server_read(reader, msg_sender); tokio::pin!(reading); tokio::pin!(writing); @@ -190,42 +143,132 @@ async fn client_sync(reader: &mut Read) -> Result<(), t } } -async fn client_handle_connection( - port: u16, - writer: mpsc::Sender, - connections: ConnectionTable, - socket: &mut TcpStream, -) { - let mut connections = connections; - let (send_connected, connected) = oneshot::channel(); - let (send_data, mut data) = mpsc::channel(32); - let channel = connections.alloc(send_connected, send_data); +/// Handle an incoming client connection, by forwarding it to the SOCKS5 +/// server at the specified port. +/// +/// 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<()> { + debug!("Handling connection!"); - if let Ok(_) = writer.send(Message::Connect(channel, port)).await { - if let Ok(_) = connected.await { - let mut writer = writer.clone(); - connection::process(channel, socket, &mut data, &mut writer).await; - } else { - error!("Failed to connect to remote"); - } + let dest_addr = SocketAddrV4::new(Ipv4Addr::LOCALHOST, socks_port); + let mut dest_socket = TcpStream::connect(dest_addr).await?; + + debug!("Connected, sending handshake request"); + let packet: [u8; 3] = [ + 0x05, // v5 + 0x01, // 1 auth method + 0x00, // my one auth method is no auth + ]; + dest_socket.write_all(&packet[..]).await?; + debug!("Initial handshake sent. Awaiting handshake response"); + + let mut response: [u8; 2] = [0; 2]; + dest_socket.read_exact(&mut response).await?; + if response[0] != 0x05 { + bail!("SOCKS incorrect response version {}", response[0]); } + if response[1] == 0xFF { + bail!("SOCKS server says no acceptable auth"); + } + if response[1] != 0x00 { + bail!("SOCKS server chose something wild? {}", response[1]); + } + + debug!("Handshake response received, sending connect request"); + let packet: [u8; 10] = [ + 0x05, // version again :P + 0x01, // connect + 0x00, // reserved! + 0x01, // ipv4 + 127, // lo + 0, // ..ca.. + 0, // ..lho.. + 1, // ..st + ((port & 0xFF00) >> 8).try_into().unwrap(), // port (high) + ((port & 0x00FF) >> 0).try_into().unwrap(), // port (low) + ]; + dest_socket.write_all(&packet[..]).await?; + + debug!("Connect request sent, awaiting response"); + let mut response: [u8; 4] = [0; 4]; + dest_socket.read_exact(&mut response).await?; + if response[0] != 0x05 { + bail!("SOCKS5 incorrect response version again? {}", response[0]); + } + if response[1] != 0x00 { + bail!("SOCKS5 reports a connect error {}", response[1]); + } + // Now we 100% do not care about the following information but we must + // discard it so we can get to the good stuff. response[3] is the type of address... + if response[3] == 0x01 { + // IPv4 - 4 bytes. + let mut response: [u8; 4] = [0; 4]; + dest_socket.read_exact(&mut response).await?; + } else if response[3] == 0x03 { + // Domain Name + let len = dest_socket.read_u8().await?; + for _ in 0..len { + dest_socket.read_u8().await?; // So slow! + } + } else if response[3] == 0x04 { + // IPv6 - 8 bytes + let mut response: [u8; 8] = [0; 8]; + dest_socket.read_exact(&mut response).await?; + } else { + bail!( + "SOCKS5 sent me an address I don't understand {}", + response[3] + ); + } + // Finally the port number. Again, garbage, but it's in the packet we need to skip. + let mut response: [u8; 2] = [0; 2]; + dest_socket.read_exact(&mut response).await?; + + info!("Connection established on port {}", port); + + let (client_read_half, client_write_half) = socket.into_split(); + let (server_read_half, server_write_half) = dest_socket.into_split(); + let client_to_server = tokio::spawn(async move { + let mut client_read_half = client_read_half; + let mut server_write_half = server_write_half; + tokio::io::copy(&mut client_read_half, &mut server_write_half).await + }); + let server_to_client = tokio::spawn(async move { + let mut server_read_half = server_read_half; + let mut client_write_half = client_write_half; + tokio::io::copy(&mut server_read_half, &mut client_write_half).await + }); + + let client_err = client_to_server.await; + debug!("Done client -> server"); + let svr_err = server_to_client.await; + debug!("Done server -> client"); + + if let Ok(Err(e)) = client_err { + return Err(e.into()); + } else if let Ok(Err(e)) = svr_err { + return Err(e.into()); + } + + Ok(()) } -async fn client_listen( - port: u16, - writer: mpsc::Sender, - connections: ConnectionTable, -) -> Result<()> { +async fn client_listen(port: u16, socks_port: u16) -> Result<()> { loop { 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 (mut socket, _) = listener.accept().await?; + let (socket, _) = listener.accept().await?; - let (writer, connections) = (writer.clone(), connections.clone()); tokio::spawn(async move { - client_handle_connection(port, writer, connections, &mut socket).await; + if let Err(e) = client_handle_connection(socks_port, port, socket).await { + error!("Error handling connection: {:?}", e); + } else { + debug!("Done???"); + } }); } } @@ -233,7 +276,6 @@ async fn client_listen( async fn client_read( reader: &mut MessageReader, - connections: ConnectionTable, events: mpsc::Sender, ) -> Result<()> { info!("Running"); @@ -244,24 +286,6 @@ async fn client_read( use Message::*; match message { Ping => (), - Connected(channel) => { - let mut connections = connections.clone(); - tokio::spawn(async move { - connections.connected(channel); - }); - } - Close(channel) => { - let mut connections = connections.clone(); - tokio::spawn(async move { - connections.remove(channel); - }); - } - Data(channel, buf) => { - let connections = connections.clone(); - tokio::spawn(async move { - connections.receive(channel, buf).await; - }); - } Ports(ports) => { if let Err(_) = events.send(ui::UIEvent::Ports(ports)).await { // TODO: Log @@ -295,6 +319,7 @@ async fn client_pipe_stderr( } async fn client_main( + socks_port: u16, reader: &mut MessageReader, writer: &mut MessageWriter, events: mpsc::Sender, @@ -308,20 +333,13 @@ async fn client_main( bail!("Expected a hello message from the remote server"); } - let connections = ConnectionTable::new(); - // And now really get into it... let (msg_sender, mut msg_receiver) = mpsc::channel(32); - _ = events - .send(ui::UIEvent::Connected( - msg_sender.clone(), - connections.clone(), - )) - .await; + _ = events.send(ui::UIEvent::Connected(socks_port)).await; let writing = pump_write(&mut msg_receiver, writer); - let reading = client_read(reader, connections, events); + let reading = client_read(reader, events); tokio::pin!(reading); tokio::pin!(writing); @@ -384,21 +402,32 @@ pub async fn run_server() { } } -async fn spawn_ssh(server: &str) -> Result { +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() + }; + let mut cmd = process::Command::new("ssh"); - cmd.arg("-T").arg(server).arg("fwd").arg("--server"); + cmd.arg("-T") + .arg("-D") + .arg(socks_port.to_string()) + .arg(server) + .arg("fwd") + .arg("--server"); cmd.stdout(std::process::Stdio::piped()); cmd.stdin(std::process::Stdio::piped()); cmd.stderr(std::process::Stdio::piped()); - cmd.spawn() + let child = cmd.spawn()?; + Ok((child, socks_port)) } async fn client_connect_loop(remote: &str, events: mpsc::Sender) { loop { _ = events.send(ui::UIEvent::Disconnected).await; - let mut child = 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 @@ -433,7 +462,7 @@ async fn client_connect_loop(remote: &str, events: mpsc::Sender) { client_pipe_stderr(&mut stderr, sec).await; }); - if let Err(e) = client_main(&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..."); diff --git a/src/message.rs b/src/message.rs index 5da977f..b638866 100644 --- a/src/message.rs +++ b/src/message.rs @@ -25,12 +25,8 @@ pub struct PortDesc { pub enum Message { Ping, // Ignored on both sides, can be used to test connection. Hello(u8, u8, Vec), // Server info announcement: major version, minor version, headers. - Connect(u64, u16), // Request to connect on a port from client to server. - Connected(u64), // Sucessfully connected from server to client. - Close(u64), // Notify that one or the other end of a channel is closed. Refresh, // Request to refresh list of ports from client. Ports(Vec), // List of available ports from server to client. - Data(u64, Bytes), // Transmit data on a channel. } impl Message { @@ -55,19 +51,6 @@ impl Message { put_string(result, detail); } } - Connect(channel, port) => { - result.put_u8(0x02); - result.put_u64(*channel); - result.put_u16(*port); - } - Connected(channel) => { - result.put_u8(0x03); - result.put_u64(*channel); - } - Close(channel) => { - result.put_u8(0x04); - result.put_u64(*channel); - } Refresh => { result.put_u8(0x05); } @@ -83,12 +66,6 @@ impl Message { put_string(result, sliced); } } - Data(channel, bytes) => { - result.put_u8(0x07); - result.put_u64(*channel); - result.put_u16(bytes.len().try_into().expect("Payload too big")); - result.put_slice(bytes); // I hate that this copies. We should make this an async write probably, maybe? - } }; } @@ -106,19 +83,6 @@ impl Message { } Ok(Hello(major, minor, details)) } - 0x02 => { - let channel = get_u64(cursor)?; - let port = get_u16(cursor)?; - Ok(Connect(channel, port)) - } - 0x03 => { - let channel = get_u64(cursor)?; - Ok(Connected(channel)) - } - 0x04 => { - let channel = get_u64(cursor)?; - Ok(Close(channel)) - } 0x05 => Ok(Refresh), 0x06 => { let count = get_u16(cursor)?; @@ -130,12 +94,6 @@ impl Message { } Ok(Ports(ports)) } - 0x07 => { - let channel = get_u64(cursor)?; - let length = get_u16(cursor)?; - let data = get_bytes(cursor, length.into())?; - Ok(Data(channel, data)) - } b => Err(MessageError::Unknown(b).into()), } } @@ -155,13 +113,6 @@ fn get_u16(cursor: &mut Cursor<&[u8]>) -> Result { Ok(cursor.get_u16()) } -fn get_u64(cursor: &mut Cursor<&[u8]>) -> Result { - if cursor.remaining() < 8 { - return Err(MessageError::Incomplete); - } - Ok(cursor.get_u64()) -} - fn get_bytes(cursor: &mut Cursor<&[u8]>, length: usize) -> Result { if cursor.remaining() < length { return Err(MessageError::Incomplete); @@ -281,9 +232,6 @@ mod message_tests { vec!["One".to_string(), "Two".to_string(), "Three".to_string()], )); assert_round_trip(Hello(0x00, 0x01, vec![])); - assert_round_trip(Connect(0x1234567890123456, 0x1234)); - assert_round_trip(Connected(0x1234567890123456)); - assert_round_trip(Close(0x1234567890123456)); assert_round_trip(Refresh); assert_round_trip(Ports(vec![])); assert_round_trip(Ports(vec![ @@ -296,8 +244,6 @@ mod message_tests { desc: "metadata-library".to_string(), }, ])); - assert_round_trip(Data(0x1234567890123456, vec![1, 2, 3, 4].into())); - assert_round_trip(Data(0x123, vec![0; u16::max_value().into()].into())); } #[test] diff --git a/src/ui.rs b/src/ui.rs index 0808018..cbabf1b 100644 --- a/src/ui.rs +++ b/src/ui.rs @@ -1,8 +1,5 @@ use crate::client_listen; -use crate::{ - message::{Message, PortDesc}, - ConnectionTable, -}; +use crate::message::PortDesc; use anyhow::Result; use crossterm::{ cursor::MoveTo, @@ -24,7 +21,7 @@ use tokio::sync::oneshot; use tokio_stream::StreamExt; pub enum UIEvent { - Connected(mpsc::Sender, ConnectionTable), + Connected(u16), Disconnected, ServerLine(String), LogLine(log::Level, String), @@ -61,18 +58,16 @@ impl log::Log for Logger { pub struct UI { events: mpsc::Receiver, - writer: Option>, - connections: Option, listeners: HashMap>, + socks_port: u16, } impl UI { pub fn new(events: mpsc::Receiver) -> UI { UI { events, - writer: None, - connections: None, listeners: HashMap::new(), + socks_port: 0, } } @@ -146,13 +141,12 @@ impl UI { ev = self.events.recv() => { match ev { Some(UIEvent::Disconnected) => { - self.writer = None; - self.connections = None; + self.socks_port = 0; connected = false; } - Some(UIEvent::Connected(w,t)) => { - self.writer = Some(w); - self.connections = Some(t); + 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)) => { @@ -282,7 +276,7 @@ impl UI { } fn enable_disable_port(&mut self, port: u16) { - if let (Some(writer), Some(connections)) = (&self.writer, &self.connections) { + if self.socks_port != 0 { if let Some(_) = self.listeners.remove(&port) { return; // We disabled the listener. } @@ -291,10 +285,10 @@ impl UI { let (l, stop) = oneshot::channel(); self.listeners.insert(port, l); - let (writer, connections) = (writer.clone(), connections.clone()); + let socks_port = self.socks_port; tokio::spawn(async move { let result = tokio::select! { - r = client_listen(port, writer, connections) => r, + r = client_listen(port, socks_port) => r, _ = stop => Ok(()), }; if let Err(e) = result {