<commit_before>#[macro_use] extern crate nom;
#[macro_use] extern crate clap;
#[macro_use] extern crate serde_derive;
#[macro_use] extern crate lazy_static;
extern crate mio;
extern crate mio_uds;
extern crate serde_json;
extern crate time;
extern crate libc;
extern crate slab;
extern crate rand;
extern crate nix;
extern crate tempfile;
extern crate futures;
extern crate regex;
#[cfg(feature = "jemallocator")]
extern crate jemallocator;
#[macro_use] extern crate sozu_lib as sozu;
#[macro_use] extern crate sozu_command_lib as sozu_command;

#[cfg(target_os = "linux")]
extern crate num_cpus;
#[cfg(target_os="linux")]
use regex::Regex;

#[cfg(feature = "jemallocator")]
#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

mod command;
mod worker;
mod logging;
mod upgrade;
mod util;
mod cli;

use std::env;
use std::panic;
use sozu_command::config::Config;
use clap::ArgMatches;

#[cfg(target_os = "linux")]
use libc::{cpu_set_t,pid_t};

use command::Worker;
use worker::{start_workers,get_executable_path};
use sozu::metrics::METRICS;

enum StartupError {
  ConfigurationFileNotSpecified,
  ConfigurationFileLoadError(std::io::Error),
  #[allow(dead_code)]
  TooManyAllowedConnections(String),
  #[allow(dead_code)]
  TooManyAllowedConnectionsForWorker(String),
  WorkersSpawnFail(nix::Error),
  PIDFileNotWritable(String)
}

fn main() {
  register_panic_hook();

  // Init parsing of arguments
  let matches = cli::init();
  // Check if we are upgrading workers or master
  let upgrade = cli::upgrade_worker(&matches).or_else(|| cli::upgrade_master(&matches));

  // If we are not, then we want to start sozu
  if upgrade == None {
    let start = get_config_file_path(&matches)
    .and_then(|config_file| load_configuration(config_file))
    .and_then(|config| {
      util::write_pid_file(&config)
        .map(|()| config)
        .map_err(|err| StartupError::PIDFileNotWritable(err))
    })
    .map(|config| {
      util::setup_logging(&config);
      info!("Starting up");
      util::setup_metrics(&config);

      config
    })
    .and_then(|config| check_process_limits(&config).map(|()| config))
    .and_then(|config| init_workers(&config).map(|workers| (config, workers)))
    .map(|(config, workers)| {
      if config.handle_process_affinity {
        set_workers_affinity(&workers);
      }
      let command_socket_path = config.command_socket_path();
      command::start(config, command_socket_path, workers);
    });

    match start {
      Ok(_) => info!("master process stopped"), // Ok() is only called when the proxy exits
      Err(StartupError::ConfigurationFileNotSpecified) => {
        error!("Configuration file hasn't been specified. Either use -c with the start command \
               or use the SOZU_CONFIG environment variable when building sozu.");
        std::process::exit(1);
      },
      Err(StartupError::ConfigurationFileLoadError(err)) => {
        error!("Invalid configuration file. Error: {:?}", err);
        std::process::exit(1);
      },
      Err(StartupError::TooManyAllowedConnections(err)) | Err(StartupError::TooManyAllowedConnectionsForWorker(err)) => {
        error!("{}", err);
        std::process::exit(1);
      },
      Err(StartupError::WorkersSpawnFail(err)) => {
        error!("At least one worker failed to spawn. Error: {:?}", err);
        std::process::exit(1);
      },
      Err(StartupError::PIDFileNotWritable(err)) => {
        error!("{}", err);
        std::process::exit(1);
      }
    }
  }
}

fn init_workers(config: &Config) -> Result<Vec<Worker>, StartupError> {
  let path = unsafe { get_executable_path() };
  match start_workers(path, &config) {
    Ok(workers) => {
      info!("created workers: {:?}", workers);
      Ok(workers)
    },
    Err(e) => Err(StartupError::WorkersSpawnFail(e))
  }
}

fn get_config_file_path<'a>(matches: &'a ArgMatches<'a>) -> Result<&'a str, StartupError> {
  let start_matches = matches.subcommand_matches("start").expect("unknown subcommand");
  match start_matches.value_of("config") {
    Some(config_file) => Ok(config_file),
    None => option_env!("SOZU_CONFIG").ok_or(StartupError::ConfigurationFileNotSpecified)
  }
}

fn load_configuration(config_file: &str) -> Result<Config, StartupError> {
  match Config::load_from_path(config_file) {
    Ok(config) => Ok(config),
    Err(e) => Err(StartupError::ConfigurationFileLoadError(e))
  }
}

/// Set workers process affinity, see man sched_setaffinity
/// Bind each worker (including the master) process to a CPU core.
/// Can bind multiple processes to a CPU core if there are more processes
/// than CPU cores. Only works on Linux.
#[cfg(target_os = "linux")]
fn set_workers_affinity(workers: &Vec<Worker>) {
  let mut cpu_count = 0;
  let max_cpu = num_cpus::get();

  // +1 for the master process that will also be bound to its CPU core
  if (workers.len() + 1) > max_cpu {
    warn!("There are more workers than available CPU cores, \
          multiple workers will be bound to the same CPU core. \
          This may impact performances");
  }

  let master_pid = unsafe { libc::getpid() };
  set_process_affinity(master_pid, cpu_count);
  cpu_count = cpu_count + 1;

  for ref worker in workers {
    if cpu_count >= max_cpu {
      cpu_count = 0;
    }

    set_process_affinity(worker.pid, cpu_count);

    cpu_count = cpu_count + 1;
  }
}

/// Set workers process affinity, see man sched_setaffinity
/// Bind each worker (including the master) process to a CPU core.
/// Can bind multiple processes to a CPU core if there are more processes
/// than CPU cores. Only works on Linux.
#[cfg(not(target_os = "linux"))]
fn set_workers_affinity(_: &Vec<Worker>) {
}

/// Set a specific process to run onto a specific CPU core
#[cfg(target_os = "linux")]
use std::mem;
#[cfg(target_os = "linux")]
fn set_process_affinity(pid: pid_t, cpu: usize) {
  unsafe {
    let mut cpu_set: cpu_set_t = mem::zeroed();
    let size_cpu_set = mem::size_of::<cpu_set_t>();
    libc::CPU_SET(cpu, &mut cpu_set);
    libc::sched_setaffinity(pid, size_cpu_set, &mut cpu_set);

    debug!("Worker {} bound to CPU core {}", pid, cpu);
  };
}

#[cfg(target_os="linux")]
// We check the hard_limit. The soft_limit can be changed at runtime
// by the process or any user. hard_limit can only be changed by root
fn check_process_limits(config: &Config) -> Result<(), StartupError> {
  let process_limits_file = sozu_command::config::Config::load_file("/proc/self/limits")
    .expect("Couldn't read /proc/self/limits to determine max open file descriptors limit");

  let re = Regex::new(r"Max open files\s*\d*\s*(\d*)\s*files").unwrap();

  if let Some(hard_limit) = re.captures(&process_limits_file).and_then(|c| c.get(1))
    .and_then(|m| m.as_str().parse::<usize>().ok()) {
    if config.max_connections > hard_limit {
      let error = format!("At least one worker can't have that many connections. \
              Current max file descriptor hard limit is: {}", hard_limit);
      return Err(StartupError::TooManyAllowedConnectionsForWorker(error));
    }
  }

  let f = sozu_command::config::Config::load_file("/proc/sys/fs/file-max")
    .expect("Couldn't read /proc/sys/fs/file-max");

  let re_max = Regex::new(r"(\d*)").unwrap();

  let system_max_fd = re_max.captures(&f).and_then(|c| c.get(1))
    .and_then(|m| m.as_str().parse::<usize>().ok())
    .expect("Couldn't parse /proc/sys/fs/file-max");

  if config.max_connections > system_max_fd {
    let error = format!("Proxies total max_connections can't be higher than system's file-max limit. \
            Current limit: {}, current value: {}", system_max_fd, config.max_connections);
    return Err(StartupError::TooManyAllowedConnections(error))
  }

  Ok(())
}

#[cfg(not(target_os = "linux"))]
fn check_process_limits(_: &Config) -> Result<(), StartupError> {
  Ok(())
}

fn register_panic_hook() {
  // We save the original panic hook so we can call it later
  // to have the original behavior
  let original_panic_hook = panic::take_hook();

  panic::set_hook(Box::new(move |panic_info| {
    incr!("panic");
    METRICS.with(|metrics| {
      (*metrics.borrow_mut()).send_data();
    });

    (*original_panic_hook)(panic_info)
  }));
}
<commit_msg>check the soft max open files limit too<commit_after>#[macro_use] extern crate nom;
#[macro_use] extern crate clap;
#[macro_use] extern crate serde_derive;
#[macro_use] extern crate lazy_static;
extern crate mio;
extern crate mio_uds;
extern crate serde_json;
extern crate time;
extern crate libc;
extern crate slab;
extern crate rand;
extern crate nix;
extern crate tempfile;
extern crate futures;
extern crate regex;
#[cfg(feature = "jemallocator")]
extern crate jemallocator;
#[macro_use] extern crate sozu_lib as sozu;
#[macro_use] extern crate sozu_command_lib as sozu_command;

#[cfg(target_os = "linux")]
extern crate num_cpus;
#[cfg(target_os="linux")]
use regex::Regex;

#[cfg(feature = "jemallocator")]
#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

mod command;
mod worker;
mod logging;
mod upgrade;
mod util;
mod cli;

use std::env;
use std::panic;
use sozu_command::config::Config;
use clap::ArgMatches;

#[cfg(target_os = "linux")]
use libc::{cpu_set_t,pid_t};

use command::Worker;
use worker::{start_workers,get_executable_path};
use sozu::metrics::METRICS;

enum StartupError {
  ConfigurationFileNotSpecified,
  ConfigurationFileLoadError(std::io::Error),
  #[allow(dead_code)]
  TooManyAllowedConnections(String),
  #[allow(dead_code)]
  TooManyAllowedConnectionsForWorker(String),
  WorkersSpawnFail(nix::Error),
  PIDFileNotWritable(String)
}

fn main() {
  register_panic_hook();

  // Init parsing of arguments
  let matches = cli::init();
  // Check if we are upgrading workers or master
  let upgrade = cli::upgrade_worker(&matches).or_else(|| cli::upgrade_master(&matches));

  // If we are not, then we want to start sozu
  if upgrade == None {
    let start = get_config_file_path(&matches)
    .and_then(|config_file| load_configuration(config_file))
    .and_then(|config| {
      util::write_pid_file(&config)
        .map(|()| config)
        .map_err(|err| StartupError::PIDFileNotWritable(err))
    })
    .map(|config| {
      util::setup_logging(&config);
      info!("Starting up");
      util::setup_metrics(&config);

      config
    })
    .and_then(|config| check_process_limits(&config).map(|()| config))
    .and_then(|config| init_workers(&config).map(|workers| (config, workers)))
    .map(|(config, workers)| {
      if config.handle_process_affinity {
        set_workers_affinity(&workers);
      }
      let command_socket_path = config.command_socket_path();
      command::start(config, command_socket_path, workers);
    });

    match start {
      Ok(_) => info!("master process stopped"), // Ok() is only called when the proxy exits
      Err(StartupError::ConfigurationFileNotSpecified) => {
        error!("Configuration file hasn't been specified. Either use -c with the start command \
               or use the SOZU_CONFIG environment variable when building sozu.");
        std::process::exit(1);
      },
      Err(StartupError::ConfigurationFileLoadError(err)) => {
        error!("Invalid configuration file. Error: {:?}", err);
        std::process::exit(1);
      },
      Err(StartupError::TooManyAllowedConnections(err)) | Err(StartupError::TooManyAllowedConnectionsForWorker(err)) => {
        error!("{}", err);
        std::process::exit(1);
      },
      Err(StartupError::WorkersSpawnFail(err)) => {
        error!("At least one worker failed to spawn. Error: {:?}", err);
        std::process::exit(1);
      },
      Err(StartupError::PIDFileNotWritable(err)) => {
        error!("{}", err);
        std::process::exit(1);
      }
    }
  }
}

