Skip to content

34-add-protection-against-connection-breaks #35

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JKearnsl merged 8 commits into from
May 9, 2025
Merged

34-add-protection-against-connection-breaks #35

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
e230294
feat: update TUN setup to handle IPv4 and IPv6 addresses; enhance Han…
JKearnsl May 9, 2025
6c0a485
feat: implement runtime state management with watch channel; refactor…
JKearnsl May 9, 2025
972b4e5
feat: fix connector executor
JKearnsl May 9, 2025
7d95033
bump ver
JKearnsl May 9, 2025
5f6e58d
bug fix
JKearnsl May 9, 2025
9c5bc69
feat: refactor state management in data, transport, and tun modules
JKearnsl May 9, 2025
35f2278
bug fix
JKearnsl May 9, 2025
7b97502
rename executable to `client`
JKearnsl May 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ members = [
]

[workspace.package]
version = "0.3.0"
version = "0.4.0"
edition = "2024"

[workspace.dependencies]
Expand Down
3 changes: 2 additions & 1 deletion crates/client/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ ws = ["tokio-tungstenite", "futures"]
cli = ["clap", "ctrlc", "anstyle"]

[[bin]]
name = "client-cli"
name = "client"
path = "src/bin/mod.rs"
required-features = ["cli"]

Expand All @@ -29,6 +29,7 @@ anstyle = { version = "1.0", optional = true }
tokio-tungstenite = { workspace = true, optional = true }
futures = { workspace = true, optional = true }
ipnetwork = "0.21.1"
console-subscriber = "0.4.1"

# Crypto
snow = "0.9"
Expand Down
112 changes: 98 additions & 14 deletions crates/client/src/bin/command/connect.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,17 @@
use clap::Args;
use std::path::PathBuf;
use std::{process, thread};
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;
use ctrlc::set_handler;
use tokio::sync::watch;
use tracing::{debug, error, info};
use tun_rs::AsyncDevice;
use client::network::RouteState;
use client::runtime::error::RuntimeError;
use client::runtime::Runtime;
use client::runtime::state::RuntimeState;
use shared::connection_config::{ConnectionConfig, InterfaceConfig, RuntimeConfig};
use shared::network::find_available_ifname;

Expand Down Expand Up @@ -42,7 +48,7 @@ impl ConnectCmd {
None => unreachable!("config or key is required should protected by clap")
}
};

if config.runtime.is_none() {
config.runtime = Some(RuntimeConfig::default());
}
Expand All @@ -61,33 +67,63 @@ impl ConnectCmd {
}
}

let sock_addr = match config.general.host.parse() {
Ok(addr) => SocketAddr::new(addr, config.general.port),
Err(err) => {
eprintln!("failed to resolve host: {}", err);
process::exit(1);
}
};

let iface_config = config.interface.unwrap_or_default();
let tun = match shared::tun::setup_tun(
iface_config.name.clone(),
iface_config.mtu,
false,
).await {
Ok(tun) => Arc::new(tun),
Err(err) => {
eprintln!("failed to setup tun: {}", err);
process::exit(1);
}
};

let routes = match RouteState::new(sock_addr.ip(), iface_config.name).build()
{
Ok(routes) => Arc::new(routes),
Err(err) => {
eprintln!("failed to setup routes: {}", err);
process::exit(1);
}
};
let routes_clone = routes.clone();

let runtime = Runtime::new(
match config.general.host.parse() {
Ok(addr) => addr,
Err(err) => {
eprintln!("failed to resolve host: {}", err);
process::exit(1);
}
},
config.general.port,
sock_addr,
tun.clone(),
config.general.alg,
config.credentials,
config.runtime.unwrap_or_default(),
config.interface.unwrap_or_default(),
config.runtime.unwrap_or_default()
);

let stop_tx = runtime.stop_tx.clone();
let state_tx = runtime.state_tx.clone();

tokio::spawn(tun_service(
state_tx.clone(),
tun.clone()
));

set_handler(move || {
println!("Ctrl-C received, stopping runtime...");
match stop_tx.send(RuntimeError::StopSignal) {
match state_tx.send(RuntimeState::Error(RuntimeError::StopSignal)) {
Ok(_) => {
debug!("stop signal sent from Ctrl-C handler");
}
Err(err) => {
debug!("stop signal not sent from Ctrl-C handler: {}", err);
}
}
routes.restore();
thread::sleep(Duration::from_secs(2));
process::exit(0);
}).expect("error setting Ctrl-C handler");
Expand All @@ -98,9 +134,57 @@ impl ConnectCmd {
info!("runtime stopped");
}
_ => {
routes_clone.restore();
error!("{}", error);
}
}
}
}
}
}


