Use SOCKS5 proxy instead of multiplexing myself
This got a lot simpler.
This commit is contained in:
parent
3844200118
commit
c412e81b26
4 changed files with 147 additions and 501 deletions
|
|
@ -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<T: AsyncRead + Unpin>(
|
||||
channel: u64,
|
||||
read: &mut T,
|
||||
writer: &mut mpsc::Sender<Message>,
|
||||
) -> 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<T: AsyncWrite + Unpin>(
|
||||
data: &mut mpsc::Receiver<Bytes>,
|
||||
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<Bytes>,
|
||||
writer: &mut mpsc::Sender<Message>,
|
||||
) {
|
||||
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<oneshot::Sender<()>>,
|
||||
|
||||
/// The channel where the connection receives [Bytes] to be written to
|
||||
/// the socket.
|
||||
data: mpsc::Sender<Bytes>,
|
||||
}
|
||||
|
||||
struct ConnectionTableState {
|
||||
next_id: u64,
|
||||
connections: HashMap<u64, Connection>,
|
||||
}
|
||||
|
||||
/// 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<Mutex<ConnectionTableState>>,
|
||||
}
|
||||
|
||||
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<Bytes>,
|
||||
) -> 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<Bytes>) {
|
||||
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();
|
||||
}
|
||||
}
|
||||
241
src/lib.rs
241
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<T: AsyncWrite + Unpin>(
|
|||
// ----------------------------------------------------------------------------
|
||||
// Server
|
||||
|
||||
async fn server_handle_connection(
|
||||
channel: u64,
|
||||
port: u16,
|
||||
writer: mpsc::Sender<Message>,
|
||||
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<T: AsyncRead + Unpin>(
|
||||
reader: &mut MessageReader<T>,
|
||||
writer: mpsc::Sender<Message>,
|
||||
connections: ConnectionTable,
|
||||
) -> Result<()> {
|
||||
// info!("< Processing packets...");
|
||||
loop {
|
||||
|
|
@ -82,28 +59,6 @@ async fn server_read<T: AsyncRead + Unpin>(
|
|||
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: AsyncRead + Unpin, Writer: AsyncWrite + Unpin>(
|
|||
reader: &mut MessageReader<Reader>,
|
||||
writer: &mut MessageWriter<Writer>,
|
||||
) -> 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<Read: AsyncRead + Unpin>(reader: &mut Read) -> Result<(), t
|
|||
}
|
||||
}
|
||||
|
||||
async fn client_handle_connection(
|
||||
port: u16,
|
||||
writer: mpsc::Sender<Message>,
|
||||
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;
|
||||
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 {
|
||||
error!("Failed to connect to remote");
|
||||
}
|
||||
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());
|
||||
}
|
||||
|
||||
async fn client_listen(
|
||||
port: u16,
|
||||
writer: mpsc::Sender<Message>,
|
||||
connections: ConnectionTable,
|
||||
) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
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<T: AsyncRead + Unpin>(
|
||||
reader: &mut MessageReader<T>,
|
||||
connections: ConnectionTable,
|
||||
events: mpsc::Sender<ui::UIEvent>,
|
||||
) -> Result<()> {
|
||||
info!("Running");
|
||||
|
|
@ -244,24 +286,6 @@ async fn client_read<T: AsyncRead + Unpin>(
|
|||
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<Debug: AsyncBufRead + Unpin>(
|
|||
}
|
||||
|
||||
async fn client_main<Reader: AsyncRead + Unpin, Writer: AsyncWrite + Unpin>(
|
||||
socks_port: u16,
|
||||
reader: &mut MessageReader<Reader>,
|
||||
writer: &mut MessageWriter<Writer>,
|
||||
events: mpsc::Sender<ui::UIEvent>,
|
||||
|
|
@ -308,20 +333,13 @@ async fn client_main<Reader: AsyncRead + Unpin, Writer: AsyncWrite + Unpin>(
|
|||
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<tokio::process::Child, 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()
|
||||
};
|
||||
|
||||
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<ui::UIEvent>) {
|
||||
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<ui::UIEvent>) {
|
|||
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...");
|
||||
|
|
|
|||
|
|
@ -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<String>), // 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<PortDesc>), // 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<u16, MessageError> {
|
|||
Ok(cursor.get_u16())
|
||||
}
|
||||
|
||||
fn get_u64(cursor: &mut Cursor<&[u8]>) -> Result<u64, MessageError> {
|
||||
if cursor.remaining() < 8 {
|
||||
return Err(MessageError::Incomplete);
|
||||
}
|
||||
Ok(cursor.get_u64())
|
||||
}
|
||||
|
||||
fn get_bytes(cursor: &mut Cursor<&[u8]>, length: usize) -> Result<Bytes, MessageError> {
|
||||
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]
|
||||
|
|
|
|||
28
src/ui.rs
28
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<Message>, 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<UIEvent>,
|
||||
writer: Option<mpsc::Sender<Message>>,
|
||||
connections: Option<ConnectionTable>,
|
||||
listeners: HashMap<u16, oneshot::Sender<()>>,
|
||||
socks_port: u16,
|
||||
}
|
||||
|
||||
impl UI {
|
||||
pub fn new(events: mpsc::Receiver<UIEvent>) -> 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 {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue