<commit_before>extern crate clap;
use clap::{Arg, App, AppSettings, SubCommand};
use slog::{LevelFilter, Level, DrainExt, Logger};
use yaml_rust::YamlLoader;
use std::io::prelude::*;
use std::fs::File;
use yaml_rust::{Yaml, YamlEmitter};
use std::error::Error;
use std::path::Path;
use errors::DotcopterError;

#[macro_use]
extern crate slog;
extern crate slog_term;
extern crate slog_stdlog;
extern crate yaml_rust;
extern crate crypto;
extern crate regex;

#[cfg(test)]
extern crate spectral;

mod model;
mod config;
mod checksum;
mod files;
mod mutate;
mod import;
mod errors;

fn main() {
  let matches: clap::ArgMatches = create_app().get_matches();
  let stream = slog_term::streamer().full().build();
  let verbose: bool = matches.is_present("verbose");
  let force: bool = matches.is_present("force");
  let log = if verbose {
    slog::Logger::root(stream.fuse(), o!())
  } else {
    slog::Logger::root(LevelFilter::new(stream, Level::Info).fuse(), o!())
  };

  let config_file = matches.value_of("config_file").unwrap();
  info!(log, "Starting engine"; "config_file" => config_file);

  let config = match load_config_file(&log, config_file) {
    Ok(content) => content,
    Err(e) => {
      error!(log, "Failed to load config file."; "error" => e.description());
      panic!(1)
    }
  };

  let mut yaml_documents = match YamlLoader::load_from_str(&config) {
    Ok(yaml) => yaml,
    Err(e) => {
      error!(log, "Failed to parse config file."; "error" => e.description());
      panic!(2)
    }
  };
  if yaml_documents.is_empty() {
    let s = "files:";
    yaml_documents = YamlLoader::load_from_str(s).unwrap();
  }

  let maybe_ln_matches = matches.subcommand_matches("ln");
  let maybe_cp_matches = matches.subcommand_matches("cp");
  let maybe_apply_matches = matches.subcommand_matches("apply");
  let maybe_import_matches = matches.subcommand_matches("import");
  if maybe_apply_matches.is_some() {
    let yaml_config = &yaml_documents[0];
    let dot_files: &Yaml = &yaml_config["files"];
    info!(log, "Liftoff! Applying configuration.");
    files::process_dot_files(&log, dot_files, force);
  } else if let Some(ln_matches) = maybe_ln_matches {
    let yaml_config = &yaml_documents[0];
    let link_target = ln_matches.value_of("link_target").unwrap();
    let link_name = ln_matches.value_of("link_name").unwrap();
    let log = log.new(o!("link_target" => link_target.to_string(), "link_name" => link_name.to_string()));
    info!(log, "Liftoff! Adding new link to configuration");
    let new_config = mutate::add_dotfiles_to_config(&log,
                                                    yaml_config,
                                                    &[model::DotFile {
                                                        target: link_name.to_string(),
                                                        source: link_target.to_string(),
                                                        dot_file_type: model::DotFileType::LINK,
                                                      }]);
    write_new_yaml(&log, &new_config, config_file);
  } else if let Some(cp_matches) = maybe_cp_matches {
    let yaml_config = &yaml_documents[0];
    let target = cp_matches.value_of("target").unwrap();
    let source = cp_matches.value_of("source").unwrap();
    let log = log.new(o!("target" => target.to_string(), "source" => source.to_string()));
    info!(log, "Liftoff! Adding new copy to configuration");
    let new_config = mutate::add_dotfiles_to_config(&log,
                                                    yaml_config,
                                                    &[model::DotFile {
                                                        target: target.to_string(),
                                                        source: source.to_string(),
                                                        dot_file_type: model::DotFileType::COPY,
                                                      }]);
    write_new_yaml(&log, &new_config, config_file);
  } else if let Some(import_matches) = maybe_import_matches {
    let dir = import_matches.value_of("dir").unwrap();
    let log = log.new(o!("import_directory" => dir.to_string()));
    let yaml_config = &yaml_documents[0];
    info!(log, "Liftoff! Importing to configuration");
    let dot_files = import::scan_dir(&log, dir);
    if !dot_files.is_empty() {
      let new_config = mutate::add_dotfiles_to_config(&log, yaml_config, &dot_files);
      write_new_yaml(&log, &new_config, config_file);
    }
  }
}

fn write_new_yaml(log: &Logger, document: &Yaml, config_file: &str) {
  let mut out_str = String::new();
  {
    let mut emitter = YamlEmitter::new(&mut out_str);
    match emitter.dump(document) {
      Ok(_) => {}
      Err(e) => {
        error!(log, "Failed to write config file."; "error" => format!("{:?}", e));
        return;
      }
    }
  }
  out_str.push_str("\n");
  match write_config_file(config_file, &out_str) {
    Ok(_) => info!(log, "Successfully written configuration"),
    Err(e) => error!(log, "Failed to write config file."; "error" => e.description()),
  };
}

fn write_config_file(file: &str, content: &str) -> Result<(), DotcopterError> {
  let mut file = try!(File::create(file));
  try!(file.write_all(content.as_bytes()));
  Ok(())
}

fn load_config_file(log: &Logger, file: &str) -> Result<String, DotcopterError> {
  let path = Path::new(file);
  if path.exists() {
    let mut file = try!(File::open(file));
    let mut content = String::new();
    try!(file.read_to_string(&mut content));
    Ok(content)
  } else {
    warn!(log, "Configuration doesn't exit"; "file" => file);
    Ok("".to_string())
  }
}

fn create_app<'a>() -> App<'a, 'a> {
  App::new("dotcopter")
    .version("0.1")
    .setting(AppSettings::ColoredHelp)
    .author("Patrick Haun <bomgar85@googlemail.com>")
    .about("manages dotfiles installation")
    .arg(Arg::with_name("force")
           .short("f")
           .long("force")
           .takes_value(false))
    .arg(Arg::with_name("verbose")
           .long("verbose")
           .short("v")
           .help("debug output")
           .required(false)
           .takes_value(false))
    .arg(Arg::with_name("config_file").required(true))
    .subcommand(SubCommand::with_name("apply").about("applies a dotfile configuration"))
    .subcommand(SubCommand::with_name("ln")
                  .about("adds new link to configuration")
                  .arg(Arg::with_name("link_target").required(true))
                  .arg(Arg::with_name("link_name").required(true)))
    .subcommand(SubCommand::with_name("cp")
                  .about("adds new copy to configuration")
                  .arg(Arg::with_name("source").required(true))
                  .arg(Arg::with_name("target").required(true)))
    .subcommand(SubCommand::with_name("import")
                  .about("imports dotfiles from a folder into the configuration")
                  .arg(Arg::with_name("dir").required(true)))
}
<commit_msg>exit on config file errors<commit_after>extern crate clap;
use clap::{Arg, App, AppSettings, SubCommand};
use slog::{LevelFilter, Level, DrainExt, Logger};
use yaml_rust::YamlLoader;
use std::io::prelude::*;
use std::fs::File;
use yaml_rust::{Yaml, YamlEmitter};
use std::error::Error;
use std::path::Path;
use std::process::exit;
use errors::DotcopterError;

#[macro_use]
extern crate slog;
extern crate slog_term;
extern crate slog_stdlog;
extern crate yaml_rust;
extern crate crypto;
extern crate regex;

#[cfg(test)]
extern crate spectral;

mod model;
mod config;
mod checksum;
mod files;
mod mutate;
mod import;
mod errors;