fn init_workers(config: &Config) -> Result<Vec<Worker>, StartupError> {
  let path = unsafe { get_executable_path() };
  match start_workers(path, &config) {
    Ok(workers) => {
      info!("created workers: {:?}", workers);
      Ok(workers)
    },
    Err(e) => Err(StartupError::WorkersSpawnFail(e))
  }
}

fn get_config_file_path<'a>(matches: &'a ArgMatches<'a>) -> Result<&'a str, StartupError> {
  let start_matches = matches.subcommand_matches("start").expect("unknown subcommand");
  match start_matches.value_of("config") {
    Some(config_file) => Ok(config_file),
    None => option_env!("SOZU_CONFIG").ok_or(StartupError::ConfigurationFileNotSpecified)
  }
}

fn load_configuration(config_file: &str) -> Result<Config, StartupError> {
  match Config::load_from_path(config_file) {
    Ok(config) => Ok(config),
    Err(e) => Err(StartupError::ConfigurationFileLoadError(e))
  }
}

/// Set workers process affinity, see man sched_setaffinity
/// Bind each worker (including the master) process to a CPU core.
/// Can bind multiple processes to a CPU core if there are more processes
/// than CPU cores. Only works on Linux.
#[cfg(target_os = "linux")]
fn set_workers_affinity(workers: &Vec<Worker>) {
  let mut cpu_count = 0;
  let max_cpu = num_cpus::get();

  // +1 for the master process that will also be bound to its CPU core
  if (workers.len() + 1) > max_cpu {
    warn!("There are more workers than available CPU cores, \
          multiple workers will be bound to the same CPU core. \
          This may impact performances");
  }

  let master_pid = unsafe { libc::getpid() };
  set_process_affinity(master_pid, cpu_count);
  cpu_count = cpu_count + 1;

  for ref worker in workers {
    if cpu_count >= max_cpu {
      cpu_count = 0;
    }

    set_process_affinity(worker.pid, cpu_count);

    cpu_count = cpu_count + 1;
  }
}

/// Set workers process affinity, see man sched_setaffinity
/// Bind each worker (including the master) process to a CPU core.
/// Can bind multiple processes to a CPU core if there are more processes
/// than CPU cores. Only works on Linux.
#[cfg(not(target_os = "linux"))]
fn set_workers_affinity(_: &Vec<Worker>) {
}

/// Set a specific process to run onto a specific CPU core
#[cfg(target_os = "linux")]
use std::mem;
#[cfg(target_os = "linux")]
fn set_process_affinity(pid: pid_t, cpu: usize) {
  unsafe {
    let mut cpu_set: cpu_set_t = mem::zeroed();
    let size_cpu_set = mem::size_of::<cpu_set_t>();
    libc::CPU_SET(cpu, &mut cpu_set);
    libc::sched_setaffinity(pid, size_cpu_set, &mut cpu_set);

    debug!("Worker {} bound to CPU core {}", pid, cpu);
  };
}

#[cfg(target_os="linux")]
// We check the hard_limit. The soft_limit can be changed at runtime
// by the process or any user. hard_limit can only be changed by root
fn check_process_limits(config: &Config) -> Result<(), StartupError> {
  let process_limits_file = sozu_command::config::Config::load_file("/proc/self/limits")
    .expect("Couldn't read /proc/self/limits to determine max open file descriptors limit");

  let re = Regex::new(r"Max open files\s*(\d*)\s*(\d*)\s*files").unwrap();
  let res = re.captures(&process_limits_file);

  if let Some(soft_limit) = res.as_ref().and_then(|c| c.get(1))
    .and_then(|m| m.as_str().parse::<usize>().ok()) {
    if config.max_connections * 2 > soft_limit {
      let error = format!("at least one worker can't have that many connections. \
              current max file descriptor soft limit is: {}, \
              configured max_connections is {} (the worker needs two file descriptors \
              per client connection)", soft_limit, config.max_connections);
      return Err(StartupError::TooManyAllowedConnectionsForWorker(error));
    }
  }
  if let Some(hard_limit) = res.as_ref().and_then(|c| c.get(2))
    .and_then(|m| m.as_str().parse::<usize>().ok()) {
    if config.max_connections * 2 > hard_limit {
      let error = format!("at least one worker can't have that many connections. \
              current max file descriptor hard limit is: {}, \
              configured max_connections is {} (the worker needs two file descriptors \
              per client connection)", hard_limit, config.max_connections);
      return Err(StartupError::TooManyAllowedConnectionsForWorker(error));
    }
  }

  let f = sozu_command::config::Config::load_file("/proc/sys/fs/file-max")
    .expect("Couldn't read /proc/sys/fs/file-max");

  let re_max = Regex::new(r"(\d*)").unwrap();

  let system_max_fd = re_max.captures(&f).and_then(|c| c.get(1))
    .and_then(|m| m.as_str().parse::<usize>().ok())
    .expect("Couldn't parse /proc/sys/fs/file-max");

  if config.max_connections > system_max_fd {
    let error = format!("Proxies total max_connections can't be higher than system's file-max limit. \
            Current limit: {}, current value: {}", system_max_fd, config.max_connections);
    return Err(StartupError::TooManyAllowedConnections(error))
  }

  Ok(())
}

#[cfg(not(target_os = "linux"))]
fn check_process_limits(_: &Config) -> Result<(), StartupError> {
  Ok(())
}

fn register_panic_hook() {
  // We save the original panic hook so we can call it later
  // to have the original behavior
  let original_panic_hook = panic::take_hook();

  panic::set_hook(Box::new(move |panic_info| {
    incr!("panic");
    METRICS.with(|metrics| {
      (*metrics.borrow_mut()).send_data();
    });

    (*original_panic_hook)(panic_info)
  }));
}
<|endoftext|>


<commit_before>use crate::back::write::create_informational_target_machine;
use crate::llvm;
use libc::c_int;
use rustc_codegen_ssa::target_features::supported_target_features;
use rustc_data_structures::fx::FxHashSet;
use rustc_middle::bug;
use rustc_session::config::PrintRequest;
use rustc_session::Session;
use rustc_span::symbol::Symbol;
use rustc_target::spec::{MergeFunctions, PanicStrategy};
use std::ffi::{CStr, CString};

use std::slice;
use std::str;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Once;

static POISONED: AtomicBool = AtomicBool::new(false);
static INIT: Once = Once::new();

pub(crate) fn init(sess: &Session) {
    unsafe {
        // Before we touch LLVM, make sure that multithreading is enabled.
        INIT.call_once(|| {
            if llvm::LLVMStartMultithreaded() != 1 {
                // use an extra bool to make sure that all future usage of LLVM
                // cannot proceed despite the Once not running more than once.
                POISONED.store(true, Ordering::SeqCst);
            }

            configure_llvm(sess);
        });

        if POISONED.load(Ordering::SeqCst) {
            bug!("couldn't enable multi-threaded LLVM");
        }
    }
}

fn require_inited() {
    INIT.call_once(|| bug!("llvm is not initialized"));
    if POISONED.load(Ordering::SeqCst) {
        bug!("couldn't enable multi-threaded LLVM");
    }
}

unsafe fn configure_llvm(sess: &Session) {
    let n_args = sess.opts.cg.llvm_args.len() + sess.target.llvm_args.len();
    let mut llvm_c_strs = Vec::with_capacity(n_args + 1);
    let mut llvm_args = Vec::with_capacity(n_args + 1);

    llvm::LLVMRustInstallFatalErrorHandler();

    fn llvm_arg_to_arg_name(full_arg: &str) -> &str {
        full_arg.trim().split(|c: char| c == '=' || c.is_whitespace()).next().unwrap_or("")
    }

    let cg_opts = sess.opts.cg.llvm_args.iter();
    let tg_opts = sess.target.llvm_args.iter();
    let sess_args = cg_opts.chain(tg_opts);

    let user_specified_args: FxHashSet<_> =
        sess_args.clone().map(|s| llvm_arg_to_arg_name(s)).filter(|s| !s.is_empty()).collect();

    {
        // This adds the given argument to LLVM. Unless `force` is true
        // user specified arguments are *not* overridden.
        let mut add = |arg: &str, force: bool| {
            if force || !user_specified_args.contains(llvm_arg_to_arg_name(arg)) {
                let s = CString::new(arg).unwrap();
                llvm_args.push(s.as_ptr());
                llvm_c_strs.push(s);
            }
        };
        // Set the llvm "program name" to make usage and invalid argument messages more clear.
        add("rustc -Cllvm-args=\"...\" with", true);
        if sess.time_llvm_passes() {
            add("-time-passes", false);
        }
        if sess.print_llvm_passes() {
            add("-debug-pass=Structure", false);
        }
        if !sess.opts.debugging_opts.no_generate_arange_section {
            add("-generate-arange-section", false);
        }
        match sess.opts.debugging_opts.merge_functions.unwrap_or(sess.target.merge_functions) {
            MergeFunctions::Disabled | MergeFunctions::Trampolines => {}
            MergeFunctions::Aliases => {
                add("-mergefunc-use-aliases", false);
            }
        }

        if sess.target.os == "emscripten" && sess.panic_strategy() == PanicStrategy::Unwind {
            add("-enable-emscripten-cxx-exceptions", false);
        }

        // HACK(eddyb) LLVM inserts `llvm.assume` calls to preserve align attributes
        // during inlining. Unfortunately these may block other optimizations.
        add("-preserve-alignment-assumptions-during-inlining=false", false);

        for arg in sess_args {
            add(&(*arg), true);
        }
    }

    if sess.opts.debugging_opts.llvm_time_trace {
        // time-trace is not thread safe and running it in parallel will cause seg faults.
        if !sess.opts.debugging_opts.no_parallel_llvm {
            bug!("`-Z llvm-time-trace` requires `-Z no-parallel-llvm")
        }

        llvm::LLVMTimeTraceProfilerInitialize();
    }

    llvm::LLVMInitializePasses();

    rustc_llvm::initialize_available_targets();

    llvm::LLVMRustSetLLVMOptions(llvm_args.len() as c_int, llvm_args.as_ptr());
}

pub fn time_trace_profiler_finish(file_name: &str) {
    unsafe {
        let file_name = CString::new(file_name).unwrap();
        llvm::LLVMTimeTraceProfilerFinish(file_name.as_ptr());
    }
}

// WARNING: the features after applying `to_llvm_feature` must be known
// to LLVM or the feature detection code will walk past the end of the feature
// array, leading to crashes.
// To find a list of LLVM's names, check llvm-project/llvm/include/llvm/Support/*TargetParser.def
// where the * matches the architecture's name
// Beware to not use the llvm github project for this, but check the git submodule
// found in src/llvm-project
// Though note that Rust can also be build with an external precompiled version of LLVM
// which might lead to failures if the oldest tested / supported LLVM version
// doesn't yet support the relevant intrinsics
pub fn to_llvm_feature<'a>(sess: &Session, s: &'a str) -> &'a str {
    let arch = if sess.target.arch == "x86_64" { "x86" } else { &*sess.target.arch };
    match (arch, s) {
        ("x86", "pclmulqdq") => "pclmul",
        ("x86", "rdrand") => "rdrnd",
        ("x86", "bmi1") => "bmi",
        ("x86", "cmpxchg16b") => "cx16",
        ("x86", "avx512vaes") => "vaes",
        ("x86", "avx512gfni") => "gfni",
        ("x86", "avx512vpclmulqdq") => "vpclmulqdq",
        ("aarch64", "fp") => "fp-armv8",
        ("aarch64", "fp16") => "fullfp16",
        (_, s) => s,
    }
}

