<commit_before>// Copyright 2016 The xi-editor Authors.
//
// 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.

#[macro_use]
extern crate log;
extern crate chrono;
extern crate fern;

extern crate dirs;

extern crate xi_core_lib;
extern crate xi_rpc;

use std::collections::HashMap;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

use xi_core_lib::XiCore;
use xi_rpc::RpcLoop;

const XI_LOG_DIR: &str = "xi-core";
const XI_LOG_FILE: &str = "xi-core.log";

fn get_logging_directory_path<P: AsRef<Path>>(directory: P) -> Result<PathBuf, io::Error> {
    match dirs::data_local_dir() {
        Some(mut log_dir) => {
            log_dir.push(directory);
            Ok(log_dir)
        }
        None => Err(io::Error::new(
            io::ErrorKind::NotFound,
            "No standard logging directory known for this platform",
        )),
    }
}

/// This function tries to create the parent directories for a file
///
/// It wraps around the `parent()` function of `Path` which returns an `Option<&Path>` and
/// `fs::create_dir_all` which returns an `io::Result<()`.
///
/// This allows you to use `?`/`try!()` to create the dir and you recive the additional custom error for when `parent()`
/// returns nothing.
///
/// # Errors
/// This can return an `io::Error` if `fs::create_dir_all` fails or if `parent()` returns `None`.
/// See `Path`'s `parent()` function for more details.
/// # Examples
/// ```
/// use std::path::PathBuf;
/// let path_with_file = PathBuf::from("/some/directory/then/file");
/// assert_eq!(Some(OsStr::new("file")), path_with_file.file_name());
/// assert_eq!(create_log_directory(path_with_file).is_ok(), true);
///
/// let path_with_other_file = PathBuf::from("/other_file");
/// assert_eq!(Some(OsStr::new("other_file")), path_with_other_file.file_name());
/// assert_eq!(create_log_directory(path_with_file).is_ok(), true);
///
/// // Path that is just the root or prefix:
/// let path_without_file = PathBuf::from("/");
/// assert_eq!(None, path_with_file.file_name());
/// assert_eq!(create_log_directory(path_with_file).is_ok(), false);
/// ```
fn create_log_directory(path_with_file: &PathBuf) -> Result<(), io::Error> {
    let log_dir = path_with_file.parent().ok_or(io::Error::new(
        io::ErrorKind::InvalidInput,
        format!(
            "Unable to get the parent of the following Path: {}",
            path_with_file.display(),
        ),
    ))?;
    // Try to create the directory.
    fs::create_dir_all(log_dir)?;
    Ok(())
}

fn setup_logging(logging_path_result: Option<PathBuf>) -> Result<(), fern::InitError> {
    let level_filter = match std::env::var("XI_LOG") {
        Ok(level) => match level.to_lowercase().as_ref() {
            "trace" => log::LevelFilter::Trace,
            "debug" => log::LevelFilter::Debug,
            _ => log::LevelFilter::Info,
        },
        // Default to info
        Err(_) => log::LevelFilter::Info,
    };

    let mut fern_dispatch = 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,
            ))
        }).level(level_filter)
        .chain(io::stderr());

    if let Some(logging_file_path) = &logging_path_result {
        // Try to create the parent directories for the logging file in the logging file path
        create_log_directory(&logging_file_path)?;

        // Attach it to fern
        fern_dispatch = fern_dispatch.chain(fern::log_file(logging_file_path)?);
    };

    // Start fern
    fern_dispatch.apply()?;
    info!("Logging with fern is setup");

    // Log details of the logging_file_path result using fern/log
    // Either logging the path fern is outputting to or the error from obtaining the path
    match &logging_path_result {
        Some(logging_file_path) => info!(
            "Logging to the following file: {}",
            logging_file_path.display()
        ),
        None => warn!("There was no path supplied for the log file, falling back to stderr."),
    }
    Ok(())
}

fn generate_logging_path(logfile_config: LogfileConfig) -> Result<PathBuf, io::Error> {
    // Use the file name set in logfile_config or fallback to the default
    let logfile_file_name = match logfile_config.file {
        Some(file_name) => file_name,
        None => PathBuf::from(XI_LOG_FILE),
    };
    if logfile_file_name.eq(Path::new("")) {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "A blank file name was supplied",
        ));
    };
    // Use the directory name set in logfile_config or fallback to the default
    let logfile_directory_name = match logfile_config.directory {
        Some(dir) => dir,
        None => PathBuf::from(XI_LOG_DIR),
    };

    let mut logging_directory_path = get_logging_directory_path(logfile_directory_name)?;

    // Add the file name & return the full path
    logging_directory_path.push(logfile_file_name);
    Ok(logging_directory_path)
}

fn get_flags() -> HashMap<String, Option<String>> {
    let mut flags: HashMap<String, Option<String>> = HashMap::new();

    let flag_prefix = "-";
    let mut args_iterator = std::env::args().peekable();
    while let Some(arg) = args_iterator.next() {
        if arg.starts_with(flag_prefix) {
            let key = arg.trim_left_matches(flag_prefix).to_string();

            // Check the next argument doesn't start with the flag prefix
            // map_or accounts for peek returning an Option
            let next_arg_not_a_flag: bool = args_iterator
                .peek()
                .map_or(false, |val| !val.starts_with(flag_prefix));
            if next_arg_not_a_flag {
                flags.insert(key, args_iterator.next());
            }
        }
    }
    flags
}

struct EnvFlagConfig {
    env_name: &'static str,
    flag_name: &'static str,
}

/// Extracts a value from the flags and the env.
///
/// In this order: `String` from the flags, then `String` from the env, then `None`
fn extract_env_or_flag(
    flags: &HashMap<String, Option<String>>,
    conf: EnvFlagConfig,
) -> Option<String> {
    flags
        .get(conf.flag_name)
        .cloned()
        .unwrap_or(std::env::var(conf.env_name).ok())
}

struct LogfileConfig {
    directory: Option<PathBuf>,
    file: Option<PathBuf>,
}

fn generate_logfile_config(flags: &HashMap<String, Option<String>>) -> LogfileConfig {
    // If the key is set, get the Option within
    let log_dir_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_DIR",
        flag_name: "log-dir",
    };
    let log_file_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_FILE",
        flag_name: "log-file",
    };
    let log_dir_flag_option = extract_env_or_flag(&flags, log_dir_env_flag).map(PathBuf::from);

    let log_file_flag_option = extract_env_or_flag(&flags, log_file_env_flag).map(PathBuf::from);

    LogfileConfig {
        directory: log_dir_flag_option,
        file: log_file_flag_option,
    }
}

fn main() {
    let mut state = XiCore::new();
    let stdin = io::stdin();
    let stdout = io::stdout();
    let mut rpc_looper = RpcLoop::new(stdout);

    let flags: HashMap<String, Option<String>> = get_flags();

    let logfile_config = generate_logfile_config(&flags);

    let logging_path_result: Result<PathBuf, io::Error> = generate_logging_path(logfile_config);

    let mut logging_path: Option<PathBuf> = None;
    let mut logging_error: Option<io::Error> = None;
    match logging_path_result {
        Ok(val) => logging_path = Some(val),
        Err(e) => logging_error = Some(e),
    };
    if let Err(e) = setup_logging(logging_path) {
        eprintln!(
            "[ERROR] setup_logging returned error, logging not enabled: {:?}",
            e
        );
    }
    if let Some(e) = logging_error {
        warn!("Unable to successfully generate the logging path: {}", e)
    }

    match rpc_looper.mainloop(|| stdin.lock(), &mut state) {
        Ok(_) => (),
        Err(err) => error!("xi-core exited with error:\n{:?}", err),
    }
}
<commit_msg>Argument isn't a result anymore. Changed logging message to be less stiff.<commit_after>// Copyright 2016 The xi-editor Authors.
//
// 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.

#[macro_use]
extern crate log;
extern crate chrono;
extern crate fern;

extern crate dirs;

extern crate xi_core_lib;
extern crate xi_rpc;

use std::collections::HashMap;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

use xi_core_lib::XiCore;
use xi_rpc::RpcLoop;

const XI_LOG_DIR: &str = "xi-core";
const XI_LOG_FILE: &str = "xi-core.log";

fn get_logging_directory_path<P: AsRef<Path>>(directory: P) -> Result<PathBuf, io::Error> {
    match dirs::data_local_dir() {
        Some(mut log_dir) => {
            log_dir.push(directory);
            Ok(log_dir)
        }
        None => Err(io::Error::new(
            io::ErrorKind::NotFound,
            "No standard logging directory known for this platform",
        )),
    }
}