fn main() {
  let matches: clap::ArgMatches = create_app().get_matches();
  let stream = slog_term::streamer().full().build();
  let verbose: bool = matches.is_present("verbose");
  let force: bool = matches.is_present("force");
  let log = if verbose {
    slog::Logger::root(stream.fuse(), o!())
  } else {
    slog::Logger::root(LevelFilter::new(stream, Level::Info).fuse(), o!())
  };

  let config_file = matches.value_of("config_file").unwrap();
  info!(log, "Starting engine"; "config_file" => config_file);

  let config = match load_config_file(&log, config_file) {
    Ok(content) => content,
    Err(e) => {
      error!(log, "Failed to load config file."; "error" => e.description());
      exit(1)
    }
  };

  let mut yaml_documents = match YamlLoader::load_from_str(&config) {
    Ok(yaml) => yaml,
    Err(e) => {
      error!(log, "Failed to parse config file."; "error" => e.description());
      exit(2)
    }
  };
  if yaml_documents.is_empty() {
    let s = "files:";
    yaml_documents = YamlLoader::load_from_str(s).unwrap();
  }

  let maybe_ln_matches = matches.subcommand_matches("ln");
  let maybe_cp_matches = matches.subcommand_matches("cp");
  let maybe_apply_matches = matches.subcommand_matches("apply");
  let maybe_import_matches = matches.subcommand_matches("import");
  if maybe_apply_matches.is_some() {
    let yaml_config = &yaml_documents[0];
    let dot_files: &Yaml = &yaml_config["files"];
    info!(log, "Liftoff! Applying configuration.");
    files::process_dot_files(&log, dot_files, force);
  } else if let Some(ln_matches) = maybe_ln_matches {
    let yaml_config = &yaml_documents[0];
    let link_target = ln_matches.value_of("link_target").unwrap();
    let link_name = ln_matches.value_of("link_name").unwrap();
    let log = log.new(o!("link_target" => link_target.to_string(), "link_name" => link_name.to_string()));
    info!(log, "Liftoff! Adding new link to configuration");
    let new_config = mutate::add_dotfiles_to_config(&log,
                                                    yaml_config,
                                                    &[model::DotFile {
                                                        target: link_name.to_string(),
                                                        source: link_target.to_string(),
                                                        dot_file_type: model::DotFileType::LINK,
                                                      }]);
    write_new_yaml(&log, &new_config, config_file);
  } else if let Some(cp_matches) = maybe_cp_matches {
    let yaml_config = &yaml_documents[0];
    let target = cp_matches.value_of("target").unwrap();
    let source = cp_matches.value_of("source").unwrap();
    let log = log.new(o!("target" => target.to_string(), "source" => source.to_string()));
    info!(log, "Liftoff! Adding new copy to configuration");
    let new_config = mutate::add_dotfiles_to_config(&log,
                                                    yaml_config,
                                                    &[model::DotFile {
                                                        target: target.to_string(),
                                                        source: source.to_string(),
                                                        dot_file_type: model::DotFileType::COPY,
                                                      }]);
    write_new_yaml(&log, &new_config, config_file);
  } else if let Some(import_matches) = maybe_import_matches {
    let dir = import_matches.value_of("dir").unwrap();
    let log = log.new(o!("import_directory" => dir.to_string()));
    let yaml_config = &yaml_documents[0];
    info!(log, "Liftoff! Importing to configuration");
    let dot_files = import::scan_dir(&log, dir);
    if !dot_files.is_empty() {
      let new_config = mutate::add_dotfiles_to_config(&log, yaml_config, &dot_files);
      write_new_yaml(&log, &new_config, config_file);
    }
  }
}

fn write_new_yaml(log: &Logger, document: &Yaml, config_file: &str) {
  let mut out_str = String::new();
  {
    let mut emitter = YamlEmitter::new(&mut out_str);
    match emitter.dump(document) {
      Ok(_) => {}
      Err(e) => {
        error!(log, "Failed to write config file."; "error" => format!("{:?}", e));
        exit(3);
      }
    }
  }
  out_str.push_str("\n");
  match write_config_file(config_file, &out_str) {
    Ok(_) => info!(log, "Successfully written configuration"),
    Err(e) => {
      error!(log, "Failed to write config file."; "error" => e.description());
      exit(4);
    }
  };
}

fn write_config_file(file: &str, content: &str) -> Result<(), DotcopterError> {
  let mut file = try!(File::create(file));
  try!(file.write_all(content.as_bytes()));
  Ok(())
}

fn load_config_file(log: &Logger, file: &str) -> Result<String, DotcopterError> {
  let path = Path::new(file);
  if path.exists() {
    let mut file = try!(File::open(file));
    let mut content = String::new();
    try!(file.read_to_string(&mut content));
    Ok(content)
  } else {
    warn!(log, "Configuration doesn't exit"; "file" => file);
    Ok("".to_string())
  }
}