pub fn target_features(sess: &Session) -> Vec<Symbol> {
    let target_machine = create_informational_target_machine(sess);
    supported_target_features(sess)
        .iter()
        .filter_map(
            |&(feature, gate)| {
                if sess.is_nightly_build() || gate.is_none() { Some(feature) } else { None }
            },
        )
        .filter(|feature| {
            let llvm_feature = to_llvm_feature(sess, feature);
            let cstr = CString::new(llvm_feature).unwrap();
            unsafe { llvm::LLVMRustHasFeature(target_machine, cstr.as_ptr()) }
        })
        .map(|feature| Symbol::intern(feature))
        .collect()
}

pub fn print_version() {
    // Can be called without initializing LLVM
    unsafe {
        println!("LLVM version: {}.{}", llvm::LLVMRustVersionMajor(), llvm::LLVMRustVersionMinor());
    }
}

pub fn get_major_version() -> u32 {
    unsafe { llvm::LLVMRustVersionMajor() }
}

pub fn print_passes() {
    // Can be called without initializing LLVM
    unsafe {
        llvm::LLVMRustPrintPasses();
    }
}

pub(crate) fn print(req: PrintRequest, sess: &Session) {
    require_inited();
    let tm = create_informational_target_machine(sess);
    unsafe {
        match req {
            PrintRequest::TargetCPUs => llvm::LLVMRustPrintTargetCPUs(tm),
            PrintRequest::TargetFeatures => llvm::LLVMRustPrintTargetFeatures(tm),
            _ => bug!("rustc_codegen_llvm can't handle print request: {:?}", req),
        }
    }
}

fn handle_native(name: &str) -> &str {
    if name != "native" {
        return name;
    }

    unsafe {
        let mut len = 0;
        let ptr = llvm::LLVMRustGetHostCPUName(&mut len);
        str::from_utf8(slice::from_raw_parts(ptr as *const u8, len)).unwrap()
    }
}

pub fn target_cpu(sess: &Session) -> &str {
    let name = match sess.opts.cg.target_cpu {
        Some(ref s) => &**s,
        None => &*sess.target.cpu,
    };

    handle_native(name)
}

pub fn handle_native_features(sess: &Session) -> Vec<String> {
    match sess.opts.cg.target_cpu {
        Some(ref s) => {
            if s != "native" {
                return vec![];
            }

            let ptr = unsafe { llvm::LLVMGetHostCPUFeatures() };
            let str = unsafe { CStr::from_ptr(ptr).to_string_lossy() };

            let features = str.split(",").map(|s| s.to_owned()).collect();

            unsafe { llvm::LLVMDisposeMessage(ptr) };

            features
        }
        None => vec![],
    }
}

pub fn tune_cpu(sess: &Session) -> Option<&str> {
    match sess.opts.debugging_opts.tune_cpu {
        Some(ref s) => Some(handle_native(&**s)),
        None => None,
    }
}
<commit_msg>Check if the pointer is null/string is not utf8<commit_after>use crate::back::write::create_informational_target_machine;
use crate::llvm;
use libc::c_int;
use rustc_codegen_ssa::target_features::supported_target_features;
use rustc_data_structures::fx::FxHashSet;
use rustc_middle::bug;
use rustc_session::config::PrintRequest;
use rustc_session::Session;
use rustc_span::symbol::Symbol;
use rustc_target::spec::{MergeFunctions, PanicStrategy};
use std::ffi::{CStr, CString};

use std::slice;
use std::str;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Once;

static POISONED: AtomicBool = AtomicBool::new(false);
static INIT: Once = Once::new();

pub(crate) fn init(sess: &Session) {
    unsafe {
        // Before we touch LLVM, make sure that multithreading is enabled.
        INIT.call_once(|| {
            if llvm::LLVMStartMultithreaded() != 1 {
                // use an extra bool to make sure that all future usage of LLVM
                // cannot proceed despite the Once not running more than once.
                POISONED.store(true, Ordering::SeqCst);
            }

            configure_llvm(sess);
        });

        if POISONED.load(Ordering::SeqCst) {
            bug!("couldn't enable multi-threaded LLVM");
        }
    }
}

fn require_inited() {
    INIT.call_once(|| bug!("llvm is not initialized"));
    if POISONED.load(Ordering::SeqCst) {
        bug!("couldn't enable multi-threaded LLVM");
    }
}

unsafe fn configure_llvm(sess: &Session) {
    let n_args = sess.opts.cg.llvm_args.len() + sess.target.llvm_args.len();
    let mut llvm_c_strs = Vec::with_capacity(n_args + 1);
    let mut llvm_args = Vec::with_capacity(n_args + 1);

    llvm::LLVMRustInstallFatalErrorHandler();

    fn llvm_arg_to_arg_name(full_arg: &str) -> &str {
        full_arg.trim().split(|c: char| c == '=' || c.is_whitespace()).next().unwrap_or("")
    }

    let cg_opts = sess.opts.cg.llvm_args.iter();
    let tg_opts = sess.target.llvm_args.iter();
    let sess_args = cg_opts.chain(tg_opts);

    let user_specified_args: FxHashSet<_> =
        sess_args.clone().map(|s| llvm_arg_to_arg_name(s)).filter(|s| !s.is_empty()).collect();

    {
        // This adds the given argument to LLVM. Unless `force` is true
        // user specified arguments are *not* overridden.
        let mut add = |arg: &str, force: bool| {
            if force || !user_specified_args.contains(llvm_arg_to_arg_name(arg)) {
                let s = CString::new(arg).unwrap();
                llvm_args.push(s.as_ptr());
                llvm_c_strs.push(s);
            }
        };
        // Set the llvm "program name" to make usage and invalid argument messages more clear.
        add("rustc -Cllvm-args=\"...\" with", true);
        if sess.time_llvm_passes() {
            add("-time-passes", false);
        }
        if sess.print_llvm_passes() {
            add("-debug-pass=Structure", false);
        }
        if !sess.opts.debugging_opts.no_generate_arange_section {
            add("-generate-arange-section", false);
        }
        match sess.opts.debugging_opts.merge_functions.unwrap_or(sess.target.merge_functions) {
            MergeFunctions::Disabled | MergeFunctions::Trampolines => {}
            MergeFunctions::Aliases => {
                add("-mergefunc-use-aliases", false);
            }
        }

        if sess.target.os == "emscripten" && sess.panic_strategy() == PanicStrategy::Unwind {
            add("-enable-emscripten-cxx-exceptions", false);
        }

        // HACK(eddyb) LLVM inserts `llvm.assume` calls to preserve align attributes
        // during inlining. Unfortunately these may block other optimizations.
        add("-preserve-alignment-assumptions-during-inlining=false", false);

        for arg in sess_args {
            add(&(*arg), true);
        }
    }

    if sess.opts.debugging_opts.llvm_time_trace {
        // time-trace is not thread safe and running it in parallel will cause seg faults.
        if !sess.opts.debugging_opts.no_parallel_llvm {
            bug!("`-Z llvm-time-trace` requires `-Z no-parallel-llvm")
        }

        llvm::LLVMTimeTraceProfilerInitialize();
    }

    llvm::LLVMInitializePasses();

    rustc_llvm::initialize_available_targets();

    llvm::LLVMRustSetLLVMOptions(llvm_args.len() as c_int, llvm_args.as_ptr());
}

pub fn time_trace_profiler_finish(file_name: &str) {
    unsafe {
        let file_name = CString::new(file_name).unwrap();
        llvm::LLVMTimeTraceProfilerFinish(file_name.as_ptr());
    }
}

// WARNING: the features after applying `to_llvm_feature` must be known
// to LLVM or the feature detection code will walk past the end of the feature
// array, leading to crashes.
// To find a list of LLVM's names, check llvm-project/llvm/include/llvm/Support/*TargetParser.def
// where the * matches the architecture's name
// Beware to not use the llvm github project for this, but check the git submodule
// found in src/llvm-project
// Though note that Rust can also be build with an external precompiled version of LLVM
// which might lead to failures if the oldest tested / supported LLVM version
// doesn't yet support the relevant intrinsics
pub fn to_llvm_feature<'a>(sess: &Session, s: &'a str) -> &'a str {
    let arch = if sess.target.arch == "x86_64" { "x86" } else { &*sess.target.arch };
    match (arch, s) {
        ("x86", "pclmulqdq") => "pclmul",
        ("x86", "rdrand") => "rdrnd",
        ("x86", "bmi1") => "bmi",
        ("x86", "cmpxchg16b") => "cx16",
        ("x86", "avx512vaes") => "vaes",
        ("x86", "avx512gfni") => "gfni",
        ("x86", "avx512vpclmulqdq") => "vpclmulqdq",
        ("aarch64", "fp") => "fp-armv8",
        ("aarch64", "fp16") => "fullfp16",
        (_, s) => s,
    }
}

pub fn target_features(sess: &Session) -> Vec<Symbol> {
    let target_machine = create_informational_target_machine(sess);
    supported_target_features(sess)
        .iter()
        .filter_map(
            |&(feature, gate)| {
                if sess.is_nightly_build() || gate.is_none() { Some(feature) } else { None }
            },
        )
        .filter(|feature| {
            let llvm_feature = to_llvm_feature(sess, feature);
            let cstr = CString::new(llvm_feature).unwrap();
            unsafe { llvm::LLVMRustHasFeature(target_machine, cstr.as_ptr()) }
        })
        .map(|feature| Symbol::intern(feature))
        .collect()
}

pub fn print_version() {
    // Can be called without initializing LLVM
    unsafe {
        println!("LLVM version: {}.{}", llvm::LLVMRustVersionMajor(), llvm::LLVMRustVersionMinor());
    }
}

pub fn get_major_version() -> u32 {
    unsafe { llvm::LLVMRustVersionMajor() }
}

pub fn print_passes() {
    // Can be called without initializing LLVM
    unsafe {
        llvm::LLVMRustPrintPasses();
    }
}

pub(crate) fn print(req: PrintRequest, sess: &Session) {
    require_inited();
    let tm = create_informational_target_machine(sess);
    unsafe {
        match req {
            PrintRequest::TargetCPUs => llvm::LLVMRustPrintTargetCPUs(tm),
            PrintRequest::TargetFeatures => llvm::LLVMRustPrintTargetFeatures(tm),
            _ => bug!("rustc_codegen_llvm can't handle print request: {:?}", req),
        }
    }
}

fn handle_native(name: &str) -> &str {
    if name != "native" {
        return name;
    }

    unsafe {
        let mut len = 0;
        let ptr = llvm::LLVMRustGetHostCPUName(&mut len);
        str::from_utf8(slice::from_raw_parts(ptr as *const u8, len)).unwrap()
    }
}

pub fn target_cpu(sess: &Session) -> &str {
    let name = match sess.opts.cg.target_cpu {
        Some(ref s) => &**s,
        None => &*sess.target.cpu,
    };

    handle_native(name)
}

pub fn handle_native_features(sess: &Session) -> Vec<String> {
    match sess.opts.cg.target_cpu {
        Some(ref s) => {
            if s != "native" {
                return vec![];
            }

            let features_string = unsafe {
                let ptr = llvm::LLVMGetHostCPUFeatures();
                let features_string = if !ptr.is_null() {
                    CStr::from_ptr(ptr)
                        .to_str()
                        .unwrap_or_else(|e| {
                            bug!("LLVM returned a non-utf8 features string: {}", e);
                        })
                        .to_owned()
                } else {
                    bug!("could not allocate host CPU features, LLVM returned a `null` string");
                };

                llvm::LLVMDisposeMessage(ptr);

                features_string
            };

            features_string.split(",").map(|s| s.to_owned()).collect()
        }
        None => vec![],
    }
}

pub fn tune_cpu(sess: &Session) -> Option<&str> {
    match sess.opts.debugging_opts.tune_cpu {
        Some(ref s) => Some(handle_native(&**s)),
        None => None,
    }
}
<|endoftext|>


