<commit_before>extern crate byteorder;
use std::io::Cursor;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

macro_rules! define_instructions {
    ($($variant:ident, $value:expr, $name:expr, $num_operands:expr;)*) => (
        #[derive(Clone, Copy, Debug, Eq, PartialEq)]
        #[repr(u8)]
        enum Inst {
            $($variant = $value),*
        }

        impl Inst {
            fn from_u8(inst: u8) -> Option<Inst> {
                match inst {
                    $($value => Some(Inst::$variant),)*
                    _ => None
                }
            }

            fn from_str(inst: &str) -> Option<Inst> {
                match inst {
                    $($name => Some(Inst::$variant),)*
                    _ => None
                }
            }

            fn num_operands(&self) -> u8 {
                match *self {
                    $(Inst::$variant => $num_operands),*
                }
            }
        }
    )
}

// Bytecode instruction opcodes. The values of these opcodes should never change, to remain
// compatible with existing bytecode programs.
define_instructions! {
    Nop,   0, "nop",   0;
    Push,  1, "push",  1;
    Dup,   2, "dup",   0;
    Pop,   3, "pop",   0;
    Swap,  4, "swap",  0;
    Add,   5, "add",   0;
    Print, 6, "print", 0;
    Halt,  7, "halt",  0;
    Jump,  8, "jump",  1;
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum VmError {
    InvalidOpcode,
    UnexpectedProgramEnd, // Hit end of program while reading operand.
    StackOverflow,
    StackUnderflow,
}

fn execute(program: &[u8],
           stack: &mut [i32],
           mut stack_idx: usize) -> Result<usize, VmError> {
    use VmError::*;

    let stack_size = stack.len();
    let mut opcodes = Cursor::new(program);

    while let Ok(opcode) = opcodes.read_u8() {
        let inst = try!(Inst::from_u8(opcode).ok_or(InvalidOpcode));
        match inst {
            Inst::Nop => {},

            Inst::Push => {
                let val = try!(opcodes.read_i32::<LittleEndian>().or(Err(UnexpectedProgramEnd)));
                let stack_top = try!(stack.get_mut(stack_idx).ok_or(StackOverflow));
                *stack_top = val;
                stack_idx += 1;
            },

            Inst::Dup => {
                if stack_idx < 1 { return Err(StackUnderflow); }
                if stack_idx >= stack_size { return Err(StackOverflow); }
                unsafe {
                    *stack.get_unchecked_mut(stack_idx) = *stack.get_unchecked(stack_idx - 1);
                }
                stack_idx += 1;
            },

            Inst::Pop => {
                if stack_idx < 1 { return Err(StackUnderflow); }
                stack_idx -= 1;
            },

            Inst::Swap => {
                if stack_idx < 2 { return Err(StackUnderflow); }
                unsafe {
                    let tmp = *stack.get_unchecked(stack_idx - 1);
                    *stack.get_unchecked_mut(stack_idx - 1) = *stack.get_unchecked(stack_idx - 2);
                    *stack.get_unchecked_mut(stack_idx - 2) = tmp;
                }
            },

            Inst::Add => {
                if stack_idx < 2 { return Err(StackUnderflow); }
                unsafe {
                    *stack.get_unchecked_mut(stack_idx - 2) += *stack.get_unchecked(stack_idx - 1);
                }
                stack_idx -= 1;
            },

            Inst::Print => {
                let val = *try!(stack.get(stack_idx - 1).ok_or(StackUnderflow));
                println!("{}", val);
                stack_idx -= 1;
            },

            Inst::Halt => {
                break;
            },

            Inst::Jump => {
                let delta = try!(opcodes.read_i32::<LittleEndian>().or(Err(UnexpectedProgramEnd)));
                let operand_size = std::mem::size_of::<i32>() as i64;
                let addr = (opcodes.position() as i64 + delta as i64 - operand_size) as u64;
                opcodes.set_position(addr);
            },
        }
    }

    Ok(stack_idx)
}

fn assemble(source: &str) -> Vec<u8> {
    use std::collections::HashMap;

    let mut program: Vec<u8> = Vec::new();
    let mut label_definitions: HashMap<&str, usize> = HashMap::new();
    let mut label_uses: Vec<(&str, usize)> = Vec::new();

    for line in source.lines() {
        let mut tokens = line.split(char::is_whitespace).filter(|s| !s.is_empty());
        let mut first_token = tokens.next();

        // Parse an optional label at the start of the line.
        if let Some(label) = first_token {
            if label.chars().next_back() == Some(':') {
                let label_name = &label[..label.len() - 1];
                if label_definitions.insert(label_name, program.len()).is_some() {
                    panic!("Attempted to redefine label '{}'", label_name);
                }
                first_token = tokens.next()
            }
        }

        // Parse the rest of the line if it's not blank.
        if let Some(opcode) = first_token {
            // Parse the instruction name.
            let inst = Inst::from_str(opcode).unwrap_or_else(|| {
                panic!("Unrecognized instruction '{}'", opcode)
            });
            program.push(inst as u8);

            // Parse the operands.
            for _ in 0..inst.num_operands() {
                let operand = tokens.next().unwrap_or_else(|| {
                    panic!("Missing one or more operands after '{}'", opcode)
                });

                if operand.chars().next() == Some('@') {
                    let label_name = &operand[1..];
                    label_uses.push((label_name, program.len()));

                    // Push four zero bytes to be overwritten by the label location later.
                    for _ in 0..4 { program.push(0); }
                } else if let Ok(number) = operand.parse::<i32>() {
                    let operand_index = program.len();
                    for _ in 0..4 { program.push(0); }
                    (&mut program[operand_index..]).write_i32::<LittleEndian>(number).unwrap();
                } else {
                    panic!("Expected label or valid 32-bit signed integer after '{}', not '{}'",
                          opcode, operand);
                }
            }
        }
    }

    // Resolve label references and fill in their values in the bytecode.
    for (label_name, use_index) in label_uses {
        let target_index = *label_definitions.get(label_name).unwrap_or_else(|| {
            panic!("Reference to undefined label '{}'", label_name);
        });
        let delta = target_index as i32 - use_index as i32;
        (&mut program[use_index..]).write_i32::<LittleEndian>(delta).unwrap();
    }

    program
}

fn main() {
    let source = r"
        push 1
        label: push 2
        add
        dup
        print
        jump @label
    ";

    let program = assemble(source);

    execute(&program, &mut [0; 256], 0).unwrap();
}
<commit_msg>Add many new instructions.<commit_after>extern crate byteorder;
use std::io::Cursor;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

macro_rules! define_instructions {
    ($($variant:ident, $value:expr, $name:expr, $num_operands:expr;)*) => (
        #[derive(Clone, Copy, Debug, Eq, PartialEq)]
        #[repr(u8)]
        enum Inst {
            $($variant = $value),*
        }

        impl Inst {
            fn from_u8(inst: u8) -> Option<Inst> {
                match inst {
                    $($value => Some(Inst::$variant),)*
                    _ => None
                }
            }

            fn from_str(inst: &str) -> Option<Inst> {
                match inst {
                    $($name => Some(Inst::$variant),)*
                    _ => None
                }
            }

            fn num_operands(&self) -> u8 {
                match *self {
                    $(Inst::$variant => $num_operands),*
                }
            }
        }
    )
}

// Bytecode instruction opcodes. The values of these opcodes should never change, to remain
// compatible with existing bytecode programs.
define_instructions! {
    Nop,   0,  "nop",   0;
    Print, 1,  "print", 0;
    Halt,  2,  "halt",  0;
    Push,  3,  "push",  1;
    Dup,   4,  "dup",   0;
    Pop,   5,  "pop",   0;
    Swap,  6,  "swap",  0;
    Add,   7,  "add",   0;
    Sub,   8,  "sub",   0;
    Mul,   9,  "mul",   0;
    Div,   10, "div",   0;
    Mod,   11, "mod",   0;
    Eq,    12, "eq",    0;
    Lt,    13, "lt",    0;
    Lte,   14, "lte",   0;
    Gt,    15, "gt",    0;
    Gte,   16, "gte",   0;
    Jz,    17, "jz",    1;
    Jnz,   18, "jnz",   1;
    Jump,  19, "jump",  1;
    Call,  20, "call",  1;
    Ret,   21, "ret",   0;
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum VmError {
    InvalidOpcode,
    UnexpectedProgramEnd, // Hit end of program while reading operand.
    StackOverflow,
    StackUnderflow,
    ControlStackOverflow,
    ControlStackUnderflow,
}

#[derive(Debug, Eq, PartialEq)]
struct Vm {
    stack: Box<[i32]>,
    stack_idx: usize,
    control_stack: Box<[i32]>,
    control_stack_idx: usize,
}

impl Vm {
    fn execute(&mut self, program: &[u8]) -> Result<(), VmError> {
        use VmError::*;

        let stack = &mut self.stack;
        let mut stack_idx = self.stack_idx;
        let mut opcodes = Cursor::new(program);

        while let Ok(opcode) = opcodes.read_u8() {
            let inst = try!(Inst::from_u8(opcode).ok_or(InvalidOpcode));
            match inst {
                Inst::Nop => {},

                Inst::Print => {
                    let val = *try!(stack.get(stack_idx - 1).ok_or(StackUnderflow));
                    println!("{}", val);
                    stack_idx -= 1;
                },

                Inst::Halt => {
                    break;
                },

                Inst::Push => {
                    let val = try!(opcodes.read_i32::<LittleEndian>().or(Err(UnexpectedProgramEnd)));
                    let stack_top = try!(stack.get_mut(stack_idx).ok_or(StackOverflow));
                    *stack_top = val;
                    stack_idx += 1;
                },

                Inst::Dup => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    if stack_idx >= stack.len() { return Err(StackOverflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx) = *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx += 1;
                },

                Inst::Pop => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    stack_idx -= 1;
                },

                Inst::Swap => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let tmp = *stack.get_unchecked(stack_idx - 1);
                        *stack.get_unchecked_mut(stack_idx - 1) = *stack.get_unchecked(stack_idx - 2);
                        *stack.get_unchecked_mut(stack_idx - 2) = tmp;
                    }
                },

                Inst::Add => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) += *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Sub => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) -= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Mul => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) *= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Div => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) /= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Mod => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        *stack.get_unchecked_mut(stack_idx - 2) %= *stack.get_unchecked(stack_idx - 1);
                    }
                    stack_idx -= 1;
                },

                Inst::Eq => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 == val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Lt => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 < val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Lte => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 <= val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Gt => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 > val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Gte => {
                    if stack_idx < 2 { return Err(StackUnderflow); }
                    unsafe {
                        let val1 = *stack.get_unchecked_mut(stack_idx - 1);
                        let ptr2 = stack.get_unchecked_mut(stack_idx - 2);
                        *ptr2 = (*ptr2 >= val1) as i32;
                    }
                    stack_idx -= 1;
                },

                Inst::Jz => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    let condition = unsafe { *stack.get_unchecked(stack_idx - 1) };
                    if condition == 0 { try!(jump(&mut opcodes)); }
                    stack_idx -= 1;
                },

                Inst::Jnz => {
                    if stack_idx < 1 { return Err(StackUnderflow); }
                    let condition = unsafe { *stack.get_unchecked(stack_idx - 1) };
                    if condition != 0 { try!(jump(&mut opcodes)); }
                    stack_idx -= 1;
                },

                Inst::Jump => {
                    try!(jump(&mut opcodes));
                },

                Inst::Call => {
                    let control_stack_top = try!(self.control_stack.get_mut(self.control_stack_idx).ok_or(ControlStackOverflow));
                    *control_stack_top = opcodes.position() as i32 + 4;
                    try!(jump(&mut opcodes));
                    self.control_stack_idx += 1;
                },

                Inst::Ret => {
                    let addr = *try!(self.control_stack.get_mut(self.control_stack_idx - 1).ok_or(ControlStackOverflow));
                    opcodes.set_position(addr as u64);
                    self.control_stack_idx -= 1;
                },
            }
        }

        fn jump(opcodes: &mut Cursor<&[u8]>) -> Result<(), VmError> {
            let delta = try!(opcodes.read_i32::<LittleEndian>().or(Err(UnexpectedProgramEnd)));
            let operand_size = std::mem::size_of::<i32>() as i64;
            let addr = (opcodes.position() as i64 + delta as i64 - operand_size) as u64;
            opcodes.set_position(addr);
            Ok(())
        }

        self.stack_idx = stack_idx;
        Ok(())
    }
}