fn create_app<'a>() -> App<'a, 'a> {
  App::new("dotcopter")
    .version("0.1")
    .setting(AppSettings::ColoredHelp)
    .author("Patrick Haun <bomgar85@googlemail.com>")
    .about("manages dotfiles installation")
    .arg(Arg::with_name("force")
           .short("f")
           .long("force")
           .takes_value(false))
    .arg(Arg::with_name("verbose")
           .long("verbose")
           .short("v")
           .help("debug output")
           .required(false)
           .takes_value(false))
    .arg(Arg::with_name("config_file").required(true))
    .subcommand(SubCommand::with_name("apply").about("applies a dotfile configuration"))
    .subcommand(SubCommand::with_name("ln")
                  .about("adds new link to configuration")
                  .arg(Arg::with_name("link_target").required(true))
                  .arg(Arg::with_name("link_name").required(true)))
    .subcommand(SubCommand::with_name("cp")
                  .about("adds new copy to configuration")
                  .arg(Arg::with_name("source").required(true))
                  .arg(Arg::with_name("target").required(true)))
    .subcommand(SubCommand::with_name("import")
                  .about("imports dotfiles from a folder into the configuration")
                  .arg(Arg::with_name("dir").required(true)))
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serilized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn cancel_serilized_frames(&mut self, error: Error) -> Result<()> {
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                // Due to a bug in epoll/mio, it doesn't seem like we can trust this, it's sometimes missing when it should be there
                /*
                if event.is_readable() {
                    self.can_read = true;
                }
                */
                self.can_read = true;
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup();
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.connection.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.socket.flush()?;

        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }

            self.socket.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                } else {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<commit_msg>IoLoop: handle error before read/write closed<commit_after>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serilized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn cancel_serilized_frames(&mut self, error: Error) -> Result<()> {
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                // Due to a bug in epoll/mio, it doesn't seem like we can trust this, it's sometimes missing when it should be there
                /*
                if event.is_readable() {
                    self.can_read = true;
                }
                */
                self.can_read = true;
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup();
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.connection.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.socket.flush()?;

        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }

            self.socket.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                } else {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<|endoftext|>


<commit_before>// Playing with paths.

#![warn(bare_trait_objects)]

use chrono::Local;
use std::{collections::BTreeMap, path::Path};
use structopt::StructOpt;
use tempdir::TempDir;

use rsure::{log_init, parse_store, show_tree, Store, StoreTags, StoreVersion, Version};

// For now, just use the crate's error type.
pub use rsure::Result;

#[derive(StructOpt)]
#[structopt(name = "rsure", about = "File integrity")]
struct Opt {
    #[structopt(short = "f", long = "file", default_value = "2sure.dat.gz")]
    /// Base of file name, default 2sure, will get .dat.gz appended
    file: String,
    #[structopt(short = "d", long = "dir", default_value = ".")]
    /// Directory to scan, defaults to "."
    dir: String,
    #[structopt(long = "tag")]
    /// key=value to associate with scan
    tag: Vec<String>,
    #[structopt(short = "v", long = "version")]
    version: Option<String>,
    #[structopt(subcommand)]
    command: Command,
}

#[derive(StructOpt)]
enum Command {
    #[structopt(name = "scan")]
    /// Scan a directory for the first time
    Scan,
    #[structopt(name = "update")]
    /// Update the scan using the dat/weave file
    Update,
    #[structopt(name = "check")]
    /// Compare the directory with the dat/weave file
    Check {
        #[structopt(short = "i", long = "ignore")]
        /// Tag to ignore when comparing.
        ignore: Vec<String>,
    },
    #[structopt(name = "signoff")]
    /// Compare dat with bak file, or last two versions in weave file
    Signoff {
        #[structopt(short = "i", long = "ignore")]
        /// Tag to ignore when comparing.
        ignore: Vec<String>,
    },
    #[structopt(name = "show")]
    /// Pretty print the dat file
    Show,
    #[structopt(name = "list")]
    /// List revisions in a given sure store
    List,
}

#[allow(dead_code)]
fn main() {
    log_init();

    let opt = Opt::from_args();

    let store = parse_store(&opt.file).unwrap();

    let mut tags = decode_tags(Some(opt.tag.iter().map(|x| x.as_str())));

    add_name_tag(&mut tags, &opt.dir);

    // Note that only the "check" command uses the version tag.
    let latest = match opt.version {
        None => Version::Latest,
        Some(ref x) => Version::Tagged(x.to_string()),
    };

    match &opt.command {
        Command::Scan => {
            rsure::update(&opt.dir, &*store, false, &tags).unwrap();
        }
        Command::Update => {
            rsure::update(&opt.dir, &*store, true, &tags).unwrap();
        }
        Command::Check { ignore } => {
            let ignore: Vec<_> = ignore.iter().map(|x| x.as_str()).collect();
            run_check(&*store, &opt, latest, &ignore).unwrap();
        }
        Command::Signoff { ignore } => {
            let ignore: Vec<_> = ignore.iter().map(|x| x.as_str()).collect();
            let old_tree = store.load_iter(Version::Prior).unwrap();
            let new_tree = store.load_iter(Version::Latest).unwrap();
            println!("signoff {}", opt.file);
            rsure::compare_trees(old_tree, new_tree, &Path::new(&opt.dir), &ignore).unwrap();
        }
        Command::Show => {
            println!("show {}", opt.file);
            show_tree(&*store).unwrap();
        }
        Command::List => {
            let version = store.get_versions().unwrap();
            dump_versions(&version);
        }
    }
}

fn run_check(store: &dyn Store, opt: &Opt, latest: Version, ignore: &[&str]) -> Result<()> {
    // Perform a full scan to a temp store.
    let tdir = TempDir::new("rsure")?;
    let tpath = tdir.path().join("check.dat.gz");
    let tstore = parse_store(tpath.to_str().unwrap())?;
    let mut tags = BTreeMap::new();
    add_name_tag(&mut tags, &opt.dir);
    println!("Scanning");
    rsure::update(&opt.dir, &*tstore, false, &tags)?;

    let old_tree = store.load_iter(latest)?;
    let new_tree = tstore.load_iter(Version::Latest)?;
    println!("Check {}", opt.file);
    rsure::compare_trees(old_tree, new_tree, &Path::new(&opt.dir), ignore)?;
    Ok(())
}

/// Decode the command-line tags.  Tags should be of the form key=value, and multiple can be
/// specified, terminated by the command.  It is also possible to specify --tag multiple times.
fn decode_tags<'a, I>(tags: Option<I>) -> StoreTags
where
    I: Iterator<Item = &'a str>,
{
    match tags {
        None => BTreeMap::new(),
        Some(tags) => tags.map(|x| decode_tag(x)).collect(),
    }
}

fn decode_tag(tag: &str) -> (String, String) {
    let fields: Vec<_> = tag.splitn(2, '=').collect();
    if fields.len() != 2 {
        panic!("Tag must be key=value");
    }
    (fields[0].to_string(), fields[1].to_string())
}

/// If the caller doesn't specify a 'name=' tag, generate one based on the current timestamp.
/// Also will add a 'dir' attribute for where the tree was captured.
fn add_name_tag<P: AsRef<Path>>(tags: &mut StoreTags, dir: P) {
    if !tags.contains_key("name") {
        tags.insert("name".to_string(), Local::now().to_rfc3339());
    }

    if !tags.contains_key("dir") {
        tags.insert(
            "dir".to_string(),
            dir.as_ref()
                .canonicalize()
                .unwrap_or_else(|_| Path::new("invalid").to_owned())
                .to_string_lossy()
                .into_owned(),
        );
    }
}

fn dump_versions(versions: &[StoreVersion]) {
    println!("vers | Time captured       | name");
    println!("-----+---------------------+------------------");
    for v in versions {
        let vers = match v.version {
            Version::Latest => "tip",
            Version::Prior => "prev",
            Version::Tagged(ref v) => v,
        };
        println!(
            "{:>4} | {} | {}",
            vers,
            v.time.with_timezone(&Local).format("%Y-%m-%d %H:%M:%S"),
            v.name
        );
    }
}
<commit_msg>Change main to return Result<commit_after>// Playing with paths.

#![warn(bare_trait_objects)]

use chrono::Local;
use std::{collections::BTreeMap, path::Path};
use structopt::StructOpt;
use tempdir::TempDir;

use rsure::{log_init, parse_store, show_tree, Store, StoreTags, StoreVersion, Version};

// For now, just use the crate's error type.
pub use rsure::Result;

#[derive(StructOpt)]
#[structopt(name = "rsure", about = "File integrity")]
struct Opt {
    #[structopt(short = "f", long = "file", default_value = "2sure.dat.gz")]
    /// Base of file name, default 2sure, will get .dat.gz appended
    file: String,
    #[structopt(short = "d", long = "dir", default_value = ".")]
    /// Directory to scan, defaults to "."
    dir: String,
    #[structopt(long = "tag")]
    /// key=value to associate with scan
    tag: Vec<String>,
    #[structopt(short = "v", long = "version")]
    version: Option<String>,
    #[structopt(subcommand)]
    command: Command,
}

#[derive(StructOpt)]
enum Command {
    #[structopt(name = "scan")]
    /// Scan a directory for the first time
    Scan,
    #[structopt(name = "update")]
    /// Update the scan using the dat/weave file
    Update,
    #[structopt(name = "check")]
    /// Compare the directory with the dat/weave file
    Check {
        #[structopt(short = "i", long = "ignore")]
        /// Tag to ignore when comparing.
        ignore: Vec<String>,
    },
    #[structopt(name = "signoff")]
    /// Compare dat with bak file, or last two versions in weave file
    Signoff {
        #[structopt(short = "i", long = "ignore")]
        /// Tag to ignore when comparing.
        ignore: Vec<String>,
    },
    #[structopt(name = "show")]
    /// Pretty print the dat file
    Show,
    #[structopt(name = "list")]
    /// List revisions in a given sure store
    List,
}

#[allow(dead_code)]
fn main() -> Result<()> {
    log_init();

    let opt = Opt::from_args();

    let store = parse_store(&opt.file)?;

    let mut tags = decode_tags(Some(opt.tag.iter().map(|x| x.as_str())));

    add_name_tag(&mut tags, &opt.dir);

    // Note that only the "check" command uses the version tag.
    let latest = match opt.version {
        None => Version::Latest,
        Some(ref x) => Version::Tagged(x.to_string()),
    };

    match &opt.command {
        Command::Scan => {
            rsure::update(&opt.dir, &*store, false, &tags)?;
        }
        Command::Update => {
            rsure::update(&opt.dir, &*store, true, &tags)?;
        }
        Command::Check { ignore } => {
            let ignore: Vec<_> = ignore.iter().map(|x| x.as_str()).collect();
            run_check(&*store, &opt, latest, &ignore)?;
        }
        Command::Signoff { ignore } => {
            let ignore: Vec<_> = ignore.iter().map(|x| x.as_str()).collect();
            let old_tree = store.load_iter(Version::Prior)?;
            let new_tree = store.load_iter(Version::Latest)?;
            println!("signoff {}", opt.file);
            rsure::compare_trees(old_tree, new_tree, &Path::new(&opt.dir), &ignore)?;
        }
        Command::Show => {
            println!("show {}", opt.file);
            show_tree(&*store)?;
        }
        Command::List => {
            let version = store.get_versions()?;
            dump_versions(&version);
        }
    }

    Ok(())
}

fn run_check(store: &dyn Store, opt: &Opt, latest: Version, ignore: &[&str]) -> Result<()> {
    // Perform a full scan to a temp store.
    let tdir = TempDir::new("rsure")?;
    let tpath = tdir.path().join("check.dat.gz");
    let tstore = parse_store(tpath.to_str().unwrap())?;
    let mut tags = BTreeMap::new();
    add_name_tag(&mut tags, &opt.dir);
    println!("Scanning");
    rsure::update(&opt.dir, &*tstore, false, &tags)?;

    let old_tree = store.load_iter(latest)?;
    let new_tree = tstore.load_iter(Version::Latest)?;
    println!("Check {}", opt.file);
    rsure::compare_trees(old_tree, new_tree, &Path::new(&opt.dir), ignore)?;
    Ok(())
}

/// Decode the command-line tags.  Tags should be of the form key=value, and multiple can be
/// specified, terminated by the command.  It is also possible to specify --tag multiple times.
fn decode_tags<'a, I>(tags: Option<I>) -> StoreTags
where
    I: Iterator<Item = &'a str>,
{
    match tags {
        None => BTreeMap::new(),
        Some(tags) => tags.map(|x| decode_tag(x)).collect(),
    }
}

fn decode_tag(tag: &str) -> (String, String) {
    let fields: Vec<_> = tag.splitn(2, '=').collect();
    if fields.len() != 2 {
        panic!("Tag must be key=value");
    }
    (fields[0].to_string(), fields[1].to_string())
}

/// If the caller doesn't specify a 'name=' tag, generate one based on the current timestamp.
/// Also will add a 'dir' attribute for where the tree was captured.
fn add_name_tag<P: AsRef<Path>>(tags: &mut StoreTags, dir: P) {
    if !tags.contains_key("name") {
        tags.insert("name".to_string(), Local::now().to_rfc3339());
    }

    if !tags.contains_key("dir") {
        tags.insert(
            "dir".to_string(),
            dir.as_ref()
                .canonicalize()
                .unwrap_or_else(|_| Path::new("invalid").to_owned())
                .to_string_lossy()
                .into_owned(),
        );
    }
}

fn dump_versions(versions: &[StoreVersion]) {
    println!("vers | Time captured       | name");
    println!("-----+---------------------+------------------");
    for v in versions {
        let vers = match v.version {
            Version::Latest => "tip",
            Version::Prior => "prev",
            Version::Tagged(ref v) => v,
        };
        println!(
            "{:>4} | {} | {}",
            vers,
            v.time.with_timezone(&Local).format("%Y-%m-%d %H:%M:%S"),
            v.name
        );
    }
}
<|endoftext|>


<commit_before>extern crate libc;

use std::io;
use std::thread;
use std::time::Duration;
use std::fs::OpenOptions;
use std::os::unix::fs::OpenOptionsExt;
use std::os::unix::io::AsRawFd;
use std::ptr::{ self, read_volatile, write_volatile };

const GPIO_ADDR : libc::off_t      = 0x3F200000; // BCM2836 GPIO address
const MEM_BLK_SIZE : libc::size_t  = 4096;       // 4 * 1024

const GPFSEL2_OFFSET_ADDR : isize  = 0x02;       // 0x08 / 4
const GPSET0_OFFSET_ADDR : isize   = 0x07;       // 0x1C / 4
const GPCLR0_OFFSET_ADDR : isize   = 0x0A;       // 0x28 / 4
const GPLEV0_OFFSET_ADDR : isize   = 0x0D;       // 0x34 / 4

const FSEL27_BIT : isize    = 21;
const SET27_BIT : isize     = 27;
const CLR27_BIT : isize     = 27;
const LEV27_BIT : isize     = 27;

const GPFSEL_INPUT : u32    = 0;
const GPFSEL_OUTPUT : u32   = 1;
const FSEL27_MASK : u32     = 0xFF1FFFFF;

const GPLEV_HIGH : u32      = 1;
const ON : u32              = 1;

const SLEEP_DELAY : u64     = 1;

fn main() {
    let mut blink_cnt : i32 = 0;

    let gpio_ptr : *mut u32 = map_gpio().expect( "failed to gpio mapping" );
    let fsel2_reg : *mut u32 = gpio_ptr.wrapping_offset( GPFSEL2_OFFSET_ADDR );
    let clr0_reg : *mut u32 = gpio_ptr.wrapping_offset( GPCLR0_OFFSET_ADDR );

    unsafe {
        write_volatile( fsel2_reg, ( *fsel2_reg & FSEL27_MASK ) | ( GPFSEL_OUTPUT << FSEL27_BIT ) );
    }

    while blink_cnt < 10 {
        blink_led( gpio_ptr );
        thread::sleep( Duration::from_secs( SLEEP_DELAY ) );
        blink_cnt = blink_cnt + 1;
    }

    unsafe {
        write_volatile( clr0_reg, *clr0_reg | ( ( ON << CLR27_BIT ) ) );
        write_volatile( fsel2_reg, ( *fsel2_reg & FSEL27_MASK ) | ( GPFSEL_INPUT << FSEL27_BIT ) );
    }
}

fn map_gpio() -> io::Result<*mut u32>  {
    let mem_file = OpenOptions::new()
                .read( true )
                .write( true )
                .custom_flags( libc::O_SYNC )
                .open( "/dev/gpiomem" )
                .expect( "can't open /dev/gpiomem" );

    unsafe {
        let gpio_ptr = libc::mmap( ptr::null_mut(),
                                   MEM_BLK_SIZE,
                                   libc::PROT_READ | libc::PROT_WRITE,
                                   libc::MAP_SHARED,
                                   mem_file.as_raw_fd(),
                                   GPIO_ADDR );

        if gpio_ptr == libc::MAP_FAILED {
             Err( io::Error::last_os_error() )
        }
        else
        {
            Ok( gpio_ptr as *mut u32 )
        }
    }
}

fn blink_led( gpio_ptr : *mut u32 ) {
    let set0_reg : *mut u32 = gpio_ptr.wrapping_offset( GPSET0_OFFSET_ADDR );
    let clr0_reg : *mut u32 = gpio_ptr.wrapping_offset( GPCLR0_OFFSET_ADDR );
    let lev0_reg : *mut u32 = gpio_ptr.wrapping_offset( GPLEV0_OFFSET_ADDR );
    let level : u32;

    unsafe {
       level = ( read_volatile( lev0_reg ) >> LEV27_BIT ) & 0x00000001;
    }

    if level == GPLEV_HIGH {
        unsafe {
            write_volatile( clr0_reg, *clr0_reg | ( ON << CLR27_BIT ) );
        }
    }
    else
    {
        unsafe {
            write_volatile( set0_reg, *set0_reg | ( ON << SET27_BIT ) );
        }
    }
}
<commit_msg>Change memory size for memory mapping<commit_after>extern crate libc;

use std::io;
use std::thread;
use std::time::Duration;
use std::fs::OpenOptions;
use std::os::unix::fs::OpenOptionsExt;
use std::os::unix::io::AsRawFd;
use std::ptr::{ self, read_volatile, write_volatile };

const GPIO_ADDR : libc::off_t      = 0x3F200000; // BCM2836 GPIO address
const MEM_BLK_SIZE : libc::size_t  = 0xB0;       // GPIO register total size

const GPFSEL2_OFFSET_ADDR : isize  = 0x02;       // 0x08 / 4
const GPSET0_OFFSET_ADDR : isize   = 0x07;       // 0x1C / 4
const GPCLR0_OFFSET_ADDR : isize   = 0x0A;       // 0x28 / 4
const GPLEV0_OFFSET_ADDR : isize   = 0x0D;       // 0x34 / 4

const FSEL27_BIT : isize    = 21;
const SET27_BIT : isize     = 27;
const CLR27_BIT : isize     = 27;
const LEV27_BIT : isize     = 27;

const GPFSEL_INPUT : u32    = 0;
const GPFSEL_OUTPUT : u32   = 1;
const FSEL27_MASK : u32     = 0xFF1FFFFF;

const GPLEV_HIGH : u32      = 1;
const ON : u32              = 1;

const SLEEP_DELAY : u64     = 1;

fn main() {
    let mut blink_cnt : i32 = 0;

    let gpio_ptr : *mut u32 = map_gpio().expect( "failed to gpio mapping" );
    let fsel2_reg : *mut u32 = gpio_ptr.wrapping_offset( GPFSEL2_OFFSET_ADDR );
    let clr0_reg : *mut u32 = gpio_ptr.wrapping_offset( GPCLR0_OFFSET_ADDR );

    unsafe {
        write_volatile( fsel2_reg, ( *fsel2_reg & FSEL27_MASK ) | ( GPFSEL_OUTPUT << FSEL27_BIT ) );
    }

    while blink_cnt < 10 {
        blink_led( gpio_ptr );
        thread::sleep( Duration::from_secs( SLEEP_DELAY ) );
        blink_cnt = blink_cnt + 1;
    }

    unsafe {
        write_volatile( clr0_reg, *clr0_reg | ( ( ON << CLR27_BIT ) ) );
        write_volatile( fsel2_reg, ( *fsel2_reg & FSEL27_MASK ) | ( GPFSEL_INPUT << FSEL27_BIT ) );
    }
}

fn map_gpio() -> io::Result<*mut u32>  {
    let mem_file = OpenOptions::new()
                .read( true )
                .write( true )
                .custom_flags( libc::O_SYNC )
                .open( "/dev/gpiomem" )
                .expect( "can't open /dev/gpiomem" );

    unsafe {
        let gpio_ptr = libc::mmap( ptr::null_mut(),
                                   MEM_BLK_SIZE,
                                   libc::PROT_READ | libc::PROT_WRITE,
                                   libc::MAP_SHARED,
                                   mem_file.as_raw_fd(),
                                   GPIO_ADDR );

        if gpio_ptr == libc::MAP_FAILED {
             Err( io::Error::last_os_error() )
        }
        else
        {
            Ok( gpio_ptr as *mut u32 )
        }
    }
}

fn blink_led( gpio_ptr : *mut u32 ) {
    let set0_reg : *mut u32 = gpio_ptr.wrapping_offset( GPSET0_OFFSET_ADDR );
    let clr0_reg : *mut u32 = gpio_ptr.wrapping_offset( GPCLR0_OFFSET_ADDR );
    let lev0_reg : *mut u32 = gpio_ptr.wrapping_offset( GPLEV0_OFFSET_ADDR );
    let level : u32;

    unsafe {
       level = ( read_volatile( lev0_reg ) >> LEV27_BIT ) & 0x00000001;
    }

    if level == GPLEV_HIGH {
        unsafe {
            write_volatile( clr0_reg, *clr0_reg | ( ON << CLR27_BIT ) );
        }
    }
    else
    {
        unsafe {
            write_volatile( set0_reg, *set0_reg | ( ON << SET27_BIT ) );
        }
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(test, allow(dead_code))]
#![feature(collections, os, io, std_misc, libc)]
#![cfg_attr(not(windows), feature(path))]

extern crate libc;

mod args;
mod matching;
mod screen;
#[cfg(not(windows))] mod ansi;

use std::os;
use std::old_io;
use std::cmp::min;
use screen::Screen;
use screen::Key;
use screen::Key::*;
use self::SearchState::*;

fn main() {
    let args = match args::parse_args() {
        Some(args) => args,
        None => {
            os::set_exit_status(1);
            return;
        },
    };

    if args.help { return; }

    let choices = read_choices();
    let initial_search = args.initial_search.clone();
    let choices = choices.iter().map(|x| &x[]).collect::<Vec<&str>>();
    if args.use_first {
        let matches = matching::compute_matches(&choices, &initial_search);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        event_loop(&choices, &initial_search);
    }
}

fn event_loop(choices: &[&str], initial_search: &str) {
    let mut search = Search::new(choices, initial_search.to_string());
    let mut screen = Screen::open_screen();

    loop {
        search.recompute_matches();

        match search.state {
            InProgress => draw_screen(&mut screen, &search),
            _ => break,
        }

        let keys = screen.get_buffered_keys();
        for key in keys.iter() {
            handle_key(&mut search, key, screen.visible_choices);
        }
    }

    screen.move_cursor_to_bottom();
    println!("{}", search.get_selection());
}

fn handle_key(search: &mut Search, key: &Key, visible_choices: u16) {
    match *key {
        Char(x) => search.append(x),
        Backspace => search.backspace(),
        Control('h') => search.backspace(),
        Control('w') => search.delete_word(),
        Control('u') => search.clear_query(),
        Control('c') => search.cancel(),
        Control('g') => search.cancel(),
        Control('n') => search.down(visible_choices),
        Control('p') => search.up(),
        Enter => search.done(),
        _ => {}
    }
}

struct Search<'a> {
    choices: &'a [&'a str],
    query: String,
    matches: Vec<&'a str>,
    stale: bool,
    index: usize,
    state: SearchState,
}

