durnitisp/src/main.rs

// Copyright 2020 Jeremy Wall
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;
use std::sync::RwLock;

use gflags;
use log::{debug, error, info};
use nursery::thread;
use nursery::{Nursery, Waitable};
use prometheus;
use prometheus::{CounterVec, Encoder, IntGaugeVec, Opts, Registry, TextEncoder};
use stderrlog;
use tiny_http;

mod icmp;
mod stun;
mod util;

gflags::define! {
    /// Print this help text.
    -h, --help = false
}

gflags::define! {
    /// Port to listen on for exporting variables prometheus style.
    --listenHost = "0.0.0.0:8080"
}

gflags::define! {
    /// Enable debug logging
    --debug = false
}

fn main() -> anyhow::Result<()> {
    let default_stun_servers: Vec<&'static str> = vec![
        "stun.l.google.com:19302",
        "stun.ekiga.net:3478",
        "stun.xten.com:3478",
        "stun.ideasip.com:3478",
        "stun.rixtelecom.se:3478",
        "stun.schlund.de:3478",
        "stun.softjoys.com:3478",
        "stun.stunprotocol.org:3478",
        "stun.voiparound.com:3478",
        "stun.voipbuster.com:3478",
        "stun.voipstunt.com:3478",
        "stun1.noc.ams-ix.net:3478",
    ];
    let mut stun_servers = gflags::parse();

    let default_ping_hosts: Vec<&'static str> = vec!["google.com"];

    if HELP.flag {
        println!("durnitisp <options> <list of hostname:port>");
        println!("");
        println!("The hostname and port are expected to be for a valid stun server.");
        println!("You can put as many of them as you want after the options.");
        println!("");
        println!("FLAGS:");
        gflags::print_help_and_exit(0);
    }

    let level = if DEBUG.flag || cfg!(debug_assertions) {
        3
    } else {
        2
    };

    stderrlog::new()
        .verbosity(level)
        .timestamp(stderrlog::Timestamp::Millisecond)
        .init()?;

    if stun_servers.is_empty() {
        stun_servers = default_stun_servers;
    }
    // FIXME(jwall): allow them to override ping hosts
    let ping_hosts = default_ping_hosts;
    let counter_opts = Opts::new(
        "stun_attempt_counter",
        "Counter for the good, bad, and total attempts to connect to stun server.",
    );
    let gauge_opts = Opts::new(
        "stun_attempt_latency_ms",
        "Latency guage in millis per stun domain.",
    );

    let stop_signal = Arc::new(RwLock::new(false));

    // Create a Registry and register metrics.
    let r = Registry::new();
    let stun_counter_vec = CounterVec::new(counter_opts, &["result", "domain"]).unwrap();
    let stun_success_vec = IntGaugeVec::new(
        Opts::new("stun_success", "Stun probe successes"),
        &["domain"],
    )
    .unwrap();
    r.register(Box::new(stun_counter_vec.clone()))
        .expect("Failed to register stun connection counter");
    let stun_latency_vec = IntGaugeVec::new(gauge_opts, &["domain"]).unwrap();
    r.register(Box::new(stun_latency_vec.clone()))
        .expect("Failed to register stun latency guage");
    r.register(Box::new(stun_success_vec.clone()))
        .expect("Failed to register stun success gauge");
    let stun_socket_addrs = util::resolve_addrs(&stun_servers).unwrap();
    let stun_servers = Arc::new(stun_servers);
    let ping_addrs = util::resolve_ip_addrs(&ping_hosts).unwrap();
    let ping_hosts = Arc::new(ping_hosts);

    let mut parent = Nursery::new();
    // First we start the render thread.
    {
        // Introduce a new scope for our Arc to clone before moving it into the thread.
        let stop_signal = stop_signal.clone();
        // thread::Handle starts the thread immediately so the render thread will usually start first.
        let render_thread = thread::Handle::new(move || {
            debug!("attempting to start server on {}", LISTENHOST.flag);
            let server = match tiny_http::Server::http(LISTENHOST.flag) {
                Ok(server) => server,
                Err(err) => {
                    let mut signal = stop_signal.write().unwrap();
                    *signal = true;
                    error!("Error starting render thread {}", err);
                    error!("Shutting down all threads...");
                    std::process::exit(1);
                }
            };
            info!("Listening for metrics request on {}", LISTENHOST.flag);
            loop {
                info!("Waiting for request");
                match server.recv() {
                    Ok(req) => {
                        let mut buffer = vec![];
                        // Gather the metrics.
                        let encoder = TextEncoder::new();
                        let metric_families = r.gather();
                        encoder.encode(&metric_families, &mut buffer).unwrap();

                        let response = tiny_http::Response::from_data(buffer).with_status_code(200);
                        if let Err(e) = req.respond(response) {
                            error!("Error responding to request {}", e);
                        }
                    }
                    Err(e) => {
                        info!("Invalid http request! {}", e);
                    }
                }
            }
        });
        parent.adopt(Box::new(render_thread));
    }
    for (i, addr) in ping_addrs.iter().cloned().enumerate() {
        // TODO(Prometheus stats)
        let ping_hosts_copy = ping_hosts.clone();
        if let Some(addr) = dbg!(addr) {
            let domain_name = *ping_hosts_copy.get(i).unwrap();
            debug!("Pinging {}", domain_name);
            let stop_signal = stop_signal.clone();
            let ping_thread = thread::Pending::new(move || {
                icmp::start_echo_loop(domain_name, stop_signal, addr, i as u16);
            });
            parent.schedule(Box::new(ping_thread));
        }
    }
    // Then we attempt to start connections to each stun server.
    for (i, s) in stun_socket_addrs.iter().enumerate() {
        let stun_servers_copy = stun_servers.clone();
        let stun_counter_vec_copy = stun_counter_vec.clone();
        let stun_latency_vec_copy = stun_latency_vec.clone();
        let stun_success_vec_copy = stun_success_vec.clone();
        if let Some(s) = s.clone() {
            let domain_name = *stun_servers_copy.get(i).unwrap();
            let stop_signal = stop_signal.clone();
            let connect_thread = thread::Pending::new(move || {
                stun::start_listen_thread(
                    domain_name,
                    stop_signal,
                    s,
                    stun_counter_vec_copy,
                    stun_latency_vec_copy,
                    stun_success_vec_copy,
                )
            });
            parent.schedule(Box::new(connect_thread));
            // Spread the probe threads out so they're somewhat uniformly distributed.
            std::thread::sleep(std::time::Duration::from_micros(
                stun::delay_secs() * 1000000 / (stun_socket_addrs.len() as u64),
            ))
        };
    }
    // Blocks forever
    parent.wait();
    Ok(())
}