fn assemble(source: &str) -> Vec<u8> {
    use std::collections::HashMap;

    let mut program: Vec<u8> = Vec::new();
    let mut label_definitions: HashMap<&str, usize> = HashMap::new();
    let mut label_uses: Vec<(&str, usize)> = Vec::new();

    for line in source.split(|c| c == '\n' || c == ';') {
        let mut tokens = line.split(char::is_whitespace).filter(|s| !s.is_empty());
        let mut first_token = tokens.next();

        // Parse an optional label at the start of the line.
        if let Some(label) = first_token {
            if label.chars().next_back() == Some(':') {
                let label_name = &label[..label.len() - 1];
                if label_definitions.insert(label_name, program.len()).is_some() {
                    panic!("Attempted to redefine label '{}'", label_name);
                }
                first_token = tokens.next()
            }
        }

        // Parse the rest of the line if it's not blank.
        if let Some(opcode) = first_token {
            // Parse the instruction name.
            let inst = Inst::from_str(opcode).unwrap_or_else(|| {
                panic!("Unrecognized instruction '{}'", opcode)
            });
            program.push(inst as u8);

            // Parse the operands.
            for _ in 0..inst.num_operands() {
                let operand = tokens.next().unwrap_or_else(|| {
                    panic!("Missing one or more operands after '{}'", opcode)
                });

                if operand.chars().next() == Some('@') {
                    let label_name = &operand[1..];
                    label_uses.push((label_name, program.len()));

                    // Push four zero bytes to be overwritten by the label location later.
                    for _ in 0..4 { program.push(0); }
                } else if let Ok(number) = operand.parse::<i32>() {
                    let operand_index = program.len();
                    for _ in 0..4 { program.push(0); }
                    (&mut program[operand_index..]).write_i32::<LittleEndian>(number).unwrap();
                } else {
                    panic!("Expected label or valid 32-bit signed integer after '{}', not '{}'",
                           opcode, operand);
                }
            }
        }
    }

    // Resolve label references and fill in their values in the bytecode.
    for (label_name, use_index) in label_uses {
        let target_index = *label_definitions.get(label_name).unwrap_or_else(|| {
            panic!("Reference to undefined label '{}'", label_name);
        });
        let delta = target_index as i32 - use_index as i32;
        (&mut program[use_index..]).write_i32::<LittleEndian>(delta).unwrap();
    }

    program
}

fn main() {
    let source = r"
        push 1
loop:   call @plus2
        dup; print
        jump @loop

plus2:  push 2
        add
        ret
    ";

    // let source = r"
    //     push 1
    //     push 10
// loop:   dup; push 0; eq; jnz @done
    //     jump @loop
// done:   print
    // ";

    let program = assemble(source);

    let mut vm = Vm {
        stack: Box::new([0; 256]),
        stack_idx: 0,
        control_stack: Box::new([0; 256]),
        control_stack_idx: 0,
    };
    vm.execute(&program).unwrap();
}
<|endoftext|>


<commit_before>extern crate alpm;
#[macro_use]
extern crate clap;
extern crate curl;
extern crate env_logger;
extern crate itertools;
#[macro_use]
extern crate log;
extern crate serde_json;

use clap::App;
use curl::easy::Easy;
use itertools::Itertools;
use serde_json::Value;
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::collections::btree_map::Entry::{Occupied, Vacant};
use std::default::Default;
use std::process::exit;
use std::str;

mod avg;
mod enums;

struct Options {
    format: Option<String>,
    quiet: u64,
    upgradable_only: bool,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            format: None,
            quiet: 0,
            upgradable_only: false,
        }
    }
}

fn main() {
    env_logger::init().expect("env_logger failed");

    let yaml = load_yaml!("cli.yml");
    let args = App::from_yaml(yaml).get_matches();

    let options = Options {
        format: {
            match args.value_of("format") {
                Some(f) => Some(f.to_string()),
                None => None,
            }
        },
        quiet: args.occurrences_of("quiet"),
        upgradable_only: args.is_present("upgradable"),
    };

    let mut avgs = String::new();
    {
        info!("Downloading AVGs...");
        let avgs_url = "https://security.archlinux.org/issues/all.json";

        let mut easy = Easy::new();
        easy.url(avgs_url).expect("curl::Easy::url failed");
        let mut transfer = easy.transfer();
        transfer
            .write_function(|data| {
                avgs.push_str(str::from_utf8(data).expect("str conversion failed"));
                Ok(data.len())
            })
            .expect("write_function failed");
        match transfer.perform() {
            Ok(_) => {}
            Err(_) => {
                println!("Cannot fetch data, please check your network connection!");
                exit(1)
            }
        };
    }

    let pacman = match args.value_of("dbpath") {
        Some(path) => {
            alpm::Alpm::with_dbpath(path.to_string()).expect("alpm::Alpm::with_dbpath failed")
        }
        None => alpm::Alpm::new().expect("alpm::Alpm::new failed"),
    };

    let mut cves: BTreeMap<String, Vec<_>> = BTreeMap::new();
    {
        let json: Value = serde_json::from_str(&avgs).expect("serde_json::from_str failed");

        for avg in json.as_array().expect("Value::as_array failed") {
            let packages = avg["packages"]
                .as_array()
                .expect("Value::as_array failed")
                .iter()
                .map(|s| s.as_str().expect("Value::as_str failed").to_string())
                .collect::<Vec<_>>();

            if !package_is_installed(&pacman, &packages) {
                continue;
            }

            let info = to_avg(avg);

            if info.status != enums::Status::NotAffected {
                for p in packages {
                    match cves.entry(p) {
                        Occupied(c) => c.into_mut(),
                        Vacant(c) => c.insert(vec![]),
                    }.push(info.clone());
                }
            }
        }
    }

    let mut affected_avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();
    for (pkg, avgs) in cves {
        for avg in &avgs {
            if system_is_affected(&pacman, &pkg, avg) {
                match affected_avgs.entry(pkg.clone()) {
                    Occupied(c) => c.into_mut(),
                    Vacant(c) => c.insert(vec![]),
                }.push(avg.clone());
            }
        }
    }

    let merged = merge_avgs(&pacman, &affected_avgs);
    print_avgs(&options, &merged);
}

/// Converts a JSON to an `avg::AVG`
fn to_avg(data: &Value) -> avg::AVG {
    avg::AVG {
        issues: data["issues"]
            .as_array()
            .expect("Value::as_array failed")
            .iter()
            .map(|s| s.as_str().expect("Value::as_str failed").to_string())
            .collect(),
        fixed: match data["fixed"].as_str() {
            Some(s) => Some(s.to_string()),
            None => None,
        },
        severity: data["severity"]
            .as_str()
            .expect("Value::as_str failed")
            .to_string()
            .parse::<enums::Severity>()
            .expect("parse::<Severity> failed"),
        status: data["status"]
            .as_str()
            .expect("Value::as_str failed")
            .to_string()
            .parse::<enums::Status>()
            .expect("parse::<Status> failed"),
    }
}

#[test]
fn test_to_avg() {
    let json: Value = serde_json::from_str(
        "{\"issues\": [\"CVE-1\", \"CVE-2\"], \"fixed\": \"1.0\", \
                               \"severity\": \"High\", \"status\": \"Not affected\"}",
    ).expect("serde_json::from_str failed");

    let avg1 = to_avg(&json);
    assert_eq!(2, avg1.issues.len());
    assert_eq!(Some("1.0".to_string()), avg1.fixed);
    assert_eq!(enums::Severity::High, avg1.severity);
    assert_eq!(enums::Status::NotAffected, avg1.status);

    let json: Value = serde_json::from_str(
        "{\"issues\": [\"CVE-1\"], \"fixed\": null, \
                               \"severity\": \"Low\", \"status\": \"Vulnerable\"}",
    ).expect("serde_json::from_str failed");

    let avg2 = to_avg(&json);
    assert_eq!(1, avg2.issues.len());
    assert_eq!(None, avg2.fixed);
    assert_eq!(enums::Severity::Low, avg2.severity);
    assert_eq!(enums::Status::Vulnerable, avg2.status);
}

/// Given a package and an `avg::AVG`, returns true if the system is affected
fn system_is_affected(pacman: &alpm::Alpm, pkg: &str, avg: &avg::AVG) -> bool {
    match pacman.query_package_version(pkg.clone()) {
        Ok(v) => {
            info!("Found installed version {} for package {}", v, pkg);
            match avg.fixed {
                Some(ref version) => {
                    info!("Comparing with fixed version {}", version);
                    let cmp = pacman.vercmp(v.clone(), version.clone()).expect(
                        "Alpm::vercmp failed",
                    );
                    if let Ordering::Less = cmp {
                        return true;
                    }
                }
                None => return true,
            };
        }
        Err(_) => debug!("Package {} not installed", pkg),
    }

    false
}