#[derive(PartialEq, Eq)]
enum SearchState {
    InProgress,
    Done,
    Canceled,
}

impl<'a> Search<'a> {
    fn new(choices: &'a [&'a str], initial_search: String) -> Search<'a> {
        let matches = choices.to_vec();
        Search {
            choices: choices,
            query: initial_search,
            matches: matches,
            stale: true,
            index: 0,
            state: InProgress,
        }
    }

    fn up(&mut self) {
        self.index = if self.index == 0 { 0 } else { self.index - 1 };
    }

    fn down(&mut self, visible_choices: u16) {
        let limit = min(visible_choices as usize - 1, self.matches.len() - 1);
        self.index = min(self.index + 1, limit);
    }

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.index = 0;
        self.matches = self.choices.to_vec();
    }

    fn delete_word(&mut self) {
        self.stale = true;
        delete_last_word(&mut self.query);
        self.matches = self.choices.to_vec();
    }

    fn append(&mut self, c: char) {
        self.query.push(c);
        self.stale = true;
        self.index = 0;
    }

    fn clear_query(&mut self) {
        self.query.clear();
        self.matches = self.choices.to_vec();
    }

    fn recompute_matches(&mut self) {
        if self.stale {
            self.matches = matching::compute_matches(&self.matches, &self.query);
            self.stale = false;
        }
    }

    fn get_selection(&mut self) -> String {
        if self.state == Canceled {
            "".to_string()
        } else {
            self.recompute_matches();
            self.matches.get(self.index).unwrap_or(&"").to_string()
        }
    }

    fn cancel(&mut self) {
        self.state = Canceled
    }

    fn done(&mut self) {
        self.state = Done
    }
}