/// This function tries to create the parent directories for a file
///
/// It wraps around the `parent()` function of `Path` which returns an `Option<&Path>` and
/// `fs::create_dir_all` which returns an `io::Result<()`.
///
/// This allows you to use `?`/`try!()` to create the dir and you recive the additional custom error for when `parent()`
/// returns nothing.
///
/// # Errors
/// This can return an `io::Error` if `fs::create_dir_all` fails or if `parent()` returns `None`.
/// See `Path`'s `parent()` function for more details.
/// # Examples
/// ```
/// use std::path::PathBuf;
/// let path_with_file = PathBuf::from("/some/directory/then/file");
/// assert_eq!(Some(OsStr::new("file")), path_with_file.file_name());
/// assert_eq!(create_log_directory(path_with_file).is_ok(), true);
///
/// let path_with_other_file = PathBuf::from("/other_file");
/// assert_eq!(Some(OsStr::new("other_file")), path_with_other_file.file_name());
/// assert_eq!(create_log_directory(path_with_file).is_ok(), true);
///
/// // Path that is just the root or prefix:
/// let path_without_file = PathBuf::from("/");
/// assert_eq!(None, path_with_file.file_name());
/// assert_eq!(create_log_directory(path_with_file).is_ok(), false);
/// ```
fn create_log_directory(path_with_file: &PathBuf) -> Result<(), io::Error> {
    let log_dir = path_with_file.parent().ok_or(io::Error::new(
        io::ErrorKind::InvalidInput,
        format!(
            "Unable to get the parent of the following Path: {}",
            path_with_file.display(),
        ),
    ))?;
    // Try to create the directory.
    fs::create_dir_all(log_dir)?;
    Ok(())
}

fn setup_logging(logging_path: Option<PathBuf>) -> Result<(), fern::InitError> {
    let level_filter = match std::env::var("XI_LOG") {
        Ok(level) => match level.to_lowercase().as_ref() {
            "trace" => log::LevelFilter::Trace,
            "debug" => log::LevelFilter::Debug,
            _ => log::LevelFilter::Info,
        },
        // Default to info
        Err(_) => log::LevelFilter::Info,
    };

    let mut fern_dispatch = 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,
            ))
        }).level(level_filter)
        .chain(io::stderr());

    if let Some(logging_file_path) = &logging_path {
        // Try to create the parent directories for the logging file in the logging file path
        create_log_directory(&logging_file_path)?;

        // Attach it to fern
        fern_dispatch = fern_dispatch.chain(fern::log_file(logging_file_path)?);
    };

    // Start fern
    fern_dispatch.apply()?;
    info!("Logging with fern is setup");

    // Log details of the logging_file_path result using fern/log
    // Either logging the path fern is outputting to or the error from obtaining the path
    match &logging_path {
        Some(logging_file_path) => info!("Writing logs to: {}", logging_file_path.display()),
        None => warn!("There was no path supplied for the log file, falling back to stderr."),
    }
    Ok(())
}

fn generate_logging_path(logfile_config: LogfileConfig) -> Result<PathBuf, io::Error> {
    // Use the file name set in logfile_config or fallback to the default
    let logfile_file_name = match logfile_config.file {
        Some(file_name) => file_name,
        None => PathBuf::from(XI_LOG_FILE),
    };
    if logfile_file_name.eq(Path::new("")) {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "A blank file name was supplied",
        ));
    };
    // Use the directory name set in logfile_config or fallback to the default
    let logfile_directory_name = match logfile_config.directory {
        Some(dir) => dir,
        None => PathBuf::from(XI_LOG_DIR),
    };

    let mut logging_directory_path = get_logging_directory_path(logfile_directory_name)?;

    // Add the file name & return the full path
    logging_directory_path.push(logfile_file_name);
    Ok(logging_directory_path)
}

fn get_flags() -> HashMap<String, Option<String>> {
    let mut flags: HashMap<String, Option<String>> = HashMap::new();

    let flag_prefix = "-";
    let mut args_iterator = std::env::args().peekable();
    while let Some(arg) = args_iterator.next() {
        if arg.starts_with(flag_prefix) {
            let key = arg.trim_left_matches(flag_prefix).to_string();

            // Check the next argument doesn't start with the flag prefix
            // map_or accounts for peek returning an Option
            let next_arg_not_a_flag: bool = args_iterator
                .peek()
                .map_or(false, |val| !val.starts_with(flag_prefix));
            if next_arg_not_a_flag {
                flags.insert(key, args_iterator.next());
            }
        }
    }
    flags
}

struct EnvFlagConfig {
    env_name: &'static str,
    flag_name: &'static str,
}

/// Extracts a value from the flags and the env.
///
/// In this order: `String` from the flags, then `String` from the env, then `None`
fn extract_env_or_flag(
    flags: &HashMap<String, Option<String>>,
    conf: EnvFlagConfig,
) -> Option<String> {
    flags
        .get(conf.flag_name)
        .cloned()
        .unwrap_or(std::env::var(conf.env_name).ok())
}

struct LogfileConfig {
    directory: Option<PathBuf>,
    file: Option<PathBuf>,
}

fn generate_logfile_config(flags: &HashMap<String, Option<String>>) -> LogfileConfig {
    // If the key is set, get the Option within
    let log_dir_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_DIR",
        flag_name: "log-dir",
    };
    let log_file_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_FILE",
        flag_name: "log-file",
    };
    let log_dir_flag_option = extract_env_or_flag(&flags, log_dir_env_flag).map(PathBuf::from);

    let log_file_flag_option = extract_env_or_flag(&flags, log_file_env_flag).map(PathBuf::from);

    LogfileConfig {
        directory: log_dir_flag_option,
        file: log_file_flag_option,
    }
}

fn main() {
    let mut state = XiCore::new();
    let stdin = io::stdin();
    let stdout = io::stdout();
    let mut rpc_looper = RpcLoop::new(stdout);

    let flags: HashMap<String, Option<String>> = get_flags();

    let logfile_config = generate_logfile_config(&flags);

    let logging_path_result: Result<PathBuf, io::Error> = generate_logging_path(logfile_config);

    let mut logging_path: Option<PathBuf> = None;
    let mut logging_error: Option<io::Error> = None;
    match logging_path_result {
        Ok(val) => logging_path = Some(val),
        Err(e) => logging_error = Some(e),
    };
    if let Err(e) = setup_logging(logging_path) {
        eprintln!(
            "[ERROR] setup_logging returned error, logging not enabled: {:?}",
            e
        );
    }
    if let Some(e) = logging_error {
        warn!("Unable to successfully generate the logging path: {}", e)
    }

    match rpc_looper.mainloop(|| stdin.lock(), &mut state) {
        Ok(_) => (),
        Err(err) => error!("xi-core exited with error:\n{:?}", err),
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 The xi-editor Authors.
//
// 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.

#[macro_use]
extern crate log;
extern crate chrono;
extern crate fern;

extern crate dirs;

extern crate xi_core_lib;
extern crate xi_rpc;

use std::collections::HashMap;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

use xi_core_lib::XiCore;
use xi_rpc::RpcLoop;

const XI_LOG_DIR: &str = "xi-core";
const XI_LOG_FILE: &str = "xi-core.log";

fn get_logging_directory_path<P: AsRef<Path>>(directory: P) -> Result<PathBuf, io::Error> {
    match dirs::data_local_dir() {
        Some(mut log_dir) => {
            log_dir.push(directory);
            Ok(log_dir)
        }
        None => Err(io::Error::new(
            io::ErrorKind::NotFound,
            "No standard logging directory known for this platform",
        )),
    }
}

/// This function tries to create the parent directories for a file
///
/// It wraps around the `parent()` function of `Path` which returns an `Option<&Path>` and
/// `fs::create_dir_all` which returns an `io::Result<()`.
///
/// This allows you to use `?`/`try!()` to create the dir and you recive the additional custom error for when `parent()`
/// returns nothing.
///
/// # Errors
/// This can return an `io::Error` if `fs::create_dir_all` fails or if `parent()` returns `None`.
/// See `Path`'s `parent()` function for more details.
/// # Examples
/// ```
/// use std::path::PathBuf;
/// let path_with_file = PathBuf::from("/some/directory/then/file");
/// assert_eq!(Some(OsStr::new("file")), path_with_file.file_name());
/// assert_eq!(create_log_directory(path_with_file).is_ok(), true);
///
/// let path_with_other_file = PathBuf::from("/other_file");
/// assert_eq!(Some(OsStr::new("other_file")), path_with_other_file.file_name());
/// assert_eq!(create_log_directory(path_with_file).is_ok(), true);
///
/// // Path that is just the root or prefix:
/// let path_without_file = PathBuf::from("/");
/// assert_eq!(None, path_with_file.file_name());
/// assert_eq!(create_log_directory(path_with_file).is_ok(), false);
/// ```
fn create_log_directory(path_with_file: &PathBuf) -> Result<(), io::Error> {
    let log_dir = path_with_file.parent().ok_or(io::Error::new(
        io::ErrorKind::InvalidInput,
        format!(
            "Unable to get the parent of the following Path: {}, Your path should contain a file name",
            path_with_file.display(),
        ),
    ))?;
    fs::create_dir_all(log_dir)?;
    Ok(())
}

fn setup_logging(logging_path: Option<PathBuf>) -> Result<(), fern::InitError> {
    let level_filter = match std::env::var("XI_LOG") {
        Ok(level) => match level.to_lowercase().as_ref() {
            "trace" => log::LevelFilter::Trace,
            "debug" => log::LevelFilter::Debug,
            _ => log::LevelFilter::Info,
        },
        // Default to info
        Err(_) => log::LevelFilter::Info,
    };

    let mut fern_dispatch = 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,
            ))
        }).level(level_filter)
        .chain(io::stderr());

    if let Some(logging_file_path) = &logging_path {
        create_log_directory(&logging_file_path)?;

        fern_dispatch = fern_dispatch.chain(fern::log_file(logging_file_path)?);
    };

    // Start fern
    fern_dispatch.apply()?;
    info!("Logging with fern is set up");

    // Log details of the logging_file_path result using fern/log
    // Either logging the path fern is outputting to or the error from obtaining the path
    match &logging_path {
        Some(logging_file_path) => info!("Writing logs to: {}", logging_file_path.display()),
        None => warn!("No path was supplied for the log file. Not saving logs to disk, falling back to just stderr"),
    }
    Ok(())
}