#[test]
fn test_system_is_affected() {
    let pacman = alpm::Alpm::new().expect("Alpm::new failed");

    let avg1 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Unknown,
    };

    assert_eq!(
        false,
        system_is_affected(&pacman, &"pacman".to_string(), &avg1)
    );

    let avg2 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("7.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Unknown,
    };

    assert!(system_is_affected(&pacman, &"pacman".to_string(), &avg2));
}

/// Given a list of package names, returns true when at least one is installed
fn package_is_installed(pacman: &alpm::Alpm, packages: &[String]) -> bool {
    for pkg in packages {
        match pacman.query_package_version(pkg.as_str()) {
            Ok(_) => {
                info!("Package {} is installed", pkg);
                return true;
            }
            Err(_) => debug!("Package {} not installed", pkg),
        }
    }
    false
}

#[test]
fn test_package_is_installed() {
    let pacman = alpm::Alpm::new().expect("Alpm::new failed");

    let packages = vec!["pacman".to_string(), "pac".to_string()];
    assert!(package_is_installed(&pacman, &packages));

    let packages = vec!["pac".to_string()];
    assert_eq!(false, package_is_installed(&pacman, &packages));
}

/// Merge a list of `avg::AVG` into a single `avg::AVG` using major version as version
fn merge_avgs(
    pacman: &alpm::Alpm,
    cves: &BTreeMap<String, Vec<avg::AVG>>,
) -> BTreeMap<String, avg::AVG> {
    let mut avgs: BTreeMap<String, avg::AVG> = BTreeMap::new();
    for (pkg, list) in cves.iter() {
        let mut avg_issues = vec![];
        let mut avg_fixed: Option<String> = None;
        let mut avg_severity = enums::Severity::Unknown;
        let mut avg_status = enums::Status::Unknown;

        for a in list.iter() {
            avg_issues.append(&mut a.issues.clone());

            match avg_fixed.clone() {
                Some(ref version) => {
                    if let Some(ref v) = a.fixed {
                        let cmp = pacman.vercmp(version.to_string(), v.to_string()).expect(
                            "Alpm::vercmp failed",
                        );
                        if let Ordering::Greater = cmp {
                            avg_fixed = a.fixed.clone();
                        }
                    }
                }
                None => avg_fixed = a.fixed.clone(),
            }

            if a.severity > avg_severity {
                avg_severity = a.severity;
            }

            // We only care about testing stuff
            if a.status == enums::Status::Testing {
                avg_status = a.status.clone();
            }
        }

        let avg = avg::AVG {
            issues: avg_issues,
            fixed: avg_fixed,
            severity: avg_severity,
            status: avg_status,
        };
        avgs.insert(pkg.to_string(), avg);
    }

    avgs
}

#[test]
fn test_merge_avgs() {
    let mut avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();

    let avg1 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Fixed,
    };

    let avg2 = avg::AVG {
        issues: vec!["CVE-4".to_string(), "CVE-10".to_string()],
        fixed: Some("0.9.8".to_string()),
        severity: enums::Severity::High,
        status: enums::Status::Testing,
    };

    assert!(enums::Severity::Critical > enums::Severity::High);

    avgs.insert("package".to_string(), vec![avg1.clone(), avg2.clone()]);

    avgs.insert("package2".to_string(), vec![avg1, avg2]);

    let pacman = alpm::Alpm::new().expect("Alpm::new failed");
    let merged = merge_avgs(&pacman, &avgs);

    assert_eq!(2, merged.len());
    assert_eq!(
        4,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .issues
            .len()
    );
    assert_eq!(
        enums::Severity::High,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .severity
    );
    assert_eq!(
        enums::Status::Testing,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .status
    );
}

/// Print a list of `avg::AVG`
fn print_avgs(options: &Options, avgs: &BTreeMap<String, avg::AVG>) {
    for (pkg, avg) in avgs {
        match avg.fixed {
            Some(ref v) => {
                if options.quiet >= 2 {
                    println!("{}", pkg);
                } else if options.quiet == 1 {
                    println!("{}>={}", pkg, v);
                } else {
                    match options.format {
                        Some(ref f) => {
                            println!(
                                "{}",
                                f.replace("%n", pkg.as_str()).replace(
                                    "%c",
                                    avg.issues
                                        .iter()
                                        .join(",")
                                        .as_str(),
                                )
                            )
                        }
                        None => {
                            print!(
                                "Package {} is affected by {}. {}!",
                                pkg,
                                avg.issues.join(", "),
                                avg.severity
                            );
                            if avg.status == enums::Status::Testing {
                                println!("Update to {} from testing repos!", v)
                            } else {
                                println!("Update to {}!", v)
                            }
                        }
                    }
                }
            }
            None => {
                if !options.upgradable_only {
                    if options.quiet > 0 {
                        println!("{}", pkg);
                    } else {
                        match options.format {
                            Some(ref f) => {
                                println!(
                                    "{}",
                                    f.replace("%n", pkg.as_str()).replace(
                                        "%c",
                                        avg.issues
                                            .iter()
                                            .join(",")
                                            .as_str(),
                                    )
                                )
                            }
                            None => {
                                println!(
                                    "Package {} is affected by {}. {}!",
                                    pkg,
                                    avg.issues.join(", "),
                                    avg.severity
                                );
                            }
                        }
                    }
                }
            }
        }
    }
}
<commit_msg>Do not report packages as fixed when they're not<commit_after>extern crate alpm;
#[macro_use]
extern crate clap;
extern crate curl;
extern crate env_logger;
extern crate itertools;
#[macro_use]
extern crate log;
extern crate serde_json;

use clap::App;
use curl::easy::Easy;
use itertools::Itertools;
use serde_json::Value;
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::collections::btree_map::Entry::{Occupied, Vacant};
use std::default::Default;
use std::process::exit;
use std::str;

mod avg;
mod enums;

struct Options {
    format: Option<String>,
    quiet: u64,
    upgradable_only: bool,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            format: None,
            quiet: 0,
            upgradable_only: false,
        }
    }
}

fn main() {
    env_logger::init().expect("env_logger failed");

    let yaml = load_yaml!("cli.yml");
    let args = App::from_yaml(yaml).get_matches();

    let options = Options {
        format: {
            match args.value_of("format") {
                Some(f) => Some(f.to_string()),
                None => None,
            }
        },
        quiet: args.occurrences_of("quiet"),
        upgradable_only: args.is_present("upgradable"),
    };

    let mut avgs = String::new();
    {
        info!("Downloading AVGs...");
        let avgs_url = "https://security.archlinux.org/issues/all.json";

        let mut easy = Easy::new();
        easy.url(avgs_url).expect("curl::Easy::url failed");
        let mut transfer = easy.transfer();
        transfer
            .write_function(|data| {
                avgs.push_str(str::from_utf8(data).expect("str conversion failed"));
                Ok(data.len())
            })
            .expect("write_function failed");
        match transfer.perform() {
            Ok(_) => {}
            Err(_) => {
                println!("Cannot fetch data, please check your network connection!");
                exit(1)
            }
        };
    }

    let pacman = match args.value_of("dbpath") {
        Some(path) => {
            alpm::Alpm::with_dbpath(path.to_string()).expect("alpm::Alpm::with_dbpath failed")
        }
        None => alpm::Alpm::new().expect("alpm::Alpm::new failed"),
    };

    let mut cves: BTreeMap<String, Vec<_>> = BTreeMap::new();
    {
        let json: Value = serde_json::from_str(&avgs).expect("serde_json::from_str failed");

        for avg in json.as_array().expect("Value::as_array failed") {
            let packages = avg["packages"]
                .as_array()
                .expect("Value::as_array failed")
                .iter()
                .map(|s| s.as_str().expect("Value::as_str failed").to_string())
                .collect::<Vec<_>>();

            if !package_is_installed(&pacman, &packages) {
                continue;
            }

            let info = to_avg(avg);

            if info.status != enums::Status::NotAffected {
                for p in packages {
                    match cves.entry(p) {
                        Occupied(c) => c.into_mut(),
                        Vacant(c) => c.insert(vec![]),
                    }.push(info.clone());
                }
            }
        }
    }

    let mut affected_avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();
    for (pkg, avgs) in cves {
        for avg in &avgs {
            if system_is_affected(&pacman, &pkg, avg) {
                match affected_avgs.entry(pkg.clone()) {
                    Occupied(c) => c.into_mut(),
                    Vacant(c) => c.insert(vec![]),
                }.push(avg.clone());
            }
        }
    }

    let merged = merge_avgs(&pacman, &affected_avgs);
    print_avgs(&options, &merged);
}

/// Converts a JSON to an `avg::AVG`
fn to_avg(data: &Value) -> avg::AVG {
    avg::AVG {
        issues: data["issues"]
            .as_array()
            .expect("Value::as_array failed")
            .iter()
            .map(|s| s.as_str().expect("Value::as_str failed").to_string())
            .collect(),
        fixed: match data["fixed"].as_str() {
            Some(s) => Some(s.to_string()),
            None => None,
        },
        severity: data["severity"]
            .as_str()
            .expect("Value::as_str failed")
            .to_string()
            .parse::<enums::Severity>()
            .expect("parse::<Severity> failed"),
        status: data["status"]
            .as_str()
            .expect("Value::as_str failed")
            .to_string()
            .parse::<enums::Status>()
            .expect("parse::<Status> failed"),
    }
}

#[test]
fn test_to_avg() {
    let json: Value = serde_json::from_str(
        "{\"issues\": [\"CVE-1\", \"CVE-2\"], \"fixed\": \"1.0\", \
                               \"severity\": \"High\", \"status\": \"Not affected\"}",
    ).expect("serde_json::from_str failed");

    let avg1 = to_avg(&json);
    assert_eq!(2, avg1.issues.len());
    assert_eq!(Some("1.0".to_string()), avg1.fixed);
    assert_eq!(enums::Severity::High, avg1.severity);
    assert_eq!(enums::Status::NotAffected, avg1.status);

    let json: Value = serde_json::from_str(
        "{\"issues\": [\"CVE-1\"], \"fixed\": null, \
                               \"severity\": \"Low\", \"status\": \"Vulnerable\"}",
    ).expect("serde_json::from_str failed");

    let avg2 = to_avg(&json);
    assert_eq!(1, avg2.issues.len());
    assert_eq!(None, avg2.fixed);
    assert_eq!(enums::Severity::Low, avg2.severity);
    assert_eq!(enums::Status::Vulnerable, avg2.status);
}

