// Copyright 2017, CZ.NIC z.s.p.o. (http://www.nic.cz/)
//
// This file is part of the pakon system.
//
// Pakon is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
//
// Pakon is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with Pakon.  If not, see <http://www.gnu.org/licenses/>.

//! The data source taking information from the guts daemon.
//!
//! This connects to the guts daemon on a unix domain socket and converts its messages into
//! updates.

#![deny(missing_docs)]

extern crate eui48;
extern crate futures;
#[macro_use]
extern crate slog;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
#[cfg(test)]
extern crate tempdir;
extern crate tokio_io;
extern crate tokio_core;
#[macro_use]
extern crate tokio_jsonrpc;
extern crate tokio_retry;
extern crate tokio_uds;
extern crate void;

extern crate dsrc;
extern crate libdata;
extern crate libflow;
extern crate libutils;
#[cfg(test)]
extern crate test_help;

mod server;
mod snapshot;

use std::io::Error as IoError;
use std::net::IpAddr;
use std::time::Duration;
use std::path::PathBuf;

use eui48::MacAddress;
use futures::Future;
use futures::future::{self, Loop};
use slog::Logger;
use tokio_core::reactor::Handle;
use tokio_retry::{Error as RetryError, Retry};
use tokio_retry::strategy::{self, FibonacciBackoff};
use tokio_uds::UnixStream;
use void::Void;

use dsrc::{Dsrc, Factory};
use libdata::column;
use libdata::flow::{Direction, IpProto, IpProtoRaw, Names, Port};
use libflow::update::{Cork, UpdateSender};
use libutils::LocalBoxFuture;

/// The guts data source.
///
/// Note that the actual logic is implemented inside the futures that trigger each other. This is
/// just that we have something to return and can implement the trait on. The tick is ignored on
/// Guts, since it has push-only semantics.
struct GutsDsrc;

impl Dsrc for GutsDsrc {
    fn tick(&mut self, _cork: Cork) {
        // We don't push anything as a result of a tick. Therefore we don't need to mark the end of
        // sending by a cork and therefore we can drop it (signal it) right away.
    }
}

/// Connect to the socket and handle the connection.
///
/// This tries (repeatedly, if it fails at first) to connect and once it succeeds, it reads the
/// records from there and handles them. The returned future terminates if the connection is lost.
///
/// It terminates with a success if the other side simply closes the connection.
fn upstream_connect(logger: Logger,
                    handle: Handle,
                    socket: PathBuf,
                    snapshot_sender: UpdateSender)
    -> LocalBoxFuture<(), IoError>
{
    // Try connecting right away and keep retrying if it doesn't help, in longer and longer
    // intervals
    let retry = FibonacciBackoff::from_millis(25)
        .max_delay(Duration::from_secs(5))
        .map(strategy::jitter);
    let convert_handle = handle.clone();
    let convert_logger = logger.clone();
    let connection_done = Retry::spawn(handle.clone(), retry, move || {
        // Keep trying until you get a connection
        debug!(logger, "Connecting to upstream");
        let result = UnixStream::connect(&socket, &handle);
        if result.is_err() {
            warn!(logger, "Connection to upstream failed: {}", result.as_ref().unwrap_err());
        }
        result
    }).map_err(|retry_err| match retry_err {
        // The retry strategy distinguishes the errors for where they come from. We don't care, so
        // we flatten them (they are the same type).
        RetryError::OperationError(e) |
        RetryError::TimerError(e) => e,
    }).and_then(move |connection| {
        // Once we have the connection...
        server::convert(&convert_logger, &convert_handle, connection, snapshot_sender)
    });
    Box::new(connection_done)
}

/// Keep running upstream all over and over.
fn upstream_forever(logger: Logger,
                    handle: Handle,
                    socket: PathBuf,
                    snapshot_sender: UpdateSender)
    -> LocalBoxFuture<Void, IoError>
{
    let forever = future::loop_fn((), move |()| {
        // We need to provide a separate logger to each, as they may live for unknown time and we
        // need to keep ours for future iterations
        let logger_clone = logger.clone();
        // Perform one full connection and once it terminates, signal the loop_fn we want to try
        // again. Forever.
        upstream_connect(logger.clone(), handle.clone(), socket.clone(), snapshot_sender.clone())
            .then(move |res| {
                match res {
                    Ok(()) => warn!(logger_clone, "Upstream connection lost"),
                    Err(e) => warn!(logger_clone, "Upstream connection lost: {}", e),
                }
                Ok(Loop::Continue(()))
            })
    });
    Box::new(forever)
}

/// The factory for the guts daemon data source.
pub struct Guts {
    /// The path to the unix domain socket where the guts daemon lives.
    path: PathBuf,
}

impl Guts {
    /// Creates a new guts factory.
    ///
    /// # Params
    ///
    /// * `path`: The path to the unix domain socket where to connect.
    pub fn new(path: PathBuf) -> Self {
        Self {
            path
        }
    }
}

