Reorganization ssh
This commit is contained in:
parent
95c1ec2b96
commit
2bd26f1db2
|
@ -1 +1,3 @@
|
|||
pub mod animals;
|
||||
pub mod ssh;
|
||||
pub mod tunnel;
|
||||
|
|
242
src/main.rs
242
src/main.rs
|
@ -1,10 +1,4 @@
|
|||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
net::SocketAddr,
|
||||
path::Path,
|
||||
sync::Arc,
|
||||
time::Duration,
|
||||
};
|
||||
use std::{net::SocketAddr, path::Path};
|
||||
|
||||
use bytes::Bytes;
|
||||
use http_body_util::{BodyExt, Full, combinators::BoxBody};
|
||||
|
@ -17,20 +11,10 @@ use hyper::{
|
|||
};
|
||||
use hyper_util::rt::TokioIo;
|
||||
use rand::rngs::OsRng;
|
||||
use russh::{
|
||||
ChannelId,
|
||||
server::{self, Handle, Server as _},
|
||||
};
|
||||
use tokio::{
|
||||
net::TcpListener,
|
||||
sync::{
|
||||
RwLock,
|
||||
mpsc::{self, UnboundedReceiver, UnboundedSender},
|
||||
},
|
||||
};
|
||||
use tracing::{debug, error, trace, warn};
|
||||
use tokio::net::TcpListener;
|
||||
use tracing::{debug, trace, warn};
|
||||
use tracing_subscriber::{EnvFilter, Registry, layer::SubscriberExt, util::SubscriberInitExt};
|
||||
use tunnel_rs::animals::get_animal_name;
|
||||
use tunnel_rs::ssh::Server;
|
||||
|
||||
fn full<T: Into<Bytes>>(chunk: T) -> BoxBody<Bytes, hyper::Error> {
|
||||
Full::new(chunk.into())
|
||||
|
@ -53,23 +37,9 @@ async fn main() {
|
|||
russh::keys::PrivateKey::random(&mut OsRng, russh::keys::Algorithm::Ed25519).unwrap()
|
||||
};
|
||||
|
||||
let config = russh::server::Config {
|
||||
inactivity_timeout: Some(Duration::from_secs(3600)),
|
||||
auth_rejection_time: Duration::from_secs(3),
|
||||
auth_rejection_time_initial: Some(Duration::from_secs(0)),
|
||||
keys: vec![key],
|
||||
preferred: russh::Preferred {
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let config = Arc::new(config);
|
||||
|
||||
let mut sh = Server::new();
|
||||
|
||||
let tunnels = sh.tunnels.clone();
|
||||
tokio::spawn(async move { sh.run_on_address(config, ("0.0.0.0", 2222)).await });
|
||||
let mut ssh = Server::new();
|
||||
let tunnels = ssh.tunnels();
|
||||
tokio::spawn(async move { ssh.run(key, ("0.0.0.0", 2222)).await });
|
||||
|
||||
let service = service_fn(move |req: Request<_>| {
|
||||
let tunnels = tunnels.clone();
|
||||
|
@ -108,16 +78,7 @@ async fn main() {
|
|||
};
|
||||
|
||||
debug!("Opening channel");
|
||||
let channel = match tunnel
|
||||
.handle
|
||||
.channel_open_forwarded_tcpip(
|
||||
&tunnel.address,
|
||||
tunnel.port,
|
||||
&tunnel.address,
|
||||
tunnel.port,
|
||||
)
|
||||
.await
|
||||
{
|
||||
let channel = match tunnel.open_tunnel().await {
|
||||
Ok(channel) => channel,
|
||||
Err(err) => {
|
||||
warn!("Failed to tunnel: {err}");
|
||||
|
@ -167,190 +128,3 @@ async fn main() {
|
|||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct Tunnel {
|
||||
handle: Handle,
|
||||
address: String,
|
||||
port: u32,
|
||||
}
|
||||
|
||||
type Tunnels = Arc<RwLock<HashMap<String, Tunnel>>>;
|
||||
|
||||
struct Server {
|
||||
tunnels: Tunnels,
|
||||
}
|
||||
|
||||
impl Server {
|
||||
fn new() -> Self {
|
||||
Server {
|
||||
tunnels: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl server::Server for Server {
|
||||
type Handler = Handler;
|
||||
|
||||
fn new_client(&mut self, _peer_addr: Option<std::net::SocketAddr>) -> Self::Handler {
|
||||
let (tx, rx) = mpsc::unbounded_channel::<Vec<u8>>();
|
||||
|
||||
Handler {
|
||||
tx,
|
||||
rx: Some(rx),
|
||||
all_tunnels: self.tunnels.clone(),
|
||||
tunnels: HashSet::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_session_error(&mut self, error: <Self::Handler as server::Handler>::Error) {
|
||||
error!("Session error: {error:#?}");
|
||||
}
|
||||
}
|
||||
|
||||
struct Handler {
|
||||
tx: UnboundedSender<Vec<u8>>,
|
||||
rx: Option<UnboundedReceiver<Vec<u8>>>,
|
||||
|
||||
all_tunnels: Tunnels,
|
||||
tunnels: HashSet<String>,
|
||||
}
|
||||
|
||||
impl Handler {
|
||||
fn send(&self, data: &str) {
|
||||
let _ = self.tx.send(data.as_bytes().to_vec());
|
||||
}
|
||||
|
||||
async fn full_address(&self, address: &str) -> Option<String> {
|
||||
let all_tunnels = self.all_tunnels.read().await;
|
||||
|
||||
let address = if address == "localhost" {
|
||||
loop {
|
||||
let address = get_animal_name();
|
||||
if !all_tunnels.contains_key(address) {
|
||||
break address;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if all_tunnels.contains_key(address) {
|
||||
return None;
|
||||
}
|
||||
address
|
||||
};
|
||||
|
||||
Some(format!("{address}.tunnel.huizinga.dev"))
|
||||
}
|
||||
}
|
||||
|
||||
impl server::Handler for Handler {
|
||||
type Error = russh::Error;
|
||||
|
||||
async fn channel_open_session(
|
||||
&mut self,
|
||||
channel: russh::Channel<server::Msg>,
|
||||
_session: &mut server::Session,
|
||||
) -> Result<bool, Self::Error> {
|
||||
debug!("channel_open_session");
|
||||
|
||||
let Some(mut rx) = self.rx.take() else {
|
||||
return Err(russh::Error::Inconsistent);
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
debug!("Waiting for message to send to client...");
|
||||
loop {
|
||||
let message = rx.recv().await;
|
||||
debug!("Message!");
|
||||
|
||||
let Some(message) = message else { break };
|
||||
|
||||
if channel.data(message.as_ref()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
debug!("Ending receive task");
|
||||
});
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn auth_publickey(
|
||||
&mut self,
|
||||
user: &str,
|
||||
_public_key: &russh::keys::ssh_key::PublicKey,
|
||||
) -> Result<server::Auth, Self::Error> {
|
||||
debug!("Login from {user}");
|
||||
|
||||
// TODO: Get ssh keys associated with user from ldap
|
||||
Ok(server::Auth::Accept)
|
||||
}
|
||||
|
||||
async fn data(
|
||||
&mut self,
|
||||
_channel: ChannelId,
|
||||
data: &[u8],
|
||||
_session: &mut server::Session,
|
||||
) -> Result<(), Self::Error> {
|
||||
if data == [3] {
|
||||
return Err(russh::Error::Disconnect);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn exec_request(
|
||||
&mut self,
|
||||
_channel: ChannelId,
|
||||
data: &[u8],
|
||||
_session: &mut server::Session,
|
||||
) -> Result<(), Self::Error> {
|
||||
debug!("exec_request data {data:?}");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn tcpip_forward(
|
||||
&mut self,
|
||||
address: &str,
|
||||
port: &mut u32,
|
||||
session: &mut server::Session,
|
||||
) -> Result<bool, Self::Error> {
|
||||
debug!("{address}:{port}");
|
||||
|
||||
let Some(full_address) = self.full_address(address).await else {
|
||||
self.send(&format!("{port} => FAILED ({address} already in use)\r\n"));
|
||||
return Ok(false);
|
||||
};
|
||||
|
||||
self.tunnels.insert(full_address.clone());
|
||||
self.all_tunnels.write().await.insert(
|
||||
full_address.clone(),
|
||||
Tunnel {
|
||||
handle: session.handle(),
|
||||
address: address.into(),
|
||||
port: *port,
|
||||
},
|
||||
);
|
||||
|
||||
self.send(&format!("{port} => https://{full_address}\r\n"));
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Handler {
|
||||
fn drop(&mut self) {
|
||||
let tunnels = self.tunnels.clone();
|
||||
let all_tunnels = self.all_tunnels.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut all_tunnels = all_tunnels.write().await;
|
||||
for tunnel in tunnels {
|
||||
all_tunnels.remove(&tunnel);
|
||||
}
|
||||
|
||||
debug!("{all_tunnels:?}");
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
221
src/ssh.rs
Normal file
221
src/ssh.rs
Normal file
|
@ -0,0 +1,221 @@
|
|||
use std::{collections::HashSet, net::SocketAddr, sync::Arc, time::Duration};
|
||||
|
||||
use russh::{
|
||||
ChannelId,
|
||||
keys::PrivateKey,
|
||||
server::{Auth, Msg, Server as _, Session},
|
||||
};
|
||||
use tokio::{
|
||||
net::ToSocketAddrs,
|
||||
sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel},
|
||||
};
|
||||
use tracing::{debug, error};
|
||||
|
||||
use crate::{
|
||||
animals::get_animal_name,
|
||||
tunnel::{self, Tunnel, Tunnels},
|
||||
};
|
||||
|
||||
pub struct Handler {
|
||||
tx: UnboundedSender<Vec<u8>>,
|
||||
rx: Option<UnboundedReceiver<Vec<u8>>>,
|
||||
|
||||
all_tunnels: Tunnels,
|
||||
tunnels: HashSet<String>,
|
||||
}
|
||||
|
||||
impl Handler {
|
||||
fn send(&self, data: &str) {
|
||||
let _ = self.tx.send(data.as_bytes().to_vec());
|
||||
}
|
||||
|
||||
async fn full_address(&self, address: &str) -> Option<String> {
|
||||
let all_tunnels = self.all_tunnels.read().await;
|
||||
|
||||
let address = if address == "localhost" {
|
||||
loop {
|
||||
let address = get_animal_name();
|
||||
if !all_tunnels.contains_key(address) {
|
||||
break address;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if all_tunnels.contains_key(address) {
|
||||
return None;
|
||||
}
|
||||
address
|
||||
};
|
||||
|
||||
Some(format!("{address}.tunnel.huizinga.dev"))
|
||||
}
|
||||
}
|
||||
|
||||
impl russh::server::Handler for Handler {
|
||||
type Error = russh::Error;
|
||||
|
||||
async fn channel_open_session(
|
||||
&mut self,
|
||||
channel: russh::Channel<Msg>,
|
||||
_session: &mut Session,
|
||||
) -> Result<bool, Self::Error> {
|
||||
debug!("channel_open_session");
|
||||
|
||||
let Some(mut rx) = self.rx.take() else {
|
||||
return Err(russh::Error::Inconsistent);
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
debug!("Waiting for message to send to client...");
|
||||
loop {
|
||||
let message = rx.recv().await;
|
||||
debug!("Message!");
|
||||
|
||||
let Some(message) = message else { break };
|
||||
|
||||
if channel.data(message.as_ref()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
debug!("Ending receive task");
|
||||
});
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn auth_publickey(
|
||||
&mut self,
|
||||
user: &str,
|
||||
_public_key: &russh::keys::ssh_key::PublicKey,
|
||||
) -> Result<Auth, Self::Error> {
|
||||
debug!("Login from {user}");
|
||||
|
||||
// TODO: Get ssh keys associated with user from ldap
|
||||
Ok(Auth::Accept)
|
||||
}
|
||||
|
||||
async fn data(
|
||||
&mut self,
|
||||
_channel: ChannelId,
|
||||
data: &[u8],
|
||||
_session: &mut Session,
|
||||
) -> Result<(), Self::Error> {
|
||||
if data == [3] {
|
||||
return Err(russh::Error::Disconnect);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn exec_request(
|
||||
&mut self,
|
||||
_channel: ChannelId,
|
||||
data: &[u8],
|
||||
_session: &mut Session,
|
||||
) -> Result<(), Self::Error> {
|
||||
debug!("exec_request data {data:?}");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn tcpip_forward(
|
||||
&mut self,
|
||||
address: &str,
|
||||
port: &mut u32,
|
||||
session: &mut Session,
|
||||
) -> Result<bool, Self::Error> {
|
||||
debug!("{address}:{port}");
|
||||
|
||||
let Some(full_address) = self.full_address(address).await else {
|
||||
self.send(&format!("{port} => FAILED ({address} already in use)\r\n"));
|
||||
return Ok(false);
|
||||
};
|
||||
|
||||
self.tunnels.insert(full_address.clone());
|
||||
self.all_tunnels.write().await.insert(
|
||||
full_address.clone(),
|
||||
Tunnel::new(session.handle(), address, *port),
|
||||
);
|
||||
|
||||
self.send(&format!("{port} => https://{full_address}\r\n"));
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Handler {
|
||||
fn drop(&mut self) {
|
||||
let tunnels = self.tunnels.clone();
|
||||
let all_tunnels = self.all_tunnels.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut all_tunnels = all_tunnels.write().await;
|
||||
for tunnel in tunnels {
|
||||
all_tunnels.remove(&tunnel);
|
||||
}
|
||||
|
||||
debug!("{all_tunnels:?}");
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Server {
|
||||
tunnels: Tunnels,
|
||||
}
|
||||
|
||||
impl Server {
|
||||
pub fn new() -> Self {
|
||||
Server {
|
||||
tunnels: tunnel::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn tunnels(&self) -> Tunnels {
|
||||
self.tunnels.clone()
|
||||
}
|
||||
|
||||
pub fn run(
|
||||
&mut self,
|
||||
key: PrivateKey,
|
||||
addr: impl ToSocketAddrs + Send,
|
||||
) -> impl Future<Output = Result<(), std::io::Error>> + Send {
|
||||
let config = russh::server::Config {
|
||||
inactivity_timeout: Some(Duration::from_secs(3600)),
|
||||
auth_rejection_time: Duration::from_secs(3),
|
||||
auth_rejection_time_initial: Some(Duration::from_secs(0)),
|
||||
keys: vec![key],
|
||||
preferred: russh::Preferred {
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let config = Arc::new(config);
|
||||
|
||||
async move { self.run_on_address(config, addr).await }
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for Server {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl russh::server::Server for Server {
|
||||
type Handler = Handler;
|
||||
|
||||
fn new_client(&mut self, _peer_addr: Option<SocketAddr>) -> Self::Handler {
|
||||
let (tx, rx) = unbounded_channel::<Vec<u8>>();
|
||||
|
||||
Handler {
|
||||
tx,
|
||||
rx: Some(rx),
|
||||
all_tunnels: self.tunnels.clone(),
|
||||
tunnels: HashSet::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_session_error(&mut self, error: <Self::Handler as russh::server::Handler>::Error) {
|
||||
error!("Session error: {error:#?}");
|
||||
}
|
||||
}
|
36
src/tunnel.rs
Normal file
36
src/tunnel.rs
Normal file
|
@ -0,0 +1,36 @@
|
|||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use russh::{
|
||||
Channel,
|
||||
server::{Handle, Msg},
|
||||
};
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Tunnel {
|
||||
handle: Handle,
|
||||
address: String,
|
||||
port: u32,
|
||||
}
|
||||
|
||||
impl Tunnel {
|
||||
pub fn new(handle: Handle, address: impl Into<String>, port: u32) -> Self {
|
||||
Self {
|
||||
handle,
|
||||
address: address.into(),
|
||||
port,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn open_tunnel(&self) -> Result<Channel<Msg>, russh::Error> {
|
||||
self.handle
|
||||
.channel_open_forwarded_tcpip(&self.address, self.port, &self.address, self.port)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
pub type Tunnels = Arc<RwLock<HashMap<String, Tunnel>>>;
|
||||
|
||||
pub fn new() -> Tunnels {
|
||||
Arc::new(RwLock::new(HashMap::new()))
|
||||
}
|
Loading…
Reference in New Issue
Block a user