fn generate_logging_path(logfile_config: LogfileConfig) -> Result<PathBuf, io::Error> {
    // Use the file name set in logfile_config or fallback to the default
    let logfile_file_name = match logfile_config.file {
        Some(file_name) => file_name,
        None => PathBuf::from(XI_LOG_FILE),
    };
    if logfile_file_name.eq(Path::new("")) {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "A blank file name was supplied",
        ));
    };
    // Use the directory name set in logfile_config or fallback to the default
    let logfile_directory_name = match logfile_config.directory {
        Some(dir) => dir,
        None => PathBuf::from(XI_LOG_DIR),
    };

    let mut logging_directory_path = get_logging_directory_path(logfile_directory_name)?;

    // Add the file name & return the full path
    logging_directory_path.push(logfile_file_name);
    Ok(logging_directory_path)
}

fn get_flags() -> HashMap<String, Option<String>> {
    let mut flags: HashMap<String, Option<String>> = HashMap::new();

    let flag_prefix = "-";
    let mut args_iterator = std::env::args().peekable();
    while let Some(arg) = args_iterator.next() {
        if arg.starts_with(flag_prefix) {
            let key = arg.trim_left_matches(flag_prefix).to_string();

            // Check the next argument doesn't start with the flag prefix
            // map_or accounts for peek returning an Option
            let next_arg_not_a_flag: bool = args_iterator
                .peek()
                .map_or(false, |val| !val.starts_with(flag_prefix));
            if next_arg_not_a_flag {
                flags.insert(key, args_iterator.next());
            }
        }
    }
    flags
}

struct EnvFlagConfig {
    env_name: &'static str,
    flag_name: &'static str,
}

/// Extracts a value from the flags and the env.
///
/// In this order: `String` from the flags, then `String` from the env, then `None`
fn extract_env_or_flag(
    flags: &HashMap<String, Option<String>>,
    conf: EnvFlagConfig,
) -> Option<String> {
    flags
        .get(conf.flag_name)
        .cloned()
        .unwrap_or(std::env::var(conf.env_name).ok())
}

struct LogfileConfig {
    directory: Option<PathBuf>,
    file: Option<PathBuf>,
}

fn generate_logfile_config(flags: &HashMap<String, Option<String>>) -> LogfileConfig {
    // If the key is set, get the Option within
    let log_dir_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_DIR",
        flag_name: "log-dir",
    };
    let log_file_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_FILE",
        flag_name: "log-file",
    };
    let log_dir_flag_option = extract_env_or_flag(&flags, log_dir_env_flag).map(PathBuf::from);

    let log_file_flag_option = extract_env_or_flag(&flags, log_file_env_flag).map(PathBuf::from);

    LogfileConfig {
        directory: log_dir_flag_option,
        file: log_file_flag_option,
    }
}

fn main() {
    let mut state = XiCore::new();
    let stdin = io::stdin();
    let stdout = io::stdout();
    let mut rpc_looper = RpcLoop::new(stdout);

    let flags = get_flags();

    let logfile_config = generate_logfile_config(&flags);

    let logging_path_result = generate_logging_path(logfile_config);

    let mut logging_path: Option<PathBuf> = None;
    let mut logging_error: Option<io::Error> = None;
    match logging_path_result {
        Ok(val) => logging_path = Some(val),
        Err(e) => logging_error = Some(e),
    };
    if let Err(e) = setup_logging(logging_path) {
        eprintln!(
            "[ERROR] setup_logging returned error, logging not enabled: {:?}",
            e
        );
    }
    if let Some(e) = logging_error {
        warn!(
            "Unable to generate the logging path to pass to set up: {}",
            e
        )
    }

    match rpc_looper.mainloop(|| stdin.lock(), &mut state) {
        Ok(_) => (),
        Err(err) => error!("xi-core exited with error:\n{:?}", err),
    }
}
<commit_msg>Swapped to passing `setup_logging` an `Option<&Path>`. Fixed `create_log_directory` to take the new path. Fixed some issues in the doc string.<commit_after>// Copyright 2016 The xi-editor Authors.
//
// 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.

#[macro_use]
extern crate log;
extern crate chrono;
extern crate fern;

extern crate dirs;

extern crate xi_core_lib;
extern crate xi_rpc;

use std::collections::HashMap;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

use xi_core_lib::XiCore;
use xi_rpc::RpcLoop;

const XI_LOG_DIR: &str = "xi-core";
const XI_LOG_FILE: &str = "xi-core.log";

fn get_logging_directory_path<P: AsRef<Path>>(directory: P) -> Result<PathBuf, io::Error> {
    match dirs::data_local_dir() {
        Some(mut log_dir) => {
            log_dir.push(directory);
            Ok(log_dir)
        }
        None => Err(io::Error::new(
            io::ErrorKind::NotFound,
            "No standard logging directory known for this platform",
        )),
    }
}

/// This function tries to create the parent directories for a file
///
/// It wraps around the `parent()` function of `Path` which returns an `Option<&Path>` and
/// `fs::create_dir_all` which returns an `io::Result<()`.
///
/// This allows you to use `?`/`try!()` to create the dir and you recive the additional custom error for when `parent()`
/// returns nothing.
///
/// # Errors
/// This can return an `io::Error` if `fs::create_dir_all` fails or if `parent()` returns `None`.
/// See `Path`'s `parent()` function for more details.
/// # Examples
/// ```
/// use std::path::Path;
/// use std::ffi::OsStr;
///
/// let path_with_file = Path::new("/some/directory/then/file");
/// assert_eq!(Some(OsStr::new("file")), path_with_file.file_name());
/// assert_eq!(create_log_directory(path_with_file).is_ok(), true);
///
/// let path_with_other_file = Path::new("/other_file");
/// assert_eq!(Some(OsStr::new("other_file")), path_with_other_file.file_name());
/// assert_eq!(create_log_directory(path_with_file).is_ok(), true);
///
/// // Path that is just the root or prefix:
/// let path_without_file = Path::new("/");
/// assert_eq!(None, path_without_file.file_name());
/// assert_eq!(create_log_directory(path_without_file).is_ok(), false);
/// ```
fn create_log_directory(path_with_file: &Path) -> Result<(), io::Error> {
    let log_dir = path_with_file.parent().ok_or(io::Error::new(
        io::ErrorKind::InvalidInput,
        format!(
            "Unable to get the parent of the following Path: {}, Your path should contain a file name",
            path_with_file.display(),
        ),
    ))?;
    fs::create_dir_all(log_dir)?;
    Ok(())
}

fn setup_logging(logging_path: Option<&Path>) -> Result<(), fern::InitError> {
    let level_filter = match std::env::var("XI_LOG") {
        Ok(level) => match level.to_lowercase().as_ref() {
            "trace" => log::LevelFilter::Trace,
            "debug" => log::LevelFilter::Debug,
            _ => log::LevelFilter::Info,
        },
        // Default to info
        Err(_) => log::LevelFilter::Info,
    };

    let mut fern_dispatch = 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,
            ))
        }).level(level_filter)
        .chain(io::stderr());

    if let Some(logging_file_path) = logging_path {
        create_log_directory(logging_file_path)?;

        fern_dispatch = fern_dispatch.chain(fern::log_file(logging_file_path)?);
    };

    // Start fern
    fern_dispatch.apply()?;
    info!("Logging with fern is set up");

    // Log details of the logging_file_path result using fern/log
    // Either logging the path fern is outputting to or the error from obtaining the path
    match logging_path {
        Some(logging_file_path) => info!("Writing logs to: {}", logging_file_path.display()),
        None => warn!("No path was supplied for the log file. Not saving logs to disk, falling back to just stderr"),
    }
    Ok(())
}