/// Given a package and an `avg::AVG`, returns true if the system is affected
fn system_is_affected(pacman: &alpm::Alpm, pkg: &str, avg: &avg::AVG) -> bool {
    match pacman.query_package_version(pkg.clone()) {
        Ok(v) => {
            info!("Found installed version {} for package {}", v, pkg);
            match avg.fixed {
                Some(ref version) => {
                    info!("Comparing with fixed version {}", version);
                    let cmp = pacman.vercmp(v.clone(), version.clone()).expect(
                        "Alpm::vercmp failed",
                    );
                    if let Ordering::Less = cmp {
                        return true;
                    }
                }
                None => return true,
            };
        }
        Err(_) => debug!("Package {} not installed", pkg),
    }

    false
}

#[test]
fn test_system_is_affected() {
    let pacman = alpm::Alpm::new().expect("Alpm::new failed");

    let avg1 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Unknown,
    };

    assert_eq!(
        false,
        system_is_affected(&pacman, &"pacman".to_string(), &avg1)
    );

    let avg2 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("7.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Unknown,
    };

    assert!(system_is_affected(&pacman, &"pacman".to_string(), &avg2));
}

/// Given a list of package names, returns true when at least one is installed
fn package_is_installed(pacman: &alpm::Alpm, packages: &[String]) -> bool {
    for pkg in packages {
        match pacman.query_package_version(pkg.as_str()) {
            Ok(_) => {
                info!("Package {} is installed", pkg);
                return true;
            }
            Err(_) => debug!("Package {} not installed", pkg),
        }
    }
    false
}

#[test]
fn test_package_is_installed() {
    let pacman = alpm::Alpm::new().expect("Alpm::new failed");

    let packages = vec!["pacman".to_string(), "pac".to_string()];
    assert!(package_is_installed(&pacman, &packages));

    let packages = vec!["pac".to_string()];
    assert_eq!(false, package_is_installed(&pacman, &packages));
}

/// Merge a list of `avg::AVG` into a single `avg::AVG` using major version as version
fn merge_avgs(
    pacman: &alpm::Alpm,
    cves: &BTreeMap<String, Vec<avg::AVG>>,
) -> BTreeMap<String, avg::AVG> {
    let mut avgs: BTreeMap<String, avg::AVG> = BTreeMap::new();
    for (pkg, list) in cves.iter() {
        let mut avg_issues = vec![];
        let mut avg_fixed: Option<String> = None;
        let mut avg_severity = enums::Severity::Unknown;
        let mut avg_status = enums::Status::Unknown;

        for a in list.iter() {
            avg_issues.append(&mut a.issues.clone());

            match avg_fixed.clone() {
                Some(ref version) => {
                    if let Some(ref v) = a.fixed {
                        let cmp = pacman.vercmp(version.to_string(), v.to_string()).expect(
                            "Alpm::vercmp failed",
                        );
                        if let Ordering::Greater = cmp {
                            avg_fixed = a.fixed.clone();
                        }
                    }
                }
                None => avg_fixed = a.fixed.clone(),
            }

            if a.severity > avg_severity {
                avg_severity = a.severity;
            }

            // We only care about testing stuff
            if a.status == enums::Status::Testing {
                avg_status = a.status.clone();
            }
        }

        let avg = avg::AVG {
            issues: avg_issues,
            fixed: avg_fixed,
            severity: avg_severity,
            status: avg_status,
        };
        avgs.insert(pkg.to_string(), avg);
    }

    avgs
}

#[test]
fn test_merge_avgs() {
    let mut avgs: BTreeMap<String, Vec<_>> = BTreeMap::new();

    let avg1 = avg::AVG {
        issues: vec!["CVE-1".to_string(), "CVE-2".to_string()],
        fixed: Some("1.0.0".to_string()),
        severity: enums::Severity::Unknown,
        status: enums::Status::Fixed,
    };

    let avg2 = avg::AVG {
        issues: vec!["CVE-4".to_string(), "CVE-10".to_string()],
        fixed: Some("0.9.8".to_string()),
        severity: enums::Severity::High,
        status: enums::Status::Testing,
    };

    assert!(enums::Severity::Critical > enums::Severity::High);

    avgs.insert("package".to_string(), vec![avg1.clone(), avg2.clone()]);

    avgs.insert("package2".to_string(), vec![avg1, avg2]);

    let pacman = alpm::Alpm::new().expect("Alpm::new failed");
    let merged = merge_avgs(&pacman, &avgs);

    assert_eq!(2, merged.len());
    assert_eq!(
        4,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .issues
            .len()
    );
    assert_eq!(
        enums::Severity::High,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .severity
    );
    assert_eq!(
        enums::Status::Testing,
        merged
            .get(&"package".to_string())
            .expect("'package' key not found")
            .status
    );
}