pub async fn tun_service(
state_tx: watch::Sender<RuntimeState>,
tun: Arc<AsyncDevice>,
) {
let mut state_rx = state_tx.subscribe();
loop {
match state_rx.changed().await {
Ok(_) => {
debug!("tun service execute");
let state = state_rx.borrow().clone();
match state {
RuntimeState::Connected((payload, _)) => {
match payload.ipaddr {
IpAddr::V4(addr) => {
if let Err(err) = tun.set_network_address(addr, 32, None) {
state_tx.send(RuntimeState::Error(RuntimeError::IO(
format!("failed to set ipv4 network address: {}", err)
))).expect("state_tx channel broken in tun_service");
break;
}
},
IpAddr::V6(addr) => {
if let Err(err) = tun.add_address_v6(addr, 128) {
state_tx.send(RuntimeState::Error(RuntimeError::IO(
format!("failed to add ipv6 network address: {}", err)
))).expect("state_tx channel broken in tun_service");
break;
}
}
}
},
RuntimeState::Error(_) => {
debug!("tun service closed by global error state");
break;
},
_ => {}
}
}
Err(err) => {
debug!("state_tx channel error in tun service: {}", err);
break;
}
}
}
}
3 changes: 2 additions & 1 deletion crates/client/src/bin/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@ const LOG_PREFIX: &str = "client.log";
async fn main() {
let opt = opt::Opt::parse();
opt.init_logging();

// console_subscriber::init();

match opt.cmd {
Commands::Connect(cmd) => cmd.exec().await,
}
Expand Down
8 changes: 4 additions & 4 deletions crates/client/src/bin/opt.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,19 +25,19 @@ impl Opt {
} else {
tracing::Level::INFO
});

let file_appender = tracing_appender::rolling::daily(LOG_DIR, LOG_PREFIX);
let (non_blocking, _guard) = tracing_appender::non_blocking(file_appender);

let file_layer = fmt::layer()
.with_writer(non_blocking)
.with_ansi(false)
.with_filter(log_level);

let console_layer = fmt::layer()
.with_ansi(std::io::stdout().is_terminal())
.with_filter(log_level);

tracing_subscriber::registry()
.with(file_layer)
.with(console_layer)
Expand Down
8 changes: 4 additions & 4 deletions crates/client/src/network.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use std::net::{IpAddr};
use tracing::{info, warn};
use tracing::{debug, info, warn};
use std::process::Command;
use anyhow::format_err;
use std::fmt::Write;
Expand All @@ -13,7 +13,7 @@ pub struct RouteState {
use ipnetwork::{IpNetwork, NetworkSize};

impl RouteState {
pub fn new(remote: IpAddr, dev: String) -> Self {
pub fn new(remote: IpAddr, dev: String) -> Self {
Self {
dev,
default_gateway: None,
Expand All @@ -30,7 +30,7 @@ impl RouteState {
format_err!("failed to get default device: {}", e)
)?;
self.default_gateway = Some(default_gateway);
info!("default gateway: {} from dev {}", default_gateway, default_dev_name);
debug!("default gateway: {} from dev {}", default_gateway, default_dev_name);
add_route(
&IpNetwork::from_str("0.0.0.0/1")?,
None,
Expand All @@ -56,7 +56,7 @@ impl RouteState {
Ok(self)
}

pub fn restore(&mut self) {
pub fn restore(&self) {
for addr in self.exclude.iter() {
match delete_route(
addr,
Expand Down
6 changes: 0 additions & 6 deletions crates/client/src/runtime/error.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,8 @@ use thiserror::Error;
pub enum RuntimeError {
#[error("IO: {0}")]
IO(String),
#[error("Tun: {0}")]
Tun(String),
#[error("Network: {0}")]
Network(String),
#[error("Handshake: {0}")]
Handshake(String),
#[error("Disconnect: {0}")]
Disconnect(String),
#[error("Unexpected: {0}")]
Unexpected(String),
#[error("StopSignal")]
Expand Down
27 changes: 14 additions & 13 deletions ...es/client/src/runtime/worker/handshake.rs → crates/client/src/runtime/handshake.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,28 +1,29 @@
use std::sync::Arc;
use std::time::Duration;
use snow::{Builder, HandshakeState, StatelessTransportState};
use tokio::select;
use tracing::warn;
use shared::connection_config::CredentialsConfig;
use shared::handshake::{
NOISE_IK_PSK2_25519_CHACHAPOLY_BLAKE2S,
NOISE_IK_PSK2_25519_AESGCM_BLAKE2S
};
use shared::protocol::{
EncryptedHandshake,
HandshakeError,
HandshakeResponderBody,
HandshakeResponderPayload,
EncryptedHandshake,
HandshakeError,
HandshakeResponderBody,
HandshakeResponderPayload,
Packet
};
use shared::session::Alg;
use crate::runtime::transport::Transport;
use super::super::{
error::RuntimeError
use super::{
error::RuntimeError,
transport::Transport
};


fn initial(
alg: Alg,
alg: Alg,
cred: &CredentialsConfig
) -> Result<(EncryptedHandshake, HandshakeState), RuntimeError> {
let mut initiator = Builder::new(match alg {
Expand All @@ -40,7 +41,7 @@ fn initial(
}

fn complete(
handshake: &EncryptedHandshake,
handshake: &EncryptedHandshake,
mut initiator: HandshakeState
) -> Result<(HandshakeResponderBody, StatelessTransportState), RuntimeError> {
let mut buffer = [0u8; 65536];
Expand All @@ -53,7 +54,7 @@ fn complete(
}
}

pub(super) async fn handshake_step(
pub async fn handshake_step(
transport: Arc<dyn Transport>,
cred: CredentialsConfig,
alg: Alg,
Expand All @@ -64,12 +65,12 @@ pub(super) async fn handshake_step(
alg,
&cred
)?;

transport.send(&Packet::HandshakeInitial(handshake).to_bytes()).await?;

// [step 2] Server complete
let mut buffer = [0u8; 65536];
let resp = tokio::select! {
let resp = select! {
_ = tokio::time::sleep(timeout) => Err(RuntimeError::Handshake(
format!("server timeout ({:?})", timeout)
)),
Expand Down Expand Up @@ -107,4 +108,4 @@ pub(super) async fn handshake_step(
}
}
}
}
}
Loading