From 48eaf0f840b5a44b04af4eae4074d56fdcb34e85 Mon Sep 17 00:00:00 2001 From: Aram Peres Date: Sat, 25 Jun 2022 10:33:37 -0400 Subject: [PATCH] Allow onetun to be used as a library --- .github/workflows/build.yml | 5 ++ Cargo.toml | 2 + src/config.rs | 24 +++---- src/lib.rs | 122 ++++++++++++++++++++++++++++++++++ src/main.rs | 126 ++++-------------------------------- 5 files changed, 154 insertions(+), 125 deletions(-) create mode 100644 src/lib.rs diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index fb404a0..01fafed 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -26,6 +26,11 @@ jobs: with: command: check + - name: Run cargo check without default features + uses: actions-rs/cargo@v1 + with: + command: check --no-default-features + test: name: Test Suite runs-on: ubuntu-latest diff --git a/Cargo.toml b/Cargo.toml index 929cdcf..55d9eea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,3 +30,5 @@ pretty_env_logger = { version = "0.4", optional = true } pcap = [] default = [ "bin" ] bin = [ "clap", "pretty_env_logger", "pcap" ] + +[lib] diff --git a/src/config.rs b/src/config.rs index e055466..20f3093 100644 --- a/src/config.rs +++ b/src/config.rs @@ -12,19 +12,19 @@ const DEFAULT_PORT_FORWARD_SOURCE: &str = "127.0.0.1"; #[derive(Clone, Debug)] pub struct Config { - pub(crate) port_forwards: Vec, + pub port_forwards: Vec, #[allow(dead_code)] - pub(crate) remote_port_forwards: Vec, - pub(crate) private_key: Arc, - pub(crate) endpoint_public_key: Arc, - pub(crate) endpoint_addr: SocketAddr, - pub(crate) endpoint_bind_addr: SocketAddr, - pub(crate) source_peer_ip: IpAddr, - pub(crate) keepalive_seconds: Option, - pub(crate) max_transmission_unit: usize, - pub(crate) log: String, - pub(crate) warnings: Vec, - pub(crate) pcap_file: Option, + pub remote_port_forwards: Vec, + pub private_key: Arc, + pub endpoint_public_key: Arc, + pub endpoint_addr: SocketAddr, + pub endpoint_bind_addr: SocketAddr, + pub source_peer_ip: IpAddr, + pub keepalive_seconds: Option, + pub max_transmission_unit: usize, + pub log: String, + pub warnings: Vec, + pub pcap_file: Option, } impl Config { diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..57a3fc4 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,122 @@ +#[macro_use] +extern crate log; + +use std::sync::Arc; + +use anyhow::Context; + +use crate::config::{Config, PortProtocol}; +use crate::events::Bus; +use crate::tunnel::tcp::TcpPortPool; +use crate::tunnel::udp::UdpPortPool; +use crate::virtual_device::VirtualIpDevice; +use crate::virtual_iface::tcp::TcpVirtualInterface; +use crate::virtual_iface::udp::UdpVirtualInterface; +use crate::virtual_iface::VirtualInterfacePoll; +use crate::wg::WireGuardTunnel; + +pub mod config; +pub mod events; +#[cfg(feature = "pcap")] +pub mod pcap; +pub mod tunnel; +pub mod virtual_device; +pub mod virtual_iface; +pub mod wg; + +/// Starts the onetun tunnels in separate tokio tasks. +/// +/// Note: This future completes immediately. +pub async fn start_tunnels(config: Config, bus: Bus) -> anyhow::Result<()> { + // Initialize the port pool for each protocol + let tcp_port_pool = TcpPortPool::new(); + let udp_port_pool = UdpPortPool::new(); + + #[cfg(feature = "pcap")] + if let Some(pcap_file) = config.pcap_file.clone() { + // Start packet capture + let bus = bus.clone(); + tokio::spawn(async move { pcap::capture(pcap_file, bus).await }); + } + + let wg = WireGuardTunnel::new(&config, bus.clone()) + .await + .with_context(|| "Failed to initialize WireGuard tunnel")?; + let wg = Arc::new(wg); + + { + // Start routine task for WireGuard + let wg = wg.clone(); + tokio::spawn(async move { wg.routine_task().await }); + } + + { + // Start consumption task for WireGuard + let wg = wg.clone(); + tokio::spawn(async move { wg.consume_task().await }); + } + + { + // Start production task for WireGuard + let wg = wg.clone(); + tokio::spawn(async move { wg.produce_task().await }); + } + + if config + .port_forwards + .iter() + .any(|pf| pf.protocol == PortProtocol::Tcp) + { + // TCP device + let bus = bus.clone(); + let device = + VirtualIpDevice::new(PortProtocol::Tcp, bus.clone(), config.max_transmission_unit); + + // Start TCP Virtual Interface + let port_forwards = config.port_forwards.clone(); + let iface = TcpVirtualInterface::new(port_forwards, bus, config.source_peer_ip); + tokio::spawn(async move { iface.poll_loop(device).await }); + } + + if config + .port_forwards + .iter() + .any(|pf| pf.protocol == PortProtocol::Udp) + { + // UDP device + let bus = bus.clone(); + let device = + VirtualIpDevice::new(PortProtocol::Udp, bus.clone(), config.max_transmission_unit); + + // Start UDP Virtual Interface + let port_forwards = config.port_forwards.clone(); + let iface = UdpVirtualInterface::new(port_forwards, bus, config.source_peer_ip); + tokio::spawn(async move { iface.poll_loop(device).await }); + } + + { + let port_forwards = config.port_forwards; + let source_peer_ip = config.source_peer_ip; + + port_forwards + .into_iter() + .map(|pf| { + ( + pf, + wg.clone(), + tcp_port_pool.clone(), + udp_port_pool.clone(), + bus.clone(), + ) + }) + .for_each(move |(pf, wg, tcp_port_pool, udp_port_pool, bus)| { + tokio::spawn(async move { + tunnel::port_forward(pf, source_peer_ip, tcp_port_pool, udp_port_pool, wg, bus) + .await + .unwrap_or_else(|e| error!("Port-forward failed for {} : {}", pf, e)) + }); + }); + } + + Ok(()) +} diff --git a/src/main.rs b/src/main.rs index c9ad9ce..78c6419 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,32 +1,13 @@ +#[cfg(feature = "bin")] #[macro_use] extern crate log; -use std::sync::Arc; - -use anyhow::Context; - -use crate::config::{Config, PortProtocol}; -use crate::events::Bus; -use crate::tunnel::tcp::TcpPortPool; -use crate::tunnel::udp::UdpPortPool; -use crate::virtual_device::VirtualIpDevice; -use crate::virtual_iface::tcp::TcpVirtualInterface; -use crate::virtual_iface::udp::UdpVirtualInterface; -use crate::virtual_iface::VirtualInterfacePoll; -use crate::wg::WireGuardTunnel; - -pub mod config; -pub mod events; -#[cfg(feature = "pcap")] -pub mod pcap; -pub mod tunnel; -pub mod virtual_device; -pub mod virtual_iface; -pub mod wg; - #[cfg(feature = "bin")] #[tokio::main] async fn main() -> anyhow::Result<()> { + use anyhow::Context; + use onetun::{config::Config, events::Bus}; + let config = Config::from_args().with_context(|| "Failed to read config")?; init_logger(&config)?; @@ -34,102 +15,21 @@ async fn main() -> anyhow::Result<()> { warn!("{}", warning); } - // Initialize the port pool for each protocol - let tcp_port_pool = TcpPortPool::new(); - let udp_port_pool = UdpPortPool::new(); - let bus = Bus::default(); - - if let Some(pcap_file) = config.pcap_file.clone() { - // Start packet capture - let bus = bus.clone(); - tokio::spawn(async move { pcap::capture(pcap_file, bus).await }); - } - - let wg = WireGuardTunnel::new(&config, bus.clone()) - .await - .with_context(|| "Failed to initialize WireGuard tunnel")?; - let wg = Arc::new(wg); - - { - // Start routine task for WireGuard - let wg = wg.clone(); - tokio::spawn(async move { wg.routine_task().await }); - } - - { - // Start consumption task for WireGuard - let wg = wg.clone(); - tokio::spawn(async move { wg.consume_task().await }); - } - - { - // Start production task for WireGuard - let wg = wg.clone(); - tokio::spawn(async move { wg.produce_task().await }); - } - - if config - .port_forwards - .iter() - .any(|pf| pf.protocol == PortProtocol::Tcp) - { - // TCP device - let bus = bus.clone(); - let device = - VirtualIpDevice::new(PortProtocol::Tcp, bus.clone(), config.max_transmission_unit); - - // Start TCP Virtual Interface - let port_forwards = config.port_forwards.clone(); - let iface = TcpVirtualInterface::new(port_forwards, bus, config.source_peer_ip); - tokio::spawn(async move { iface.poll_loop(device).await }); - } - - if config - .port_forwards - .iter() - .any(|pf| pf.protocol == PortProtocol::Udp) - { - // UDP device - let bus = bus.clone(); - let device = - VirtualIpDevice::new(PortProtocol::Udp, bus.clone(), config.max_transmission_unit); - - // Start UDP Virtual Interface - let port_forwards = config.port_forwards.clone(); - let iface = UdpVirtualInterface::new(port_forwards, bus, config.source_peer_ip); - tokio::spawn(async move { iface.poll_loop(device).await }); - } - - { - let port_forwards = config.port_forwards; - let source_peer_ip = config.source_peer_ip; - - port_forwards - .into_iter() - .map(|pf| { - ( - pf, - wg.clone(), - tcp_port_pool.clone(), - udp_port_pool.clone(), - bus.clone(), - ) - }) - .for_each(move |(pf, wg, tcp_port_pool, udp_port_pool, bus)| { - tokio::spawn(async move { - tunnel::port_forward(pf, source_peer_ip, tcp_port_pool, udp_port_pool, wg, bus) - .await - .unwrap_or_else(|e| error!("Port-forward failed for {} : {}", pf, e)) - }); - }); - } + onetun::start_tunnels(config, bus).await?; futures::future::pending().await } +#[cfg(not(feature = "bin"))] +fn main() -> anyhow::Result<()> { + Err(anyhow::anyhow!("Binary compiled without 'bin' feature")) +} + #[cfg(feature = "bin")] -fn init_logger(config: &Config) -> anyhow::Result<()> { +fn init_logger(config: &onetun::config::Config) -> anyhow::Result<()> { + use anyhow::Context; + let mut builder = pretty_env_logger::formatted_timed_builder(); builder.parse_filters(&config.log); builder