/// Print a list of `avg::AVG`
fn print_avgs(options: &Options, avgs: &BTreeMap<String, avg::AVG>) {
    for (pkg, avg) in avgs {
        match avg.fixed {
            Some(ref v) => {
                if options.quiet >= 2 {
                    println!("{}", pkg);
                } else if options.quiet == 1 {
                    println!("{}>={}", pkg, v);
                } else {
                    match options.format {
                        Some(ref f) => {
                            println!(
                                "{}",
                                f.replace("%n", pkg.as_str()).replace(
                                    "%c",
                                    avg.issues
                                        .iter()
                                        .join(",")
                                        .as_str(),
                                )
                            )
                        }
                        None => {
                            let msg = format!(
                                "Package {} is affected by {}. {}!",
                                pkg,
                                avg.issues.join(", "),
                                avg.severity
                            );

                            if avg.status == enums::Status::Testing {
                                println!("{}. Update to {} from testing repos!", msg, v)
                            } else if avg.status == enums::Status::Fixed {
                                println!("{}. Update to {}!", msg, v)
                            } else {
                                println!("{}", msg)
                            }
                        }
                    }
                }
            }
            None => {
                if !options.upgradable_only {
                    if options.quiet > 0 {
                        println!("{}", pkg);
                    } else {
                        match options.format {
                            Some(ref f) => {
                                println!(
                                    "{}",
                                    f.replace("%n", pkg.as_str()).replace(
                                        "%c",
                                        avg.issues
                                            .iter()
                                            .join(",")
                                            .as_str(),
                                    )
                                )
                            }
                            None => {
                                println!(
                                    "Package {} is affected by {}. {}!",
                                    pkg,
                                    avg.issues.join(", "),
                                    avg.severity
                                );
                            }
                        }
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>#![feature(env, fs, io, old_io, path)]
use std::fs;
use std::io;
use std::io::ReadExt;
use std::old_io::stdio;
use std::path::Path;

const MEMSIZE: usize = 4096;

enum ParserResult {
    Something(Op),
    Nothing,
    EOF
}

use ParserResult::*;

type Ops = Vec<Op>;
enum Op {
    Add(u8),
    Mov(usize),
    In,
    Out,
    Loop(Ops)
}

use Op::*;

struct State {
    index: usize,
    memory: [u8; MEMSIZE]
}

impl State {
    fn new() -> State {
        State { index: 0, memory: [0; MEMSIZE] }
    }

    fn peek(&self) -> u8 {
        self.memory[self.index % MEMSIZE]
    }

    fn poke(&mut self, value: u8) {
        self.memory[self.index % MEMSIZE] = value;
    }

    fn step(&mut self, op: &Op) {
        match op {
            &Add(i) => { let x = self.peek(); self.poke(x + i); },
            &Mov(n) => self.index += n,
            &In => {
                self.poke(stdio::stdin().read_u8().unwrap());
            },
            &Out => {
                stdio::stdout().write_u8(self.peek()).unwrap();
            },
            &Loop(ref ops) => {
                while self.peek() != 0 {
                    self.run(ops);
                }
            }
        }
    }

    fn run(&mut self, ops: &[Op]) {
        for op in ops {
            self.step(op);
        }
    }
}

#[test]
fn test_state_peek() {
    let mut state = State::new();
    state.memory[state.index] = 23;
    assert_eq!(23, state.peek());
    state.index = 5;
    state.memory[state.index] = 42;
    assert_eq!(42, state.peek());
}

#[test]
fn test_state_poke() {
    let mut state = State::new();
    state.poke(23);
    assert_eq!(23, state.memory[state.index]);
    state.index = 5;
    state.poke(42);
    assert_eq!(42, state.memory[state.index]);
}

fn main() {
    let filenames: Vec<String> = std::iter::FromIterator::from_iter(std::env::args());

    for filename in &filenames[1..] {
        let reader = reader(&Path::new(filename));
        let mut state = State::new();
        let mut chars = reader.chars().peekable();

        loop {
            match parse(&mut chars) {
                Something(op) => state.step(&op),
                EOF => break,
                _ => {}
            }
        }
    }
}

fn parse<T: io::Read>(mut chars: &mut std::iter::Peekable<io::Chars<T>>) -> ParserResult {
    match chars.next() {
        Some(Ok(c)) => match c {
            '+' => Something(Add(1)),
            '-' => Something(Add(-1)),
            '<' => Something(Mov(-1)),
            '>' => Something(Mov(1)),
            ',' => Something(In),
            '.' => Something(Out),
            '[' => {
                let mut children: Ops = Vec::new();
                while chars.peek() != Some(&Ok(']')) {
                    match parse(&mut chars) {
                        Something(op) => children.push(op),
                        Nothing => {},
                        EOF => panic!()
                    }
                }
                chars.next();
                Something(Loop(children))
            },
            _ => Nothing  // other characters
        },
        _ => EOF
    }
}

fn reader(path: &Path) -> io::BufReader<fs::File> {
    match fs::File::open(path) {
        Ok(f) => io::BufReader::new(f),
        Err(e) => panic!(e)
    }
}
<commit_msg>refactor: don't reimplement unwrap<commit_after>#![feature(env, fs, io, old_io, path)]
use std::fs;
use std::io;
use std::io::ReadExt;
use std::old_io::stdio;
use std::path::Path;

const MEMSIZE: usize = 4096;

enum ParserResult {
    Something(Op),
    Nothing,
    EOF
}

use ParserResult::*;

type Ops = Vec<Op>;
enum Op {
    Add(u8),
    Mov(usize),
    In,
    Out,
    Loop(Ops)
}

use Op::*;

struct State {
    index: usize,
    memory: [u8; MEMSIZE]
}

impl State {
    fn new() -> State {
        State { index: 0, memory: [0; MEMSIZE] }
    }

    fn peek(&self) -> u8 {
        self.memory[self.index % MEMSIZE]
    }

    fn poke(&mut self, value: u8) {
        self.memory[self.index % MEMSIZE] = value;
    }

    fn step(&mut self, op: &Op) {
        match op {
            &Add(i) => { let x = self.peek(); self.poke(x + i); },
            &Mov(n) => self.index += n,
            &In => {
                self.poke(stdio::stdin().read_u8().unwrap());
            },
            &Out => {
                stdio::stdout().write_u8(self.peek()).unwrap();
            },
            &Loop(ref ops) => {
                while self.peek() != 0 {
                    self.run(ops);
                }
            }
        }
    }

    fn run(&mut self, ops: &[Op]) {
        for op in ops {
            self.step(op);
        }
    }
}

#[test]
fn test_state_peek() {
    let mut state = State::new();
    state.memory[state.index] = 23;
    assert_eq!(23, state.peek());
    state.index = 5;
    state.memory[state.index] = 42;
    assert_eq!(42, state.peek());
}

#[test]
fn test_state_poke() {
    let mut state = State::new();
    state.poke(23);
    assert_eq!(23, state.memory[state.index]);
    state.index = 5;
    state.poke(42);
    assert_eq!(42, state.memory[state.index]);
}

fn main() {
    let filenames: Vec<String> = std::iter::FromIterator::from_iter(std::env::args());

    for filename in &filenames[1..] {
        let reader = reader(&Path::new(filename));
        let mut state = State::new();
        let mut chars = reader.chars().peekable();

        loop {
            match parse(&mut chars) {
                Something(op) => state.step(&op),
                EOF => break,
                _ => {}
            }
        }
    }
}

fn parse<T: io::Read>(mut chars: &mut std::iter::Peekable<io::Chars<T>>) -> ParserResult {
    match chars.next() {
        Some(Ok(c)) => match c {
            '+' => Something(Add(1)),
            '-' => Something(Add(-1)),
            '<' => Something(Mov(-1)),
            '>' => Something(Mov(1)),
            ',' => Something(In),
            '.' => Something(Out),
            '[' => {
                let mut children: Ops = Vec::new();
                while chars.peek() != Some(&Ok(']')) {
                    match parse(&mut chars) {
                        Something(op) => children.push(op),
                        Nothing => {},
                        EOF => panic!()
                    }
                }
                chars.next();
                Something(Loop(children))
            },
            _ => Nothing  // other characters
        },
        _ => EOF
    }
}

fn reader(path: &Path) -> io::BufReader<fs::File> {
    io::BufReader::new(fs::File::open(path).unwrap())
}
<|endoftext|>


<commit_before>extern crate mio;
extern crate bytes;
extern crate log4rs;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;

use mio::tcp::*;
use mio::EventLoop;
use mio::{TryRead,TryWrite};
use mio::util::Slab;
use std::collections::VecDeque;
use clap::{Arg, App};

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
    NewConnection(TcpStream),
}

struct MiChat {
    connections: Slab<EventHandler>,
}

impl MiChat {
    fn new() -> MiChat {
        MiChat {
            connections: Slab::new(1024),
        }
    }

    fn listen(&mut self, event_loop: &mut mio::EventLoop<MiChat>, listener: TcpListener) {
        let l = EventHandler::Listener(Listener::new(listener));
        let token = self.connections.insert(l).expect("insert listener");
        &self.connections[token].register(event_loop, token);
    }

    fn process_action(&mut self, msg: MiChatCommand, event_loop: &mut mio::EventLoop<MiChat>) {
        trace!("{:p}; got {:?}", self, msg);
        match msg {
            MiChatCommand::Broadcast(s) => {
                for c in self.connections.iter_mut() {
                    if let &mut EventHandler::Conn(ref mut c) = c {
                        c.enqueue(&s);
                    }
                }
            },

            MiChatCommand::NewConnection(socket) => {
                let token = self.connections
                    .insert_with(|token| EventHandler::Conn(Connection::new(socket, token)))
                    .expect("token insert");
                &self.connections[token].register(event_loop, token);
            }
        }
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::EventSet,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    failed: bool,
    write_buf: Vec<u8>,
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::EventSet::none(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
            failed: false,
        }
    }

    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        event_loop.register_opt(
                &self.socket,
                token,
                mio::EventSet::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot())
            .expect("event loop register");
    }

    // Event updates arrive here
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        self.sock_status.insert(events);
        info!("Connection::handle_event: {:?}; this time: {:?}; now: {:?}",
                self.socket.peer_addr(), events, self.sock_status);
    }

    // actions are processed here on down.

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>) {
        if self.sock_status.is_readable() {
            self.read();
            self.sock_status.remove(mio::EventSet::readable());
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.write();
            self.sock_status.remove(mio::EventSet::writable());
        }

        if !self.is_closed() {
            self.reregister(event_loop)
        }
    }

    fn process_buffer(&mut self, to_parent: &mut VecDeque<MiChatCommand>) {
        let mut prev = 0;
        info!("{:?}: Read buffer: {:?}", self.socket.peer_addr(), self.read_buf);
        for n in self.read_buf.iter().enumerate()
                .filter_map(|(i, e)| if *e == '\n' as u8 { Some(i) } else { None } ) {
            info!("{:?}: Pos: {:?}; chunk: {:?}", self.socket.peer_addr(), n, &self.read_buf[prev..n]);
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent.push_back(cmd);
            prev = n+1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!("{:?}: read Remainder: {}", self.socket.peer_addr(), remainder.len());
        self.read_buf = remainder;
    }

    fn read(&mut self) {
        let mut abuf = Vec::new();
        match self.socket.try_read_buf(&mut abuf) {
            Ok(Some(0)) => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true
            },
            Ok(Some(n)) => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(abuf);
            },
            Ok(None) => {
                info!("{:?}: Noop!", self.socket.peer_addr());
            },
            Err(e) => {
                error!("got an error trying to read; err={:?}", e);
                self.failed =true;
            }
        }
    }

    fn write(&mut self) {
        match self.socket.try_write(&mut self.write_buf) {
            Ok(Some(n)) => {
                info!("{:?}: Wrote {} of {} in buffer", self.socket.peer_addr(), n,
                    self.write_buf.len());
                self.write_buf = self.write_buf[n..].to_vec();
                info!("{:?}: Now {:?}b", self.socket.peer_addr(), self.write_buf.len());
            },
            Ok(None) => {
                info!("Write unready");
            },
            Err(e) => {
                error!("got an error trying to write; err={:?}", e);
                self.failed = true;
            }
        }
    }

    fn reregister(&mut self, event_loop: &mut mio::EventLoop<MiChat>) {
        let mut flags = mio::EventSet::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::EventSet::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop.reregister(
                &self.socket,
                self.token,
                flags,
                mio::PollOpt::oneshot()).expect("EventLoop#reregister")
    }

    fn is_closed(&self) -> bool {
        self.failed || (self.read_eof && self.write_buf.is_empty())
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

#[derive(Debug)]
struct Listener {
    listener: TcpListener,
    sock_status: mio::EventSet,
}

impl Listener {
    fn new(listener: TcpListener) -> Listener {
        Listener {
            listener: listener,
            sock_status: mio::EventSet::none(),
        }
    }

    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        event_loop.register(&self.listener, token).expect("Register listener");
    }

    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        assert!(events.is_readable());
        self.sock_status.insert(events);
        info!("Listener::handle_event: {:?}; this time: {:?}; now: {:?}",
                self.listener.local_addr(), events, self.sock_status);
    }

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
            to_parent: &mut VecDeque<MiChatCommand>) {
        info!("the listener socket is ready to accept a connection");
        match self.listener.accept() {
            Ok(Some(socket)) => {
                let cmd = MiChatCommand::NewConnection(socket);
                to_parent.push_back(cmd);
            }
            Ok(None) => {
                info!("the listener socket wasn't actually ready");
            }
            Err(e) => {
                info!("listener.accept() errored: {}", e);
                event_loop.shutdown();
            }
        }
    }

    fn is_closed(&self) -> bool {
        false
    }
}

#[derive(Debug)]
enum EventHandler {
    Listener (Listener),
    Conn (Connection),
}

impl EventHandler {
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn.handle_event(_event_loop, events),
            &mut EventHandler::Listener(ref mut listener) => listener.handle_event(_event_loop, events)
        }
    }

    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        match self {
            &EventHandler::Conn(ref conn) => conn.register(event_loop, token),
            &EventHandler::Listener(ref listener) => listener.register(event_loop, token)
        }
    }

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>) {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn.process_rules(event_loop, to_parent),
            &mut EventHandler::Listener(ref mut listener) => listener.process_rules(event_loop, to_parent)
        }
    }

    fn is_closed(&self) -> bool {
        match self {
            &EventHandler::Conn(ref conn) => conn.is_closed(),
            &EventHandler::Listener(ref listener) => listener.is_closed()
        }
    }
}

impl mio::Handler for MiChat {
    type Timeout = ();
    type Message = ();
    fn ready(&mut self, event_loop: &mut mio::EventLoop<Self>, token: mio::Token, events: mio::EventSet) {
        info!("{:?}: {:?}", token, events);
        match token {
            _ => {
                self.connections[token].handle_event(event_loop, events);
                if self.connections[token].is_closed() {
                    info!("Removing; {:?}", token);
                    self.connections.remove(token);
                }
            }
        }

        let mut parent_actions = VecDeque::new();
        loop {
            for conn in self.connections.iter_mut() {
                conn.process_rules(event_loop, &mut parent_actions);
            }
            // Anything left to process?
            if parent_actions.is_empty() { break; }

            for action in parent_actions.drain(..) {
                self.process_action(action, event_loop);
            }
        }
    }
}

const LOG_FILE: &'static str = "log.toml";

fn main() {
    if let Err(e) = log4rs::init_file(LOG_FILE, Default::default()) {
        panic!("Could not init logger from file {}: {}", LOG_FILE, e);
    }
    let matches = App::new("michat")
        .arg(Arg::with_name("bind").short("l").takes_value(true).required(true))
        .get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::EventLoop::new().expect("Create event loop");
    let mut service = MiChat::new();

    service.listen(&mut event_loop, listener);
    info!("running michat listener at: {:?}", address);
    event_loop.run(&mut service).expect("Run loop");
}
<commit_msg>Remove pointless match.<commit_after>extern crate mio;
extern crate bytes;
extern crate log4rs;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;

