Allow onetun to be used as a library

This commit is contained in:
Aram 🍐 2022-06-25 10:33:37 -04:00
parent f85692950f
commit 48eaf0f840
5 changed files with 154 additions and 125 deletions

View file

@ -26,6 +26,11 @@ jobs:
with: with:
command: check command: check
- name: Run cargo check without default features
uses: actions-rs/cargo@v1
with:
command: check --no-default-features
test: test:
name: Test Suite name: Test Suite
runs-on: ubuntu-latest runs-on: ubuntu-latest

View file

@ -30,3 +30,5 @@ pretty_env_logger = { version = "0.4", optional = true }
pcap = [] pcap = []
default = [ "bin" ] default = [ "bin" ]
bin = [ "clap", "pretty_env_logger", "pcap" ] bin = [ "clap", "pretty_env_logger", "pcap" ]
[lib]

View file

@ -12,19 +12,19 @@ const DEFAULT_PORT_FORWARD_SOURCE: &str = "127.0.0.1";
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct Config { pub struct Config {
pub(crate) port_forwards: Vec<PortForwardConfig>, pub port_forwards: Vec<PortForwardConfig>,
#[allow(dead_code)] #[allow(dead_code)]
pub(crate) remote_port_forwards: Vec<PortForwardConfig>, pub remote_port_forwards: Vec<PortForwardConfig>,
pub(crate) private_key: Arc<X25519SecretKey>, pub private_key: Arc<X25519SecretKey>,
pub(crate) endpoint_public_key: Arc<X25519PublicKey>, pub endpoint_public_key: Arc<X25519PublicKey>,
pub(crate) endpoint_addr: SocketAddr, pub endpoint_addr: SocketAddr,
pub(crate) endpoint_bind_addr: SocketAddr, pub endpoint_bind_addr: SocketAddr,
pub(crate) source_peer_ip: IpAddr, pub source_peer_ip: IpAddr,
pub(crate) keepalive_seconds: Option<u16>, pub keepalive_seconds: Option<u16>,
pub(crate) max_transmission_unit: usize, pub max_transmission_unit: usize,
pub(crate) log: String, pub log: String,
pub(crate) warnings: Vec<String>, pub warnings: Vec<String>,
pub(crate) pcap_file: Option<String>, pub pcap_file: Option<String>,
} }
impl Config { impl Config {

122
src/lib.rs Normal file
View file

@ -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(())
}

View file

@ -1,32 +1,13 @@
#[cfg(feature = "bin")]
#[macro_use] #[macro_use]
extern crate log; 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")] #[cfg(feature = "bin")]
#[tokio::main] #[tokio::main]
async fn main() -> anyhow::Result<()> { 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")?; let config = Config::from_args().with_context(|| "Failed to read config")?;
init_logger(&config)?; init_logger(&config)?;
@ -34,102 +15,21 @@ async fn main() -> anyhow::Result<()> {
warn!("{}", warning); 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(); let bus = Bus::default();
onetun::start_tunnels(config, bus).await?;
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))
});
});
}
futures::future::pending().await futures::future::pending().await
} }
#[cfg(not(feature = "bin"))]
fn main() -> anyhow::Result<()> {
Err(anyhow::anyhow!("Binary compiled without 'bin' feature"))
}
#[cfg(feature = "bin")] #[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(); let mut builder = pretty_env_logger::formatted_timed_builder();
builder.parse_filters(&config.log); builder.parse_filters(&config.log);
builder builder