rustocat/src/main.rs

mod config;
mod listener;
mod shutdown;
mod tcp;
mod udp;

#[cfg(feature = "consul")]
mod consul;

use config::ConfigProvider;
use log::{debug, error, info, warn};
use std::collections::{HashMap, HashSet};
use std::error::Error;
use std::sync::Arc;
use tokio::sync::{broadcast, RwLock};

pub type Result<T> = std::result::Result<T, Box<dyn Error>>;

// #[derive(Debug)]
struct ActiveListener {
    notify_shutdown: broadcast::Sender<()>,
    targets: Arc<RwLock<Vec<String>>>,
    notify_targets: broadcast::Sender<()>,
}

#[tokio::main]
async fn main() -> Result<()> {
    fern::Dispatch::new()
        .format(|out, message, record| {
            out.finish(format_args!(
                "{}[{}][{}] {}",
                chrono::Local::now().format("[%Y-%m-%d][%H:%M:%S]"),
                record.target(),
                record.level(),
                message
            ))
        })
        // Add blanket level filter -
        .level(log::LevelFilter::Info)
        .level_for("rustocat", log::LevelFilter::Debug)
        .chain(std::io::stdout())
        .apply()?;

    let listeners: HashMap<String, ActiveListener> = HashMap::new();
    let text_config_provider = config::FileConfigProvider::new();

    let target_provider: Box<dyn config::ConfigProvider> = {
        #[cfg(feature = "consul")]
        {
            let cfg = text_config_provider.load_config()?;
            info!("Loaded yaml config");
            if cfg.consul.is_some() && cfg.consul.unwrap() {
                info!("Using consul config provider");
                let consul_config_provider = consul::ConsulConfigProvider::new(Some(&cfg));
                info!("Loaded consul config provider");
                Box::new(consul_config_provider)
            } else {
                Box::new(text_config_provider)
            }
        }

        #[cfg(not(feature = "consul"))]
        {
            let cfg = Box::new(text_config_provider);
            info!("Loaded yaml config");
            cfg
        }
    };

    match run_loop(target_provider, listeners).await {
        Ok(_) => {}
        Err(e) => {
            error!("Error in run loop: {}", e);
            info!("Exiting");
        }
    }

    Ok(())
}

async fn run_loop(
    mut target_provider: Box<dyn ConfigProvider>,
    mut listeners: HashMap<String, ActiveListener>,
) -> Result<()> {
    loop {
        let mappings = target_provider.get_targets().await?;
        let mut required_listeners: HashSet<String> = HashSet::new();

        for target in mappings {
            let mut source_str = "".to_owned();
            if target.udp == None || target.udp == Some(false) {
                source_str.push_str("udp:");
            } else {
                source_str.push_str("tcp:");
            }
            source_str.push_str(&target.source);

            required_listeners.insert(source_str.clone());
            if let Some(listener) = listeners.get(&source_str) {
                let mut invalid = false;

                let targets = listener.targets.read().await;
                for t in &target.targets {
                    if !targets.iter().any(|e| *e == *t) {
                        invalid = true;
                        break;
                    }
                }

                if !invalid {
                    for t in targets.iter() {
                        if !target.targets.iter().any(|e| *e == *t) {
                            invalid = true;
                            break;
                        }
                    }
                }

                if invalid {
                    warn!("Found invalid targets! Adjusting!");
                    let mut w = listener.targets.write().await;
                    *w = target.targets;

                    _ = listener.notify_targets.send(())?; //TODO: Maybe ignore this error
                }
            } else {
                let (notify_shutdown, _) = broadcast::channel(1);
                let (notify_targets, _) = broadcast::channel(1);

                let listener = ActiveListener {
                    notify_shutdown,
                    targets: Arc::new(RwLock::new(target.targets)),
                    notify_targets,
                };

                let l = listener::Listener {
                    shutdown: shutdown::ShutdownReceiver::new(listener.notify_shutdown.subscribe()),
                    source: target.source.clone(),
                    targets: listener.targets.clone(),
                    targets_changed: listener.notify_targets.subscribe(),
                };

                tokio::spawn(async move {
                    if target.udp == None || target.udp == Some(false) {
                        if let Err(err) = tcp::start_tcp_listener(l).await {
                            error!("tcp listener error: {}", err);
                        }
                    } else {
                        if let Err(err) = udp::start_udp_listener(l).await {
                            error!("udp listener error {}: {}", target.source, err);
                        }
                    }
                });

                listeners.insert(source_str, listener);
            }
        }

        let to_delete: Vec<_> = listeners
            .keys()
            .filter(|x| required_listeners.get(*x).is_none())
            .cloned() // Clone the result to make listeners mutable again!
            .collect();

        for del_key in to_delete {
            if let Some(listener) = listeners.get(&del_key) {
                let res = listener.notify_shutdown.send(()); //Errors are irrelevant here. I guess....
                match res {
                    Ok(_) => {
                        info!("Sent shutdown signal!");
                    }
                    Err(_) => {
                        warn!("Failed to send shutdown signal!");
                    }
                }

                // TODO: Maybe Wait for the listener to shut down

                debug!("Removing listener!");
                listeners.remove(&del_key);
            }
        }

        target_provider.wait_for_change().await?;
        info!("Reloading config!");
    }
}