use mio::tcp::*;
use mio::EventLoop;
use mio::{TryRead,TryWrite};
use mio::util::Slab;
use std::collections::VecDeque;
use clap::{Arg, App};

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
    NewConnection(TcpStream),
}

struct MiChat {
    connections: Slab<EventHandler>,
}

impl MiChat {
    fn new() -> MiChat {
        MiChat {
            connections: Slab::new(1024),
        }
    }

    fn listen(&mut self, event_loop: &mut mio::EventLoop<MiChat>, listener: TcpListener) {
        let l = EventHandler::Listener(Listener::new(listener));
        let token = self.connections.insert(l).expect("insert listener");
        &self.connections[token].register(event_loop, token);
    }

    fn process_action(&mut self, msg: MiChatCommand, event_loop: &mut mio::EventLoop<MiChat>) {
        trace!("{:p}; got {:?}", self, msg);
        match msg {
            MiChatCommand::Broadcast(s) => {
                for c in self.connections.iter_mut() {
                    if let &mut EventHandler::Conn(ref mut c) = c {
                        c.enqueue(&s);
                    }
                }
            },

            MiChatCommand::NewConnection(socket) => {
                let token = self.connections
                    .insert_with(|token| EventHandler::Conn(Connection::new(socket, token)))
                    .expect("token insert");
                &self.connections[token].register(event_loop, token);
            }
        }
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::EventSet,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    failed: bool,
    write_buf: Vec<u8>,
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::EventSet::none(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
            failed: false,
        }
    }

    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        event_loop.register_opt(
                &self.socket,
                token,
                mio::EventSet::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot())
            .expect("event loop register");
    }

    // Event updates arrive here
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        self.sock_status.insert(events);
        info!("Connection::handle_event: {:?}; this time: {:?}; now: {:?}",
                self.socket.peer_addr(), events, self.sock_status);
    }

    // actions are processed here on down.

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>) {
        if self.sock_status.is_readable() {
            self.read();
            self.sock_status.remove(mio::EventSet::readable());
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.write();
            self.sock_status.remove(mio::EventSet::writable());
        }

        if !self.is_closed() {
            self.reregister(event_loop)
        }
    }

    fn process_buffer(&mut self, to_parent: &mut VecDeque<MiChatCommand>) {
        let mut prev = 0;
        info!("{:?}: Read buffer: {:?}", self.socket.peer_addr(), self.read_buf);
        for n in self.read_buf.iter().enumerate()
                .filter_map(|(i, e)| if *e == '\n' as u8 { Some(i) } else { None } ) {
            info!("{:?}: Pos: {:?}; chunk: {:?}", self.socket.peer_addr(), n, &self.read_buf[prev..n]);
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent.push_back(cmd);
            prev = n+1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!("{:?}: read Remainder: {}", self.socket.peer_addr(), remainder.len());
        self.read_buf = remainder;
    }

    fn read(&mut self) {
        let mut abuf = Vec::new();
        match self.socket.try_read_buf(&mut abuf) {
            Ok(Some(0)) => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true
            },
            Ok(Some(n)) => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(abuf);
            },
            Ok(None) => {
                info!("{:?}: Noop!", self.socket.peer_addr());
            },
            Err(e) => {
                error!("got an error trying to read; err={:?}", e);
                self.failed =true;
            }
        }
    }

    fn write(&mut self) {
        match self.socket.try_write(&mut self.write_buf) {
            Ok(Some(n)) => {
                info!("{:?}: Wrote {} of {} in buffer", self.socket.peer_addr(), n,
                    self.write_buf.len());
                self.write_buf = self.write_buf[n..].to_vec();
                info!("{:?}: Now {:?}b", self.socket.peer_addr(), self.write_buf.len());
            },
            Ok(None) => {
                info!("Write unready");
            },
            Err(e) => {
                error!("got an error trying to write; err={:?}", e);
                self.failed = true;
            }
        }
    }

    fn reregister(&mut self, event_loop: &mut mio::EventLoop<MiChat>) {
        let mut flags = mio::EventSet::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::EventSet::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop.reregister(
                &self.socket,
                self.token,
                flags,
                mio::PollOpt::oneshot()).expect("EventLoop#reregister")
    }

    fn is_closed(&self) -> bool {
        self.failed || (self.read_eof && self.write_buf.is_empty())
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

#[derive(Debug)]
struct Listener {
    listener: TcpListener,
    sock_status: mio::EventSet,
}

impl Listener {
    fn new(listener: TcpListener) -> Listener {
        Listener {
            listener: listener,
            sock_status: mio::EventSet::none(),
        }
    }

    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        event_loop.register(&self.listener, token).expect("Register listener");
    }

    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        assert!(events.is_readable());
        self.sock_status.insert(events);
        info!("Listener::handle_event: {:?}; this time: {:?}; now: {:?}",
                self.listener.local_addr(), events, self.sock_status);
    }

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
            to_parent: &mut VecDeque<MiChatCommand>) {
        info!("the listener socket is ready to accept a connection");
        match self.listener.accept() {
            Ok(Some(socket)) => {
                let cmd = MiChatCommand::NewConnection(socket);
                to_parent.push_back(cmd);
            }
            Ok(None) => {
                info!("the listener socket wasn't actually ready");
            }
            Err(e) => {
                info!("listener.accept() errored: {}", e);
                event_loop.shutdown();
            }
        }
    }

    fn is_closed(&self) -> bool {
        false
    }
}

#[derive(Debug)]
enum EventHandler {
    Listener (Listener),
    Conn (Connection),
}

impl EventHandler {
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn.handle_event(_event_loop, events),
            &mut EventHandler::Listener(ref mut listener) => listener.handle_event(_event_loop, events)
        }
    }

    fn register(&self, event_loop: &mut mio::EventLoop<MiChat>, token: mio::Token) {
        match self {
            &EventHandler::Conn(ref conn) => conn.register(event_loop, token),
            &EventHandler::Listener(ref listener) => listener.register(event_loop, token)
        }
    }

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>) {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn.process_rules(event_loop, to_parent),
            &mut EventHandler::Listener(ref mut listener) => listener.process_rules(event_loop, to_parent)
        }
    }

    fn is_closed(&self) -> bool {
        match self {
            &EventHandler::Conn(ref conn) => conn.is_closed(),
            &EventHandler::Listener(ref listener) => listener.is_closed()
        }
    }
}

impl mio::Handler for MiChat {
    type Timeout = ();
    type Message = ();
    fn ready(&mut self, event_loop: &mut mio::EventLoop<Self>, token: mio::Token, events: mio::EventSet) {
        info!("{:?}: {:?}", token, events);
        self.connections[token].handle_event(event_loop, events);
        if self.connections[token].is_closed() {
            info!("Removing; {:?}", token);
            self.connections.remove(token);
        }

        let mut parent_actions = VecDeque::new();
        loop {
            for conn in self.connections.iter_mut() {
                conn.process_rules(event_loop, &mut parent_actions);
            }
            // Anything left to process?
            if parent_actions.is_empty() { break; }

            for action in parent_actions.drain(..) {
                self.process_action(action, event_loop);
            }
        }
    }
}

const LOG_FILE: &'static str = "log.toml";

fn main() {
    if let Err(e) = log4rs::init_file(LOG_FILE, Default::default()) {
        panic!("Could not init logger from file {}: {}", LOG_FILE, e);
    }
    let matches = App::new("michat")
        .arg(Arg::with_name("bind").short("l").takes_value(true).required(true))
        .get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::EventLoop::new().expect("Create event loop");
    let mut service = MiChat::new();

    service.listen(&mut event_loop, listener);
    info!("running michat listener at: {:?}", address);
    event_loop.run(&mut service).expect("Run loop");
}
<|endoftext|>


<commit_before>extern crate iron;
extern crate router;

use iron::prelude::*;
use iron::status;
use router::{Router};

fn main() {
    let mut router = Router::new();

    router.get("/", roothandler);
    router.get("/:query", roothandler);

    Iron::new(router).http("localhost:3000").unwrap();

    fn roothandler(req: &mut Request) -> IronResult<Response> {
        let ref query = req.extensions.get::<Router>()
            .unwrap().find("query").unwrap_or("/");

        Ok(Response::with((status::Ok, *query)))
    }
}
<commit_msg>Redirect handling<commit_after>#![feature(collections)]

extern crate toml;
extern crate iron;
extern crate router;

use iron::prelude::*;
use iron::headers::*;
use iron::{status, headers};

use router::{Router};

fn main() {
    let mut router = Router::new();

    //TODO: fetch it from toml config.
    router.get("/wiki", wikiredirect);
    router.get("/wiki/*", wikihandler);

    router.get("/", roothandler);
    router.get("/:query", roothandler);

    Iron::new(router).http("localhost:3000").unwrap();

    fn roothandler(req: &mut Request) -> IronResult<Response> {
        let ref query = req.extensions.get::<Router>()
            .unwrap().find("query").unwrap_or("/");

        Ok(Response::with((status::Ok, *query)))
    }

    fn wikiredirect(req: &mut Request) -> IronResult<Response> {
        let mut url = req.url.path.clone();
        let mut header = Headers::new();
        let mut newpath = String::new();

        //TODO: customizable root page.
        url.push(String::from_str("FrontPage"));

        for i in url.iter() {
            newpath.push_str("/");
            newpath.push_str(i);
        }

        println!("{}", newpath);
        header.set(Location(newpath));

        let mut resp: Response = Response::with((status::MovedPermanently));
        resp.headers = header;

        Ok(resp)
    }

    fn wikihandler(req: &mut Request) -> IronResult<Response> {
        let ref requrl = req.url.clone();
        let mut urlpath = requrl.path.clone();

        for i in urlpath.iter() {
            println!("{}\n", i);
        }

        Ok(Response::with((status::Ok, "Watch console.")))
    }
}
<|endoftext|>


<commit_before>//! SEGIMAP is an IMAP server implementation.
#![feature(macro_rules)]
#![deny(non_camel_case_types)]
#![feature(phase)]
#![feature(default_type_params)]

extern crate "rust-crypto" as crypto;
extern crate regex;
#[phase(plugin)] extern crate regex_macros;
#[phase(plugin, link)] extern crate log;
#[phase(plugin)] extern crate peg_syntax_ext;
extern crate serialize;
extern crate time;
extern crate toml;

pub use config::Config;
pub use email::Email;
pub use lmtp::Lmtp;
pub use login::LoginData;
pub use message::Message;
pub use server::Server;
pub use session::Session;
pub use user::User;

