use anyhow::Result;
use std::sync::Arc;
use parking_lot::RwLock;
use std::collections::HashMap;
use crate::config::{PortMapping, Protocol};
use super::tcp::TcpForwarder;
use super::udp::UdpForwarder;
use crate::metrics::collector::MetricsCollector;
#[derive(Debug, Default, Clone)]
pub struct ForwarderStats {
pub tcp: super::tcp::TcpStats,
pub udp: super::udp::UdpStats,
}
pub struct ForwarderManager {
tcp: Arc<TcpForwarder>,
udp: Arc<UdpForwarder>,
active_forwards: Arc<RwLock<HashMap<String, Protocol>>>,
metrics: Arc<MetricsCollector>,
}
impl ForwarderManager {
pub fn new() -> Self {
Self::with_metrics(Arc::new(MetricsCollector::new()))
}
pub fn with_metrics(metrics: Arc<MetricsCollector>) -> Self {
Self {
tcp: Arc::new(TcpForwarder::new()),
udp: Arc::new(UdpForwarder::new()),
active_forwards: Arc::new(RwLock::new(HashMap::new())),
metrics,
}
}
pub async fn start_forward(&self, mapping: &PortMapping) -> Result<()> {
if !mapping.enabled {
tracing::info!("Skipping disabled port: {}", mapping.id);
return Ok(());
}
match mapping.protocol {
Protocol::Tcp => {
self.tcp.start_forwarding(mapping).await?;
}
Protocol::Udp => {
self.udp.start_forwarding(mapping).await?;
}
}
let mut forwards = self.active_forwards.write();
forwards.insert(mapping.id.clone(), mapping.protocol.clone());
Ok(())
}
pub async fn stop_forward(&self, mapping_id: &str) -> Result<()> {
let mut forwards = self.active_forwards.write();
forwards.remove(mapping_id);
tracing::info!("Stopped forwarding: {}", mapping_id);
Ok(())
}
pub async fn stats(&self) -> ForwarderStats {
ForwarderStats {
tcp: self.tcp.stats(),
udp: self.udp.stats().await,
}
}
}
impl Default for ForwarderManager {
fn default() -> Self {
Self::new()
}
}