<commit_before>use proc_macro2::TokenStream;
use quote::format_ident;

use internals::ast::{Container, Data, Field, Style, Variant};

// Suppress dead_code warnings that would otherwise appear when using a remote
// derive. Other than this pretend code, a struct annotated with remote derive
// never has its fields referenced and an enum annotated with remote derive
// never has its variants constructed.
//
//     warning: field is never used: `i`
//      --> src/main.rs:4:20
//       |
//     4 | struct StructDef { i: i32 }
//       |                    ^^^^^^
//
//     warning: variant is never constructed: `V`
//      --> src/main.rs:8:16
//       |
//     8 | enum EnumDef { V }
//       |                ^
//
pub fn pretend_used(cont: &Container, is_packed: bool) -> TokenStream {
    let pretend_fields = pretend_fields_used(cont, is_packed);
    let pretend_variants = pretend_variants_used(cont);

    quote! {
        #pretend_fields
        #pretend_variants
    }
}

// For structs with named fields, expands to:
//
//     match None::<&T> {
//         Some(T { a: __v0, b: __v1 }) => {}
//         _ => {}
//     }
//
// For packed structs on sufficiently new rustc, expands to:
//
//     match None::<&T> {
//         Some(__v @ T { a: _, b: _ }) => {
//             let _ = addr_of!(__v.a);
//             let _ = addr_of!(__v.b);
//         }
//         _ => {}
//     }
//
// For packed structs on older rustc, we assume Sized and !Drop, and expand to:
//
//     match None::<T> {
//         Some(T { a: __v0, b: __v1 }) => {}
//         _ => {}
//     }
//
// For enums, expands to the following but only including struct variants:
//
//     match None::<&T> {
//         Some(T::A { a: __v0 }) => {}
//         Some(T::B { b: __v0 }) => {}
//         _ => {}
//     }
//
fn pretend_fields_used(cont: &Container, is_packed: bool) -> TokenStream {
    match &cont.data {
        Data::Enum(variants) => pretend_fields_used_enum(cont, variants),
        Data::Struct(Style::Struct, fields) => if is_packed {
            pretend_fields_used_struct_packed(cont, fields)
        } else {
            pretend_fields_used_struct(cont, fields)
        },
        Data::Struct(_, _) => quote!(),
    }
}

fn pretend_fields_used_struct(cont: &Container, fields: &[Field]) -> TokenStream {
    let type_ident = &cont.ident;
    let (_, ty_generics, _) = cont.generics.split_for_impl();

    let members = fields.iter().map(|field| &field.member);
    let placeholders = (0usize..).map(|i| format_ident!("__v{}", i));

    quote! {
        match _serde::__private::None::<&#type_ident #ty_generics> {
            _serde::__private::Some(#type_ident { #(#members: #placeholders),* }) => {}
            _ => {}
        }
    }
}

fn pretend_fields_used_struct_packed(cont: &Container, fields: &[Field]) -> TokenStream {
    let type_ident = &cont.ident;
    let (_, ty_generics, _) = cont.generics.split_for_impl();

    let members = fields.iter().map(|field| &field.member).collect::<Vec<_>>();

    #[cfg(ptr_addr_of)]
    {
        quote! {
            match _serde::__private::None::<&#type_ident #ty_generics> {
                _serde::__private::Some(__v @ #type_ident { #(#members: _),* }) => {
                    #(
                        let _ = _serde::__private::ptr::addr_of!(__v.#members);
                    )*
                }
                _ => {}
            }
        }
    }

    #[cfg(not(ptr_addr_of))]
    {
        let placeholders = (0usize..).map(|i| format_ident!("__v{}", i));

        quote! {
            match _serde::__private::None::<#type_ident #ty_generics> {
                _serde::__private::Some(#type_ident { #(#members: #placeholders),* }) => {}
                _ => {}
            }
        }
    }
}