use std::io::{Listener, Acceptor, BufferedStream};
use std::sync::Arc;

mod auth;
mod command;
mod config;
mod email;
mod error;
mod folder;
mod lmtp;
mod login;
mod message;
mod parser;
mod server;
mod session;
mod user;
mod util;

fn main() {
    // Load configuration.
    let config = Config::new();

    // This function only needs to be run once, really.
    // create_default_users(config.users.clone());

    // Create the server. We wrap it so that it is atomically reference
    // counted. This allows us to safely share it across threads
    let serv = Arc::new(Server::new(config));

    // Spawn a separate thread for listening for LMTP connections
    match serv.lmtp_listener() {
        Err(e) => {
            error!("Error listening on LMTP port: {}", e);
        }
        Ok(v) => {
            let lmtp_serv = serv.clone();
            spawn(proc() {
                let mut acceptor = v.listen();
                // We spawn a separate thread for each LMTP session
                for stream in acceptor.incoming() {
                    match stream {
                        Err(e) => {
                            error!("Error accepting incoming LMTP connection: {}",
                                   e);
                        }
                        Ok(stream) => {
                            let session_serv = lmtp_serv.clone();
                            spawn(proc() {
                                let mut lmtp = Lmtp::new(session_serv);
                                lmtp.handle(BufferedStream::new(stream));
                            });
                        }
                    }
                }
            });
        }
    }

    // The main thread handles listening for IMAP connections
    match serv.imap_listener() {
        Err(e) => {
            error!("Error listening on IMAP port: {}", e);
        }
        Ok(v) => {
            let mut acceptor = v.listen();
            // For each IMAP session, we spawn a separate thread.
            for stream in acceptor.incoming() {
                match stream {
                    Err(e) => {
                        error!("Error accepting incoming IMAP connection: {}",
                               e)
                    }
                    Ok(stream) => {
                        let session_serv = serv.clone();
                        spawn(proc() {
                            let mut session = Session::new(
                                               BufferedStream::new(stream),
                                               session_serv);
                            session.handle();
                        });
                    }
                }
            }
            drop(acceptor);
        }
    }
}

/// Function to create our default users.json for testing
#[allow(dead_code)]
fn create_default_users(filename: String) {
    let mut users = Vec::new();
    users.push(User::new(Email::new("will".to_string(), "xqz.ca".to_string()),
                         "54321".to_string(), "./maildir".to_string()));
    users.push(User::new(Email::new("nikitapekin".to_string(),
                                    "gmail.com".to_string()),
                         "12345".to_string(), "./maildir".to_string()));
    user::save_users(filename, users);
}
<commit_msg>Fix `proc()` to `move ||` syntax<commit_after>//! SEGIMAP is an IMAP server implementation.
#![feature(macro_rules)]
#![deny(non_camel_case_types)]
#![feature(phase)]
#![feature(default_type_params)]

extern crate "rust-crypto" as crypto;
extern crate regex;
#[phase(plugin)] extern crate regex_macros;
#[phase(plugin, link)] extern crate log;
#[phase(plugin)] extern crate peg_syntax_ext;
extern crate serialize;
extern crate time;
extern crate toml;

pub use config::Config;
pub use email::Email;
pub use lmtp::Lmtp;
pub use login::LoginData;
pub use message::Message;
pub use server::Server;
pub use session::Session;
pub use user::User;

use std::io::{Listener, Acceptor, BufferedStream};
use std::sync::Arc;

mod auth;
mod command;
mod config;
mod email;
mod error;
mod folder;
mod lmtp;
mod login;
mod message;
mod parser;
mod server;
mod session;
mod user;
mod util;

fn main() {
    // Load configuration.
    let config = Config::new();

    // This function only needs to be run once, really.
    // create_default_users(config.users.clone());

    // Create the server. We wrap it so that it is atomically reference
    // counted. This allows us to safely share it across threads
    let serv = Arc::new(Server::new(config));

    // Spawn a separate thread for listening for LMTP connections
    match serv.lmtp_listener() {
        Err(e) => {
            error!("Error listening on LMTP port: {}", e);
        }
        Ok(v) => {
            let lmtp_serv = serv.clone();
            spawn(move || {
                let mut acceptor = v.listen();
                // We spawn a separate thread for each LMTP session
                for stream in acceptor.incoming() {
                    match stream {
                        Err(e) => {
                            error!("Error accepting incoming LMTP connection: {}",
                                   e);
                        }
                        Ok(stream) => {
                            let session_serv = lmtp_serv.clone();
                            spawn(move || {
                                let mut lmtp = Lmtp::new(session_serv);
                                lmtp.handle(BufferedStream::new(stream));
                            });
                        }
                    }
                }
            });
        }
    }

    // The main thread handles listening for IMAP connections
    match serv.imap_listener() {
        Err(e) => {
            error!("Error listening on IMAP port: {}", e);
        }
        Ok(v) => {
            let mut acceptor = v.listen();
            // For each IMAP session, we spawn a separate thread.
            for stream in acceptor.incoming() {
                match stream {
                    Err(e) => {
                        error!("Error accepting incoming IMAP connection: {}",
                               e)
                    }
                    Ok(stream) => {
                        let session_serv = serv.clone();
                        spawn(move || {
                            let mut session = Session::new(
                                               BufferedStream::new(stream),
                                               session_serv);
                            session.handle();
                        });
                    }
                }
            }
            drop(acceptor);
        }
    }
}

/// Function to create our default users.json for testing
#[allow(dead_code)]
fn create_default_users(filename: String) {
    let mut users = Vec::new();
    users.push(User::new(Email::new("will".to_string(), "xqz.ca".to_string()),
                         "54321".to_string(), "./maildir".to_string()));
    users.push(User::new(Email::new("nikitapekin".to_string(),
                                    "gmail.com".to_string()),
                         "12345".to_string(), "./maildir".to_string()));
    user::save_users(filename, users);
}
<|endoftext|>


<commit_before>#![recursion_limit = "256"]
#![feature(proc_macro_hygiene, decl_macro)]

#[macro_use]
extern crate diesel;
#[macro_use]
extern crate diesel_migrations;
#[cfg(feature = "profile-index")]
#[macro_use]
extern crate flamer;

#[cfg(unix)]
use log::error;
#[cfg(unix)]
use sd_notify::{self, NotifyState};
#[cfg(unix)]
use std::fs::File;
#[cfg(unix)]
use std::io::prelude::*;
#[cfg(unix)]
use unix_daemonize::{daemonize_redirect, ChdirMode};

use anyhow::*;
use core::ops::Deref;
use getopts::Options;
use log::info;
use simplelog::{LevelFilter, SimpleLogger, TermLogger, TerminalMode};
use std::path::Path;
use std::sync::Arc;

mod api;
#[cfg(test)]
mod api_tests;
mod config;
mod db;
mod ddns;
mod index;
mod lastfm;
mod metadata;
mod playlist;
mod serve;
mod server;
mod swagger;
#[cfg(test)]
mod test;
mod thumbnails;
mod ui;
mod user;
mod utils;
mod vfs;
mod web;

fn log_config() -> simplelog::Config {
	simplelog::ConfigBuilder::new()
		.set_location_level(LevelFilter::Error)
		.build()
}

#[cfg(unix)]
fn daemonize(options: &getopts::Matches) -> Result<()> {
	if options.opt_present("f") {
		return Ok(());
	}
	let mut log_file = utils::get_data_root()?;
	log_file.push("polaris.log");
	let pid = match daemonize_redirect(Some(&log_file), Some(&log_file), ChdirMode::NoChdir) {
		Ok(p) => p,
		Err(e) => bail!("Daemonize error: {#?}", e),
	};
	let mut pid_path = utils::get_data_root()?;
	pid_path.push("polaris.pid");
	let mut file = File::create(pid_path)?;
	file.write_all(pid.to_string().as_bytes())?;
	Ok(())
}

#[cfg(unix)]
fn init_log(log_level: LevelFilter, options: &getopts::Matches) -> Result<()> {
	if options.opt_present("f") {
		if let Err(e) = TermLogger::init(log_level, log_config(), TerminalMode::Stdout) {
			println!("Error starting terminal logger: {}", e);
		} else {
			return Ok(());
		}
	}

	if let Err(e) = SimpleLogger::init(log_level, log_config()) {
		bail!("Error starting simple logger: {}", e);
	}
	Ok(())
}

#[cfg(windows)]
fn init_log(log_level: LevelFilter, _: &getopts::Matches) -> Result<()> {
	if TermLogger::init(log_level, log_config(), TerminalMode::Stdout).is_err() {
		if let Err(e) = SimpleLogger::init(log_level, log_config()) {
			bail!("Error starting simple logger: {}", e);
		}
	};
	Ok(())
}

#[cfg(unix)]
fn notify_ready() {
	if let Ok(true) = sd_notify::booted() {
		if let Err(e) = sd_notify::notify(true, &[NotifyState::Ready]) {
			error!("Unable to send ready notification: {}", e);
		}
	}
}

#[cfg(not(unix))]
fn notify_ready() {}