fn generate_logging_path(logfile_config: LogfileConfig) -> Result<PathBuf, io::Error> {
    // Use the file name set in logfile_config or fallback to the default
    let logfile_file_name = match logfile_config.file {
        Some(file_name) => file_name,
        None => PathBuf::from(XI_LOG_FILE),
    };
    if logfile_file_name.eq(Path::new("")) {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "A blank file name was supplied",
        ));
    };
    // Use the directory name set in logfile_config or fallback to the default
    let logfile_directory_name = match logfile_config.directory {
        Some(dir) => dir,
        None => PathBuf::from(XI_LOG_DIR),
    };

    let mut logging_directory_path = get_logging_directory_path(logfile_directory_name)?;

    // Add the file name & return the full path
    logging_directory_path.push(logfile_file_name);
    Ok(logging_directory_path)
}

fn get_flags() -> HashMap<String, Option<String>> {
    let mut flags: HashMap<String, Option<String>> = HashMap::new();

    let flag_prefix = "-";
    let mut args_iterator = std::env::args().peekable();
    while let Some(arg) = args_iterator.next() {
        if arg.starts_with(flag_prefix) {
            let key = arg.trim_left_matches(flag_prefix).to_string();

            // Check the next argument doesn't start with the flag prefix
            // map_or accounts for peek returning an Option
            let next_arg_not_a_flag: bool = args_iterator
                .peek()
                .map_or(false, |val| !val.starts_with(flag_prefix));
            if next_arg_not_a_flag {
                flags.insert(key, args_iterator.next());
            }
        }
    }
    flags
}

struct EnvFlagConfig {
    env_name: &'static str,
    flag_name: &'static str,
}

/// Extracts a value from the flags and the env.
///
/// In this order: `String` from the flags, then `String` from the env, then `None`
fn extract_env_or_flag(
    flags: &HashMap<String, Option<String>>,
    conf: EnvFlagConfig,
) -> Option<String> {
    flags
        .get(conf.flag_name)
        .cloned()
        .unwrap_or(std::env::var(conf.env_name).ok())
}

struct LogfileConfig {
    directory: Option<PathBuf>,
    file: Option<PathBuf>,
}

fn generate_logfile_config(flags: &HashMap<String, Option<String>>) -> LogfileConfig {
    // If the key is set, get the Option within
    let log_dir_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_DIR",
        flag_name: "log-dir",
    };
    let log_file_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_FILE",
        flag_name: "log-file",
    };
    let log_dir_flag_option = extract_env_or_flag(&flags, log_dir_env_flag).map(PathBuf::from);

    let log_file_flag_option = extract_env_or_flag(&flags, log_file_env_flag).map(PathBuf::from);

    LogfileConfig {
        directory: log_dir_flag_option,
        file: log_file_flag_option,
    }
}

fn main() {
    let mut state = XiCore::new();
    let stdin = io::stdin();
    let stdout = io::stdout();
    let mut rpc_looper = RpcLoop::new(stdout);

    let flags = get_flags();

    let logfile_config = generate_logfile_config(&flags);

    let logging_path_result = generate_logging_path(logfile_config);

    let logging_path = logging_path_result
        .as_ref()
        .map(|p: &PathBuf| -> &Path { p.as_path() })
        .ok();

    if let Err(e) = setup_logging(logging_path) {
        eprintln!(
            "[ERROR] setup_logging returned error, logging not enabled: {:?}",
            e
        );
    }
    if let Err(e) = logging_path_result.as_ref() {
        warn!(
            "Unable to generate the logging path to pass to set up: {}",
            e
        )
    }

    match rpc_looper.mainloop(|| stdin.lock(), &mut state) {
        Ok(_) => (),
        Err(err) => error!("xi-core exited with error:\n{:?}", err),
    }
}
<|endoftext|>


<commit_before>extern crate image2emblem;

fn main() {
}
<commit_msg>Adds seconds since year 2000 functions<commit_after>extern crate image2emblem;
extern crate time;
extern crate chrono;

use chrono::*;

fn python_total_seconds(microseconds: i64) -> f64 {
    microseconds as f64 / 10i64.pow(6) as f64
}

fn seconds_since_2000(now: chrono::datetime::DateTime<UTC>) -> f64 {
    let year_2000 = chrono::UTC.ymd(2000, 1, 1).and_hms(0, 0, 0);
    let duration = now - year_2000;
    let microseconds = duration.num_microseconds().unwrap();

    python_total_seconds(microseconds)
}