fn pretend_fields_used_enum(cont: &Container, variants: &[Variant]) -> TokenStream {
    let type_ident = &cont.ident;
    let (_, ty_generics, _) = cont.generics.split_for_impl();

    let patterns = variants
        .iter()
        .filter_map(|variant| match variant.style {
            Style::Struct => {
                let variant_ident = &variant.ident;
                let members = variant.fields.iter().map(|field| &field.member);
                let placeholders = (0usize..).map(|i| format_ident!("__v{}", i));
                Some(quote!(#type_ident::#variant_ident { #(#members: #placeholders),* }))
            }
            _ => None,
        })
        .collect::<Vec<_>>();

    quote! {
        match _serde::__private::None::<&#type_ident #ty_generics> {
            #(
                _serde::__private::Some(#patterns) => {}
            )*
            _ => {}
        }
    }
}

// Expands to one of these per enum variant:
//
//     match None {
//         Some((__v0, __v1,)) => {
//             let _ = E::V { a: __v0, b: __v1 };
//         }
//         _ => {}
//     }
//
fn pretend_variants_used(cont: &Container) -> TokenStream {
    let variants = match &cont.data {
        Data::Enum(variants) => variants,
        Data::Struct(_, _) => {
            return quote!();
        }
    };

    let type_ident = &cont.ident;
    let (_, ty_generics, _) = cont.generics.split_for_impl();
    let turbofish = ty_generics.as_turbofish();

    let cases = variants.iter().map(|variant| {
        let variant_ident = &variant.ident;
        let placeholders = &(0..variant.fields.len())
            .map(|i| format_ident!("__v{}", i))
            .collect::<Vec<_>>();

        let pat = match variant.style {
            Style::Struct => {
                let members = variant.fields.iter().map(|field| &field.member);
                quote!({ #(#members: #placeholders),* })
            }
            Style::Tuple | Style::Newtype => quote!(( #(#placeholders),* )),
            Style::Unit => quote!(),
        };

        quote! {
            match _serde::__private::None {
                _serde::__private::Some((#(#placeholders,)*)) => {
                    let _ = #type_ident::#variant_ident #turbofish #pat;
                }
                _ => {}
            }
        }
    });

    quote!(#(#cases)*)
}
<commit_msg>Format with rustfmt 2021-09-29<commit_after>use proc_macro2::TokenStream;
use quote::format_ident;

use internals::ast::{Container, Data, Field, Style, Variant};

// Suppress dead_code warnings that would otherwise appear when using a remote
// derive. Other than this pretend code, a struct annotated with remote derive
// never has its fields referenced and an enum annotated with remote derive
// never has its variants constructed.
//
//     warning: field is never used: `i`
//      --> src/main.rs:4:20
//       |
//     4 | struct StructDef { i: i32 }
//       |                    ^^^^^^
//
//     warning: variant is never constructed: `V`
//      --> src/main.rs:8:16
//       |
//     8 | enum EnumDef { V }
//       |                ^
//
pub fn pretend_used(cont: &Container, is_packed: bool) -> TokenStream {
    let pretend_fields = pretend_fields_used(cont, is_packed);
    let pretend_variants = pretend_variants_used(cont);

    quote! {
        #pretend_fields
        #pretend_variants
    }
}

// For structs with named fields, expands to:
//
//     match None::<&T> {
//         Some(T { a: __v0, b: __v1 }) => {}
//         _ => {}
//     }
//
// For packed structs on sufficiently new rustc, expands to:
//
//     match None::<&T> {
//         Some(__v @ T { a: _, b: _ }) => {
//             let _ = addr_of!(__v.a);
//             let _ = addr_of!(__v.b);
//         }
//         _ => {}
//     }
//
// For packed structs on older rustc, we assume Sized and !Drop, and expand to:
//
//     match None::<T> {
//         Some(T { a: __v0, b: __v1 }) => {}
//         _ => {}
//     }
//
// For enums, expands to the following but only including struct variants:
//
//     match None::<&T> {
//         Some(T::A { a: __v0 }) => {}
//         Some(T::B { b: __v0 }) => {}
//         _ => {}
//     }
//
fn pretend_fields_used(cont: &Container, is_packed: bool) -> TokenStream {
    match &cont.data {
        Data::Enum(variants) => pretend_fields_used_enum(cont, variants),
        Data::Struct(Style::Struct, fields) => {
            if is_packed {
                pretend_fields_used_struct_packed(cont, fields)
            } else {
                pretend_fields_used_struct(cont, fields)
            }
        }
        Data::Struct(_, _) => quote!(),
    }
}

fn pretend_fields_used_struct(cont: &Container, fields: &[Field]) -> TokenStream {
    let type_ident = &cont.ident;
    let (_, ty_generics, _) = cont.generics.split_for_impl();

    let members = fields.iter().map(|field| &field.member);
    let placeholders = (0usize..).map(|i| format_ident!("__v{}", i));

    quote! {
        match _serde::__private::None::<&#type_ident #ty_generics> {
            _serde::__private::Some(#type_ident { #(#members: #placeholders),* }) => {}
            _ => {}
        }
    }
}

fn pretend_fields_used_struct_packed(cont: &Container, fields: &[Field]) -> TokenStream {
    let type_ident = &cont.ident;
    let (_, ty_generics, _) = cont.generics.split_for_impl();

    let members = fields.iter().map(|field| &field.member).collect::<Vec<_>>();

    #[cfg(ptr_addr_of)]
    {
        quote! {
            match _serde::__private::None::<&#type_ident #ty_generics> {
                _serde::__private::Some(__v @ #type_ident { #(#members: _),* }) => {
                    #(
                        let _ = _serde::__private::ptr::addr_of!(__v.#members);
                    )*
                }
                _ => {}
            }
        }
    }

    #[cfg(not(ptr_addr_of))]
    {
        let placeholders = (0usize..).map(|i| format_ident!("__v{}", i));

        quote! {
            match _serde::__private::None::<#type_ident #ty_generics> {
                _serde::__private::Some(#type_ident { #(#members: #placeholders),* }) => {}
                _ => {}
            }
        }
    }
}

fn pretend_fields_used_enum(cont: &Container, variants: &[Variant]) -> TokenStream {
    let type_ident = &cont.ident;
    let (_, ty_generics, _) = cont.generics.split_for_impl();

    let patterns = variants
        .iter()
        .filter_map(|variant| match variant.style {
            Style::Struct => {
                let variant_ident = &variant.ident;
                let members = variant.fields.iter().map(|field| &field.member);
                let placeholders = (0usize..).map(|i| format_ident!("__v{}", i));
                Some(quote!(#type_ident::#variant_ident { #(#members: #placeholders),* }))
            }
            _ => None,
        })
        .collect::<Vec<_>>();

    quote! {
        match _serde::__private::None::<&#type_ident #ty_generics> {
            #(
                _serde::__private::Some(#patterns) => {}
            )*
            _ => {}
        }
    }
}

// Expands to one of these per enum variant:
//
//     match None {
//         Some((__v0, __v1,)) => {
//             let _ = E::V { a: __v0, b: __v1 };
//         }
//         _ => {}
//     }
//
fn pretend_variants_used(cont: &Container) -> TokenStream {
    let variants = match &cont.data {
        Data::Enum(variants) => variants,
        Data::Struct(_, _) => {
            return quote!();
        }
    };

    let type_ident = &cont.ident;
    let (_, ty_generics, _) = cont.generics.split_for_impl();
    let turbofish = ty_generics.as_turbofish();

    let cases = variants.iter().map(|variant| {
        let variant_ident = &variant.ident;
        let placeholders = &(0..variant.fields.len())
            .map(|i| format_ident!("__v{}", i))
            .collect::<Vec<_>>();

        let pat = match variant.style {
            Style::Struct => {
                let members = variant.fields.iter().map(|field| &field.member);
                quote!({ #(#members: #placeholders),* })
            }
            Style::Tuple | Style::Newtype => quote!(( #(#placeholders),* )),
            Style::Unit => quote!(),
        };

        quote! {
            match _serde::__private::None {
                _serde::__private::Some((#(#placeholders,)*)) => {
                    let _ = #type_ident::#variant_ident #turbofish #pat;
                }
                _ => {}
            }
        }
    });

    quote!(#(#cases)*)
}
<|endoftext|>


<commit_before>use diesel;
use diesel::prelude::*;

use rocket_contrib::{JSON, SerdeError};

use helpers::db::DB;
use validation::user::{UserSerializer, UserSettingsSerializer};

use schema::{users, pods, queues};
use schema::users::dsl::*;

use models::user::{UserModel, NewUser, ChangedUser};
use models::pod::{PodModel, NewPod};
use models::queue::{QueueModel, NewQueue};

use responses::{
    APIResponse,
    ok,
    created,
    conflict,
    bad_request,
    forbidden,
    unauthorized,
};

#[get("/info")]
pub fn info(current_user: UserModel) -> APIResponse {
    ok().data(json!(current_user))
}

#[post("/register", data = "<user_data>", format = "application/json")]
pub fn register(user_data: Result<JSON<UserSerializer>, SerdeError>, db: DB) -> APIResponse {

    // Return specific error if invalid JSON has been sent.
    if user_data.is_err() {
        return bad_request().message(format!("{}", user_data.err().unwrap()).as_str());
    }

    let user_data = user_data.unwrap();

    // Check for existing user email
    let results = users.filter(email.eq(user_data.email.clone()))
        .first::<UserModel>(&*db);
    if results.is_ok() {
        return conflict().message("User with this email already exists.");
    }

    // Check for existing nickname
    let results = users.filter(nickname.eq(user_data.nickname.clone()))
        .first::<UserModel>(&*db);
    if results.is_ok() {
        return conflict().message("Nickname already taken.");
    }

    // Create new password hash
    let new_password_hash = UserModel::make_password_hash(user_data.password.as_str());

    // New user model for table insertion
    let new_user = NewUser {
        nickname: user_data.nickname.clone(),
        email: user_data.email.clone(),
        password_hash: new_password_hash,
    };

    // Insert user to get id for pod
    let user = diesel::insert(&new_user)
        .into(users::table)
        .get_result::<UserModel>(&*db)
        .expect("Error saving new user");

    // New pod
    let new_pod = NewPod {
        name: format!("{}'s Pod", user.nickname.clone()),
        user_id: user.id.clone(),
    };

    let pod = diesel::insert(&new_pod)
        .into(pods::table)
        .get_result::<PodModel>(&*db)
        .expect("Error creating pod");

    let new_queue = NewQueue {
        slots: 2,
        pod_id: Some(pod.id.clone()),
        base_id: None,
    };

    diesel::insert(&new_queue)
        .into(queues::table)
        .get_result::<QueueModel>(&*db)
        .expect("Error creating pod's queue");

    created().message("User created.").data(json!(&user))
}

#[post("/settings", data = "<user_data>", format = "application/json")]
pub fn settings(current_user: UserModel, user_data: Result<JSON<UserSettingsSerializer>, SerdeError>, db: DB) -> APIResponse {

    // Return specific error if invalid JSON has been sent.
    if user_data.is_err() {
        return bad_request().message(format!("{}", user_data.err().unwrap()).as_str());
    }

    let user_data = user_data.unwrap();

    let mut new_password_hash: Option<String> = None;

    if user_data.new_password.is_some() {
        if user_data.password.is_none() {
            return forbidden().message("The current passwords needs to be specified, if you want to change your password.");
        }
        if !current_user.verify_password(user_data.password.as_ref().unwrap().as_str()) {
            return unauthorized().message("Password incorrect.");
        }

        // Create new password hash
        new_password_hash = Some(UserModel::make_password_hash(user_data.password.as_ref().unwrap().as_str()));

    }

    let changed_user = ChangedUser {
        nickname: user_data.nickname.clone(),
        email: user_data.email.clone(),
        password_hash: new_password_hash,
    };

    let user = diesel::update(users.filter(id.eq(current_user.id)))
        .set(&changed_user)
        .get_result::<UserModel>(&*db)
        .expect("Failed to update user.");

    ok().message("User data changed.").data(json!(&user))
}
<commit_msg>Uniform api error messages<commit_after>use diesel;
use diesel::prelude::*;

use rocket_contrib::{JSON, SerdeError};

use helpers::db::DB;
use validation::user::{UserSerializer, UserSettingsSerializer};

use schema::{users, pods, queues};
use schema::users::dsl::*;

use models::user::{UserModel, NewUser, ChangedUser};
use models::pod::{PodModel, NewPod};
use models::queue::{QueueModel, NewQueue};

use responses::{
    APIResponse,
    ok,
    created,
    conflict,
    bad_request,
    forbidden,
    unauthorized,
};

#[get("/info")]
pub fn info(current_user: UserModel) -> APIResponse {
    ok().data(json!(current_user))
}

#[post("/register", data = "<user_data>", format = "application/json")]
pub fn register(user_data: Result<JSON<UserSerializer>, SerdeError>, db: DB) -> APIResponse {

    // Return specific error if invalid JSON has been sent.
    if user_data.is_err() {
        return bad_request().message(format!("{}", user_data.err().unwrap()).as_str());
    }

    let user_data = user_data.unwrap();

    // Check for existing user email
    let results = users.filter(email.eq(user_data.email.clone()))
        .first::<UserModel>(&*db);
    if results.is_ok() {
        return conflict().message("Email already taken.");
    }

    // Check for existing nickname
    let results = users.filter(nickname.eq(user_data.nickname.clone()))
        .first::<UserModel>(&*db);
    if results.is_ok() {
        return conflict().message("Nickname already taken.");
    }

    // Create new password hash
    let new_password_hash = UserModel::make_password_hash(user_data.password.as_str());

    // New user model for table insertion
    let new_user = NewUser {
        nickname: user_data.nickname.clone(),
        email: user_data.email.clone(),
        password_hash: new_password_hash,
    };

    // Insert user to get id for pod
    let user = diesel::insert(&new_user)
        .into(users::table)
        .get_result::<UserModel>(&*db)
        .expect("Error saving new user");

    // New pod
    let new_pod = NewPod {
        name: format!("{}'s Pod", user.nickname.clone()),
        user_id: user.id.clone(),
    };

    let pod = diesel::insert(&new_pod)
        .into(pods::table)
        .get_result::<PodModel>(&*db)
        .expect("Error creating pod");

    let new_queue = NewQueue {
        slots: 2,
        pod_id: Some(pod.id.clone()),
        base_id: None,
    };

    diesel::insert(&new_queue)
        .into(queues::table)
        .get_result::<QueueModel>(&*db)
        .expect("Error creating pod's queue");

    created().message("User created.").data(json!(&user))
}

#[post("/settings", data = "<user_data>", format = "application/json")]
pub fn settings(current_user: UserModel, user_data: Result<JSON<UserSettingsSerializer>, SerdeError>, db: DB) -> APIResponse {

    // Return specific error if invalid JSON has been sent.
    if user_data.is_err() {
        return bad_request().message(format!("{}", user_data.err().unwrap()).as_str());
    }

    let user_data = user_data.unwrap();

    let mut new_password_hash: Option<String> = None;

    if user_data.new_password.is_some() {
        if user_data.password.is_none() {
            return forbidden().message("The current passwords needs to be specified, if you want to change your password.");
        }
        if !current_user.verify_password(user_data.password.as_ref().unwrap().as_str()) {
            return unauthorized().message("Incorrect password.");
        }

        // Create new password hash
        new_password_hash = Some(UserModel::make_password_hash(user_data.password.as_ref().unwrap().as_str()));

    }

    let changed_user = ChangedUser {
        nickname: user_data.nickname.clone(),
        email: user_data.email.clone(),
        password_hash: new_password_hash,
    };

    let user = diesel::update(users.filter(id.eq(current_user.id)))
        .set(&changed_user)
        .get_result::<UserModel>(&*db)
        .expect("Failed to update user.");

    ok().message("User data changed.").data(json!(&user))
}
<|endoftext|>


<commit_before> #[feature(struct_variant)];
extern mod extra;
use extra::getopts;
use std::rand;
use std::rc::RcMut;
use std::fmt;

fn randrange(min : uint, max : uint) -> uint {
    min + rand::random::<uint>() % (max - min)
}

struct ServerID(uint);

#[deriving(Ord)]
struct Term(uint);

#[deriving(Ord)]
struct Index(uint);

#[deriving(Ord)]
struct Time(uint);

struct Environment {
    clock: Time,
    num_servers: uint,
    network: ~[Message],
}

impl Environment {
    fn new(num_servers: uint) -> Environment {
        Environment {
            clock: Time(0),
            num_servers: num_servers,
            network: ~[],
        }
    }
    fn make_time(&self, min: uint, max: uint) -> Time {
        Time(*self.clock + randrange(min, max))
    }
}

struct Message {
    from: ServerID,
    to: ServerID,
    term: Term,
    body: MessageBody,
}

enum MessageBody {
    RequestVoteRequest {
        lastLogTerm: Term,
        lastLogIndex: Index,
    },
    RequestVoteResponse {
        granted: bool,
    },
}

enum ServerState {
    Follower  { timer: Time },
    Candidate { timer: Time, votes: uint },
    Leader,
}

impl fmt::Default for Time {
    fn fmt(time: &Time, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **time)
    }
}

impl fmt::Default for ServerState {
    fn fmt(state: &ServerState, f: &mut fmt::Formatter) {
        match *state {
            Follower{timer} => write!(f.buf, "Follower({})", timer),
            Candidate{timer, votes} => write!(f.buf, "Candidate({}, {})", timer, votes),
            Leader => write!(f.buf, "Leader"),
        }
    }
}

struct LogEntry {
    term: Term,
    // command: ~str,
}

struct Log(~[LogEntry]);

impl Log {
    fn new() -> Log {
        Log(~[])
    }
}

struct Server {
    env: RcMut<Environment>,
    id: ServerID,
    term: Term,
    state: ServerState,
    log: Log,
}

impl Server {
    fn new(id : ServerID, env: &RcMut<Environment>) -> Server {
        Server {
            env: env.clone(),
            id: id,
            term: Term(0),
            state: Follower { timer: env.with_borrow(|e| e.make_time(150, 300)) },
            log: Log::new(),
        }
    }
    /*
    fn reset_timer(&mut self) {
        self.set_timer(Time(*self.env.clock + randrange(150, 300)));
    }
    fn set_timer(&mut self, time: Time) {
        match self.state {
            Follower  {timer : ref mut timer, _} |
            Candidate {timer : ref mut timer, _} =>
                { *timer = time },
            Leader => {},
        }
    }
    */
    fn tick(&mut self) {
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} => {
                do self.env.with_borrow |env| {
                    if timer <= env.clock {
                        self.term = Term(*self.term + 1);
                        self.state = Candidate {
                            timer: env.make_time(150, 300),
                            votes: 1,
                        };
                    }
                }
            },
            Leader => {},
        }
    }
    fn try_become_leader(&mut self) {
        match self.state {
            Follower {_} => {},
            Candidate {votes, _} => {
                do self.env.with_borrow |env| {
                    if votes > env.num_servers / 2 {
                        self.state = Leader;
                    }
                }
            },
            Leader => {},
        }
    }
    /*
    fn next_event_time(&self) -> Time {
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} =>
                timer,
            Leader => Time(!0)
        }
    }
    */
}

impl ToStr for Server {
    fn to_str(&self) -> ~str {
        format!("Server {}: {}", *self.id, self.state)
    }
}

struct Servers(~[Server]);

impl Servers {
    fn new(env: &RcMut<Environment>) -> Servers {
        let mut servers = ~[];
        do env.with_borrow |e| {
            for i in range(0, e.num_servers) {
                let s = Server::new(ServerID(i + 1), env);
                servers.push(s);
            }
        }
        return Servers(servers);
    }
}

fn main() {
    let args = std::os::args();
    let opts = ~[
        getopts::optflag("h"),
        getopts::optflag("help"),
        getopts::optopt("servers"),
    ];
    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m },
        Err(f) => { fail!(f.to_err_msg()) },
    };
    let usage = || {
        println!("Usage: {} [options]", args[0]);
        println("-h, --help   Print this help message");
        println("--servers=N  Simulate cluster with N servers");
    };
    if matches.opt_present("h") || matches.opt_present("help") {
        usage();
        return;
    }
    if !matches.free.is_empty() {
        usage();
        fail!(format!("Extra arguments: {}", matches.free.len()));
    }
    let num_servers : uint = match matches.opt_str("servers") {
        Some(s) => { match std::from_str::from_str(s) {
            Some(i) => { i },
            None => {
                usage();
                fail!(format!("Couldn't parse number of servers from '{}'", s));
            },
        }},
        None => { 5 },
    };
    let env = RcMut::new(Environment::new(num_servers));
    let mut servers = Servers::new(&env);
    loop {
        for server in servers.iter() {
            println(server.to_str());
        }
        do env.with_mut_borrow |e| {
            e.clock = Time(*e.clock + 1);
        }
        for server in servers.mut_iter() {
            server.tick();
        }
    }

    do env.with_mut_borrow |e| {
        e.network.push(Message {
            from: ServerID(1),
            to: ServerID(2),
            term: Term(0),
            body: RequestVoteRequest {
                lastLogTerm: Term(0),
                lastLogIndex: Index(0),
            },
        });
    }
    do env.with_borrow |e| {
        for message in e.network.iter() {
            println(format!("{} to {}", *message.from, *message.to));
        }
    }
}

<commit_msg>drop ref counted environment<commit_after> #[feature(struct_variant)];
extern mod extra;
use extra::getopts;
use std::rand;
use std::fmt;

fn randrange(min : uint, max : uint) -> uint {
    min + rand::random::<uint>() % (max - min)
}

struct ServerID(uint);

#[deriving(Ord)]
struct Term(uint);

#[deriving(Ord)]
struct Index(uint);

#[deriving(Ord)]
struct Time(uint);

struct Environment {
    clock: Time,
    num_servers: uint,
    network: ~[Message],
}

impl Environment {
    fn new(num_servers: uint) -> Environment {
        Environment {
            clock: Time(0),
            num_servers: num_servers,
            network: ~[],
        }
    }
    fn make_time(&self, min: uint, max: uint) -> Time {
        Time(*self.clock + randrange(min, max))
    }
}

struct Message {
    from: ServerID,
    to: ServerID,
    term: Term,
    body: MessageBody,
}

enum MessageBody {
    RequestVoteRequest {
        lastLogTerm: Term,
        lastLogIndex: Index,
    },
    RequestVoteResponse {
        granted: bool,
    },
}

enum ServerState {
    Follower  { timer: Time },
    Candidate { timer: Time, votes: uint },
    Leader,
}

impl fmt::Default for Time {
    fn fmt(time: &Time, f: &mut fmt::Formatter) {
        write!(f.buf, "{}", **time)
    }
}

impl fmt::Default for ServerState {
    fn fmt(state: &ServerState, f: &mut fmt::Formatter) {
        match *state {
            Follower{timer} => write!(f.buf, "Follower({})", timer),
            Candidate{timer, votes} => write!(f.buf, "Candidate({}, {})", timer, votes),
            Leader => write!(f.buf, "Leader"),
        }
    }
}

struct LogEntry {
    term: Term,
    // command: ~str,
}

struct Log(~[LogEntry]);

impl Log {
    fn new() -> Log {
        Log(~[])
    }
}

struct Server {
    id: ServerID,
    term: Term,
    state: ServerState,
    log: Log,
}

impl Server {
    fn new(id : ServerID, env: &Environment) -> Server {
        Server {
            id: id,
            term: Term(0),
            state: Follower { timer: env.make_time(150, 300) },
            log: Log::new(),
        }
    }
    /*
    fn reset_timer(&mut self) {
        self.set_timer(Time(*self.env.clock + randrange(150, 300)));
    }
    fn set_timer(&mut self, time: Time) {
        match self.state {
            Follower  {timer : ref mut timer, _} |
            Candidate {timer : ref mut timer, _} =>
                { *timer = time },
            Leader => {},
        }
    }
    */
    fn tick(&mut self, env: &Environment) {
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} => {
                if timer <= env.clock {
                    self.term = Term(*self.term + 1);
                    self.state = Candidate {
                        timer: env.make_time(150, 300),
                        votes: 1,
                    };
                }
            },
            Leader => {},
        }
    }
    fn try_become_leader(&mut self, env: &Environment) {
        match self.state {
            Follower {_} => {},
            Candidate {votes, _} => {
                if votes > env.num_servers / 2 {
                    self.state = Leader;
                }
            },
            Leader => {},
        }
    }
    /*
    fn next_event_time(&self) -> Time {
        match self.state {
            Follower  {timer, _} |
            Candidate {timer, _} =>
                timer,
            Leader => Time(!0)
        }
    }
    */
}

impl ToStr for Server {
    fn to_str(&self) -> ~str {
        format!("Server {}: {}", *self.id, self.state)
    }
}

struct Servers(~[Server]);

impl Servers {
    fn new(env: &Environment) -> Servers {
        let mut servers = ~[];
        for i in range(0, env.num_servers) {
            let s = Server::new(ServerID(i + 1), env);
            servers.push(s);
        }
        return Servers(servers);
    }
}

fn main() {
    let args = std::os::args();
    let opts = ~[
        getopts::optflag("h"),
        getopts::optflag("help"),
        getopts::optopt("servers"),
    ];
    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m },
        Err(f) => { fail!(f.to_err_msg()) },
    };
    let usage = || {
        println!("Usage: {} [options]", args[0]);
        println("-h, --help   Print this help message");
        println("--servers=N  Simulate cluster with N servers");
    };
    if matches.opt_present("h") || matches.opt_present("help") {
        usage();
        return;
    }
    if !matches.free.is_empty() {
        usage();
        fail!(format!("Extra arguments: {}", matches.free.len()));
    }
    let num_servers : uint = match matches.opt_str("servers") {
        Some(s) => { match std::from_str::from_str(s) {
            Some(i) => { i },
            None => {
                usage();
                fail!(format!("Couldn't parse number of servers from '{}'", s));
            },
        }},
        None => { 5 },
    };
    let env = &mut Environment::new(num_servers);
    let mut servers = Servers::new(env);
    loop {
        for server in servers.iter() {
            println(server.to_str());
        }
        env.clock = Time(*env.clock + 1);
        for server in servers.mut_iter() {
            server.tick(env);
        }
    }

    env.network.push(Message {
        from: ServerID(1),
        to: ServerID(2),
        term: Term(0),
        body: RequestVoteRequest {
            lastLogTerm: Term(0),
            lastLogIndex: Index(0),
        },
    });

    for message in env.network.iter() {
        println(format!("{} to {}", *message.from, *message.to));
    }
}

<|endoftext|>


<commit_before>use crate::{build::BuildRoot,
            error::{Error,
                    Result},
            naming::Naming,
            util,
            Credentials};
use failure::SyncFailure;
use habitat_common::ui::{Status,
                         UIWriter,
                         UI};
use habitat_core::package::PackageIdent;
use handlebars::Handlebars;
use serde_json;
use std::{fs,
          path::{Path,
                 PathBuf},
          str::FromStr};

// This code makes heavy use of `#[cfg(unix)]` and `#[cfg(windows)]`. This should potentially be
// changed to use the various target feature flags.

/// The `Dockerfile` template.
#[cfg(unix)]
const DOCKERFILE: &str = include_str!("../defaults/Dockerfile.hbs");
#[cfg(windows)]
const DOCKERFILE: &str = include_str!("../defaults/Dockerfile_win.hbs");
/// The build report template.
const BUILD_REPORT: &str = include_str!("../defaults/last_docker_export.env.hbs");

// TODO (CM): public temporarily
pub(crate) trait Identified {
    /// The base name of an image.
    fn name(&self) -> String;

    /// The possibly-empty list of tags for an image.
    fn tags(&self) -> Vec<String>;

    /// Returns a non-empty collection of names this image is known
    /// by.
    ///
    /// If an image has no tags, it includes just the name. If it
    /// *does* have tags, it includes the tags prepended with the
    /// name.
    ///
    /// Thus, you could get as little as:
    ///
    /// core/redis
    ///
    /// or as much as:
    ///
    /// core/redis:latest
    /// core/redis:4.0.14
    /// core/redis:4.0.14-20190319155852
    /// core/redis:latest
    /// core/redis:my-custom-tag
    fn expanded_identifiers(&self) -> Vec<String> {
        let mut ids = vec![];

        let tags = self.tags();
        let name = self.name();

        if tags.is_empty() {
            ids.push(name);
        } else {
            for tag in tags {
                ids.push(format!("{}:{}", name, tag));
            }
        }

        ids
    }
}

/// A builder used to create a Docker image.
pub struct DockerBuilder {
    /// The base workdir which hosts the root file system.
    workdir: PathBuf,
    /// The name for the image.
    name:    String,
    /// A list of tags for the image.
    tags:    Vec<String>,
    /// Optional memory limit to pass to pass to the docker build
    memory:  Option<String>,
}

impl Identified for DockerBuilder {
    fn name(&self) -> String { self.name.clone() }

    fn tags(&self) -> Vec<String> { self.tags.clone() }
}

impl DockerBuilder {
    fn new(workdir: &Path, name: &str) -> Self {
        DockerBuilder { workdir: workdir.to_path_buf(),
                        name:    name.to_string(),
                        tags:    Vec::new(),
                        memory:  None, }
    }

    /// Adds a tag for the Docker image.
    pub fn tag(mut self, tag: String) -> Self {
        self.tags.push(tag);
        self
    }

    /// Specifies an amount of memory to allocate to build
    pub fn memory(mut self, memory: &str) -> Self {
        self.memory = Some(memory.to_string());
        self
    }

    /// Builds the Docker image locally and returns the corresponding `DockerImage`.
    ///
    /// # Errors
    ///
    /// * If building the Docker image fails
    pub fn build(self) -> Result<DockerImage> {
        let mut cmd = util::docker_cmd();
        cmd.current_dir(&self.workdir)
           .arg("build")
           .arg("--force-rm");
        if let Some(ref mem) = self.memory {
            cmd.arg("--memory").arg(mem);
        }
        for identifier in &self.expanded_identifiers() {
            cmd.arg("--tag").arg(identifier);
        }
        cmd.arg(".");
        debug!("Running: {:?}", &cmd);
        let exit_status = cmd.spawn()?.wait()?;
        if !exit_status.success() {
            return Err(Error::BuildFailed(exit_status).into());
        }

        let id = match self.tags.first() {
            Some(tag) => self.image_id(&format!("{}:{}", &self.name, tag))?,
            None => self.image_id(&self.name)?,
        };

        Ok(DockerImage { id,
                         name: self.name,
                         tags: self.tags,
                         workdir: self.workdir.to_owned() })
    }

    fn image_id(&self, image_tag: &str) -> Result<String> {
        let mut cmd = util::docker_cmd();
        cmd.arg("images").arg("-q").arg(image_tag);
        debug!("Running: {:?}", &cmd);
        let output = cmd.output()?;
        let stdout = String::from_utf8_lossy(&output.stdout);

        match stdout.lines().next() {
            Some(id) => Ok(id.to_string()),
            None => Err(Error::DockerImageIdNotFound(image_tag.to_string()).into()),
        }
    }
}

/// A built Docker image which exists locally.
pub struct DockerImage {
    /// The image ID for this image.
    id:      String,
    /// The name of this image.
    name:    String,
    /// The list of tags for this image.
    tags:    Vec<String>,
    /// The base workdir which hosts the root file system.
    workdir: PathBuf,
}

impl Identified for DockerImage {
    fn name(&self) -> String { self.name.clone() }

    fn tags(&self) -> Vec<String> { self.tags.clone() }
}

impl DockerImage {
    /// Pushes the Docker image, with all tags, to a remote registry using the provided
    /// `Credentials`.
    ///
    /// # Errors
    ///
    /// * If a registry login is not successful
    /// * If a pushing one or more of the image tags fails
    /// * If a registry logout is not successful
    pub fn push(&self,
                ui: &mut UI,
                credentials: &Credentials,
                registry_url: Option<&str>)
                -> Result<()> {
        ui.begin(format!("Pushing Docker image '{}' with all tags to remote registry",
                         self.name))?;
        self.create_docker_config_file(credentials, registry_url)
            .unwrap();

        for image_tag in self.expanded_identifiers() {
            ui.status(Status::Uploading,
                      format!("image '{}' to remote registry", image_tag))?;
            let mut cmd = util::docker_cmd();
            cmd.arg("--config");
            cmd.arg(self.workdir.to_str().unwrap());
            cmd.arg("push").arg(&image_tag);
            debug!("Running: {:?}", &cmd);
            let exit_status = cmd.spawn()?.wait()?;
            if !exit_status.success() {
                return Err(Error::PushImageFailed(exit_status).into());
            }
            ui.status(Status::Uploaded, format!("image '{}'", &image_tag))?;
        }

        ui.end(format!("Docker image '{}' published with tags: {}",
                       self.name,
                       self.tags.join(", "),))?;

        Ok(())
    }

    /// Removes the image from the local Docker engine along with all tags.
    ///
    /// # Errors
    ///
    /// * If one or more of the image tags cannot be removed
    pub fn rm(self, ui: &mut UI) -> Result<()> {
        ui.begin(format!("Cleaning up local Docker image '{}' with all tags",
                         self.name))?;

        for image_tag in self.expanded_identifiers() {
            ui.status(Status::Deleting, format!("local image '{}'", image_tag))?;
            let mut cmd = util::docker_cmd();
            cmd.arg("rmi").arg(image_tag);
            debug!("Running: {:?}", &cmd);
            let exit_status = cmd.spawn()?.wait()?;
            if !exit_status.success() {
                return Err(Error::RemoveImageFailed(exit_status).into());
            }
        }

        ui.end(format!("Local Docker image '{}' with tags: {} cleaned up",
                       self.name,
                       self.tags.join(", "),))?;
        Ok(())
    }

    /// Create a build report with image metadata in the given path.
    ///
    /// # Errors
    ///
    /// * If the destination directory cannot be created
    /// * If the report file cannot be written
    pub fn create_report<P: AsRef<Path>>(&self, ui: &mut UI, dst: P) -> Result<()> {
        let report = dst.as_ref().join("last_docker_export.env");
        ui.status(Status::Creating,
                  format!("build report {}", report.display()))?;
        fs::create_dir_all(&dst)?;
        let name_tags: Vec<_> = self.tags
                                    .iter()
                                    .map(|t| format!("{}:{}", &self.name, t))
                                    .collect();
        let json = json!({
            "id": &self.id,
            "name": &self.name,
            "tags": self.tags.join(","),
            "name_tags": name_tags.join(","),
        });
        util::write_file(&report,
                         &Handlebars::new().template_render(BUILD_REPORT, &json)
                                           .map_err(SyncFailure::new)?)?;
        Ok(())
    }

    pub fn create_docker_config_file(&self,
                                     credentials: &Credentials,
                                     registry_url: Option<&str>)
                                     -> Result<()> {
        let config = self.workdir.join("config.json");
        fs::create_dir_all(&self.workdir)?;
        let registry = match registry_url {
            Some(url) => url,
            None => "https://index.docker.io/v1/",
        };
        debug!("Using registry: {:?}", registry);
        let json = json!({
            "auths": {
                registry: {
                    "auth": credentials.token
                }
            }
        });
        util::write_file(&config, &serde_json::to_string(&json).unwrap())?;
        Ok(())
    }
}

/// A temporary file system build root for building a Docker image, based on Habitat packages.
pub struct DockerBuildRoot(BuildRoot);

impl DockerBuildRoot {
    /// Builds a completed Docker build root from a `BuildRoot`, performing any final tasks on the
    /// root file system.
    ///
    /// # Errors
    ///
    /// * If any remaining tasks cannot be performed in the build root
    #[cfg(unix)]
    pub fn from_build_root(build_root: BuildRoot, ui: &mut UI) -> Result<Self> {
        let root = DockerBuildRoot(build_root);
        root.add_users_and_groups(ui)?;
        root.create_entrypoint(ui)?;
        root.create_dockerfile(ui)?;

        Ok(root)
    }

    #[cfg(windows)]
    pub fn from_build_root(build_root: BuildRoot, ui: &mut UI) -> Result<Self> {
        let root = DockerBuildRoot(build_root);
        root.create_dockerfile(ui)?;

        Ok(root)
    }

    /// Destroys the temporary build root.
    ///
    /// Note that the build root will automatically destroy itself when it falls out of scope, so
    /// a call to this method is not required, but calling this will provide more user-facing
    /// progress and error reporting.
    ///
    /// # Errors
    ///
    /// * If the temporary work directory cannot be removed
    pub fn destroy(self, ui: &mut UI) -> Result<()> { self.0.destroy(ui) }

    #[cfg(unix)]
    fn add_users_and_groups(&self, ui: &mut UI) -> Result<()> {
        use std::{fs::OpenOptions,
                  io::Write};

        let ctx = self.0.ctx();
        let (users, groups) = ctx.svc_users_and_groups()?;
        {
            let file = "etc/passwd";
            let mut f = OpenOptions::new().append(true)
                                          .open(ctx.rootfs().join(&file))?;
            for user in users {
                ui.status(Status::Creating,
                          format!("user '{}' in /{}", user.name, &file))?;
                writeln!(f, "{}", user)?;
            }
        }
        {
            let file = "etc/group";
            let mut f = OpenOptions::new().append(true)
                                          .open(ctx.rootfs().join(&file))?;
            for group in groups {
                ui.status(Status::Creating,
                          format!("group '{}' in /{}", group.name, &file))?;
                writeln!(f, "{}", group)?;
            }
        }
        Ok(())
    }

    #[cfg(unix)]
    fn create_entrypoint(&self, ui: &mut UI) -> Result<()> {
        use habitat_core::util::posix_perm;

        /// The entrypoint script template.
        const INIT_SH: &str = include_str!("../defaults/init.sh.hbs");

        ui.status(Status::Creating, "entrypoint script")?;
        let ctx = self.0.ctx();
        let busybox_shell =
            util::pkg_path_for(&util::busybox_ident()?, ctx.rootfs())?.join("bin/sh");
        let json = json!({
            "busybox_shell": busybox_shell,
            "path": ctx.env_path(),
            "sup_bin": format!("{} sup", ctx.bin_path().join("hab").display()),
            "primary_svc_ident": ctx.primary_svc_ident().to_string(),
        });
        let init = ctx.rootfs().join("init.sh");
        util::write_file(&init,
                         &Handlebars::new().template_render(INIT_SH, &json)
                                           .map_err(SyncFailure::new)?)?;
        posix_perm::set_permissions(init.to_string_lossy().as_ref(), 0o0755)?;
        Ok(())
    }

    fn create_dockerfile(&self, ui: &mut UI) -> Result<()> {
        ui.status(Status::Creating, "image Dockerfile")?;
        let ctx = self.0.ctx();
        let json = json!({
            "base_image": ctx.base_image(),
            "rootfs": ctx.rootfs().file_name().expect("file_name exists")
                .to_string_lossy()
                .as_ref(),
            "path": ctx.env_path(),
            "hab_path": util::pkg_path_for(
                &PackageIdent::from_str("core/hab")?,
                ctx.rootfs())?.join("bin/hab")
                .to_string_lossy()
                .replace("\\", "/"),
            "exposes": ctx.svc_exposes().join(" "),
            "multi_layer": ctx.multi_layer(),
            "primary_svc_ident": ctx.primary_svc_ident().to_string(),
            "installed_primary_svc_ident": ctx.installed_primary_svc_ident()?.to_string(),
            "environment": ctx.environment,
            "packages": self.0.graph().reverse_topological_sort().iter().map(ToString::to_string).collect::<Vec<_>>(),
        });
        util::write_file(self.0.workdir().join("Dockerfile"),
                         &Handlebars::new().template_render(DOCKERFILE, &json)
                                           .map_err(SyncFailure::new)?)?;
        Ok(())
    }

    /// Build the Docker image locally using the provided naming policy.
    pub fn export(&self,
                  ui: &mut UI,
                  naming: &Naming,
                  memory: Option<&str>)
                  -> Result<DockerImage> {
        ui.status(Status::Creating, "Docker image")?;
        let ident = self.0.ctx().installed_primary_svc_ident()?;
        let channel = self.0.ctx().channel();

        // TODO (CM): Ideally, we'd toss this error much earlier,
        // since this error would be based on user input errors
        let (image_name, tags) = naming.image_identifiers(&ident, &channel)?;

        let mut builder = DockerBuilder::new(self.0.workdir(), &image_name);
        for tag in tags {
            builder = builder.tag(tag);
        }
        if let Some(memory) = memory {
            builder = builder.memory(memory);
        }
        builder.build()
    }
}
<commit_msg>Rename DockerBuilder to ImageBuilder<commit_after>use crate::{build::BuildRoot,
            error::{Error,
                    Result},
            naming::Naming,
            util,
            Credentials};
use failure::SyncFailure;
use habitat_common::ui::{Status,
                         UIWriter,
                         UI};
use habitat_core::package::PackageIdent;
use handlebars::Handlebars;
use serde_json;
use std::{fs,
          path::{Path,
                 PathBuf},
          str::FromStr};

// This code makes heavy use of `#[cfg(unix)]` and `#[cfg(windows)]`. This should potentially be
// changed to use the various target feature flags.

/// The `Dockerfile` template.
#[cfg(unix)]
const DOCKERFILE: &str = include_str!("../defaults/Dockerfile.hbs");
#[cfg(windows)]
const DOCKERFILE: &str = include_str!("../defaults/Dockerfile_win.hbs");
/// The build report template.
const BUILD_REPORT: &str = include_str!("../defaults/last_docker_export.env.hbs");

// TODO (CM): public temporarily
pub(crate) trait Identified {
    /// The base name of an image.
    fn name(&self) -> String;

    /// The possibly-empty list of tags for an image.
    fn tags(&self) -> Vec<String>;

    /// Returns a non-empty collection of names this image is known
    /// by.
    ///
    /// If an image has no tags, it includes just the name. If it
    /// *does* have tags, it includes the tags prepended with the
    /// name.
    ///
    /// Thus, you could get as little as:
    ///
    /// core/redis
    ///
    /// or as much as:
    ///
    /// core/redis:latest
    /// core/redis:4.0.14
    /// core/redis:4.0.14-20190319155852
    /// core/redis:latest
    /// core/redis:my-custom-tag
    fn expanded_identifiers(&self) -> Vec<String> {
        let mut ids = vec![];

        let tags = self.tags();
        let name = self.name();

        if tags.is_empty() {
            ids.push(name);
        } else {
            for tag in tags {
                ids.push(format!("{}:{}", name, tag));
            }
        }

        ids
    }
}

/// A builder used to create a Docker image.
pub struct ImageBuilder {
    /// The base workdir which hosts the root file system.
    workdir: PathBuf,
    /// The name for the image.
    name:    String,
    /// A list of tags for the image.
    tags:    Vec<String>,
    /// Optional memory limit to pass to pass to the docker build
    memory:  Option<String>,
}

impl Identified for ImageBuilder {
    fn name(&self) -> String { self.name.clone() }

    fn tags(&self) -> Vec<String> { self.tags.clone() }
}

impl ImageBuilder {
    fn new(workdir: &Path, name: &str) -> Self {
        ImageBuilder { workdir: workdir.to_path_buf(),
                       name:    name.to_string(),
                       tags:    Vec::new(),
                       memory:  None, }
    }

    /// Adds a tag for the Docker image.
    pub fn tag(mut self, tag: String) -> Self {
        self.tags.push(tag);
        self
    }

    /// Specifies an amount of memory to allocate to build
    pub fn memory(mut self, memory: &str) -> Self {
        self.memory = Some(memory.to_string());
        self
    }

    /// Builds the Docker image locally and returns the corresponding `DockerImage`.
    ///
    /// # Errors
    ///
    /// * If building the Docker image fails
    pub fn build(self) -> Result<DockerImage> {
        let mut cmd = util::docker_cmd();
        cmd.current_dir(&self.workdir)
           .arg("build")
           .arg("--force-rm");
        if let Some(ref mem) = self.memory {
            cmd.arg("--memory").arg(mem);
        }
        for identifier in &self.expanded_identifiers() {
            cmd.arg("--tag").arg(identifier);
        }
        cmd.arg(".");
        debug!("Running: {:?}", &cmd);
        let exit_status = cmd.spawn()?.wait()?;
        if !exit_status.success() {
            return Err(Error::BuildFailed(exit_status).into());
        }

        let id = match self.tags.first() {
            Some(tag) => self.image_id(&format!("{}:{}", &self.name, tag))?,
            None => self.image_id(&self.name)?,
        };

        Ok(DockerImage { id,
                         name: self.name,
                         tags: self.tags,
                         workdir: self.workdir.to_owned() })
    }

    fn image_id(&self, image_tag: &str) -> Result<String> {
        let mut cmd = util::docker_cmd();
        cmd.arg("images").arg("-q").arg(image_tag);
        debug!("Running: {:?}", &cmd);
        let output = cmd.output()?;
        let stdout = String::from_utf8_lossy(&output.stdout);

        match stdout.lines().next() {
            Some(id) => Ok(id.to_string()),
            None => Err(Error::DockerImageIdNotFound(image_tag.to_string()).into()),
        }
    }
}

/// A built Docker image which exists locally.
pub struct DockerImage {
    /// The image ID for this image.
    id:      String,
    /// The name of this image.
    name:    String,
    /// The list of tags for this image.
    tags:    Vec<String>,
    /// The base workdir which hosts the root file system.
    workdir: PathBuf,
}

impl Identified for DockerImage {
    fn name(&self) -> String { self.name.clone() }

    fn tags(&self) -> Vec<String> { self.tags.clone() }
}

impl DockerImage {
    /// Pushes the Docker image, with all tags, to a remote registry using the provided
    /// `Credentials`.
    ///
    /// # Errors
    ///
    /// * If a registry login is not successful
    /// * If a pushing one or more of the image tags fails
    /// * If a registry logout is not successful
    pub fn push(&self,
                ui: &mut UI,
                credentials: &Credentials,
                registry_url: Option<&str>)
                -> Result<()> {
        ui.begin(format!("Pushing Docker image '{}' with all tags to remote registry",
                         self.name))?;
        self.create_docker_config_file(credentials, registry_url)
            .unwrap();

        for image_tag in self.expanded_identifiers() {
            ui.status(Status::Uploading,
                      format!("image '{}' to remote registry", image_tag))?;
            let mut cmd = util::docker_cmd();
            cmd.arg("--config");
            cmd.arg(self.workdir.to_str().unwrap());
            cmd.arg("push").arg(&image_tag);
            debug!("Running: {:?}", &cmd);
            let exit_status = cmd.spawn()?.wait()?;
            if !exit_status.success() {
                return Err(Error::PushImageFailed(exit_status).into());
            }
            ui.status(Status::Uploaded, format!("image '{}'", &image_tag))?;
        }

        ui.end(format!("Docker image '{}' published with tags: {}",
                       self.name,
                       self.tags.join(", "),))?;

        Ok(())
    }

    /// Removes the image from the local Docker engine along with all tags.
    ///
    /// # Errors
    ///
    /// * If one or more of the image tags cannot be removed
    pub fn rm(self, ui: &mut UI) -> Result<()> {
        ui.begin(format!("Cleaning up local Docker image '{}' with all tags",
                         self.name))?;

        for image_tag in self.expanded_identifiers() {
            ui.status(Status::Deleting, format!("local image '{}'", image_tag))?;
            let mut cmd = util::docker_cmd();
            cmd.arg("rmi").arg(image_tag);
            debug!("Running: {:?}", &cmd);
            let exit_status = cmd.spawn()?.wait()?;
            if !exit_status.success() {
                return Err(Error::RemoveImageFailed(exit_status).into());
            }
        }

        ui.end(format!("Local Docker image '{}' with tags: {} cleaned up",
                       self.name,
                       self.tags.join(", "),))?;
        Ok(())
    }

    /// Create a build report with image metadata in the given path.
    ///
    /// # Errors
    ///
    /// * If the destination directory cannot be created
    /// * If the report file cannot be written
    pub fn create_report<P: AsRef<Path>>(&self, ui: &mut UI, dst: P) -> Result<()> {
        let report = dst.as_ref().join("last_docker_export.env");
        ui.status(Status::Creating,
                  format!("build report {}", report.display()))?;
        fs::create_dir_all(&dst)?;
        let name_tags: Vec<_> = self.tags
                                    .iter()
                                    .map(|t| format!("{}:{}", &self.name, t))
                                    .collect();
        let json = json!({
            "id": &self.id,
            "name": &self.name,
            "tags": self.tags.join(","),
            "name_tags": name_tags.join(","),
        });
        util::write_file(&report,
                         &Handlebars::new().template_render(BUILD_REPORT, &json)
                                           .map_err(SyncFailure::new)?)?;
        Ok(())
    }

    pub fn create_docker_config_file(&self,
                                     credentials: &Credentials,
                                     registry_url: Option<&str>)
                                     -> Result<()> {
        let config = self.workdir.join("config.json");
        fs::create_dir_all(&self.workdir)?;
        let registry = match registry_url {
            Some(url) => url,
            None => "https://index.docker.io/v1/",
        };
        debug!("Using registry: {:?}", registry);
        let json = json!({
            "auths": {
                registry: {
                    "auth": credentials.token
                }
            }
        });
        util::write_file(&config, &serde_json::to_string(&json).unwrap())?;
        Ok(())
    }
}

/// A temporary file system build root for building a Docker image, based on Habitat packages.
pub struct DockerBuildRoot(BuildRoot);

impl DockerBuildRoot {
    /// Builds a completed Docker build root from a `BuildRoot`, performing any final tasks on the
    /// root file system.
    ///
    /// # Errors
    ///
    /// * If any remaining tasks cannot be performed in the build root
    #[cfg(unix)]
    pub fn from_build_root(build_root: BuildRoot, ui: &mut UI) -> Result<Self> {
        let root = DockerBuildRoot(build_root);
        root.add_users_and_groups(ui)?;
        root.create_entrypoint(ui)?;
        root.create_dockerfile(ui)?;

        Ok(root)
    }

    #[cfg(windows)]
    pub fn from_build_root(build_root: BuildRoot, ui: &mut UI) -> Result<Self> {
        let root = DockerBuildRoot(build_root);
        root.create_dockerfile(ui)?;

        Ok(root)
    }

    /// Destroys the temporary build root.
    ///
    /// Note that the build root will automatically destroy itself when it falls out of scope, so
    /// a call to this method is not required, but calling this will provide more user-facing
    /// progress and error reporting.
    ///
    /// # Errors
    ///
    /// * If the temporary work directory cannot be removed
    pub fn destroy(self, ui: &mut UI) -> Result<()> { self.0.destroy(ui) }

    #[cfg(unix)]
    fn add_users_and_groups(&self, ui: &mut UI) -> Result<()> {
        use std::{fs::OpenOptions,
                  io::Write};

        let ctx = self.0.ctx();
        let (users, groups) = ctx.svc_users_and_groups()?;
        {
            let file = "etc/passwd";
            let mut f = OpenOptions::new().append(true)
                                          .open(ctx.rootfs().join(&file))?;
            for user in users {
                ui.status(Status::Creating,
                          format!("user '{}' in /{}", user.name, &file))?;
                writeln!(f, "{}", user)?;
            }
        }
        {
            let file = "etc/group";
            let mut f = OpenOptions::new().append(true)
                                          .open(ctx.rootfs().join(&file))?;
            for group in groups {
                ui.status(Status::Creating,
                          format!("group '{}' in /{}", group.name, &file))?;
                writeln!(f, "{}", group)?;
            }
        }
        Ok(())
    }

    #[cfg(unix)]
    fn create_entrypoint(&self, ui: &mut UI) -> Result<()> {
        use habitat_core::util::posix_perm;

        /// The entrypoint script template.
        const INIT_SH: &str = include_str!("../defaults/init.sh.hbs");

        ui.status(Status::Creating, "entrypoint script")?;
        let ctx = self.0.ctx();
        let busybox_shell =
            util::pkg_path_for(&util::busybox_ident()?, ctx.rootfs())?.join("bin/sh");
        let json = json!({
            "busybox_shell": busybox_shell,
            "path": ctx.env_path(),
            "sup_bin": format!("{} sup", ctx.bin_path().join("hab").display()),
            "primary_svc_ident": ctx.primary_svc_ident().to_string(),
        });
        let init = ctx.rootfs().join("init.sh");
        util::write_file(&init,
                         &Handlebars::new().template_render(INIT_SH, &json)
                                           .map_err(SyncFailure::new)?)?;
        posix_perm::set_permissions(init.to_string_lossy().as_ref(), 0o0755)?;
        Ok(())
    }

    fn create_dockerfile(&self, ui: &mut UI) -> Result<()> {
        ui.status(Status::Creating, "image Dockerfile")?;
        let ctx = self.0.ctx();
        let json = json!({
            "base_image": ctx.base_image(),
            "rootfs": ctx.rootfs().file_name().expect("file_name exists")
                .to_string_lossy()
                .as_ref(),
            "path": ctx.env_path(),
            "hab_path": util::pkg_path_for(
                &PackageIdent::from_str("core/hab")?,
                ctx.rootfs())?.join("bin/hab")
                .to_string_lossy()
                .replace("\\", "/"),
            "exposes": ctx.svc_exposes().join(" "),
            "multi_layer": ctx.multi_layer(),
            "primary_svc_ident": ctx.primary_svc_ident().to_string(),
            "installed_primary_svc_ident": ctx.installed_primary_svc_ident()?.to_string(),
            "environment": ctx.environment,
            "packages": self.0.graph().reverse_topological_sort().iter().map(ToString::to_string).collect::<Vec<_>>(),
        });
        util::write_file(self.0.workdir().join("Dockerfile"),
                         &Handlebars::new().template_render(DOCKERFILE, &json)
                                           .map_err(SyncFailure::new)?)?;
        Ok(())
    }

    /// Build the Docker image locally using the provided naming policy.
    pub fn export(&self,
                  ui: &mut UI,
                  naming: &Naming,
                  memory: Option<&str>)
                  -> Result<DockerImage> {
        ui.status(Status::Creating, "Docker image")?;
        let ident = self.0.ctx().installed_primary_svc_ident()?;
        let channel = self.0.ctx().channel();

        // TODO (CM): Ideally, we'd toss this error much earlier,
        // since this error would be based on user input errors
        let (image_name, tags) = naming.image_identifiers(&ident, &channel)?;

        let mut builder = ImageBuilder::new(self.0.workdir(), &image_name);
        for tag in tags {
            builder = builder.tag(tag);
        }
        if let Some(memory) = memory {
            builder = builder.memory(memory);
        }
        builder.build()
    }
}
<|endoftext|>