fn draw_screen(screen: &mut Screen, search: &Search) {
    screen.hide_cursor();
    screen.blank_screen();
    screen.write(&format!("> {} ({} choices)\n", search.query, search.choices.len()));

    print_matches(screen, &search.matches, search.index);

    let start = screen.start_line;
    screen.move_cursor(start, 2 + search.query.len() as u16);
    screen.show_cursor();
}

fn print_matches(screen: &mut Screen, matches: &Vec<&str>, index: usize) {
    let mut i = 1;
    for choice in matches.iter() {
        if i == index + 1 {
            screen.write_inverted(&choice);
        } else {
            screen.write(&choice);
        }
        if i >= screen.visible_choices as usize {
            return;
        } else {
            screen.write("\n");
        }
        i += 1;
    }
}

fn read_choices() -> Vec<String> {
    let mut stdin = old_io::stdio::stdin();
    let mut lines = Vec::new();

    while let Ok(mut s) = stdin.read_line() {
        trim(&mut s);
        lines.push(s);
    }

    lines
}

pub fn trim(s: &mut String) {
    while let Some(x) = s.pop() {
        if x != '\n' {
            s.push(x);
            return;
        }
    }
}

fn delete_last_word(s: &mut String) {
    let mut deleted_something = false;
    while let Some(x) = s.pop() {
        if x == ' ' {
            if deleted_something {
                s.push(x);
                return;
            }
        } else {
            deleted_something = true;
        }
    }
}