fn main() -> Result<()> {
	// Parse CLI options
	let args: Vec<String> = std::env::args().collect();
	let mut options = Options::new();
	options.optopt("c", "config", "set the configuration file", "FILE");
	options.optopt("p", "port", "set polaris to run on a custom port", "PORT");
	options.optopt("d", "database", "set the path to index database", "FILE");
	options.optopt("w", "web", "set the path to web client files", "DIRECTORY");
	options.optopt("s", "swagger", "set the path to swagger files", "DIRECTORY");
	options.optopt(
		"l",
		"log",
		"set the log level to a value between 0 (off) and 3 (debug)",
		"LEVEL",
	);

	#[cfg(unix)]
	options.optflag(
		"f",
		"foreground",
		"run polaris in the foreground instead of daemonizing",
	);

	options.optflag("h", "help", "print this help menu");

	let matches = options.parse(&args[1..])?;

	if matches.opt_present("h") {
		let program = args[0].clone();
		let brief = format!("Usage: {} [options]", program);
		print!("{}", options.usage(&brief));
		return Ok(());
	}

	let log_level = match matches.opt_str("l").as_ref().map(String::as_ref) {
		Some("0") => LevelFilter::Off,
		Some("1") => LevelFilter::Error,
		Some("2") => LevelFilter::Info,
		Some("3") => LevelFilter::Debug,
		_ => LevelFilter::Info,
	};

	init_log(log_level, &matches)?;

	#[cfg(unix)]
	daemonize(&matches)?;

	// Init DB
	info!("Starting up database");
	let db_path = matches.opt_str("d");
	let mut default_db_path = utils::get_data_root()?;
	default_db_path.push("db.sqlite");
	let db_path = db_path
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_db_path);
	let db = Arc::new(db::DB::new(&db_path)?);

	// Parse config
	info!("Parsing configuration");
	let config_file_name = matches.opt_str("c");
	let config_file_path = config_file_name.map(|p| Path::new(p.as_str()).to_path_buf());
	if let Some(path) = config_file_path {
		let config = config::parse_toml_file(&path)?;
		info!("Applying configuration");
		config::amend(db.deref(), &config)?;
	}
	let config = config::read(db.deref())?;
	let auth_secret = config::get_auth_secret(db.deref())?;

	// Init index
	info!("Initializing index");
	let command_sender = index::init(db.clone());

	// Trigger auto-indexing
	let db_auto_index = db.clone();
	let command_sender_auto_index = command_sender.clone();
	std::thread::spawn(move || {
		index::self_trigger(db_auto_index.deref(), &command_sender_auto_index);
	});

	// API mount target
	let prefix_url = config.prefix_url.unwrap_or_else(|| "".to_string());
	let api_url = format!("/{}api", &prefix_url);
	info!("Mounting API on {}", api_url);

	// Web client mount target
	let web_dir_name = matches.opt_str("w");
	let mut default_web_dir = utils::get_data_root()?;
	default_web_dir.push("web");
	let web_dir_path = web_dir_name
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_web_dir);
	info!("Static files location is {}", web_dir_path.display());
	let web_url = format!("/{}", &prefix_url);
	info!("Mounting web client files on {}", web_url);

	// Swagger files mount target
	let swagger_dir_name = matches.opt_str("s");
	let mut default_swagger_dir = utils::get_data_root()?;
	default_swagger_dir.push("swagger");
	let swagger_dir_path = swagger_dir_name
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_swagger_dir);
	info!("Swagger files location is {}", swagger_dir_path.display());
	let swagger_url = format!("/{}swagger", &prefix_url);
	info!("Mounting swagger files on {}", swagger_url);

	// Start server
	info!("Starting up server");
	let port: u16 = matches
		.opt_str("p")
		.unwrap_or_else(|| "5050".to_owned())
		.parse()
		.with_context(|| "Invalid port number")?;

	let server = server::get_server(
		port,
		Some(auth_secret.as_slice()),
		&api_url,
		&web_url,
		&web_dir_path,
		&swagger_url,
		&swagger_dir_path,
		db.clone(),
		command_sender,
	)?;
	std::thread::spawn(move || {
		server.launch();
	});

	// Start DDNS updates
	let db_ddns = db.clone();
	std::thread::spawn(move || {
		ddns::run(db_ddns.deref());
	});

	// Send readiness notification
	notify_ready();

	// Run UI
	ui::run();

	info!("Shutting down server");
	Ok(())
}
<commit_msg>Fixed typo<commit_after>#![recursion_limit = "256"]
#![feature(proc_macro_hygiene, decl_macro)]

#[macro_use]
extern crate diesel;
#[macro_use]
extern crate diesel_migrations;
#[cfg(feature = "profile-index")]
#[macro_use]
extern crate flamer;

#[cfg(unix)]
use log::error;
#[cfg(unix)]
use sd_notify::{self, NotifyState};
#[cfg(unix)]
use std::fs::File;
#[cfg(unix)]
use std::io::prelude::*;
#[cfg(unix)]
use unix_daemonize::{daemonize_redirect, ChdirMode};

use anyhow::*;
use core::ops::Deref;
use getopts::Options;
use log::info;
use simplelog::{LevelFilter, SimpleLogger, TermLogger, TerminalMode};
use std::path::Path;
use std::sync::Arc;

mod api;
#[cfg(test)]
mod api_tests;
mod config;
mod db;
mod ddns;
mod index;
mod lastfm;
mod metadata;
mod playlist;
mod serve;
mod server;
mod swagger;
#[cfg(test)]
mod test;
mod thumbnails;
mod ui;
mod user;
mod utils;
mod vfs;
mod web;

fn log_config() -> simplelog::Config {
	simplelog::ConfigBuilder::new()
		.set_location_level(LevelFilter::Error)
		.build()
}

#[cfg(unix)]
fn daemonize(options: &getopts::Matches) -> Result<()> {
	if options.opt_present("f") {
		return Ok(());
	}
	let mut log_file = utils::get_data_root()?;
	log_file.push("polaris.log");
	let pid = match daemonize_redirect(Some(&log_file), Some(&log_file), ChdirMode::NoChdir) {
		Ok(p) => p,
		Err(e) => bail!("Daemonize error: {:#?}", e),
	};
	let mut pid_path = utils::get_data_root()?;
	pid_path.push("polaris.pid");
	let mut file = File::create(pid_path)?;
	file.write_all(pid.to_string().as_bytes())?;
	Ok(())
}

#[cfg(unix)]
fn init_log(log_level: LevelFilter, options: &getopts::Matches) -> Result<()> {
	if options.opt_present("f") {
		if let Err(e) = TermLogger::init(log_level, log_config(), TerminalMode::Stdout) {
			println!("Error starting terminal logger: {}", e);
		} else {
			return Ok(());
		}
	}

	if let Err(e) = SimpleLogger::init(log_level, log_config()) {
		bail!("Error starting simple logger: {}", e);
	}
	Ok(())
}

#[cfg(windows)]
fn init_log(log_level: LevelFilter, _: &getopts::Matches) -> Result<()> {
	if TermLogger::init(log_level, log_config(), TerminalMode::Stdout).is_err() {
		if let Err(e) = SimpleLogger::init(log_level, log_config()) {
			bail!("Error starting simple logger: {}", e);
		}
	};
	Ok(())
}

#[cfg(unix)]
fn notify_ready() {
	if let Ok(true) = sd_notify::booted() {
		if let Err(e) = sd_notify::notify(true, &[NotifyState::Ready]) {
			error!("Unable to send ready notification: {}", e);
		}
	}
}

#[cfg(not(unix))]
fn notify_ready() {}

fn main() -> Result<()> {
	// Parse CLI options
	let args: Vec<String> = std::env::args().collect();
	let mut options = Options::new();
	options.optopt("c", "config", "set the configuration file", "FILE");
	options.optopt("p", "port", "set polaris to run on a custom port", "PORT");
	options.optopt("d", "database", "set the path to index database", "FILE");
	options.optopt("w", "web", "set the path to web client files", "DIRECTORY");
	options.optopt("s", "swagger", "set the path to swagger files", "DIRECTORY");
	options.optopt(
		"l",
		"log",
		"set the log level to a value between 0 (off) and 3 (debug)",
		"LEVEL",
	);

	#[cfg(unix)]
	options.optflag(
		"f",
		"foreground",
		"run polaris in the foreground instead of daemonizing",
	);

	options.optflag("h", "help", "print this help menu");

	let matches = options.parse(&args[1..])?;

	if matches.opt_present("h") {
		let program = args[0].clone();
		let brief = format!("Usage: {} [options]", program);
		print!("{}", options.usage(&brief));
		return Ok(());
	}

	let log_level = match matches.opt_str("l").as_ref().map(String::as_ref) {
		Some("0") => LevelFilter::Off,
		Some("1") => LevelFilter::Error,
		Some("2") => LevelFilter::Info,
		Some("3") => LevelFilter::Debug,
		_ => LevelFilter::Info,
	};

	init_log(log_level, &matches)?;

	#[cfg(unix)]
	daemonize(&matches)?;

	// Init DB
	info!("Starting up database");
	let db_path = matches.opt_str("d");
	let mut default_db_path = utils::get_data_root()?;
	default_db_path.push("db.sqlite");
	let db_path = db_path
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_db_path);
	let db = Arc::new(db::DB::new(&db_path)?);

	// Parse config
	info!("Parsing configuration");
	let config_file_name = matches.opt_str("c");
	let config_file_path = config_file_name.map(|p| Path::new(p.as_str()).to_path_buf());
	if let Some(path) = config_file_path {
		let config = config::parse_toml_file(&path)?;
		info!("Applying configuration");
		config::amend(db.deref(), &config)?;
	}
	let config = config::read(db.deref())?;
	let auth_secret = config::get_auth_secret(db.deref())?;

	// Init index
	info!("Initializing index");
	let command_sender = index::init(db.clone());

	// Trigger auto-indexing
	let db_auto_index = db.clone();
	let command_sender_auto_index = command_sender.clone();
	std::thread::spawn(move || {
		index::self_trigger(db_auto_index.deref(), &command_sender_auto_index);
	});

	// API mount target
	let prefix_url = config.prefix_url.unwrap_or_else(|| "".to_string());
	let api_url = format!("/{}api", &prefix_url);
	info!("Mounting API on {}", api_url);

	// Web client mount target
	let web_dir_name = matches.opt_str("w");
	let mut default_web_dir = utils::get_data_root()?;
	default_web_dir.push("web");
	let web_dir_path = web_dir_name
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_web_dir);
	info!("Static files location is {}", web_dir_path.display());
	let web_url = format!("/{}", &prefix_url);
	info!("Mounting web client files on {}", web_url);

	// Swagger files mount target
	let swagger_dir_name = matches.opt_str("s");
	let mut default_swagger_dir = utils::get_data_root()?;
	default_swagger_dir.push("swagger");
	let swagger_dir_path = swagger_dir_name
		.map(|n| Path::new(n.as_str()).to_path_buf())
		.unwrap_or(default_swagger_dir);
	info!("Swagger files location is {}", swagger_dir_path.display());
	let swagger_url = format!("/{}swagger", &prefix_url);
	info!("Mounting swagger files on {}", swagger_url);

	// Start server
	info!("Starting up server");
	let port: u16 = matches
		.opt_str("p")
		.unwrap_or_else(|| "5050".to_owned())
		.parse()
		.with_context(|| "Invalid port number")?;

	let server = server::get_server(
		port,
		Some(auth_secret.as_slice()),
		&api_url,
		&web_url,
		&web_dir_path,
		&swagger_url,
		&swagger_dir_path,
		db.clone(),
		command_sender,
	)?;
	std::thread::spawn(move || {
		server.launch();
	});

	// Start DDNS updates
	let db_ddns = db.clone();
	std::thread::spawn(move || {
		ddns::run(db_ddns.deref());
	});

	// Send readiness notification
	notify_ready();

	// Run UI
	ui::run();

	info!("Shutting down server");
	Ok(())
}
<|endoftext|>