fn main() {
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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::collections::HashMap;
use std::path::Path;
use std::sync::mpsc::{sync_channel, Receiver, SyncSender, TryRecvError};
use std::thread;
use std::time::Duration;

use butterfly;
use common::ui::UI;
use depot_client;
use env;
use hcore::package::{PackageIdent, PackageInstall};
use hcore::service::ServiceGroup;
use hcore::crypto::default_cache_key_path;
use hcore::fs::{CACHE_ARTIFACT_PATH, FS_ROOT_PATH};
use launcher_client::LauncherCli;
use time::{SteadyTime, Duration as TimeDuration};

use {PRODUCT, VERSION};
use error::Result;
use census::CensusRing;
use manager::service::{Service, Topology, UpdateStrategy};

static LOGKEY: &'static str = "SU";
const FREQUENCY_ENVVAR: &'static str = "HAB_UPDATE_STRATEGY_FREQUENCY_MS";
const DEFAULT_FREQUENCY: i64 = 60_000;

type UpdaterStateList = HashMap<ServiceGroup, UpdaterState>;

enum UpdaterState {
    AtOnce(Receiver<PackageInstall>),
    Rolling(RollingState),
}

enum RollingState {
    AwaitingElection,
    InElection,
    Leader(LeaderState),
    Follower(FollowerState),
}

enum LeaderState {
    Polling(Receiver<PackageInstall>),
    Waiting,
}

enum FollowerState {
    Waiting,
    Updating(Receiver<PackageInstall>),
}

pub struct ServiceUpdater {
    states: UpdaterStateList,
    butterfly: butterfly::Server,
}

impl ServiceUpdater {
    pub fn new(butterfly: butterfly::Server) -> Self {
        ServiceUpdater {
            states: UpdaterStateList::default(),
            butterfly: butterfly,
        }
    }

    pub fn add(&mut self, service: &Service) -> bool {
        match service.update_strategy {
            UpdateStrategy::None => false,
            UpdateStrategy::AtOnce => {
                self.states
                    .entry(service.service_group.clone())
                    .or_insert_with(|| {
                        let rx = Worker::new(service).start(&service.service_group, None);
                        UpdaterState::AtOnce(rx)
                    });
                true
            }
            UpdateStrategy::Rolling => {
                self.states.entry(service.service_group.clone()).or_insert(
                    UpdaterState::Rolling(RollingState::AwaitingElection),
                );
                true
            }
        }
    }

    pub fn check_for_updated_package(
        &mut self,
        service: &mut Service,
        census_ring: &CensusRing,
        launcher: &LauncherCli,
    ) -> bool {
        let mut updated = false;
        match self.states.get_mut(&service.service_group) {
            Some(&mut UpdaterState::AtOnce(ref mut rx)) => {
                match rx.try_recv() {
                    Ok(package) => {
                        service.update_package(package, launcher);
                        return true;
                    }
                    Err(TryRecvError::Empty) => return false,
                    Err(TryRecvError::Disconnected) => {}
                }
                debug!("Service Updater worker has died; restarting...");
                *rx = Worker::new(service).start(&service.service_group, None);
            }
            Some(&mut UpdaterState::Rolling(ref mut st @ RollingState::AwaitingElection)) => {
                if let Some(census_group) = census_ring.census_group_for(&service.service_group) {
                    if service.topology == Topology::Leader {
                        debug!(
                            "Rolling Update, determining proper suitability because we're in \
                                a leader topology"
                        );
                        match (census_group.me(), census_group.leader()) {
                            (Some(me), Some(leader)) => {
                                let suitability = if me.member_id == leader.member_id {
                                    u64::min_value()
                                } else {
                                    u64::max_value()
                                };
                                self.butterfly.start_update_election(
                                    service.service_group.clone(),
                                    suitability,
                                    0,
                                );
                                *st = RollingState::InElection
                            }
                            _ => return false,
                        }
                    } else {
                        debug!("Rolling update, using default suitability");
                        self.butterfly.start_update_election(
                            service.service_group.clone(),
                            0,
                            0,
                        );
                        *st = RollingState::InElection;
                    }
                }
            }
            Some(&mut UpdaterState::Rolling(ref mut st @ RollingState::InElection)) => {
                if let Some(census_group) = census_ring.census_group_for(&service.service_group) {
                    match (census_group.me(), census_group.update_leader()) {
                        (Some(me), Some(leader)) => {
                            if me.member_id == leader.member_id {
                                debug!("We're the leader");
                                // Start in waiting state to ensure all members agree with our
                                // version before attempting a new rolling upgrade.
                                *st = RollingState::Leader(LeaderState::Waiting);
                            } else {
                                debug!("We're a follower");
                                *st = RollingState::Follower(FollowerState::Waiting);
                            }
                        }
                        (Some(_), None) => return false,
                        _ => return false,
                    }
                }
            }
            Some(&mut UpdaterState::Rolling(RollingState::Leader(ref mut state))) => {
                match *state {
                    LeaderState::Polling(ref mut rx) => {
                        match rx.try_recv() {
                            Ok(package) => {
                                debug!("Rolling Update, polling found a new package");
                                service.update_package(package, launcher);
                                updated = true;
                            }
                            Err(TryRecvError::Empty) => return false,
                            Err(TryRecvError::Disconnected) => {
                                debug!("Service Updater worker has died; restarting...");
                                *rx = Worker::new(service).start(&service.service_group, None);
                            }
                        }
                    }
                    LeaderState::Waiting => {
                        match census_ring.census_group_for(&service.service_group) {
                            Some(census_group) => {
                                if census_group.members().iter().any(|cm| {
                                    cm.pkg.as_ref().unwrap() !=
                                        census_group.me().unwrap().pkg.as_ref().unwrap()
                                })
                                {
                                    debug!("Update leader still waiting for followers...");
                                    return false;
                                }
                                let rx = Worker::new(service).start(&service.service_group, None);
                                *state = LeaderState::Polling(rx);
                            }
                            None => {
                                panic!(
                                    "Expected census list to have service group '{}'!",
                                    &*service.service_group
                                )
                            }
                        }
                    }
                }
                if updated {
                    *state = LeaderState::Waiting;
                }
            }
            Some(&mut UpdaterState::Rolling(RollingState::Follower(ref mut state))) => {
                match *state {
                    FollowerState::Waiting => {
                        match census_ring.census_group_for(&service.service_group) {
                            Some(census_group) => {
                                match (
                                    census_group.update_leader(),
                                    census_group.previous_peer(),
                                    census_group.me(),
                                ) {
                                    (Some(leader), Some(peer), Some(me)) => {
                                        if leader.pkg == me.pkg {
                                            debug!("We're not in an update");
                                            return false;
                                        }
                                        if leader.pkg != peer.pkg {
                                            debug!("We're in an update but it's not our turn");
                                            return false;
                                        }
                                        debug!("We're in an update and it's our turn");
                                        let rx = Worker::new(service).start(
                                            &service.service_group,
                                            leader.pkg.clone(),
                                        );
                                        *state = FollowerState::Updating(rx);
                                    }
                                    _ => return false,
                                }
                            }
                            None => {
                                panic!(
                                    "Expected census list to have service group '{}'!",
                                    &*service.service_group
                                )
                            }
                        }
                    }
                    FollowerState::Updating(ref mut rx) => {
                        match census_ring.census_group_for(&service.service_group) {
                            Some(census_group) => {
                                match rx.try_recv() {
                                    Ok(package) => {
                                        service.update_package(package, launcher);
                                        updated = true
                                    }
                                    Err(TryRecvError::Empty) => return false,
                                    Err(TryRecvError::Disconnected) => {
                                        debug!("Service Updater worker has died; restarting...");
                                        let package =
                                            census_group.update_leader().unwrap().pkg.clone();
                                        *rx = Worker::new(service).start(
                                            &service.service_group,
                                            package,
                                        );
                                    }
                                }
                            }
                            None => {
                                panic!(
                                    "Expected census list to have service group '{}'!",
                                    &*service.service_group
                                )
                            }
                        }
                    }
                }
                if updated {
                    *state = FollowerState::Waiting;
                }
            }
            None => {}
        }
        updated
    }
}

struct Worker {
    current: PackageIdent,
    spec_ident: PackageIdent,
    depot: depot_client::Client,
    channel: String,
    update_strategy: UpdateStrategy,
    ui: UI,
}

impl Worker {
    fn new(service: &Service) -> Self {
        Worker {
            current: service.pkg.ident.clone(),
            spec_ident: service.spec_ident.clone(),
            depot: depot_client::Client::new(&service.depot_url, PRODUCT, VERSION, None).unwrap(),
            channel: service.channel.clone(),
            update_strategy: service.update_strategy.clone(),
            ui: UI::default(),
        }
    }

    /// Start a new update worker.
    ///
    /// Passing an optional package identifier will make the worker perform a run-once update to
    /// retrieve a specific version from a remote Depot. If no package identifier is specified,
    /// then the updater will poll until a newer more suitable package is found.
    fn start(mut self, sg: &ServiceGroup, ident: Option<PackageIdent>) -> Receiver<PackageInstall> {
        let (tx, rx) = sync_channel(0);
        thread::Builder::new()
            .name(format!("service-updater-{}", sg))
            .spawn(move || match ident {
                Some(latest) => self.run_once(tx, latest),
                None => self.run_poll(tx),
            })
            .expect("unable to start service-updater thread");
        rx
    }

    fn run_once(&mut self, sender: SyncSender<PackageInstall>, ident: PackageIdent) {
        outputln!("Updating from {} to {}", self.current, ident);
        loop {
            let next_check = SteadyTime::now() +
                TimeDuration::milliseconds(self.update_frequency());
            match self.install(&ident, true) {
                Ok(package) => {
                    self.current = package.ident().clone();
                    sender.send(package).expect("Main thread has gone away!");
                    break;
                }
                Err(e) => warn!("Failed to install updated package: {:?}", e),
            }
            let time_to_wait = (next_check - SteadyTime::now()).num_milliseconds();
            if time_to_wait > 0 {
                thread::sleep(Duration::from_millis(time_to_wait as u64));
            }
        }
    }

    fn run_poll(&mut self, sender: SyncSender<PackageInstall>) {
        loop {
            let next_check = SteadyTime::now() +
                TimeDuration::milliseconds(self.update_frequency());
            let mut package: Option<PackageInstall> = None;
            match self.depot.show_package(
                &self.spec_ident,
                Some(&self.channel),
            ) {
                Ok(remote) => {
                    let latest: PackageIdent = remote.get_ident().clone().into();
                    if latest > self.current {
                        outputln!("Updating from {} to {}", self.current, latest);
                        match self.install(&latest, true) {
                            Ok(pkg) => package = Some(pkg),
                            Err(e) => warn!("Failed to install updated package: {:?}", e),
                        }
                    } else {
                        info!("Package found is not newer than ours");
                    }
                }
                Err(e) => warn!("Updater failed to get latest package: {:?}", e),
            }
            if self.update_strategy == UpdateStrategy::AtOnce {
                if let Ok(cached) = PackageInstall::load(
                    &self.spec_ident,
                    Some(&Path::new(&*FS_ROOT_PATH)),
                )
                {
                    let compare = match package {
                        Some(ref pkg) => pkg.ident.clone(),
                        None => self.current.clone(),
                    };

                    if cached.ident > compare {
                        package = Some(cached);
                    }
                }
            }

            if let Some(pkg) = package {
                self.current = pkg.ident.clone();
                sender.send(pkg).expect("Main thread has gone away!");
                break;
            }
            let time_to_wait = (next_check - SteadyTime::now()).num_milliseconds();
            if time_to_wait > 0 {
                thread::sleep(Duration::from_millis(time_to_wait as u64));
            }
        }
    }

    fn install(&mut self, package: &PackageIdent, recurse: bool) -> Result<PackageInstall> {
        let package = match PackageInstall::load(package, Some(&*FS_ROOT_PATH)) {
            Ok(pkg) => pkg,
            Err(_) => self.download(package)?,
        };
        if recurse {
            for ident in package.tdeps()?.iter() {
                self.install(&ident, false)?;
            }
        }
        Ok(package)
    }

    fn download(&mut self, package: &PackageIdent) -> Result<PackageInstall> {
        outputln!("Downloading {}", package);
        let mut archive = self.depot.fetch_package(
            package,
            &Path::new(&*FS_ROOT_PATH).join(
                CACHE_ARTIFACT_PATH,
            ),
            self.ui.progress(),
        )?;
        archive.verify(&default_cache_key_path(None))?;
        outputln!("Installing {}", package);
        archive.unpack(None)?;
        let pkg = PackageInstall::load(archive.ident().as_ref().unwrap(), Some(&*FS_ROOT_PATH))?;
        Ok(pkg)
    }

    fn update_frequency(&self) -> i64 {
        match env::var(FREQUENCY_ENVVAR) {
            Ok(val) => {
                match val.parse::<i64>() {
                    Ok(num) => num,
                    Err(_) => {
                        outputln!(
                            "Unable to parse '{}' from {} as a valid integer. Falling back \
                                  to default {} MS frequency.",
                            val,
                            FREQUENCY_ENVVAR,
                            DEFAULT_FREQUENCY
                        );
                        DEFAULT_FREQUENCY
                    }
                }
            }
            Err(_) => DEFAULT_FREQUENCY,
        }
    }
}
<commit_msg>Download keys for verification when updating services<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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::collections::HashMap;
use std::path::Path;
use std::sync::mpsc::{sync_channel, Receiver, SyncSender, TryRecvError};
use std::thread;
use std::time::Duration;

use butterfly;
use common::ui::{Status, UI};
use depot_client;
use env;
use hcore::package::{PackageIdent, PackageInstall};
use hcore::service::ServiceGroup;
use hcore::crypto::{artifact, default_cache_key_path, SigKeyPair};
use hcore::crypto::keys::parse_name_with_rev;
use hcore::fs::{CACHE_ARTIFACT_PATH, FS_ROOT_PATH};
use launcher_client::LauncherCli;
use time::{SteadyTime, Duration as TimeDuration};

use {PRODUCT, VERSION};
use error::Result;
use census::CensusRing;
use manager::service::{Service, Topology, UpdateStrategy};

static LOGKEY: &'static str = "SU";
const FREQUENCY_ENVVAR: &'static str = "HAB_UPDATE_STRATEGY_FREQUENCY_MS";
const DEFAULT_FREQUENCY: i64 = 60_000;

type UpdaterStateList = HashMap<ServiceGroup, UpdaterState>;

enum UpdaterState {
    AtOnce(Receiver<PackageInstall>),
    Rolling(RollingState),
}

enum RollingState {
    AwaitingElection,
    InElection,
    Leader(LeaderState),
    Follower(FollowerState),
}

enum LeaderState {
    Polling(Receiver<PackageInstall>),
    Waiting,
}

enum FollowerState {
    Waiting,
    Updating(Receiver<PackageInstall>),
}

pub struct ServiceUpdater {
    states: UpdaterStateList,
    butterfly: butterfly::Server,
}

impl ServiceUpdater {
    pub fn new(butterfly: butterfly::Server) -> Self {
        ServiceUpdater {
            states: UpdaterStateList::default(),
            butterfly: butterfly,
        }
    }

    pub fn add(&mut self, service: &Service) -> bool {
        match service.update_strategy {
            UpdateStrategy::None => false,
            UpdateStrategy::AtOnce => {
                self.states
                    .entry(service.service_group.clone())
                    .or_insert_with(|| {
                        let rx = Worker::new(service).start(&service.service_group, None);
                        UpdaterState::AtOnce(rx)
                    });
                true
            }
            UpdateStrategy::Rolling => {
                self.states.entry(service.service_group.clone()).or_insert(
                    UpdaterState::Rolling(RollingState::AwaitingElection),
                );
                true
            }
        }
    }

    pub fn check_for_updated_package(
        &mut self,
        service: &mut Service,
        census_ring: &CensusRing,
        launcher: &LauncherCli,
    ) -> bool {
        let mut updated = false;
        match self.states.get_mut(&service.service_group) {
            Some(&mut UpdaterState::AtOnce(ref mut rx)) => {
                match rx.try_recv() {
                    Ok(package) => {
                        service.update_package(package, launcher);
                        return true;
                    }
                    Err(TryRecvError::Empty) => return false,
                    Err(TryRecvError::Disconnected) => {}
                }
                debug!("Service Updater worker has died; restarting...");
                *rx = Worker::new(service).start(&service.service_group, None);
            }
            Some(&mut UpdaterState::Rolling(ref mut st @ RollingState::AwaitingElection)) => {
                if let Some(census_group) = census_ring.census_group_for(&service.service_group) {
                    if service.topology == Topology::Leader {
                        debug!(
                            "Rolling Update, determining proper suitability because we're in \
                                a leader topology"
                        );
                        match (census_group.me(), census_group.leader()) {
                            (Some(me), Some(leader)) => {
                                let suitability = if me.member_id == leader.member_id {
                                    u64::min_value()
                                } else {
                                    u64::max_value()
                                };
                                self.butterfly.start_update_election(
                                    service.service_group.clone(),
                                    suitability,
                                    0,
                                );
                                *st = RollingState::InElection
                            }
                            _ => return false,
                        }
                    } else {
                        debug!("Rolling update, using default suitability");
                        self.butterfly.start_update_election(
                            service.service_group.clone(),
                            0,
                            0,
                        );
                        *st = RollingState::InElection;
                    }
                }
            }
            Some(&mut UpdaterState::Rolling(ref mut st @ RollingState::InElection)) => {
                if let Some(census_group) = census_ring.census_group_for(&service.service_group) {
                    match (census_group.me(), census_group.update_leader()) {
                        (Some(me), Some(leader)) => {
                            if me.member_id == leader.member_id {
                                debug!("We're the leader");
                                // Start in waiting state to ensure all members agree with our
                                // version before attempting a new rolling upgrade.
                                *st = RollingState::Leader(LeaderState::Waiting);
                            } else {
                                debug!("We're a follower");
                                *st = RollingState::Follower(FollowerState::Waiting);
                            }
                        }
                        (Some(_), None) => return false,
                        _ => return false,
                    }
                }
            }
            Some(&mut UpdaterState::Rolling(RollingState::Leader(ref mut state))) => {
                match *state {
                    LeaderState::Polling(ref mut rx) => {
                        match rx.try_recv() {
                            Ok(package) => {
                                debug!("Rolling Update, polling found a new package");
                                service.update_package(package, launcher);
                                updated = true;
                            }
                            Err(TryRecvError::Empty) => return false,
                            Err(TryRecvError::Disconnected) => {
                                debug!("Service Updater worker has died; restarting...");
                                *rx = Worker::new(service).start(&service.service_group, None);
                            }
                        }
                    }
                    LeaderState::Waiting => {
                        match census_ring.census_group_for(&service.service_group) {
                            Some(census_group) => {
                                if census_group.members().iter().any(|cm| {
                                    cm.pkg.as_ref().unwrap() !=
                                        census_group.me().unwrap().pkg.as_ref().unwrap()
                                })
                                {
                                    debug!("Update leader still waiting for followers...");
                                    return false;
                                }
                                let rx = Worker::new(service).start(&service.service_group, None);
                                *state = LeaderState::Polling(rx);
                            }
                            None => {
                                panic!(
                                    "Expected census list to have service group '{}'!",
                                    &*service.service_group
                                )
                            }
                        }
                    }
                }
                if updated {
                    *state = LeaderState::Waiting;
                }
            }
            Some(&mut UpdaterState::Rolling(RollingState::Follower(ref mut state))) => {
                match *state {
                    FollowerState::Waiting => {
                        match census_ring.census_group_for(&service.service_group) {
                            Some(census_group) => {
                                match (
                                    census_group.update_leader(),
                                    census_group.previous_peer(),
                                    census_group.me(),
                                ) {
                                    (Some(leader), Some(peer), Some(me)) => {
                                        if leader.pkg == me.pkg {
                                            debug!("We're not in an update");
                                            return false;
                                        }
                                        if leader.pkg != peer.pkg {
                                            debug!("We're in an update but it's not our turn");
                                            return false;
                                        }
                                        debug!("We're in an update and it's our turn");
                                        let rx = Worker::new(service).start(
                                            &service.service_group,
                                            leader.pkg.clone(),
                                        );
                                        *state = FollowerState::Updating(rx);
                                    }
                                    _ => return false,
                                }
                            }
                            None => {
                                panic!(
                                    "Expected census list to have service group '{}'!",
                                    &*service.service_group
                                )
                            }
                        }
                    }
                    FollowerState::Updating(ref mut rx) => {
                        match census_ring.census_group_for(&service.service_group) {
                            Some(census_group) => {
                                match rx.try_recv() {
                                    Ok(package) => {
                                        service.update_package(package, launcher);
                                        updated = true
                                    }
                                    Err(TryRecvError::Empty) => return false,
                                    Err(TryRecvError::Disconnected) => {
                                        debug!("Service Updater worker has died; restarting...");
                                        let package =
                                            census_group.update_leader().unwrap().pkg.clone();
                                        *rx = Worker::new(service).start(
                                            &service.service_group,
                                            package,
                                        );
                                    }
                                }
                            }
                            None => {
                                panic!(
                                    "Expected census list to have service group '{}'!",
                                    &*service.service_group
                                )
                            }
                        }
                    }
                }
                if updated {
                    *state = FollowerState::Waiting;
                }
            }
            None => {}
        }
        updated
    }
}

struct Worker {
    current: PackageIdent,
    spec_ident: PackageIdent,
    depot: depot_client::Client,
    channel: String,
    update_strategy: UpdateStrategy,
    ui: UI,
}

impl Worker {
    fn new(service: &Service) -> Self {
        Worker {
            current: service.pkg.ident.clone(),
            spec_ident: service.spec_ident.clone(),
            depot: depot_client::Client::new(&service.depot_url, PRODUCT, VERSION, None).unwrap(),
            channel: service.channel.clone(),
            update_strategy: service.update_strategy.clone(),
            ui: UI::default(),
        }
    }

    /// Start a new update worker.
    ///
    /// Passing an optional package identifier will make the worker perform a run-once update to
    /// retrieve a specific version from a remote Depot. If no package identifier is specified,
    /// then the updater will poll until a newer more suitable package is found.
    fn start(mut self, sg: &ServiceGroup, ident: Option<PackageIdent>) -> Receiver<PackageInstall> {
        let (tx, rx) = sync_channel(0);
        thread::Builder::new()
            .name(format!("service-updater-{}", sg))
            .spawn(move || match ident {
                Some(latest) => self.run_once(tx, latest),
                None => self.run_poll(tx),
            })
            .expect("unable to start service-updater thread");
        rx
    }

    fn run_once(&mut self, sender: SyncSender<PackageInstall>, ident: PackageIdent) {
        outputln!("Updating from {} to {}", self.current, ident);
        loop {
            let next_check = SteadyTime::now() +
                TimeDuration::milliseconds(self.update_frequency());
            match self.install(&ident, true) {
                Ok(package) => {
                    self.current = package.ident().clone();
                    sender.send(package).expect("Main thread has gone away!");
                    break;
                }
                Err(e) => warn!("Failed to install updated package: {:?}", e),
            }
            let time_to_wait = (next_check - SteadyTime::now()).num_milliseconds();
            if time_to_wait > 0 {
                thread::sleep(Duration::from_millis(time_to_wait as u64));
            }
        }
    }

    fn run_poll(&mut self, sender: SyncSender<PackageInstall>) {
        loop {
            let next_check = SteadyTime::now() +
                TimeDuration::milliseconds(self.update_frequency());
            let mut package: Option<PackageInstall> = None;
            match self.depot.show_package(
                &self.spec_ident,
                Some(&self.channel),
            ) {
                Ok(remote) => {
                    let latest: PackageIdent = remote.get_ident().clone().into();
                    if latest > self.current {
                        outputln!("Updating from {} to {}", self.current, latest);
                        match self.install(&latest, true) {
                            Ok(pkg) => package = Some(pkg),
                            Err(e) => warn!("Failed to install updated package: {:?}", e),
                        }
                    } else {
                        info!("Package found is not newer than ours");
                    }
                }
                Err(e) => warn!("Updater failed to get latest package: {:?}", e),
            }
            if self.update_strategy == UpdateStrategy::AtOnce {
                if let Ok(cached) = PackageInstall::load(
                    &self.spec_ident,
                    Some(&Path::new(&*FS_ROOT_PATH)),
                )
                {
                    let compare = match package {
                        Some(ref pkg) => pkg.ident.clone(),
                        None => self.current.clone(),
                    };

                    if cached.ident > compare {
                        package = Some(cached);
                    }
                }
            }

            if let Some(pkg) = package {
                self.current = pkg.ident.clone();
                sender.send(pkg).expect("Main thread has gone away!");
                break;
            }
            let time_to_wait = (next_check - SteadyTime::now()).num_milliseconds();
            if time_to_wait > 0 {
                thread::sleep(Duration::from_millis(time_to_wait as u64));
            }
        }
    }

    fn install(&mut self, package: &PackageIdent, recurse: bool) -> Result<PackageInstall> {
        let package = match PackageInstall::load(package, Some(&*FS_ROOT_PATH)) {
            Ok(pkg) => pkg,
            Err(_) => self.download(package)?,
        };
        if recurse {
            for ident in package.tdeps()?.iter() {
                self.install(&ident, false)?;
            }
        }
        Ok(package)
    }

    fn download(&mut self, package: &PackageIdent) -> Result<PackageInstall> {
        outputln!("Downloading {}", package);
        let mut archive = self.depot.fetch_package(
            package,
            &Path::new(&*FS_ROOT_PATH).join(
                CACHE_ARTIFACT_PATH,
            ),
            self.ui.progress(),
        )?;

        // TODO (CM): Copied (with modifications) from
        // common::command::package::install; will be properly
        // factored in an upcoming broad refactoring and consolidation
        // of our installation logic.
        let cache_key_path = &default_cache_key_path(None);
        let nwr = artifact::artifact_signer(&archive.path)?;
        if let Err(_) = SigKeyPair::get_public_key_path(&nwr, cache_key_path) {
            self.fetch_origin_key(&nwr)?;
        }

        archive.verify(cache_key_path)?;
        outputln!("Installing {}", package);
        archive.unpack(None)?;
        let pkg = PackageInstall::load(archive.ident().as_ref().unwrap(), Some(&*FS_ROOT_PATH))?;
        Ok(pkg)
    }

    fn update_frequency(&self) -> i64 {
        match env::var(FREQUENCY_ENVVAR) {
            Ok(val) => {
                match val.parse::<i64>() {
                    Ok(num) => num,
                    Err(_) => {
                        outputln!(
                            "Unable to parse '{}' from {} as a valid integer. Falling back \
                                  to default {} MS frequency.",
                            val,
                            FREQUENCY_ENVVAR,
                            DEFAULT_FREQUENCY
                        );
                        DEFAULT_FREQUENCY
                    }
                }
            }
            Err(_) => DEFAULT_FREQUENCY,
        }
    }

    // TODO (CM): Copied (with modifications) from
    // common::command::package::install; will be properly factored in
    // an upcoming broad refactoring and consolidation of our
    // installation logic.
    fn fetch_origin_key(&mut self, name_with_rev: &str) -> Result<()> {
        self.ui.status(
            Status::Downloading,
            format!("{} public origin key", &name_with_rev),
        )?;
        let (name, rev) = parse_name_with_rev(&name_with_rev)?;
        self.depot.fetch_origin_key(
            &name,
            &rev,
            &default_cache_key_path(None),
            self.ui.progress(),
        )?;
        self.ui.status(
            Status::Cached,
            format!("{} public origin key", &name_with_rev),
        )?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Aead Ciphers

use crypto::cipher::{CipherCategory, CipherResult, CipherType};

use crypto::ring::RingAeadCipher;
#[cfg(feature = "miscreant")]
use crypto::siv::MiscreantCipher;

use ring::digest::SHA1;
use ring::hkdf;
use ring::hmac::SigningKey;

use bytes::{Bytes, BytesMut};

/// Encryptor API for AEAD ciphers
pub trait AeadEncryptor {
    /// Encrypt `input` to `output` with `tag`. `output.len()` should equals to `input.len()+tag.len()`.
    fn encrypt(&mut self, input: &[u8], output: &mut [u8]);
}

/// Decryptor API for AEAD ciphers
pub trait AeadDecryptor {
    /// Decrypt `input` to `output` with `tag`. `output.len()` should equals to `input.len()+tag.len()`.
    fn decrypt(&mut self, input: &[u8], output: &mut [u8]) -> CipherResult<()>;
}

/// Variant `AeadDecryptor`
pub type BoxAeadDecryptor = Box<AeadDecryptor + Send>;

/// Variant `AeadEncryptor`
pub type BoxAeadEncryptor = Box<AeadEncryptor + Send>;

/// Generate a specific AEAD cipher encryptor
pub fn new_aead_encryptor(t: CipherType, key: &[u8], nonce: &[u8]) -> BoxAeadEncryptor {
    assert!(t.category() == CipherCategory::Aead);

    match t {
        CipherType::Aes128Gcm | CipherType::Aes256Gcm | CipherType::ChaCha20Poly1305 => {
            Box::new(RingAeadCipher::new(t, key, nonce, true))
        }

        #[cfg(feature = "miscreant")]
        CipherType::Aes128PmacSiv | CipherType::Aes256PmacSiv => Box::new(MiscreantCipher::new(t, key, nonce)),

        _ => unreachable!(),
    }
}

/// Generate a specific AEAD cipher decryptor
pub fn new_aead_decryptor(t: CipherType, key: &[u8], nonce: &[u8]) -> BoxAeadDecryptor {
    assert!(t.category() == CipherCategory::Aead);

    match t {
        CipherType::Aes128Gcm | CipherType::Aes256Gcm | CipherType::ChaCha20Poly1305 => {
            Box::new(RingAeadCipher::new(t, key, nonce, false))
        }

        #[cfg(feature = "miscreant")]
        CipherType::Aes128PmacSiv | CipherType::Aes256PmacSiv => Box::new(MiscreantCipher::new(t, key, nonce)),

        _ => unreachable!(),
    }
}

const SUBKEY_INFO: &'static [u8] = b"ss-subkey";

/// Make Session key
///
/// ## Session key (SIP007)
///
/// AEAD ciphers require a per-session subkey derived from the pre-shared master key using HKDF, and use the subkey
/// to encrypt/decrypt. Essentially it means we are moving from (M+N)-bit (PSK, nonce) pair to
/// (M+N)-bit (HKDF(PSK, salt), nonce) pair. Because HKDF is a PRF, the new construction significantly expands the
/// amount of randomness (from N to at least M where M is much greater than N), thus correcting the previously
/// mentioned design flaw.
///
/// Assuming we already have a user-supplied pre-shared master key PSK.
///
/// Function HKDF_SHA1 is a HKDF constructed using SHA1 hash. Its signature is
///
/// ```plain
/// HKDF_SHA1(secret_key, salt, info)
/// ```
///
/// The "info" string argument allows us to bind the derived subkey to a specific application context.
///
/// For AEAD ciphers, the encryption scheme is:
///
/// 1. Pick a random R-bit salt (R = max(128, len(SK)))
/// 2. Derive subkey SK = HKDF_SHA1(PSK, salt, "ss-subkey")
/// 3. Send salt
/// 4. For each chunk, encrypt and authenticate payload using SK with a counting nonce
///    (starting from 0 and increment by 1 after each use)
/// 5. Send encrypted chunk
pub fn make_skey(t: CipherType, key: &[u8], salt: &[u8]) -> Bytes {
    assert!(t.category() == CipherCategory::Aead);

    let salt = SigningKey::new(&SHA1, salt);

    let mut skey = BytesMut::with_capacity(key.len());
    unsafe {
        skey.set_len(key.len());
    }

    hkdf::extract_and_expand(&salt, key, SUBKEY_INFO, &mut skey);

    skey.freeze()
}

/// Increase nonce by 1
///
/// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk
#[cfg(feature = "sodium")]
pub fn increase_nonce(nonce: &mut [u8]) {
    use libsodium_ffi::sodium_increment;

    unsafe {
        sodium_increment(nonce.as_mut_ptr(), nonce.len());
    }
}

/// Increase nonce by 1
///
/// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk
#[cfg(not(feature = "sodium"))]
pub fn increase_nonce(nonce: &mut [u8]) {
    let mut prev: u16 = 1;
    for i in nonce {
        prev += *i as u16;
        *i = prev as u8;
        prev >>= 8;
    }
}
<commit_msg>Add missing register in new aead encryptor/decryptor<commit_after>//! Aead Ciphers

use crypto::cipher::{CipherCategory, CipherResult, CipherType};

use crypto::ring::RingAeadCipher;
#[cfg(feature = "miscreant")]
use crypto::siv::MiscreantCipher;
#[cfg(feature = "sodium")]
use crypto::sodium::SodiumAeadCipher;

use ring::digest::SHA1;
use ring::hkdf;
use ring::hmac::SigningKey;

use bytes::{Bytes, BytesMut};

/// Encryptor API for AEAD ciphers
pub trait AeadEncryptor {
    /// Encrypt `input` to `output` with `tag`. `output.len()` should equals to `input.len()+tag.len()`.
    fn encrypt(&mut self, input: &[u8], output: &mut [u8]);
}

/// Decryptor API for AEAD ciphers
pub trait AeadDecryptor {
    /// Decrypt `input` to `output` with `tag`. `output.len()` should equals to `input.len()+tag.len()`.
    fn decrypt(&mut self, input: &[u8], output: &mut [u8]) -> CipherResult<()>;
}

/// Variant `AeadDecryptor`
pub type BoxAeadDecryptor = Box<AeadDecryptor + Send>;

/// Variant `AeadEncryptor`
pub type BoxAeadEncryptor = Box<AeadEncryptor + Send>;

/// Generate a specific AEAD cipher encryptor
pub fn new_aead_encryptor(t: CipherType, key: &[u8], nonce: &[u8]) -> BoxAeadEncryptor {
    assert!(t.category() == CipherCategory::Aead);

    match t {
        CipherType::Aes128Gcm | CipherType::Aes256Gcm | CipherType::ChaCha20Poly1305 => {
            Box::new(RingAeadCipher::new(t, key, nonce, true))
        }

        #[cfg(feature = "sodium")]
        CipherType::XChaCha20Poly1305 => {
            Box::new(SodiumAeadCipher::new(t, key, nonce))
        }

        #[cfg(feature = "miscreant")]
        CipherType::Aes128PmacSiv | CipherType::Aes256PmacSiv => Box::new(MiscreantCipher::new(t, key, nonce)),

        _ => unreachable!(),
    }
}

/// Generate a specific AEAD cipher decryptor
pub fn new_aead_decryptor(t: CipherType, key: &[u8], nonce: &[u8]) -> BoxAeadDecryptor {
    assert!(t.category() == CipherCategory::Aead);

    match t {
        CipherType::Aes128Gcm | CipherType::Aes256Gcm | CipherType::ChaCha20Poly1305 => {
            Box::new(RingAeadCipher::new(t, key, nonce, false))
        }

        #[cfg(feature = "sodium")]
        CipherType::XChaCha20Poly1305 => {
            Box::new(SodiumAeadCipher::new(t, key, nonce))
        }

        #[cfg(feature = "miscreant")]
        CipherType::Aes128PmacSiv | CipherType::Aes256PmacSiv => Box::new(MiscreantCipher::new(t, key, nonce)),

        _ => unreachable!(),
    }
}

const SUBKEY_INFO: &'static [u8] = b"ss-subkey";

/// Make Session key
///
/// ## Session key (SIP007)
///
/// AEAD ciphers require a per-session subkey derived from the pre-shared master key using HKDF, and use the subkey
/// to encrypt/decrypt. Essentially it means we are moving from (M+N)-bit (PSK, nonce) pair to
/// (M+N)-bit (HKDF(PSK, salt), nonce) pair. Because HKDF is a PRF, the new construction significantly expands the
/// amount of randomness (from N to at least M where M is much greater than N), thus correcting the previously
/// mentioned design flaw.
///
/// Assuming we already have a user-supplied pre-shared master key PSK.
///
/// Function HKDF_SHA1 is a HKDF constructed using SHA1 hash. Its signature is
///
/// ```plain
/// HKDF_SHA1(secret_key, salt, info)
/// ```
///
/// The "info" string argument allows us to bind the derived subkey to a specific application context.
///
/// For AEAD ciphers, the encryption scheme is:
///
/// 1. Pick a random R-bit salt (R = max(128, len(SK)))
/// 2. Derive subkey SK = HKDF_SHA1(PSK, salt, "ss-subkey")
/// 3. Send salt
/// 4. For each chunk, encrypt and authenticate payload using SK with a counting nonce
///    (starting from 0 and increment by 1 after each use)
/// 5. Send encrypted chunk
pub fn make_skey(t: CipherType, key: &[u8], salt: &[u8]) -> Bytes {
    assert!(t.category() == CipherCategory::Aead);

    let salt = SigningKey::new(&SHA1, salt);

    let mut skey = BytesMut::with_capacity(key.len());
    unsafe {
        skey.set_len(key.len());
    }

    hkdf::extract_and_expand(&salt, key, SUBKEY_INFO, &mut skey);

    skey.freeze()
}

/// Increase nonce by 1
///
/// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk
#[cfg(feature = "sodium")]
pub fn increase_nonce(nonce: &mut [u8]) {
    use libsodium_ffi::sodium_increment;

    unsafe {
        sodium_increment(nonce.as_mut_ptr(), nonce.len());
    }
}

/// Increase nonce by 1
///
/// AEAD ciphers requires to increase nonce after encrypt/decrypt every chunk
#[cfg(not(feature = "sodium"))]
pub fn increase_nonce(nonce: &mut [u8]) {
    let mut prev: u16 = 1;
    for i in nonce {
        prev += *i as u16;
        *i = prev as u8;
        prev >>= 8;
    }
}
<|endoftext|>