#[test]
fn trim_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        trim(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("\n", "");
    should_become("\n\n", "");
    should_become("asdf", "asdf");
    should_become("asdf\n", "asdf");
    should_become("asdf\nasdf\n", "asdf\nasdf");
}

#[test]
fn delete_word_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        delete_last_word(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("a", "");
    should_become("asdf", "");
    should_become("asdf asdf asdf", "asdf asdf ");
    should_become("asdf asdf asdf ", "asdf asdf ");
    should_become("asdf asdf asdf  ", "asdf asdf ");
}
<commit_msg>Show match count in prompt<commit_after>#![cfg_attr(test, allow(dead_code))]
#![feature(collections, os, io, std_misc, libc)]
#![cfg_attr(not(windows), feature(path))]

extern crate libc;

mod args;
mod matching;
mod screen;
#[cfg(not(windows))] mod ansi;

use std::os;
use std::old_io;
use std::cmp::min;
use screen::Screen;
use screen::Key;
use screen::Key::*;
use self::SearchState::*;

fn main() {
    let args = match args::parse_args() {
        Some(args) => args,
        None => {
            os::set_exit_status(1);
            return;
        },
    };

    if args.help { return; }

    let choices = read_choices();
    let initial_search = args.initial_search.clone();
    let choices = choices.iter().map(|x| &x[]).collect::<Vec<&str>>();
    if args.use_first {
        let matches = matching::compute_matches(&choices, &initial_search);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        event_loop(&choices, &initial_search);
    }
}

fn event_loop(choices: &[&str], initial_search: &str) {
    let mut search = Search::new(choices, initial_search.to_string());
    let mut screen = Screen::open_screen();

    loop {
        search.recompute_matches();

        match search.state {
            InProgress => draw_screen(&mut screen, &search),
            _ => break,
        }

        let keys = screen.get_buffered_keys();
        for key in keys.iter() {
            handle_key(&mut search, key, screen.visible_choices);
        }
    }

    screen.move_cursor_to_bottom();
    println!("{}", search.get_selection());
}

fn handle_key(search: &mut Search, key: &Key, visible_choices: u16) {
    match *key {
        Char(x) => search.append(x),
        Backspace => search.backspace(),
        Control('h') => search.backspace(),
        Control('w') => search.delete_word(),
        Control('u') => search.clear_query(),
        Control('c') => search.cancel(),
        Control('g') => search.cancel(),
        Control('n') => search.down(visible_choices),
        Control('p') => search.up(),
        Enter => search.done(),
        _ => {}
    }
}

struct Search<'a> {
    choices: &'a [&'a str],
    query: String,
    matches: Vec<&'a str>,
    stale: bool,
    index: usize,
    state: SearchState,
}

#[derive(PartialEq, Eq)]
enum SearchState {
    InProgress,
    Done,
    Canceled,
}

impl<'a> Search<'a> {
    fn new(choices: &'a [&'a str], initial_search: String) -> Search<'a> {
        let matches = choices.to_vec();
        Search {
            choices: choices,
            query: initial_search,
            matches: matches,
            stale: true,
            index: 0,
            state: InProgress,
        }
    }

    fn up(&mut self) {
        self.index = if self.index == 0 { 0 } else { self.index - 1 };
    }

    fn down(&mut self, visible_choices: u16) {
        let limit = min(visible_choices as usize - 1, self.matches.len() - 1);
        self.index = min(self.index + 1, limit);
    }

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.index = 0;
        self.matches = self.choices.to_vec();
    }

    fn delete_word(&mut self) {
        self.stale = true;
        delete_last_word(&mut self.query);
        self.matches = self.choices.to_vec();
    }

    fn append(&mut self, c: char) {
        self.query.push(c);
        self.stale = true;
        self.index = 0;
    }

    fn clear_query(&mut self) {
        self.query.clear();
        self.matches = self.choices.to_vec();
    }

    fn recompute_matches(&mut self) {
        if self.stale {
            self.matches = matching::compute_matches(&self.matches, &self.query);
            self.stale = false;
        }
    }

    fn get_selection(&mut self) -> String {
        if self.state == Canceled {
            "".to_string()
        } else {
            self.recompute_matches();
            self.matches.get(self.index).unwrap_or(&"").to_string()
        }
    }

    fn cancel(&mut self) {
        self.state = Canceled
    }

    fn done(&mut self) {
        self.state = Done
    }
}

fn draw_screen(screen: &mut Screen, search: &Search) {
    screen.hide_cursor();
    screen.blank_screen();
    screen.write(&format!("> {} ({}/{} choices)\n", search.query, search.matches.len(), search.choices.len()));

    print_matches(screen, &search.matches, search.index);

    let start = screen.start_line;
    screen.move_cursor(start, 2 + search.query.len() as u16);
    screen.show_cursor();
}

fn print_matches(screen: &mut Screen, matches: &Vec<&str>, index: usize) {
    let mut i = 1;
    for choice in matches.iter() {
        if i == index + 1 {
            screen.write_inverted(&choice);
        } else {
            screen.write(&choice);
        }
        if i >= screen.visible_choices as usize {
            return;
        } else {
            screen.write("\n");
        }
        i += 1;
    }
}