impl Factory for Guts {
    fn create(self, logger: &Logger, handle: Handle, channel: UpdateSender)
        -> (Box<Dsrc>, LocalBoxFuture<(), ()>)
    {
        info!(logger, "Creating new guts data source");
        // Register the columns we produce
        column::register::<IpProto>();
        column::register::<IpProtoRaw>();
        column::register::<Direction>();
        column::register_endpoint::<MacAddress>();
        column::register_endpoint::<IpAddr>();
        column::register_endpoint::<Names>();
        column::register_endpoint::<Port>();
        let child_logger = logger.new(o!("context" => "guts"));
        let err_logger = child_logger.clone();
        let future = upstream_forever(child_logger, handle, self.path, channel)
            .map(|_| ())
            .map_err(move |e| error!(err_logger, "{}", e));
        (Box::new(GutsDsrc), Box::new(future))
    }
}

#[cfg(test)]
mod tests {
    use futures::Stream;
    use futures::unsync::mpsc;
    use tempdir::TempDir;
    use tokio_core::reactor::{Core, Timeout};
    use tokio_io::io;
    use tokio_uds::UnixListener;

    use super::*;
    use test_help::ONE_SNAPSHOT;

    /// Check the upstream keeps reconnecting until it succeeds.
    #[test]
    fn upstream_reconnect() {
        let mut core = Core::new().unwrap();
        let handle = core.handle();
        let dir = TempDir::new("upstream_test").unwrap();
        let path = dir.path().join("socket");
        let logger = libutils::test_logger();
        let (channel_snd, channel_rcv) = mpsc::channel(5);
        // Create the listener, but only after a while
        let path_clone = path.clone();
        let one_snapshot = ONE_SNAPSHOT.as_bytes();
        let timer = Timeout::new(Duration::from_millis(500), &handle)
            .unwrap()
            // Create the listening socket once the timeout fires
            .and_then(move |_| {
                let listener = UnixListener::bind(path_clone, &handle).unwrap();
                // And get the first connection it receives
                listener.incoming()
                    .into_future()
                    .map(|(conn, _)| conn.unwrap().0)
                    .map_err(|(e, _)| e)
            })
            // Write all the data into the connection
            .and_then(move |conn| {
                io::write_all(conn, one_snapshot)
            })
            // Flush them
            .and_then(|(conn, _)| io::flush(conn))
            // And drop the connection (closes it)
            .and_then(|conn| {
                drop(conn);
                Ok(())
            });
        // Create the upstream connection. Note that the socket doesn't exist yet.
        let connection_finished = upstream_connect(logger, core.handle(), path, channel_snd);
        // Wait for the snapshot to appear on the channel
        let snapshot = channel_rcv.into_future()
            .map(|(rep, _)| rep.unwrap())
            .map_err(|_| panic!());
        // And wait for all the things to finish
        let all = connection_finished.join3(snapshot, timer);
        core.run(all).unwrap();
    }

    /// Check we reconnect multiple times if the connection is lost.
    #[test]
    fn upstream_forever_retry() {
        let mut core = Core::new().unwrap();
        let handle = core.handle();
        let dir = TempDir::new("upstream_forever").unwrap();
        let path = dir.path().join("socket");
        let logger = libutils::test_logger();
        let (channel_snd, channel_rcv) = mpsc::channel(5);
        // Create a listener future that accepts connections, pushes a bit of data into each and
        // terminates.
        let listener = UnixListener::bind(&path, &handle).unwrap();
        let accepts = listener.incoming()
            .for_each(|(conn, _)| {
                // One complete snapshot
                let one_snapshot = ONE_SNAPSHOT.as_bytes();
                // One incomplete
                io::write_all(conn, one_snapshot)
                    .and_then(|(conn, _)| {
                        let garbage = &ONE_SNAPSHOT.as_bytes()[..5];
                        io::write_all(conn, garbage)
                    })
                    // Flush
                    .and_then(|(conn, _)| io::flush(conn))
                    // And terminate
                    .and_then(|conn| {
                        drop(conn);
                        Ok(())
                    })
            })
            .map_err(|e| panic!("error: {}", e));
        // And let the acceptor run in the background (as it runs forever)
        handle.spawn(accepts);
        // Run the conversion in the background as well (because it would never end anyway)
        let upstream = upstream_forever(logger, handle.clone(), path, channel_snd)
            .map(|_| ())
            .map_err(|e| panic!("error: {}", e));
        handle.spawn(upstream);
        // And try to receive three snapshots (which need to be generated by three different
        // reconnects internally)
        let snapshots = channel_rcv.take(3).collect();
        let snapshots_received = core.run(snapshots).unwrap();
        assert_eq!(3, snapshots_received.len());
    }
}