fn read_choices() -> Vec<String> {
    let mut stdin = old_io::stdio::stdin();
    let mut lines = Vec::new();

    while let Ok(mut s) = stdin.read_line() {
        trim(&mut s);
        lines.push(s);
    }

    lines
}

pub fn trim(s: &mut String) {
    while let Some(x) = s.pop() {
        if x != '\n' {
            s.push(x);
            return;
        }
    }
}

fn delete_last_word(s: &mut String) {
    let mut deleted_something = false;
    while let Some(x) = s.pop() {
        if x == ' ' {
            if deleted_something {
                s.push(x);
                return;
            }
        } else {
            deleted_something = true;
        }
    }
}

#[test]
fn trim_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        trim(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("\n", "");
    should_become("\n\n", "");
    should_become("asdf", "asdf");
    should_become("asdf\n", "asdf");
    should_become("asdf\nasdf\n", "asdf\nasdf");
}

#[test]
fn delete_word_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        delete_last_word(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("a", "");
    should_become("asdf", "");
    should_become("asdf asdf asdf", "asdf asdf ");
    should_become("asdf asdf asdf ", "asdf asdf ");
    should_become("asdf asdf asdf  ", "asdf asdf ");
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use foxbox_core::traits::Controller;
use foxbox_taxonomy::manager::*;
use iron::{AfterMiddleware, Chain, Handler, Iron, IronResult, Request, Response, Protocol};
use iron_cors::CORS;
use iron::error::IronError;
use iron::method::Method;
use iron::status::Status;
use mount::Mount;
use router::NoRoute;
use static_router;
use std::net::SocketAddr;
use std::sync::Arc;
use std::thread;
use taxonomy_router;

const THREAD_COUNT: usize = 8;

struct Custom404;

impl AfterMiddleware for Custom404 {
    fn catch(&self, _: &mut Request, err: IronError) -> IronResult<Response> {
        use std::io::Error as StdError;
        use std::io::ErrorKind;

        if err.error.downcast::<NoRoute>().is_some() {
            // Router error
            return Ok(Response::with((Status::NotFound, format!("Unknown resource: {}", err))));
        } else if let Some(err) = err.error.downcast::<StdError>() {
            // StaticFile error
            if err.kind() == ErrorKind::NotFound {
                return Ok(Response::with((Status::NotFound, format!("Unknown resource: {}", err))));
            }
        }

        // Just let other errors go through, like 401.
        Err(err)
    }
}

struct Ping;

impl Handler for Ping {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        Ok(Response::with(Status::NoContent))
    }
}

pub struct HttpServer<T: Controller> {
    controller: T,
}

impl<T: Controller> HttpServer<T> {
    pub fn new(controller: T) -> Self {
        HttpServer { controller: controller }
    }

    pub fn start(&mut self, adapter_api: &Arc<AdapterManager>) {
        let taxonomy_chain = taxonomy_router::create(self.controller.clone(), adapter_api);

        let users_manager = self.controller.get_users_manager();
        let mut mount = Mount::new();
        mount.mount("/", static_router::create(users_manager.clone()))
            .mount("/ping", Ping)
            .mount("/api/v1", taxonomy_chain)
            .mount("/users", users_manager.get_router_chain());

        let mut chain = Chain::new(mount);
        chain.link_after(Custom404);

        let cors = CORS::new(vec![(vec![Method::Get], "ping".to_owned()),

                                  // Taxonomy router paths. Keep in sync with taxonomy_router.rs
                                  (vec![Method::Get, Method::Post], "api/v1/services".to_owned()),
                                  (vec![Method::Post, Method::Delete],
                                   "api/v1/services/tags".to_owned()),
                                  (vec![Method::Get, Method::Post], "api/v1/channels".to_owned()),
                                  (vec![Method::Put], "api/v1/channels/get".to_owned()),
                                  (vec![Method::Put], "api/v1/channels/set".to_owned()),
                                  (vec![Method::Post, Method::Delete],
                                   "api/v1/channels/tags".to_owned())]);
        chain.link_after(cors);

        let addrs: Vec<_> = self.controller.http_as_addrs().unwrap().collect();

        if self.controller.get_tls_enabled() {
            let certificate_manager = self.controller.get_certificate_manager();
            let fingerprint = certificate_manager.get_box_certificate()
                .unwrap()
                .get_certificate_fingerprint();
            // Get the certificate record for the hostname, and use its certificate and
            // private key files.
            let host_name = format!("remote.{}.{}", fingerprint, self.controller.get_domain());
            let record = certificate_manager.get_certificate(&host_name)
                .expect("Unable to start https server without remote certificate");
            start_server(addrs,
                         chain,
                         Protocol::Https {
                             certificate: record.cert_file,
                             key: record.private_key_file,
                         });
        } else {
            start_server(addrs, chain, Protocol::Http);
        }
    }
}

fn start_server(addrs: Vec<SocketAddr>, chain: Chain, protocol: Protocol) {

    thread::Builder::new()
        .name("HttpServer".to_owned())
        .spawn(move || {
            Iron::new(chain)
                .listen_with(addrs[0], THREAD_COUNT, protocol, None)
                .unwrap();
        })
        .unwrap();
}

#[cfg(test)]
describe! ping {
    before_each {
        use mount::Mount;
        use iron::Headers;
        use iron::status::Status;
        use iron_test::request;
        use super::Ping;

        let mut mount = Mount::new();
        mount.mount("/ping", Ping);
    }

    it "should response 204 NoContent" {
        let response = request::get("http://localhost:3000/ping",
                                    Headers::new(),
                                    &mount).unwrap();
        assert_eq!(response.status.unwrap(), Status::NoContent);
    }
}

#[cfg(test)]
describe! http_server {
    before_each {
        extern crate hyper;

        use foxbox_taxonomy::manager::AdapterManager;
        use std::thread;
        use std::sync::Arc;
        use std::time::Duration;
        use stubs::controller::ControllerStub;

        let taxo_manager = Arc::new(AdapterManager::new(None));

        let mut http_server = HttpServer::new(ControllerStub::new());
        http_server.start(&taxo_manager);
        // HACK: Let some time for the http server to start.
        thread::sleep(Duration::new(3, 0));
    }

    it "should get the appropriate CORS headers" {
        use iron::headers;
        use iron::method::Method;

        let endpoints = vec![
            (vec![Method::Get], "ping".to_owned())
        ];
        let client = hyper::Client::new();
        for endpoint in endpoints {
            let (_, path) = endpoint;
            let path = "http://localhost:3000/".to_owned() +
                       &(path.replace(":", "foo"));
            let res = client.get(&path).send();
            let headers = &res.unwrap().headers;
            assert!(headers.has::<headers::AccessControlAllowOrigin>());
            assert!(headers.has::<headers::AccessControlAllowHeaders>());
            assert!(headers.has::<headers::AccessControlAllowMethods>());
        };
    }

    it "should respond with 404" {
        use iron::status::Status;
        use std::io::Read;

        let client = hyper::Client::new();
        let path = "http://localhost:3000/foo/bar".to_owned();
        let mut res = client.get(&path).send().unwrap();
        assert_eq!(res.status, Status::NotFound);
        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();
        assert_eq!(body, "Unknown resource: No such file or \
                   directory (os error 2)".to_owned());
    }
}
<commit_msg>Add HTTP headers to improve seurity. (#572)<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use foxbox_core::traits::Controller;
use foxbox_taxonomy::manager::*;
use iron::{AfterMiddleware, Chain, Handler, Iron, IronResult, Request, Response, Protocol};
use iron_cors::CORS;
use iron::error::IronError;
use iron::method::Method;
use iron::status::Status;
use mount::Mount;
use router::NoRoute;
use static_router;
use std::net::SocketAddr;
use std::sync::Arc;
use std::thread;
use taxonomy_router;

const THREAD_COUNT: usize = 8;

// 404 middleware.
struct Custom404;

impl AfterMiddleware for Custom404 {
    fn catch(&self, _: &mut Request, err: IronError) -> IronResult<Response> {
        use std::io::Error as StdError;
        use std::io::ErrorKind;

        if err.error.downcast::<NoRoute>().is_some() {
            // Router error
            return Ok(Response::with((Status::NotFound, format!("Unknown resource: {}", err))));
        } else if let Some(err) = err.error.downcast::<StdError>() {
            // StaticFile error
            if err.kind() == ErrorKind::NotFound {
                return Ok(Response::with((Status::NotFound, format!("Unknown resource: {}", err))));
            }
        }

        // Just let other errors go through, like 401.
        Err(err)
    }
}

// Middleware that adds security related headers.
// See https://wiki.mozilla.org/Security/Guidelines/Web_Security
struct SecurityHeaders;

impl AfterMiddleware for SecurityHeaders {
    fn after(&self, _: &mut Request, mut res: Response) -> IronResult<Response> {
        use iron::{headers, Set};
        use iron::modifiers::Header;

        // HSTS
        let header = headers::StrictTransportSecurity {
            include_subdomains: false,
            max_age: 31536000u64,
        };
        res.set_mut(Header(header));

        // Referrer-Policy
        res.set_mut(Header(headers::ReferrerPolicy::NoReferrer));

        // X-Frame-Options
        header! { (XFrameOptions, "X-Frame-Options") => [String] }
        res.set_mut(Header(XFrameOptions("DENY".to_owned())));

        // X-Content-Type-Options
        header! { (XContentTypeOptions, "X-Content-Type-Options") => [String] }
        res.set_mut(Header(XContentTypeOptions("nosniff".to_owned())));

        // Content-Security-Policy
        // TODO: refine the value of this header.
        header! { (Xcsp, "Content-Security-Policy") => [String] }
        res.set_mut(Header(Xcsp("default-src 'none'; frame-ancestors 'none'".to_owned())));

        // X-XSS-Protection
        header! { (XXSSProtection, "X-XSS-Protection") => [String] }
        res.set_mut(Header(XXSSProtection("1; mode=block".to_owned())));

        Ok(res)
    }
}

struct Ping;

impl Handler for Ping {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        Ok(Response::with(Status::NoContent))
    }
}

pub struct HttpServer<T: Controller> {
    controller: T,
}

impl<T: Controller> HttpServer<T> {
    pub fn new(controller: T) -> Self {
        HttpServer { controller: controller }
    }

    pub fn start(&mut self, adapter_api: &Arc<AdapterManager>) {
        let taxonomy_chain = taxonomy_router::create(self.controller.clone(), adapter_api);

        let users_manager = self.controller.get_users_manager();
        let mut mount = Mount::new();
        mount.mount("/", static_router::create(users_manager.clone()))
            .mount("/ping", Ping)
            .mount("/api/v1", taxonomy_chain)
            .mount("/users", users_manager.get_router_chain());

        let mut chain = Chain::new(mount);
        chain.link_after(Custom404);

        let cors = CORS::new(vec![(vec![Method::Get], "ping".to_owned()),

                                  // Taxonomy router paths. Keep in sync with taxonomy_router.rs
                                  (vec![Method::Get, Method::Post], "api/v1/services".to_owned()),
                                  (vec![Method::Post, Method::Delete],
                                   "api/v1/services/tags".to_owned()),
                                  (vec![Method::Get, Method::Post], "api/v1/channels".to_owned()),
                                  (vec![Method::Put], "api/v1/channels/get".to_owned()),
                                  (vec![Method::Put], "api/v1/channels/set".to_owned()),
                                  (vec![Method::Post, Method::Delete],
                                   "api/v1/channels/tags".to_owned())]);
        chain.link_after(cors);

        let addrs: Vec<_> = self.controller.http_as_addrs().unwrap().collect();

        if self.controller.get_tls_enabled() {
            // When running with TLS enabled, add the security headers.
            chain.link_after(SecurityHeaders);

            let certificate_manager = self.controller.get_certificate_manager();
            let fingerprint = certificate_manager.get_box_certificate()
                .unwrap()
                .get_certificate_fingerprint();
            // Get the certificate record for the hostname, and use its certificate and
            // private key files.
            let host_name = format!("remote.{}.{}", fingerprint, self.controller.get_domain());
            let record = certificate_manager.get_certificate(&host_name)
                .expect("Unable to start https server without remote certificate");
            start_server(addrs,
                         chain,
                         Protocol::Https {
                             certificate: record.full_chain.unwrap_or(record.cert_file),
                             key: record.private_key_file,
                         });
        } else {
            start_server(addrs, chain, Protocol::Http);
        }
    }
}

fn start_server(addrs: Vec<SocketAddr>, chain: Chain, protocol: Protocol) {

    thread::Builder::new()
        .name("HttpServer".to_owned())
        .spawn(move || {
            Iron::new(chain)
                .listen_with(addrs[0], THREAD_COUNT, protocol, None)
                .unwrap();
        })
        .unwrap();
}

#[cfg(test)]
describe! ping {
    before_each {
        use mount::Mount;
        use iron::Headers;
        use iron::status::Status;
        use iron_test::request;
        use super::Ping;

        let mut mount = Mount::new();
        mount.mount("/ping", Ping);
    }

    it "should response 204 NoContent" {
        let response = request::get("http://localhost:3000/ping",
                                    Headers::new(),
                                    &mount).unwrap();
        assert_eq!(response.status.unwrap(), Status::NoContent);
    }
}

#[cfg(test)]
describe! http_server {
    before_each {
        extern crate hyper;

        use foxbox_taxonomy::manager::AdapterManager;
        use std::thread;
        use std::sync::Arc;
        use std::time::Duration;
        use stubs::controller::ControllerStub;

        let taxo_manager = Arc::new(AdapterManager::new(None));

        let mut http_server = HttpServer::new(ControllerStub::new());
        http_server.start(&taxo_manager);
        // HACK: Let some time for the http server to start.
        thread::sleep(Duration::new(3, 0));
    }

    it "should get the appropriate CORS headers" {
        use iron::headers;
        use iron::method::Method;

        let endpoints = vec![
            (vec![Method::Get], "ping".to_owned())
        ];
        let client = hyper::Client::new();
        for endpoint in endpoints {
            let (_, path) = endpoint;
            let path = "http://localhost:3000/".to_owned() +
                       &(path.replace(":", "foo"));
            let res = client.get(&path).send();
            let headers = &res.unwrap().headers;
            assert!(headers.has::<headers::AccessControlAllowOrigin>());
            assert!(headers.has::<headers::AccessControlAllowHeaders>());
            assert!(headers.has::<headers::AccessControlAllowMethods>());
        };
    }

    it "should respond with 404" {
        use iron::status::Status;
        use std::io::Read;

        let client = hyper::Client::new();
        let path = "http://localhost:3000/foo/bar".to_owned();
        let mut res = client.get(&path).send().unwrap();
        assert_eq!(res.status, Status::NotFound);
        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();
        assert_eq!(body, "Unknown resource: No such file or \
                   directory (os error 2)".to_owned());
    }
}
<|endoftext|>