<commit_before>extern crate rs6502;

use std::io::{self, BufRead, Write};
use std::sync::mpsc::{channel, Receiver};
use std::thread;

use rs6502::{Cpu, Disassembler};

#[derive(Debug)]
pub struct MemoryMonitor {
    pub enabled: bool,
    start_addr: usize,
    end_addr: usize,
}

pub struct VirtualMachine {
    pub cpu: Cpu,
    pub monitor: MemoryMonitor,
    code_offset: u16,
    clock_rate: Option<u32>,
    receiver: Receiver<String>,
    last_command: String,
    breakpoints: [u8; 64 * 1024],
    broken: bool,
}

impl VirtualMachine {
    pub fn new<CR>(cpu: Cpu, code_offset: u16, clock_rate: CR) -> VirtualMachine
        where CR: Into<Option<u32>>
    {
        let (tx, rx) = channel();

        thread::spawn(move || {
            std::io::stdout().write(b"hakka> ").unwrap();
            std::io::stdout().flush().unwrap();

            loop {
                let mut line = String::new();
                let stdin = io::stdin();
                stdin.lock().read_line(&mut line).expect("Could not read line");
                tx.send(line).unwrap();
            }
        });

        VirtualMachine {
            cpu: cpu,
            code_offset: code_offset,
            clock_rate: clock_rate.into(),
            receiver: rx,
            monitor: MemoryMonitor {
                enabled: false,
                start_addr: 0,
                end_addr: 0,
            },
            last_command: "".into(),
            breakpoints: [0; 64 * 1024],
            broken: false,
        }
    }

    pub fn cycle(&mut self) {
        if let Some(clock_rate) = self.clock_rate {
            let mut n = 0;
            while n < clock_rate && !self.broken {
                n += self.cpu.step().expect("SEGFAULT") as u32;
                if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                    self.broken = true;
                    println!("");
                    println!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC);
                }
            }
        } else {
            self.cpu.step().expect("SEGFAULT");
            if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                self.broken = true;
                println!("");
                println!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC);
            }
        }
    }

    pub fn try_execute_command(&mut self) {
        if let Ok(input) = self.receiver.try_recv() {
            let mut input: String = input.trim().into();

            if input.len() == 0 {
                if self.monitor.enabled {
                    self.monitor.enabled = false;
                } else {
                    input = self.last_command.clone();
                }
            }

            if input == "source" {
                self.dump_disassembly();
            }

            if input == "list" {
                self.dump_local_disassembly();
            }

            if input.starts_with("monitor") {
                if self.monitor.enabled {
                    self.monitor.enabled = false;
                } else {
                    self.enable_memory_monitor(&input);
                }
            }

            if input.starts_with("memset") {
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() < 3 {
                    println!("ERR: Requires 2 arguments. Example: memset 0x00 0x01 to store 0x01 \
                              in 0x00.");
                } else if parts.len() == 3 {
                    if let Ok(dst) = u16::from_str_radix(&parts[1].replace("0x", "")[..], 16) {
                        if let Ok(src) = u8::from_str_radix(&parts[2].replace("0x", "")[..], 16) {
                            self.cpu.memory[dst as usize] = src;
                        } else {
                            println!("ERR: Unable to parse source byte value");
                        }
                    } else {
                        println!("ERR: Unable to parse destination byte value");
                    }
                } else {
                    if let Ok(mut dst) = usize::from_str_radix(&parts[1].replace("0x", "")[..],
                                                               16) {
                        for p in &parts[2..] {
                            if let Ok(byte) = u8::from_str_radix(&p.replace("0x", "")[..], 16) {
                                self.cpu.memory[dst] = byte;
                                dst += 0x01;
                            } else {
                                println!("ERR: Unable to parse source byte value");
                            }
                        }
                    } else {
                        println!("ERR: Unable to parse destination byte value");
                    }
                }
            }

            if input.starts_with("memdmp") {
                // 1 argument assumes 1 memory "page"
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() == 2 {
                    let page_number: usize =
                        parts[1].parse().expect("Unable to parse memory page number");

                    self.dump_memory_page(page_number);
                }
            }

            if input.starts_with("break") {
                // 1 argument assumes 1 memory "page"
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() == 2 {
                    if let Ok(addr) = usize::from_str_radix(&parts[1][..], 16) {
                        self.breakpoints[addr] = 1;
                    } else {
                        println!("Err: Unable to parse breakpoint address");
                    }
                } else {
                    println!("ERR: Requires 1 argument");
                }
            }

            if input == "continue" {
                self.broken = false;
                println!("Execution resumed");
            }

            std::io::stdout().write(b"hakka> ").unwrap();
            std::io::stdout().flush().unwrap();

            // Don't assign a blank command as a last command
            if input.len() > 0 {
                self.last_command = input.clone();
            }
        }

    }

    fn enable_memory_monitor(&mut self, input: &str) {
        let parts: Vec<&str> = input.split(" ").collect();
        if parts.len() < 3 {
            println!("ERR: Requires 2 arguments. Example: monitor 0x00 0xFF");
        } else {
            let start = usize::from_str_radix(&parts[1].replace("0x", "")[..], 16).unwrap();
            let end = usize::from_str_radix(&parts[2].replace("0x", "")[..], 16).unwrap();

            self.monitor.start_addr = start;
            self.monitor.end_addr = end;
            self.monitor.enabled = true;

            println!("{:?}", self.monitor);
        }
    }

    fn dump_disassembly(&self) {
        println!("");
        println!("-- Disassembly --");

        let disassembler = Disassembler::with_offset(self.code_offset);
        let pairs = disassembler.disassemble_with_addresses(self.cpu.get_code());
        let result = self.highlight_lines(self.cpu.registers.PC as usize, pairs, false).join("");
        print!("{}", result);
    }

    fn dump_local_disassembly(&self) {
        println!("");
        println!("-- Disassembly --");

        let disassembler = Disassembler::with_offset(self.code_offset);
        let pc = self.cpu.registers.PC as usize;
        let code = self.cpu.get_code();
        let pairs = disassembler.disassemble_with_addresses(code);
        let result = self.highlight_lines(pc, pairs, true).join("");
        print!("{}", result);
    }

    pub fn dump_memory_page(&self, page: usize) {
        let mut addr = page * 0x100;
        println!("-- Memory dump --");
        for chunk in self.cpu.memory[page * 0x100..(page * 0x100) + 0x100].chunks(8) {
            print!("{:04X}: ", addr);
            for b in chunk {
                print!("{:02X} ", *b);
            }
            addr += 0x08;
            println!("");
        }
    }

    pub fn dump_memory(&self) {
        println!("-- Memory dump --");
        for chunk in self.cpu.memory[self.monitor.start_addr..self.monitor.end_addr + 0x01]
            .chunks(8) {
            for b in chunk {
                print!("{:02X} ", *b);
            }
            println!("");
        }
    }

    fn highlight_lines(&self,
                       pc: usize,
                       pairs: Vec<(String, u16)>,
                       limit_results: bool)
                       -> Vec<String> {
        let mut result = Vec::new();

        let base = pc - self.code_offset as usize;

        for pair in pairs {
            if limit_results {
                let start = if base > 0x0A { base - 0x0A } else { 0 };
                if (pair.1 as usize) < start || (pair.1 as usize) > base + 0x0A {
                    continue;
                }
            }
            let breakpoint = self.breakpoints[self.code_offset as usize + pair.1 as usize] > 0x00;
            let current_line = pc as u16 == self.code_offset + pair.1;

            if breakpoint && current_line {
                result.push(format!("> * {}", pair.0));
            } else if breakpoint && !current_line {
                result.push(format!("  * {}", pair.0));
            } else {
                result.push(format!("    {}", pair.0));
            }
        }

        result
    }
}<commit_msg>Add breakpoint clearing support<commit_after>extern crate rs6502;

use std::io::{self, BufRead, Write};
use std::sync::mpsc::{channel, Receiver};
use std::thread;

use rs6502::{Cpu, Disassembler};

#[derive(Debug)]
pub struct MemoryMonitor {
    pub enabled: bool,
    start_addr: usize,
    end_addr: usize,
}

pub struct VirtualMachine {
    pub cpu: Cpu,
    pub monitor: MemoryMonitor,
    code_offset: u16,
    clock_rate: Option<u32>,
    receiver: Receiver<String>,
    last_command: String,
    breakpoints: [u8; 64 * 1024],
    broken: bool,
}

impl VirtualMachine {
    pub fn new<CR>(cpu: Cpu, code_offset: u16, clock_rate: CR) -> VirtualMachine
        where CR: Into<Option<u32>>
    {
        let (tx, rx) = channel();

        thread::spawn(move || {
            std::io::stdout().write(b"hakka> ").unwrap();
            std::io::stdout().flush().unwrap();

            loop {
                let mut line = String::new();
                let stdin = io::stdin();
                stdin.lock().read_line(&mut line).expect("Could not read line");
                tx.send(line).unwrap();
            }
        });

        VirtualMachine {
            cpu: cpu,
            code_offset: code_offset,
            clock_rate: clock_rate.into(),
            receiver: rx,
            monitor: MemoryMonitor {
                enabled: false,
                start_addr: 0,
                end_addr: 0,
            },
            last_command: "".into(),
            breakpoints: [0; 64 * 1024],
            broken: false,
        }
    }

    pub fn cycle(&mut self) {
        if let Some(clock_rate) = self.clock_rate {
            let mut n = 0;
            while n < clock_rate && !self.broken {
                n += self.cpu.step().expect("SEGFAULT") as u32;
                if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                    self.broken = true;
                    println!("");
                    println!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC);
                }
            }
        } else {
            self.cpu.step().expect("SEGFAULT");
            if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                self.broken = true;
                println!("");
                println!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC);
            }
        }
    }

    pub fn try_execute_command(&mut self) {
        if let Ok(input) = self.receiver.try_recv() {
            let mut input: String = input.trim().into();

            if input.len() == 0 {
                if self.monitor.enabled {
                    self.monitor.enabled = false;
                } else {
                    input = self.last_command.clone();
                }
            }

            if input == "source" {
                self.dump_disassembly();
            }

            if input == "list" {
                self.dump_local_disassembly();
            }

            if input.starts_with("monitor") {
                if self.monitor.enabled {
                    self.monitor.enabled = false;
                } else {
                    self.enable_memory_monitor(&input);
                }
            }

            if input.starts_with("memset") {
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() < 3 {
                    println!("ERR: Requires 2 arguments. Example: memset 0x00 0x01 to store 0x01 \
                              in 0x00.");
                } else if parts.len() == 3 {
                    if let Ok(dst) = u16::from_str_radix(&parts[1].replace("0x", "")[..], 16) {
                        if let Ok(src) = u8::from_str_radix(&parts[2].replace("0x", "")[..], 16) {
                            self.cpu.memory[dst as usize] = src;
                        } else {
                            println!("ERR: Unable to parse source byte value");
                        }
                    } else {
                        println!("ERR: Unable to parse destination byte value");
                    }
                } else {
                    if let Ok(mut dst) = usize::from_str_radix(&parts[1].replace("0x", "")[..],
                                                               16) {
                        for p in &parts[2..] {
                            if let Ok(byte) = u8::from_str_radix(&p.replace("0x", "")[..], 16) {
                                self.cpu.memory[dst] = byte;
                                dst += 0x01;
                            } else {
                                println!("ERR: Unable to parse source byte value");
                            }
                        }
                    } else {
                        println!("ERR: Unable to parse destination byte value");
                    }
                }
            }

            if input.starts_with("memdmp") {
                // 1 argument assumes 1 memory "page"
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() == 2 {
                    let page_number: usize =
                        parts[1].parse().expect("Unable to parse memory page number");

                    self.dump_memory_page(page_number);
                }
            }

            if input.starts_with("break") {
                // 1 argument assumes 1 memory "page"
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() == 2 {
                    if let Ok(addr) = usize::from_str_radix(&parts[1][..], 16) {
                        if self.breakpoints[addr] > 0 {
                            self.breakpoints[addr] = 0;
                        } else {
                            self.breakpoints[addr] = 1;
                        }
                    } else {
                        println!("Err: Unable to parse breakpoint address");
                    }
                } else {
                    println!("ERR: Requires 1 argument");
                }
            }

            if input == "continue" {
                self.broken = false;
                println!("Execution resumed");
            }

            std::io::stdout().write(b"hakka> ").unwrap();
            std::io::stdout().flush().unwrap();

            // Don't assign a blank command as a last command
            if input.len() > 0 {
                self.last_command = input.clone();
            }
        }

    }

    fn enable_memory_monitor(&mut self, input: &str) {
        let parts: Vec<&str> = input.split(" ").collect();
        if parts.len() < 3 {
            println!("ERR: Requires 2 arguments. Example: monitor 0x00 0xFF");
        } else {
            let start = usize::from_str_radix(&parts[1].replace("0x", "")[..], 16).unwrap();
            let end = usize::from_str_radix(&parts[2].replace("0x", "")[..], 16).unwrap();

            self.monitor.start_addr = start;
            self.monitor.end_addr = end;
            self.monitor.enabled = true;

            println!("{:?}", self.monitor);
        }
    }

    fn dump_disassembly(&self) {
        println!("");
        println!("-- Disassembly --");

        let disassembler = Disassembler::with_offset(self.code_offset);
        let pairs = disassembler.disassemble_with_addresses(self.cpu.get_code());
        let result = self.highlight_lines(self.cpu.registers.PC as usize, pairs, false).join("");
        print!("{}", result);
    }

    fn dump_local_disassembly(&self) {
        println!("");
        println!("-- Disassembly --");

        let disassembler = Disassembler::with_offset(self.code_offset);
        let pc = self.cpu.registers.PC as usize;
        let code = self.cpu.get_code();
        let pairs = disassembler.disassemble_with_addresses(code);
        let result = self.highlight_lines(pc, pairs, true).join("");
        print!("{}", result);
    }

    pub fn dump_memory_page(&self, page: usize) {
        let mut addr = page * 0x100;
        println!("-- Memory dump --");
        for chunk in self.cpu.memory[page * 0x100..(page * 0x100) + 0x100].chunks(8) {
            print!("{:04X}: ", addr);
            for b in chunk {
                print!("{:02X} ", *b);
            }
            addr += 0x08;
            println!("");
        }
    }

    pub fn dump_memory(&self) {
        println!("-- Memory dump --");
        for chunk in self.cpu.memory[self.monitor.start_addr..self.monitor.end_addr + 0x01]
            .chunks(8) {
            for b in chunk {
                print!("{:02X} ", *b);
            }
            println!("");
        }
    }

    fn highlight_lines(&self,
                       pc: usize,
                       pairs: Vec<(String, u16)>,
                       limit_results: bool)
                       -> Vec<String> {
        let mut result = Vec::new();

        let base = pc - self.code_offset as usize;

        for pair in pairs {
            if limit_results {
                let start = if base > 0x0A { base - 0x0A } else { 0 };
                if (pair.1 as usize) < start || (pair.1 as usize) > base + 0x0A {
                    continue;
                }
            }
            let breakpoint = self.breakpoints[self.code_offset as usize + pair.1 as usize] > 0x00;
            let current_line = pc as u16 == self.code_offset + pair.1;

            if breakpoint && current_line {
                result.push(format!("> * {}", pair.0));
            } else if breakpoint && !current_line {
                result.push(format!("  * {}", pair.0));
            } else {
                result.push(format!("    {}", pair.0));
            }
        }

        result
    }
}<|endoftext|>


<commit_before>//! This is a "monotonic HashMap": A HashMap that, when shared, can be pushed to but not
//! otherwise mutated.  We also Box items in the map. This means we can safely provide
//! shared references into existing items in the HashMap, because they will not be dropped
//! (from being removed) or moved (because they are boxed).
//! The API is is completely tailored to what `memory.rs` needs. It is still in
//! a separate file to minimize the amount of code that has to care about the unsafety.

use std::collections::hash_map::Entry;
use std::cell::RefCell;
use std::hash::Hash;
use std::borrow::Borrow;

use rustc_data_structures::fx::FxHashMap;

use crate::AllocMap;

#[derive(Debug, Clone)]
pub struct MonoHashMap<K: Hash + Eq, V>(RefCell<FxHashMap<K, Box<V>>>);

impl<K: Hash + Eq, V> MonoHashMap<K, V> {
    /// This function exists for priroda to be able to iterate over all evaluator memory
    ///
    /// The memory of constants does not show up in this list.
    pub fn iter<T>(&self, f: impl FnOnce(&mut dyn Iterator<Item=(&K, &V)>) -> T) -> T {
        f(&mut self.0.borrow().iter().map(|(k, v)| (k, &**v)))
    }
}

impl<K: Hash + Eq, V> Default for MonoHashMap<K, V> {
    fn default() -> Self {
        MonoHashMap(RefCell::new(Default::default()))
    }
}

impl<K: Hash + Eq, V> AllocMap<K, V> for MonoHashMap<K, V> {
    #[inline(always)]
    fn contains_key<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> bool
        where K: Borrow<Q>
    {
        self.0.get_mut().contains_key(k)
    }

    #[inline(always)]
    fn insert(&mut self, k: K, v: V) -> Option<V>
    {
        self.0.get_mut().insert(k, Box::new(v)).map(|x| *x)
    }

    #[inline(always)]
    fn remove<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>
    {
        self.0.get_mut().remove(k).map(|x| *x)
    }

    #[inline(always)]
    fn filter_map_collect<T>(&self, mut f: impl FnMut(&K, &V) -> Option<T>) -> Vec<T> {
        self.0.borrow()
            .iter()
            .filter_map(move |(k, v)| f(k, &*v))
            .collect()
    }

    /// The most interesting method: Providing a shared ref without
    /// holding the `RefCell` open, and inserting new data if the key
    /// is not used yet.
    /// `vacant` is called if the key is not found in the map;
    /// if it returns a reference, that is used directly, if it
    /// returns owned data, that is put into the map and returned.
    #[inline(always)]
    fn get_or<E>(
        &self,
        k: K,
        vacant: impl FnOnce() -> Result<V, E>
    ) -> Result<&V, E> {
        let val: *const V = match self.0.borrow_mut().entry(k) {
            Entry::Occupied(entry) => &**entry.get(),
            Entry::Vacant(entry) => &**entry.insert(Box::new(vacant()?)),
        };
        // This is safe because `val` points into a `Box`, that we know will not move and
        // will also not be dropped as long as the shared reference `self` is live.
        unsafe { Ok(&*val) }
    }

    #[inline(always)]
    fn get_mut_or<E>(
        &mut self,
        k: K,
        vacant: impl FnOnce() -> Result<V, E>
    ) -> Result<&mut V, E>
    {
        match self.0.get_mut().entry(k) {
            Entry::Occupied(e) => Ok(e.into_mut()),
            Entry::Vacant(e) => {
                let v = vacant()?;
                Ok(e.insert(Box::new(v)))
            }
        }
    }
}
<commit_msg>Explain Ref problems<commit_after>//! This is a "monotonic HashMap": A HashMap that, when shared, can be pushed to but not
//! otherwise mutated.  We also Box items in the map. This means we can safely provide
//! shared references into existing items in the HashMap, because they will not be dropped
//! (from being removed) or moved (because they are boxed).
//! The API is is completely tailored to what `memory.rs` needs. It is still in
//! a separate file to minimize the amount of code that has to care about the unsafety.

use std::collections::hash_map::Entry;
use std::cell::RefCell;
use std::hash::Hash;
use std::borrow::Borrow;

use rustc_data_structures::fx::FxHashMap;

use crate::AllocMap;

#[derive(Debug, Clone)]
pub struct MonoHashMap<K: Hash + Eq, V>(RefCell<FxHashMap<K, Box<V>>>);

impl<K: Hash + Eq, V> MonoHashMap<K, V> {
    /// This function exists for priroda to be able to iterate over all evaluator memory
    ///
    /// The function is somewhat roundabout with the closure argument because internally the
    /// `MonoHashMap` uses a `RefCell`. When iterating over the `HashMap` inside the `RefCell`,
    /// we need to keep a borrow to the `HashMap` inside the iterator. The borrow is only alive
    /// as long as the `Ref` returned by `RefCell::borrow()` is alive. So we can't return the
    /// iterator, as that would drop the `Ref`. We can't return both, as it's not possible in Rust
    /// to have a struct/tuple with a field that refers to another field.
    pub fn iter<T>(&self, f: impl FnOnce(&mut dyn Iterator<Item=(&K, &V)>) -> T) -> T {
        f(&mut self.0.borrow().iter().map(|(k, v)| (k, &**v)))
    }
}

impl<K: Hash + Eq, V> Default for MonoHashMap<K, V> {
    fn default() -> Self {
        MonoHashMap(RefCell::new(Default::default()))
    }
}

impl<K: Hash + Eq, V> AllocMap<K, V> for MonoHashMap<K, V> {
    #[inline(always)]
    fn contains_key<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> bool
        where K: Borrow<Q>
    {
        self.0.get_mut().contains_key(k)
    }

    #[inline(always)]
    fn insert(&mut self, k: K, v: V) -> Option<V>
    {
        self.0.get_mut().insert(k, Box::new(v)).map(|x| *x)
    }

    #[inline(always)]
    fn remove<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>
    {
        self.0.get_mut().remove(k).map(|x| *x)
    }

    #[inline(always)]
    fn filter_map_collect<T>(&self, mut f: impl FnMut(&K, &V) -> Option<T>) -> Vec<T> {
        self.0.borrow()
            .iter()
            .filter_map(move |(k, v)| f(k, &*v))
            .collect()
    }

    /// The most interesting method: Providing a shared ref without
    /// holding the `RefCell` open, and inserting new data if the key
    /// is not used yet.
    /// `vacant` is called if the key is not found in the map;
    /// if it returns a reference, that is used directly, if it
    /// returns owned data, that is put into the map and returned.
    #[inline(always)]
    fn get_or<E>(
        &self,
        k: K,
        vacant: impl FnOnce() -> Result<V, E>
    ) -> Result<&V, E> {
        let val: *const V = match self.0.borrow_mut().entry(k) {
            Entry::Occupied(entry) => &**entry.get(),
            Entry::Vacant(entry) => &**entry.insert(Box::new(vacant()?)),
        };
        // This is safe because `val` points into a `Box`, that we know will not move and
        // will also not be dropped as long as the shared reference `self` is live.
        unsafe { Ok(&*val) }
    }

    #[inline(always)]
    fn get_mut_or<E>(
        &mut self,
        k: K,
        vacant: impl FnOnce() -> Result<V, E>
    ) -> Result<&mut V, E>
    {
        match self.0.get_mut().entry(k) {
            Entry::Occupied(e) => Ok(e.into_mut()),
            Entry::Vacant(e) => {
                let v = vacant()?;
                Ok(e.insert(Box::new(v)))
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 click2stream, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

///! Network scanner for RTSP streams.

use std::io;
use std::fmt;
use std::thread;
use std::result;

use std::fs::File;
use std::sync::Arc;
use std::error::Error;
use std::collections::HashSet;
use std::io::{BufReader, BufRead};
use std::fmt::{Display, Formatter};
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};

use net::rtsp;
use net::raw::pcap;

use net::rtsp::Client as RtspClient;
use net::raw::devices::EthernetDevice;
use net::raw::ether::MacAddr;
use net::arrow::protocol::Service;
use net::raw::arp::scanner::Ipv4ArpScanner;
use net::raw::icmp::scanner::IcmpScanner;
use net::raw::tcp::scanner::{TcpPortScanner, PortCollection};
use net::rtsp::sdp::{SessionDescription, MediaType, RTPMap, FromAttribute};

static RTSP_PATH_FILE: &'static str = "/etc/arrow/rtsp-paths";

/// Discovery error.
#[derive(Debug, Clone)]
pub struct DiscoveryError {
    msg: String,
}

impl Error for DiscoveryError {
    fn description(&self) -> &str {
        &self.msg
    }
}

impl Display for DiscoveryError {
    fn fmt(&self, f: &mut Formatter) -> result::Result<(), fmt::Error> {
        f.write_str(&self.msg)
    }
}

impl From<String> for DiscoveryError {
    fn from(msg: String) -> DiscoveryError {
        DiscoveryError { msg: msg }
    }
}

impl<'a> From<&'a str> for DiscoveryError {
    fn from(msg: &'a str) -> DiscoveryError {
        DiscoveryError { msg: msg.to_string() }
    }
}

impl From<rtsp::RtspError> for DiscoveryError {
    fn from(err: rtsp::RtspError) -> DiscoveryError {
        DiscoveryError::from(format!("RTSP client error: {}",
            err.description()))
    }
}

impl From<pcap::PcapError> for DiscoveryError {
    fn from(err: pcap::PcapError) -> DiscoveryError {
        DiscoveryError::from(format!("pcap error: {}", err.description()))
    }
}

impl From<io::Error> for DiscoveryError {
    fn from(err: io::Error) -> DiscoveryError {
        DiscoveryError::from(format!("IO error: {}", err.description()))
    }
}

/// Discovery result type alias.
pub type Result<T> = result::Result<T, DiscoveryError>;

/// Discovery host type alias.
type Host = (MacAddr, Ipv4Addr);
/// Discovery service type alias.
type Socket = (MacAddr, SocketAddrV4);

/// Find all RTSP streams in all local networks.
pub fn find_rtsp_streams() -> Result<Vec<Service>> {
    let tc      = pcap::new_threading_context();
    let devices = EthernetDevice::list();

    let port_candidates = PortCollection::new()
        .add_all(&[554, 88, 81, 555, 7447, 8554, 7070, 10554, 80]);

    let mut threads = Vec::new();

    for dev in devices {
        let pc     = port_candidates.clone();
        let tc     = tc.clone();
        let handle = thread::spawn(move || {
            find_services(tc, &dev, &pc)
        });

        threads.push(handle);
    }

    let mut services = Vec::new();

    for handle in threads {
        match handle.join() {
            Err(_)  => return Err(DiscoveryError::from("port scanner thread panicked")),
            Ok(res) => services.extend(try!(res))
        }
    }

    let rtsp_services = try!(find_rtsp_services(&services));

    let mut threads = Vec::new();
    let paths       = Arc::new(try!(load_rtsp_paths(RTSP_PATH_FILE)));

    for (mac, addr) in rtsp_services {
        let addr   = SocketAddr::V4(addr);
        let paths  = paths.clone();
        let handle = thread::spawn(move || {
            find_rtsp_paths(mac, addr, &paths)
        });
        threads.push(handle);
    }

    let mut res = Vec::new();

    for handle in threads {
        match handle.join() {
            Err(_) => return Err(DiscoveryError::from(
                "path testing thread panicked")),
            Ok(svc) => res.push(try!(svc))
        }
    }

    Ok(res)
}

/// Load all known RTSP path variants from a given file.
fn load_rtsp_paths(file: &str) -> Result<Vec<String>> {
    let file      = try!(File::open(file));
    let breader   = BufReader::new(file);
    let mut paths = Vec::new();

    for line in breader.lines() {
        let path = try!(line);
        if !path.starts_with('#') {
            paths.push(path);
        }
    }

    Ok(paths)
}

/// Check if a given service is an RTSP service.
fn is_rtsp_service(addr: SocketAddr) -> Result<bool> {
    let mut client = try!(RtspClient::new(addr));
    client.set_timeout(Some(1000));
    Ok(client.options().is_ok())
}

/// Check if a given session description contains at least one H.264 or
/// a general MPEG4 video stream.
fn is_supported_service(sdp: &[u8]) -> bool {
    if let Ok(sdp) = SessionDescription::parse(sdp) {
        let mut vcodecs   = HashSet::new();
        let video_streams = sdp.media_descriptions.into_iter()
            .filter(|md| md.media_type == MediaType::Video);

        for md in video_streams {
            for attr in md.attributes {
                if let Ok(rtpmap) = RTPMap::parse(&attr) {
                    vcodecs.insert(rtpmap.encoding.to_uppercase());
                }
            }
        }

        vcodecs.contains("H264") ||
            vcodecs.contains("H264-RCDO") ||
            vcodecs.contains("H264-SVC") ||
            vcodecs.contains("MP4V-ES") ||
            vcodecs.contains("MPEG4-GENERIC")
    } else {
        false
    }
}

/// RTSP DESCRIBE status.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum DescribeStatus {
    Ok,
    Locked,
    Unsupported,
    NotFound,
    Error
}

/// Get describe status code for a given RTSP service and path.
fn get_describe_status(addr: SocketAddr, path: &str) -> Result<DescribeStatus> {
    let mut client = try!(RtspClient::new(addr));
    client.set_timeout(Some(1000));
    if let Ok(response) = client.describe(path) {
        let header = response.header;
        let hipcam = match header.get_str("Server") {
            Some("HiIpcam/V100R003 VodServer/1.0.0") => true,
            Some("Hipcam RealServer/V1.0")           => true,
            _ => false
        };

        if hipcam && path != "/11" && path != "/12" {
            Ok(DescribeStatus::NotFound)
        } else {
            match header.code {
                404 => Ok(DescribeStatus::NotFound),
                401 => Ok(DescribeStatus::Locked),
                200 if is_supported_service(&response.body) =>
                    Ok(DescribeStatus::Ok),
                200 => Ok(DescribeStatus::Unsupported),
                _   => Ok(DescribeStatus::Error)
            }
        }
    } else {
        Ok(DescribeStatus::Error)
    }
}

/// Find open ports on all available hosts within a given network and port
/// range.
fn find_services(
    pc: pcap::ThreadingContext,
    device: &EthernetDevice,
    ports: &PortCollection) -> Result<Vec<Socket>> {
    let mut hosts  = HashSet::new();
    let arp_hosts  = try!(Ipv4ArpScanner::scan_device(pc.clone(), device));
    let icmp_hosts = try!(IcmpScanner::scan_device(pc.clone(), device));

    hosts.extend(arp_hosts);
    hosts.extend(icmp_hosts);

    let res = try!(find_open_ports(pc, device,
        hosts.into_iter(), ports));

    Ok(res)
}

/// Check if any of given TCP ports is open on on any host from a given set.
fn find_open_ports<HI: Iterator<Item=Host>>(
    pc: pcap::ThreadingContext,
    device: &EthernetDevice,
    hosts: HI,
    ports: &PortCollection) -> Result<Vec<Socket>> {
    let res = try!(TcpPortScanner::scan_ipv4_hosts(pc, device, hosts, ports))
        .into_iter()
        .map(|(mac, ip, p)| (mac, SocketAddrV4::new(ip, p)))
        .collect::<Vec<_>>();

    Ok(res)
}

/// Find all RTSP services among a given set of sockets.
fn find_rtsp_services(sockets: &[Socket]) -> Result<Vec<Socket>> {
    let mut threads = Vec::new();
    let mut res     = Vec::new();

    for &(mac, addr) in sockets {
        let handle = thread::spawn(move || {
            (mac, addr, is_rtsp_service(SocketAddr::V4(addr)))
        });
        threads.push(handle);
    }

    for handle in threads {
        match handle.join() {
            Err(_) => return Err(DiscoveryError::from("RTSP service testing thread panicked")),
            Ok((mac, addr, rtsp)) => {
                if try!(rtsp) {
                    res.push((mac, addr))
                }
            }
        }
    }

    Ok(res)
}

/// Find all available RTSP paths for a given RTSP service.
fn find_rtsp_paths(
    mac: MacAddr,
    addr: SocketAddr,
    paths: &[String]) -> Result<Service> {
    let mut service = Service::UnknownRTSP(mac, addr);

    for path in paths {
        let status = try!(get_describe_status(addr, path));
        if status == DescribeStatus::Ok {
            service = Service::RTSP(mac, addr, path.to_string());
        } else if status == DescribeStatus::Unsupported {
            service = Service::UnsupportedRTSP(mac, addr, path.to_string());
        } else if status == DescribeStatus::Locked {
            service = Service::LockedRTSP(mac, addr);
        }

        match status {
            DescribeStatus::Ok     => break,
            DescribeStatus::Locked => break,
            _ => ()
        }
    }

    Ok(service)
}
<commit_msg>the network scanner should report only one RTSP service per host (in order to avoid crashing some faulty RTSP servers)<commit_after>// Copyright 2015 click2stream, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

///! Network scanner for RTSP streams.

use std::io;
use std::fmt;
use std::thread;
use std::result;

use std::fs::File;
use std::sync::Arc;
use std::error::Error;
use std::collections::HashSet;
use std::collections::HashMap;
use std::io::{BufReader, BufRead};
use std::fmt::{Display, Formatter};
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};

use net::rtsp;
use net::raw::pcap;

use net::rtsp::Client as RtspClient;
use net::raw::devices::EthernetDevice;
use net::raw::ether::MacAddr;
use net::arrow::protocol::Service;
use net::raw::arp::scanner::Ipv4ArpScanner;
use net::raw::icmp::scanner::IcmpScanner;
use net::raw::tcp::scanner::{TcpPortScanner, PortCollection};
use net::rtsp::sdp::{SessionDescription, MediaType, RTPMap, FromAttribute};

static RTSP_PATH_FILE: &'static str = "/etc/arrow/rtsp-paths";

/// Discovery error.
#[derive(Debug, Clone)]
pub struct DiscoveryError {
    msg: String,
}

impl Error for DiscoveryError {
    fn description(&self) -> &str {
        &self.msg
    }
}

impl Display for DiscoveryError {
    fn fmt(&self, f: &mut Formatter) -> result::Result<(), fmt::Error> {
        f.write_str(&self.msg)
    }
}

impl From<String> for DiscoveryError {
    fn from(msg: String) -> DiscoveryError {
        DiscoveryError { msg: msg }
    }
}

impl<'a> From<&'a str> for DiscoveryError {
    fn from(msg: &'a str) -> DiscoveryError {
        DiscoveryError { msg: msg.to_string() }
    }
}

impl From<rtsp::RtspError> for DiscoveryError {
    fn from(err: rtsp::RtspError) -> DiscoveryError {
        DiscoveryError::from(format!("RTSP client error: {}",
            err.description()))
    }
}

impl From<pcap::PcapError> for DiscoveryError {
    fn from(err: pcap::PcapError) -> DiscoveryError {
        DiscoveryError::from(format!("pcap error: {}", err.description()))
    }
}

impl From<io::Error> for DiscoveryError {
    fn from(err: io::Error) -> DiscoveryError {
        DiscoveryError::from(format!("IO error: {}", err.description()))
    }
}

/// Discovery result type alias.
pub type Result<T> = result::Result<T, DiscoveryError>;

/// Discovery host type alias.
type Host = (MacAddr, Ipv4Addr);
/// Discovery service type alias.
type Socket = (MacAddr, SocketAddrV4);

/// RTSP port candidates.
static RTSP_PORT_CANDIDATES: &'static [u16] = &[
      554,    88,    81,   555,  7447,
     8554,  7070, 10554,    80
];

/// Find all RTSP streams in all local networks.
pub fn find_rtsp_streams() -> Result<Vec<Service>> {
    let tc      = pcap::new_threading_context();
    let devices = EthernetDevice::list();

    let port_candidates = PortCollection::new()
        .add_all(RTSP_PORT_CANDIDATES);

    let mut threads = Vec::new();

    for dev in devices {
        let pc     = port_candidates.clone();
        let tc     = tc.clone();
        let handle = thread::spawn(move || {
            find_services(tc, &dev, &pc)
        });

        threads.push(handle);
    }

    let mut services = Vec::new();

    for handle in threads {
        match handle.join() {
            Err(_)  => return Err(DiscoveryError::from("port scanner thread panicked")),
            Ok(res) => services.extend(try!(res))
        }
    }

    // note: we permit only one RTSP service per host (some stupid RTSP servers
    // are accessible from more than one port and they tend to crash when they
    // are accessed from the "incorrect" one)
    let rtsp_services   = try!(find_rtsp_services(&services));
    let port_priorities = get_port_priorities(RTSP_PORT_CANDIDATES);
    let rtsp_services   = filter_duplicit_services(
        &rtsp_services,
        &port_priorities);

    let mut threads = Vec::new();
    let paths       = Arc::new(try!(load_rtsp_paths(RTSP_PATH_FILE)));

    for (mac, addr) in rtsp_services {
        let addr   = SocketAddr::V4(addr);
        let paths  = paths.clone();
        let handle = thread::spawn(move || {
            find_rtsp_paths(mac, addr, &paths)
        });
        threads.push(handle);
    }

    let mut res = Vec::new();

    for handle in threads {
        match handle.join() {
            Err(_) => return Err(DiscoveryError::from(
                "path testing thread panicked")),
            Ok(svc) => res.push(try!(svc))
        }
    }

    Ok(res)
}

/// Load all known RTSP path variants from a given file.
fn load_rtsp_paths(file: &str) -> Result<Vec<String>> {
    let file      = try!(File::open(file));
    let breader   = BufReader::new(file);
    let mut paths = Vec::new();

    for line in breader.lines() {
        let path = try!(line);
        if !path.starts_with('#') {
            paths.push(path);
        }
    }

    Ok(paths)
}

/// Check if a given service is an RTSP service.
fn is_rtsp_service(addr: SocketAddr) -> Result<bool> {
    let mut client = try!(RtspClient::new(addr));
    client.set_timeout(Some(1000));
    Ok(client.options().is_ok())
}

/// Check if a given session description contains at least one H.264 or
/// a general MPEG4 video stream.
fn is_supported_service(sdp: &[u8]) -> bool {
    if let Ok(sdp) = SessionDescription::parse(sdp) {
        let mut vcodecs   = HashSet::new();
        let video_streams = sdp.media_descriptions.into_iter()
            .filter(|md| md.media_type == MediaType::Video);

        for md in video_streams {
            for attr in md.attributes {
                if let Ok(rtpmap) = RTPMap::parse(&attr) {
                    vcodecs.insert(rtpmap.encoding.to_uppercase());
                }
            }
        }

        vcodecs.contains("H264") ||
            vcodecs.contains("H264-RCDO") ||
            vcodecs.contains("H264-SVC") ||
            vcodecs.contains("MP4V-ES") ||
            vcodecs.contains("MPEG4-GENERIC")
    } else {
        false
    }
}

/// RTSP DESCRIBE status.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum DescribeStatus {
    Ok,
    Locked,
    Unsupported,
    NotFound,
    Error
}

/// Get describe status code for a given RTSP service and path.
fn get_describe_status(addr: SocketAddr, path: &str) -> Result<DescribeStatus> {
    let mut client = try!(RtspClient::new(addr));
    client.set_timeout(Some(1000));
    if let Ok(response) = client.describe(path) {
        let header = response.header;
        let hipcam = match header.get_str("Server") {
            Some("HiIpcam/V100R003 VodServer/1.0.0") => true,
            Some("Hipcam RealServer/V1.0")           => true,
            _ => false
        };

        if hipcam && path != "/11" && path != "/12" {
            Ok(DescribeStatus::NotFound)
        } else {
            match header.code {
                404 => Ok(DescribeStatus::NotFound),
                401 => Ok(DescribeStatus::Locked),
                200 if is_supported_service(&response.body) =>
                    Ok(DescribeStatus::Ok),
                200 => Ok(DescribeStatus::Unsupported),
                _   => Ok(DescribeStatus::Error)
            }
        }
    } else {
        Ok(DescribeStatus::Error)
    }
}

/// Find open ports on all available hosts within a given network and port
/// range.
fn find_services(
    pc: pcap::ThreadingContext,
    device: &EthernetDevice,
    ports: &PortCollection) -> Result<Vec<Socket>> {
    let mut hosts  = HashSet::new();
    let arp_hosts  = try!(Ipv4ArpScanner::scan_device(pc.clone(), device));
    let icmp_hosts = try!(IcmpScanner::scan_device(pc.clone(), device));

    hosts.extend(arp_hosts);
    hosts.extend(icmp_hosts);

    let res = try!(find_open_ports(pc, device,
        hosts.into_iter(), ports));

    Ok(res)
}

/// Check if any of given TCP ports is open on on any host from a given set.
fn find_open_ports<HI: Iterator<Item=Host>>(
    pc: pcap::ThreadingContext,
    device: &EthernetDevice,
    hosts: HI,
    ports: &PortCollection) -> Result<Vec<Socket>> {
    let res = try!(TcpPortScanner::scan_ipv4_hosts(pc, device, hosts, ports))
        .into_iter()
        .map(|(mac, ip, p)| (mac, SocketAddrV4::new(ip, p)))
        .collect::<Vec<_>>();

    Ok(res)
}

/// Find all RTSP services among a given set of sockets.
fn find_rtsp_services(sockets: &[Socket]) -> Result<Vec<Socket>> {
    let mut threads = Vec::new();
    let mut res     = Vec::new();

    for &(mac, addr) in sockets {
        let handle = thread::spawn(move || {
            (mac, addr, is_rtsp_service(SocketAddr::V4(addr)))
        });
        threads.push(handle);
    }

    for handle in threads {
        match handle.join() {
            Err(_) => return Err(DiscoveryError::from("RTSP service testing thread panicked")),
            Ok((mac, addr, rtsp)) => {
                if try!(rtsp) {
                    res.push((mac, addr))
                }
            }
        }
    }

    Ok(res)
}

/// Find all available RTSP paths for a given RTSP service.
fn find_rtsp_paths(
    mac: MacAddr,
    addr: SocketAddr,
    paths: &[String]) -> Result<Service> {
    let mut service = Service::UnknownRTSP(mac, addr);

    for path in paths {
        let status = try!(get_describe_status(addr, path));
        if status == DescribeStatus::Ok {
            service = Service::RTSP(mac, addr, path.to_string());
        } else if status == DescribeStatus::Unsupported {
            service = Service::UnsupportedRTSP(mac, addr, path.to_string());
        } else if status == DescribeStatus::Locked {
            service = Service::LockedRTSP(mac, addr);
        }

        match status {
            DescribeStatus::Ok     => break,
            DescribeStatus::Locked => break,
            _ => ()
        }
    }

    Ok(service)
}

/// Assuming the given list of ports is sorted according to port priority
/// (from highest to lowest), get a map of port -> port_priority pairs.
fn get_port_priorities(ports: &[u16]) -> HashMap<u16, usize> {
    let mut res = HashMap::new();
    let len = ports.len();
    for i in 0..len {
        res.insert(ports[i], len - i);
    }

    res
}

/// Filter out duplicit services from a given list using given priorities.
fn filter_duplicit_services(
    services: &[Socket],
    port_priorities: &HashMap<u16, usize>) -> Vec<Socket> {
    let mut svc_map = HashMap::new();

    for &(ref mac, ref saddr) in services {
        let mac  = *mac;
        let ip   = *saddr.ip();
        let port = saddr.port();

        if svc_map.contains_key(&ip) {
            let old_port = svc_map.get(&ip)
                .map(|&(_, _, ref port)| *port)
                .unwrap_or(0);
            let old_priority = port_priorities.get(&old_port)
                .map(|priority| *priority)
                .unwrap_or(0);
            let new_priority = port_priorities.get(&port)
                .map(|priority| *priority)
                .unwrap_or(0);
            if new_priority > old_priority {
                svc_map.insert(ip, (mac, ip, port));
            }
        } else {
            svc_map.insert(ip, (mac, ip, port));
        }
    }

    svc_map.into_iter()
        .map(|(_, (mac, ip, port))| (mac, SocketAddrV4::new(ip, port)))
        .collect::<_>()
}

#[cfg(test)]
#[test]
/// Test the service priority filtering function.
fn test_service_filtering() {
    let ports = [554, 80];
    let mac   = MacAddr::new(0, 0, 0, 0, 0, 0);
    let ip    = Ipv4Addr::new(0, 0, 0, 0);

    let mut services = Vec::new();

    services.push((mac, SocketAddrV4::new(ip, 80)));
    services.push((mac, SocketAddrV4::new(ip, 554)));

    let port_priorities = get_port_priorities(&ports);
    let services = filter_duplicit_services(&services, &port_priorities);

    assert_eq!(services.len(), 1);
    assert_eq!(services[0].1.port(), 554);
}
<|endoftext|>


<commit_before>// Copyright 2017 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Providers
//!
//! These are the providers which coreos-metadata knows how to retrieve metadata
//! from. Internally, they handle the ins and outs of each providers metadata
//! services, and externally, they provide a function to fetch that metadata in
//! a regular format.
//!
//! To add a provider, put a `pub mod provider;` line in this file, export a
//! function to fetch the metadata, and then add a match line in the top-level
//! `fetch_metadata()` function in metadata.rs.

pub mod azure;
pub mod digitalocean;
pub mod cloudstack;
pub mod ec2;
pub mod gce;
pub mod openstack;
pub mod packet;
pub mod vagrant_virtualbox;

use std::collections::HashMap;
use std::fs::{self, File};
use std::io::prelude::*;
use std::path::Path;

use openssh_keys::PublicKey;
use update_ssh_keys::{AuthorizedKeys, AuthorizedKeyEntry};
use users;

use errors::*;
use network;

fn create_file(filename: &str) -> Result<File> {
    let file_path = Path::new(&filename);
    // create the directories if they don't exist
    let folder = file_path.parent()
        .ok_or_else(|| format!("could not get parent directory of {:?}", file_path))?;
    fs::create_dir_all(&folder)
        .chain_err(|| format!("failed to create directory {:?}", folder))?;
    // create (or truncate) the file we want to write to
    File::create(file_path)
        .chain_err(|| format!("failed to create file {:?}", file_path))
}

pub trait MetadataProvider {
    fn attributes(&self) -> Result<HashMap<String, String>>;
    fn hostname(&self) -> Result<Option<String>>;
    fn ssh_keys(&self) -> Result<Vec<PublicKey>>;
    fn networks(&self) -> Result<Vec<network::Interface>>;
    fn network_devices(&self) -> Result<Vec<network::Device>>;
    fn boot_checkin(&self) -> Result<()>;

    fn write_attributes(&self, attributes_file_path: String) -> Result<()> {
        let mut attributes_file = create_file(&attributes_file_path)?;
        for (k,v) in self.attributes()? {
            writeln!(&mut attributes_file, "COREOS_{}={}", k, v)
                .chain_err(|| format!("failed to write attributes to file {:?}", attributes_file))?;
        }
        Ok(())
    }

    fn write_ssh_keys(&self, ssh_keys_user: String) -> Result<()> {
        let ssh_keys = self.ssh_keys()?;

        if !ssh_keys.is_empty() {
            // find the ssh keys user and open their ssh authorized keys directory
            let user = users::get_user_by_name(&ssh_keys_user)
                .ok_or_else(|| format!("could not find user with username {:?}", ssh_keys_user))?;
            let mut authorized_keys_dir = AuthorizedKeys::open(user, true, None)
                .chain_err(|| format!("failed to open authorized keys directory for user '{}'", ssh_keys_user))?;

            // add the ssh keys to the directory
            let entries = ssh_keys
                .into_iter()
                .map(|key| AuthorizedKeyEntry::Valid{key})
                .collect::<Vec<_>>();
            authorized_keys_dir.add_keys("coreos-metadata", entries, true, true)?;

            // write the changes and sync the directory
            authorized_keys_dir.write()
                .chain_err(|| "failed to update authorized keys directory")?;
            authorized_keys_dir.sync()
                .chain_err(|| "failed to update authorized keys")?;
        }

        Ok(())
    }

    fn write_hostname(&self, hostname_file_path: String) -> Result<()> {
        match self.hostname()? {
            Some(ref hostname) => {
                let mut hostname_file = create_file(&hostname_file_path)?;
                writeln!(&mut hostname_file, "{}", hostname)
                    .chain_err(|| format!("failed to write hostname {:?} to file {:?}", hostname, hostname_file))
            }
            None => Ok(())
        }
    }

    fn write_network_units(&self, network_units_dir: String) -> Result<()> {
        let dir_path = Path::new(&network_units_dir);
        fs::create_dir_all(&dir_path)
            .chain_err(|| format!("failed to create directory {:?}", dir_path))?;
        for interface in &self.networks()? {
            let file_path = dir_path.join(interface.unit_name());
            let mut unit_file = File::create(&file_path)
                .chain_err(|| format!("failed to create file {:?}", file_path))?;
            write!(&mut unit_file, "{}", interface.config())
                .chain_err(|| format!("failed to write network interface unit file {:?}", unit_file))?;
        }
        for device in &self.network_devices()? {
            let file_path = dir_path.join(device.unit_name());
            let mut unit_file = File::create(&file_path)
                .chain_err(|| format!("failed to create file {:?}", file_path))?;
            write!(&mut unit_file, "{}", device.config())
                .chain_err(|| format!("failed to write network device unit file {:?}", unit_file))?;
        }
        Ok(())
    }
}
<commit_msg>providers: move update_ssh_keys wrapper to helper function<commit_after>// Copyright 2017 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Providers
//!
//! These are the providers which coreos-metadata knows how to retrieve metadata
//! from. Internally, they handle the ins and outs of each providers metadata
//! services, and externally, they provide a function to fetch that metadata in
//! a regular format.
//!
//! To add a provider, put a `pub mod provider;` line in this file, export a
//! function to fetch the metadata, and then add a match line in the top-level
//! `fetch_metadata()` function in metadata.rs.

pub mod azure;
pub mod digitalocean;
pub mod cloudstack;
pub mod ec2;
pub mod gce;
pub mod openstack;
pub mod packet;
pub mod vagrant_virtualbox;

use std::collections::HashMap;
use std::fs::{self, File};
use std::io::prelude::*;
use std::path::Path;

use openssh_keys::PublicKey;
use update_ssh_keys::{AuthorizedKeys, AuthorizedKeyEntry};
use users::{self, User};

use errors::*;
use network;

fn create_file(filename: &str) -> Result<File> {
    let file_path = Path::new(&filename);
    // create the directories if they don't exist
    let folder = file_path.parent()
        .ok_or_else(|| format!("could not get parent directory of {:?}", file_path))?;
    fs::create_dir_all(&folder)
        .chain_err(|| format!("failed to create directory {:?}", folder))?;
    // create (or truncate) the file we want to write to
    File::create(file_path)
        .chain_err(|| format!("failed to create file {:?}", file_path))
}

fn write_ssh_keys(user: User, ssh_keys: Vec<PublicKey>) -> Result<()> {
    // If we don't have any SSH keys, don't bother trying to write them as
    // update-ssh-keys will yell at us.
    if !ssh_keys.is_empty() {
        // open the user's authorized keys directory
        let user_name = user.name().to_string_lossy().into_owned();
        let mut authorized_keys_dir = AuthorizedKeys::open(user, true, None)
            .chain_err(|| format!("failed to open authorized keys directory for user '{}'", user_name))?;

        // add the ssh keys to the directory
        let entries = ssh_keys
            .into_iter()
            .map(|key| AuthorizedKeyEntry::Valid{key})
            .collect::<Vec<_>>();
        authorized_keys_dir.add_keys("coreos-metadata", entries, true, true)?;

        // write the changes and sync the directory
        authorized_keys_dir.write()
            .chain_err(|| "failed to update authorized keys directory")?;
        authorized_keys_dir.sync()
            .chain_err(|| "failed to update authorized keys")?;
    }

    Ok(())
}

pub trait MetadataProvider {
    fn attributes(&self) -> Result<HashMap<String, String>>;
    fn hostname(&self) -> Result<Option<String>>;
    fn ssh_keys(&self) -> Result<Vec<PublicKey>>;
    fn networks(&self) -> Result<Vec<network::Interface>>;
    fn network_devices(&self) -> Result<Vec<network::Device>>;
    fn boot_checkin(&self) -> Result<()>;

    fn write_attributes(&self, attributes_file_path: String) -> Result<()> {
        let mut attributes_file = create_file(&attributes_file_path)?;
        for (k,v) in self.attributes()? {
            writeln!(&mut attributes_file, "COREOS_{}={}", k, v)
                .chain_err(|| format!("failed to write attributes to file {:?}", attributes_file))?;
        }
        Ok(())
    }

    fn write_ssh_keys(&self, ssh_keys_user: String) -> Result<()> {
        let ssh_keys = self.ssh_keys()?;
        let user = users::get_user_by_name(&ssh_keys_user)
            .ok_or_else(|| format!("could not find user with username {:?}", ssh_keys_user))?;

        write_ssh_keys(user, ssh_keys)?;

        Ok(())
    }

    fn write_hostname(&self, hostname_file_path: String) -> Result<()> {
        match self.hostname()? {
            Some(ref hostname) => {
                let mut hostname_file = create_file(&hostname_file_path)?;
                writeln!(&mut hostname_file, "{}", hostname)
                    .chain_err(|| format!("failed to write hostname {:?} to file {:?}", hostname, hostname_file))
            }
            None => Ok(())
        }
    }

    fn write_network_units(&self, network_units_dir: String) -> Result<()> {
        let dir_path = Path::new(&network_units_dir);
        fs::create_dir_all(&dir_path)
            .chain_err(|| format!("failed to create directory {:?}", dir_path))?;
        for interface in &self.networks()? {
            let file_path = dir_path.join(interface.unit_name());
            let mut unit_file = File::create(&file_path)
                .chain_err(|| format!("failed to create file {:?}", file_path))?;
            write!(&mut unit_file, "{}", interface.config())
                .chain_err(|| format!("failed to write network interface unit file {:?}", unit_file))?;
        }
        for device in &self.network_devices()? {
            let file_path = dir_path.join(device.unit_name());
            let mut unit_file = File::create(&file_path)
                .chain_err(|| format!("failed to create file {:?}", file_path))?;
            write!(&mut unit_file, "{}", device.config())
                .chain_err(|| format!("failed to write network device unit file {:?}", unit_file))?;
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use backend::Backend;
use provider::error;
use provider::Provider;
use provider::Output;
use libc::uint32_t;

pub struct File<'a> {
    name: &'static str,
    backend: &'a Backend,
    provider: &'a Provider,
}

impl<'a> File<'a> {
    pub fn new(n: &'static str, b: &'a Backend, p: &'a Provider) -> File<'a> {
        File {
            name: n,
            backend: b,
            provider: p,
        }
    }

    pub fn mode(&self) -> Result<u32, error::Error> {
        self.backend
            .handle(self.provider.file.mode(self.name))
            .and_then(Output::to_u32)
    }

    pub fn size(&self) -> Result<i64, error::Error> {
        self.backend
            .handle(self.provider.file.size(self.name))
            .and_then(Output::to_i64)
    }

    pub fn is_file(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_file(self.name))
            .and_then(Output::to_bool)
    }

    pub fn exist(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.exist(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_directory(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_directory(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_block_device(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_block_device(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_character_device(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_character_device(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_pipe(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_pipe(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_socket(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_socket(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_symlink(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_symlink(self.name))
            .and_then(Output::to_bool)
    }

    pub fn contents(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.contents(self.name))
            .and_then(Output::to_string)
    }

    pub fn owner(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.owner(self.name))
            .and_then(Output::to_string)
    }

    pub fn group(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.group(self.name))
            .and_then(Output::to_string)
    }

    pub fn linked_to(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.linked_to(self.name))
            .and_then(Output::to_string)
    }

    pub fn is_readable(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_readable(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_readable_by_owner(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_readable_by_owner(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_readable_by_group(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_readable_by_group(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_readable_by_others(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_readable_by_others(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_readable_by_user(&self, user: &'static str) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_readable_by_user(self.name, user))
            .and_then(Output::to_bool)
    }

    pub fn is_writable(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_writable(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_writable_by_owner(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_writable_by_owner(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_writable_by_group(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_writable_by_group(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_writable_by_others(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_writable_by_others(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_writable_by_user(&self, user: &'static str) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_writable_by_user(self.name, user))
            .and_then(Output::to_bool)
    }

    pub fn md5sum(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.md5sum(self.name))
            .and_then(Output::to_string)
    }

    pub fn sha256sum(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.sha256sum(self.name))
            .and_then(Output::to_string)
    }
}

// Wrapper functions for FFI

#[no_mangle]
pub extern "C" fn resource_file_free(ptr: *mut File) {
    if ptr.is_null() {
        return;
    }
    unsafe {
        Box::from_raw(ptr);
    }
}

#[no_mangle]
pub extern "C" fn resource_file_mode(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    f.mode().unwrap()
}

#[no_mangle]
pub extern "C" fn resource_file_is_file(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    if f.is_file().unwrap() { 1 } else { 0 }
}

#[no_mangle]
pub extern "C" fn resource_file_exist(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    if f.exist().unwrap() { 1 } else { 0 }
}

#[no_mangle]
pub extern "C" fn resource_file_is_directory(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    if f.is_directory().unwrap() { 1 } else { 0 }
}

#[no_mangle]
pub extern "C" fn resource_file_is_block_device(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    if f.is_block_device().unwrap() { 1 } else { 0 }
}

#[no_mangle]
pub extern "C" fn resource_file_is_character_device(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    if f.is_character_device().unwrap() {
        1
    } else {
        0
    }
}
<commit_msg>Add wrapper function for file.is_pipe()<commit_after>use backend::Backend;
use provider::error;
use provider::Provider;
use provider::Output;
use libc::uint32_t;

pub struct File<'a> {
    name: &'static str,
    backend: &'a Backend,
    provider: &'a Provider,
}

impl<'a> File<'a> {
    pub fn new(n: &'static str, b: &'a Backend, p: &'a Provider) -> File<'a> {
        File {
            name: n,
            backend: b,
            provider: p,
        }
    }

    pub fn mode(&self) -> Result<u32, error::Error> {
        self.backend
            .handle(self.provider.file.mode(self.name))
            .and_then(Output::to_u32)
    }

    pub fn size(&self) -> Result<i64, error::Error> {
        self.backend
            .handle(self.provider.file.size(self.name))
            .and_then(Output::to_i64)
    }

    pub fn is_file(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_file(self.name))
            .and_then(Output::to_bool)
    }

    pub fn exist(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.exist(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_directory(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_directory(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_block_device(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_block_device(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_character_device(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_character_device(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_pipe(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_pipe(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_socket(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_socket(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_symlink(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_symlink(self.name))
            .and_then(Output::to_bool)
    }

    pub fn contents(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.contents(self.name))
            .and_then(Output::to_string)
    }

    pub fn owner(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.owner(self.name))
            .and_then(Output::to_string)
    }

    pub fn group(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.group(self.name))
            .and_then(Output::to_string)
    }

    pub fn linked_to(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.linked_to(self.name))
            .and_then(Output::to_string)
    }

    pub fn is_readable(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_readable(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_readable_by_owner(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_readable_by_owner(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_readable_by_group(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_readable_by_group(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_readable_by_others(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_readable_by_others(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_readable_by_user(&self, user: &'static str) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_readable_by_user(self.name, user))
            .and_then(Output::to_bool)
    }

    pub fn is_writable(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_writable(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_writable_by_owner(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_writable_by_owner(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_writable_by_group(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_writable_by_group(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_writable_by_others(&self) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_writable_by_others(self.name))
            .and_then(Output::to_bool)
    }

    pub fn is_writable_by_user(&self, user: &'static str) -> Result<bool, error::Error> {
        self.backend
            .handle(self.provider.file.is_writable_by_user(self.name, user))
            .and_then(Output::to_bool)
    }

    pub fn md5sum(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.md5sum(self.name))
            .and_then(Output::to_string)
    }

    pub fn sha256sum(&self) -> Result<String, error::Error> {
        self.backend
            .handle(self.provider.file.sha256sum(self.name))
            .and_then(Output::to_string)
    }
}

// Wrapper functions for FFI

#[no_mangle]
pub extern "C" fn resource_file_free(ptr: *mut File) {
    if ptr.is_null() {
        return;
    }
    unsafe {
        Box::from_raw(ptr);
    }
}

#[no_mangle]
pub extern "C" fn resource_file_mode(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    f.mode().unwrap()
}

#[no_mangle]
pub extern "C" fn resource_file_is_file(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    if f.is_file().unwrap() { 1 } else { 0 }
}

#[no_mangle]
pub extern "C" fn resource_file_exist(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    if f.exist().unwrap() { 1 } else { 0 }
}

#[no_mangle]
pub extern "C" fn resource_file_is_directory(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    if f.is_directory().unwrap() { 1 } else { 0 }
}

#[no_mangle]
pub extern "C" fn resource_file_is_block_device(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    if f.is_block_device().unwrap() { 1 } else { 0 }
}

#[no_mangle]
pub extern "C" fn resource_file_is_character_device(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    if f.is_character_device().unwrap() {
        1
    } else {
        0
    }
}

#[no_mangle]
pub extern "C" fn resource_file_is_pipe(ptr: *const File) -> uint32_t {
    let f = unsafe {
        assert!(!ptr.is_null());
        &*ptr
    };

    if f.is_pipe().unwrap() { 1 } else { 0 }
}
<|endoftext|>


<commit_before>
use std::error::Error as StdError;
use std::fmt;
use std::io::Error as IoError;

pub enum ParseError {
    Eof,
    NotFound,
    Expected(Vec<u8>),
    ExpectedType(&'static str),
    Io(IoError),
    InvalidBodyChar(u8),
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error>
    {
        match *self {
            ParseError::Expected(ref bytes) => write!(f, "{}. Expected {:?}",
                                                      self.description(), bytes),
            ParseError::ExpectedType(ref t) => write!(f, "{}. Expected {}",
                                                      self.description(), t),
            ParseError::Io(ref e) => write!(f, "{}: {}",
                                            self.description(), e),
            ParseError::InvalidBodyChar(ref c) => write!(f, "{}: {}",
                                                         self.description(), c),
            _ => write!(f, "{}", self.description()),
        }
    }
}

impl fmt::Debug for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error>
    {
        match *self {
            ParseError::Expected(ref bytes) => write!(f, "{}. Expected {:?}",
                                                      self.description(), bytes),
            ParseError::ExpectedType(ref t) => write!(f, "{}. Expected {}",
                                                      self.description(), t),
            ParseError::Io(ref e) => write!(f, "{}: {:?}",
                                             self.description(), e),
            ParseError::InvalidBodyChar(ref c) => write!(f, "{}: {}",
                                                         self.description(), c),
            _ => write!(f, "{}", self.description()),
        }
    }
}

impl StdError for ParseError {
    fn description(&self) -> &str
    {
        match *self {
            ParseError::Eof => "End of File",
            ParseError::NotFound => "Not Found",
            ParseError::Expected(_) => "Expectation Failed",
            ParseError::ExpectedType(_) => "Expectation Failed",
            ParseError::Io(_) => "I/O Error",
            ParseError::InvalidBodyChar(_) => "Invalid Body Character",
        }
    }

    fn cause(&self) -> Option<&StdError>
    {
        match *self {
            ParseError::Eof => None,
            ParseError::NotFound => None,
            ParseError::Expected(_) => None,
            ParseError::ExpectedType(_) => None,
            ParseError::Io(ref e) => Some(e),
            ParseError::InvalidBodyChar(_) => None,
        }
    }
}
<commit_msg>rfc5322: ParseError::LineTooLong<commit_after>
use std::error::Error as StdError;
use std::fmt;
use std::io::Error as IoError;

pub enum ParseError {
    Eof,
    NotFound,
    Expected(Vec<u8>),
    ExpectedType(&'static str),
    Io(IoError),
    InvalidBodyChar(u8),
    LineTooLong(usize),
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error>
    {
        match *self {
            ParseError::Expected(ref bytes) => write!(f, "{}. Expected {:?}",
                                                      self.description(), bytes),
            ParseError::ExpectedType(ref t) => write!(f, "{}. Expected {}",
                                                      self.description(), t),
            ParseError::Io(ref e) => write!(f, "{}: {}",
                                            self.description(), e),
            ParseError::InvalidBodyChar(ref c) => write!(f, "{}: {}",
                                                         self.description(), c),
            ParseError::LineTooLong(ref l) => write!(f, "Line {} is too long", l),
            _ => write!(f, "{}", self.description()),
        }
    }
}

impl fmt::Debug for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error>
    {
        match *self {
            ParseError::Expected(ref bytes) => write!(f, "{}. Expected {:?}",
                                                      self.description(), bytes),
            ParseError::ExpectedType(ref t) => write!(f, "{}. Expected {}",
                                                      self.description(), t),
            ParseError::Io(ref e) => write!(f, "{}: {:?}",
                                             self.description(), e),
            ParseError::InvalidBodyChar(ref c) => write!(f, "{}: {}",
                                                         self.description(), c),
            ParseError::LineTooLong(ref l) => write!(f, "Line {} is too long", l),
            _ => write!(f, "{}", self.description()),
        }
    }
}

impl StdError for ParseError {
    fn description(&self) -> &str
    {
        match *self {
            ParseError::Eof => "End of File",
            ParseError::NotFound => "Not Found",
            ParseError::Expected(_) => "Expectation Failed",
            ParseError::ExpectedType(_) => "Expectation Failed",
            ParseError::Io(_) => "I/O Error",
            ParseError::InvalidBodyChar(_) => "Invalid Body Character",
            ParseError::LineTooLong(_) => "Line too long",
        }
    }

    fn cause(&self) -> Option<&StdError>
    {
        match *self {
            ParseError::Eof => None,
            ParseError::NotFound => None,
            ParseError::Expected(_) => None,
            ParseError::ExpectedType(_) => None,
            ParseError::Io(ref e) => Some(e),
            ParseError::InvalidBodyChar(_) => None,
            ParseError::LineTooLong(_) => None,
        }
    }
}
<|endoftext|>


<commit_before>
use Dimension;
use {Shape, StrideShape};
use Ix;
use dimension::IntoDimension;

/// A trait for `Shape` and `D where D: Dimension` that allows
/// customizing the memory layout (strides) of an array shape.
///
/// This trait is used together with array constructor methods like
/// `Array::from_shape_vec`.
pub trait ShapeBuilder {
    type Dim: Dimension;
    type Strides;

    fn into_shape(self) -> Shape<Self::Dim>;
    fn f(self) -> Shape<Self::Dim>;
    fn set_f(self, is_f: bool) -> Shape<Self::Dim>;
    fn strides(self, strides: Self::Strides) -> StrideShape<Self::Dim>;
}

impl<T, D> From<T> for StrideShape<D>
    where D: Dimension,
          T: ShapeBuilder<Dim=D>,
{
    fn from(value: T) -> Self {
        let shape = value.into_shape();
        let d = shape.dim;
        let st = if shape.is_c { d.default_strides() } else { d.fortran_strides() };
        StrideShape {
            strides: st,
            dim: d,
            custom: false,
        }
    }
}

/*
impl<D> From<Shape<D>> for StrideShape<D>
    where D: Dimension
{
    fn from(shape: Shape<D>) -> Self {
        let d = shape.dim;
        let st = if shape.is_c { d.default_strides() } else { d.fortran_strides() };
        StrideShape {
            strides: st,
            dim: d,
            custom: false,
        }
    }
}
*/

impl<T> ShapeBuilder for T
    where T: IntoDimension
{
    type Dim = T::Dim;
    type Strides = T;
    fn into_shape(self) -> Shape<Self::Dim> {
        Shape {
            dim: self.into_dimension(),
            is_c: true,
        }
    }
    fn f(self) -> Shape<Self::Dim> { self.set_f(true) }
    fn set_f(self, is_f: bool) -> Shape<Self::Dim> {
        self.into_shape().set_f(is_f)
    }
    fn strides(self, st: T) -> StrideShape<Self::Dim> {
        self.into_shape().strides(st.into_dimension())
    }
}

impl<D> ShapeBuilder for Shape<D>
    where D: Dimension
{
    type Dim = D;
    type Strides = D;
    fn into_shape(self) -> Shape<D> { self }
    fn f(self) -> Self { self.set_f(true) }
    fn set_f(mut self, is_f: bool) -> Self {
        self.is_c = !is_f;
        self
    }
    fn strides(self, st: D) -> StrideShape<D> {
        StrideShape {
            dim: self.dim,
            strides: st,
            custom: true,
        }
    }
}

impl<D> Shape<D>
    where D: Dimension,
{
    // Return a reference to the dimension
    //pub fn dimension(&self) -> &D { &self.dim }
    /// Return the size of the shape in number of elements
    pub fn size(&self) -> usize { self.dim.size() }
}
<commit_msg>Remove unused imports in shape_builder.rs<commit_after>
use Dimension;
use {Shape, StrideShape};
use dimension::IntoDimension;

/// A trait for `Shape` and `D where D: Dimension` that allows
/// customizing the memory layout (strides) of an array shape.
///
/// This trait is used together with array constructor methods like
/// `Array::from_shape_vec`.
pub trait ShapeBuilder {
    type Dim: Dimension;
    type Strides;

    fn into_shape(self) -> Shape<Self::Dim>;
    fn f(self) -> Shape<Self::Dim>;
    fn set_f(self, is_f: bool) -> Shape<Self::Dim>;
    fn strides(self, strides: Self::Strides) -> StrideShape<Self::Dim>;
}

impl<T, D> From<T> for StrideShape<D>
    where D: Dimension,
          T: ShapeBuilder<Dim=D>,
{
    fn from(value: T) -> Self {
        let shape = value.into_shape();
        let d = shape.dim;
        let st = if shape.is_c { d.default_strides() } else { d.fortran_strides() };
        StrideShape {
            strides: st,
            dim: d,
            custom: false,
        }
    }
}

/*
impl<D> From<Shape<D>> for StrideShape<D>
    where D: Dimension
{
    fn from(shape: Shape<D>) -> Self {
        let d = shape.dim;
        let st = if shape.is_c { d.default_strides() } else { d.fortran_strides() };
        StrideShape {
            strides: st,
            dim: d,
            custom: false,
        }
    }
}
*/

impl<T> ShapeBuilder for T
    where T: IntoDimension
{
    type Dim = T::Dim;
    type Strides = T;
    fn into_shape(self) -> Shape<Self::Dim> {
        Shape {
            dim: self.into_dimension(),
            is_c: true,
        }
    }
    fn f(self) -> Shape<Self::Dim> { self.set_f(true) }
    fn set_f(self, is_f: bool) -> Shape<Self::Dim> {
        self.into_shape().set_f(is_f)
    }
    fn strides(self, st: T) -> StrideShape<Self::Dim> {
        self.into_shape().strides(st.into_dimension())
    }
}

impl<D> ShapeBuilder for Shape<D>
    where D: Dimension
{
    type Dim = D;
    type Strides = D;
    fn into_shape(self) -> Shape<D> { self }
    fn f(self) -> Self { self.set_f(true) }
    fn set_f(mut self, is_f: bool) -> Self {
        self.is_c = !is_f;
        self
    }
    fn strides(self, st: D) -> StrideShape<D> {
        StrideShape {
            dim: self.dim,
            strides: st,
            custom: true,
        }
    }
}

impl<D> Shape<D>
    where D: Dimension,
{
    // Return a reference to the dimension
    //pub fn dimension(&self) -> &D { &self.dim }
    /// Return the size of the shape in number of elements
    pub fn size(&self) -> usize { self.dim.size() }
}
<|endoftext|>


<commit_before>//! Error management
//!
//! there are two ways to handle errors in nom. The first one,
//! activated by default, uses the `nom::ErrorKind<E=u32>` enum
//! in the error branch of `IResult`. This enum can hold either
//! a parser specific error code, or a custom error type you
//! specify.
//!
//! If you need more advanced error management, you can activate
//! the "verbose-errors" compilation feature, which will give you
//! the error system available in nom 1.0. The verbose errors
//! accumulate error codes and positions as you backtrack through
//! the parser tree. From there, you can precisely identify which
//! parts of the input triggered the error case.
//!
//! Please note that the verbose error management is a bit slower
//! than the simple one.
use util::ErrorKind;
use internal::{IResult, IError};
use internal::IResult::*;

pub type Err<E=u32> = ErrorKind<E>;

impl<I,O,E> IResult<I,O,E> {
  /// Maps a `IResult<I, O, E>` to `IResult<I, O, N>` by appling a function
  /// to a contained `Error` value, leaving `Done` and `Incomplete` value
  /// untouched.
  #[inline]
  pub fn map_err<N, F>(self, f: F) -> IResult<I, O, N>
   where F: FnOnce(Err<E>) -> Err<N> {
    match self {
      Error(e)      => Error(f(e)),
      Incomplete(n) => Incomplete(n),
      Done(i, o)    => Done(i, o),
    }
  }

  /// Unwrap the contained `Error(E)` value, or panic if the `IResult` is not
  /// `Error`.
  pub fn unwrap_err(self) -> Err<E> {
    match self {
      Error(e)      => e,
      Done(_, _)    => panic!("unwrap_err() called on an IResult that is Done"),
      Incomplete(_) => panic!("unwrap_err() called on an IResult that is Incomplete"),
    }
  }

  /// Convert the IResult to a std::result::Result
  pub fn to_full_result(self) -> Result<O, IError<E>> {
    match self {
      Done(_, o)    => Ok(o),
      Incomplete(n) => Err(IError::Incomplete(n)),
      Error(e)      => Err(IError::Error(e))
    }
  }

  /// Convert the IResult to a std::result::Result
  pub fn to_result(self) -> Result<O, Err<E>> {
    match self {
      Done(_, o)    => Ok(o),
      Error(e)      => Err(e),
      Incomplete(_) => panic!("to_result() called on an IResult that is Incomplete")
    }
  }
}

#[cfg(not(feature = "core"))]
use std::any::Any;
#[cfg(not(feature = "core"))]
use std::{error,fmt};
#[cfg(not(feature = "core"))]
impl<E: fmt::Debug+Any> error::Error for Err<E> {
  fn description(&self) -> &str {
    self.description()
  }
}

#[cfg(not(feature = "core"))]
impl<E: fmt::Debug> fmt::Display for Err<E> {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "{}", self.description())
  }
}

/// translate parser result from IResult<I,O,u32> to IResult<I,O,E> with a custom type
///
/// ```
/// # #[macro_use] extern crate nom;
/// # use std::collections;
/// # use nom::IResult::Error;
/// # use nom::ErrorKind;
/// # fn main() {
///     // will add a Custom(42) error to the error chain
///     named!(err_test, add_return_error!(ErrorKind::Custom(42), tag!("abcd")));
///     // Convert to IREsult<&[u8], &[u8], &str>
///     named!(parser<&[u8], &[u8], &str>, add_return_error!(ErrorKind::Custom("custom error message"), fix_error!(&str, err_test)));
///
///     let a = &b"efghblah"[..];
///     let res_a = parser(a);
///     assert_eq!(res_a,  Error(error_node_position!( ErrorKind::Custom("custom error message"), a, Position(ErrorKind::Fix, a))));
/// # }
/// ```
#[macro_export]
macro_rules! fix_error (
  ($i:expr, $t:ty, $submac:ident!( $($args:tt)* )) => (
    {
      match $submac!($i, $($args)*) {
        $crate::IResult::Incomplete(x) => $crate::IResult::Incomplete(x),
        $crate::IResult::Done(i, o)    => $crate::IResult::Done(i, o),
        $crate::IResult::Error(_) => {
          let e: $crate::ErrorKind<$t> = $crate::ErrorKind::Fix;
          $crate::IResult::Error(e)
        }
      }
    }
  );
  ($i:expr, $t:ty, $f:expr) => (
    fix_error!($i, $t, call!($f));
  );
);

/// `flat_map!(R -> IResult<R,S>, S -> IResult<S,T>) => R -> IResult<R, T>`
///
/// combines a parser R -> IResult<R,S> and
/// a parser S -> IResult<S,T> to return another
/// parser R -> IResult<R,T>
#[macro_export]
macro_rules! flat_map(
  ($i:expr, $submac:ident!( $($args:tt)* ), $submac2:ident!( $($args2:tt)* )) => (
    {
      match $submac!($i, $($args)*) {
        $crate::IResult::Error(e)                            => $crate::IResult::Error(e),
        $crate::IResult::Incomplete($crate::Needed::Unknown) => $crate::IResult::Incomplete($crate::Needed::Unknown),
        $crate::IResult::Incomplete($crate::Needed::Size(i)) => $crate::IResult::Incomplete($crate::Needed::Size(i)),
        $crate::IResult::Done(i, o)                          => match $submac2!(o, $($args2)*) {
          $crate::IResult::Error(e)                                 => $crate::IResult::Error(e),
          $crate::IResult::Incomplete($crate::Needed::Unknown)      => $crate::IResult::Incomplete($crate::Needed::Unknown),
          $crate::IResult::Incomplete($crate::Needed::Size(ref i2)) => $crate::IResult::Incomplete($crate::Needed::Size(*i2)),
          $crate::IResult::Done(_, o2)                              => $crate::IResult::Done(i, o2)
        }
      }
    }
  );
  ($i:expr, $submac:ident!( $($args:tt)* ), $g:expr) => (
    flat_map!($i, $submac!($($args)*), call!($g));
  );
  ($i:expr, $f:expr, $g:expr) => (
    flat_map!($i, call!($f), call!($g));
  );
  ($i:expr, $f:expr, $submac:ident!( $($args:tt)* )) => (
    flat_map!($i, call!($f), $submac!($($args)*));
  );
);
<commit_msg>Clarify a potential panic<commit_after>//! Error management
//!
//! there are two ways to handle errors in nom. The first one,
//! activated by default, uses the `nom::ErrorKind<E=u32>` enum
//! in the error branch of `IResult`. This enum can hold either
//! a parser specific error code, or a custom error type you
//! specify.
//!
//! If you need more advanced error management, you can activate
//! the "verbose-errors" compilation feature, which will give you
//! the error system available in nom 1.0. The verbose errors
//! accumulate error codes and positions as you backtrack through
//! the parser tree. From there, you can precisely identify which
//! parts of the input triggered the error case.
//!
//! Please note that the verbose error management is a bit slower
//! than the simple one.
use util::ErrorKind;
use internal::{IResult, IError};
use internal::IResult::*;

pub type Err<E=u32> = ErrorKind<E>;

impl<I,O,E> IResult<I,O,E> {
  /// Maps a `IResult<I, O, E>` to `IResult<I, O, N>` by appling a function
  /// to a contained `Error` value, leaving `Done` and `Incomplete` value
  /// untouched.
  #[inline]
  pub fn map_err<N, F>(self, f: F) -> IResult<I, O, N>
   where F: FnOnce(Err<E>) -> Err<N> {
    match self {
      Error(e)      => Error(f(e)),
      Incomplete(n) => Incomplete(n),
      Done(i, o)    => Done(i, o),
    }
  }

  /// Unwrap the contained `Error(E)` value, or panic if the `IResult` is not
  /// `Error`.
  pub fn unwrap_err(self) -> Err<E> {
    match self {
      Error(e)      => e,
      Done(_, _)    => panic!("unwrap_err() called on an IResult that is Done"),
      Incomplete(_) => panic!("unwrap_err() called on an IResult that is Incomplete"),
    }
  }

  /// Convert the IResult to a std::result::Result
  pub fn to_full_result(self) -> Result<O, IError<E>> {
    match self {
      Done(_, o)    => Ok(o),
      Incomplete(n) => Err(IError::Incomplete(n)),
      Error(e)      => Err(IError::Error(e))
    }
  }

  /// Convert the IResult to a std::result::Result, or panic if the `IResult` is `Incomplete`
  pub fn to_result(self) -> Result<O, Err<E>> {
    match self {
      Done(_, o)    => Ok(o),
      Error(e)      => Err(e),
      Incomplete(_) => panic!("to_result() called on an IResult that is Incomplete")
    }
  }
}

#[cfg(not(feature = "core"))]
use std::any::Any;
#[cfg(not(feature = "core"))]
use std::{error,fmt};
#[cfg(not(feature = "core"))]
impl<E: fmt::Debug+Any> error::Error for Err<E> {
  fn description(&self) -> &str {
    self.description()
  }
}

#[cfg(not(feature = "core"))]
impl<E: fmt::Debug> fmt::Display for Err<E> {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "{}", self.description())
  }
}

/// translate parser result from IResult<I,O,u32> to IResult<I,O,E> with a custom type
///
/// ```
/// # #[macro_use] extern crate nom;
/// # use std::collections;
/// # use nom::IResult::Error;
/// # use nom::ErrorKind;
/// # fn main() {
///     // will add a Custom(42) error to the error chain
///     named!(err_test, add_return_error!(ErrorKind::Custom(42), tag!("abcd")));
///     // Convert to IREsult<&[u8], &[u8], &str>
///     named!(parser<&[u8], &[u8], &str>, add_return_error!(ErrorKind::Custom("custom error message"), fix_error!(&str, err_test)));
///
///     let a = &b"efghblah"[..];
///     let res_a = parser(a);
///     assert_eq!(res_a,  Error(error_node_position!( ErrorKind::Custom("custom error message"), a, Position(ErrorKind::Fix, a))));
/// # }
/// ```
#[macro_export]
macro_rules! fix_error (
  ($i:expr, $t:ty, $submac:ident!( $($args:tt)* )) => (
    {
      match $submac!($i, $($args)*) {
        $crate::IResult::Incomplete(x) => $crate::IResult::Incomplete(x),
        $crate::IResult::Done(i, o)    => $crate::IResult::Done(i, o),
        $crate::IResult::Error(_) => {
          let e: $crate::ErrorKind<$t> = $crate::ErrorKind::Fix;
          $crate::IResult::Error(e)
        }
      }
    }
  );
  ($i:expr, $t:ty, $f:expr) => (
    fix_error!($i, $t, call!($f));
  );
);

/// `flat_map!(R -> IResult<R,S>, S -> IResult<S,T>) => R -> IResult<R, T>`
///
/// combines a parser R -> IResult<R,S> and
/// a parser S -> IResult<S,T> to return another
/// parser R -> IResult<R,T>
#[macro_export]
macro_rules! flat_map(
  ($i:expr, $submac:ident!( $($args:tt)* ), $submac2:ident!( $($args2:tt)* )) => (
    {
      match $submac!($i, $($args)*) {
        $crate::IResult::Error(e)                            => $crate::IResult::Error(e),
        $crate::IResult::Incomplete($crate::Needed::Unknown) => $crate::IResult::Incomplete($crate::Needed::Unknown),
        $crate::IResult::Incomplete($crate::Needed::Size(i)) => $crate::IResult::Incomplete($crate::Needed::Size(i)),
        $crate::IResult::Done(i, o)                          => match $submac2!(o, $($args2)*) {
          $crate::IResult::Error(e)                                 => $crate::IResult::Error(e),
          $crate::IResult::Incomplete($crate::Needed::Unknown)      => $crate::IResult::Incomplete($crate::Needed::Unknown),
          $crate::IResult::Incomplete($crate::Needed::Size(ref i2)) => $crate::IResult::Incomplete($crate::Needed::Size(*i2)),
          $crate::IResult::Done(_, o2)                              => $crate::IResult::Done(i, o2)
        }
      }
    }
  );
  ($i:expr, $submac:ident!( $($args:tt)* ), $g:expr) => (
    flat_map!($i, $submac!($($args)*), call!($g));
  );
  ($i:expr, $f:expr, $g:expr) => (
    flat_map!($i, call!($f), call!($g));
  );
  ($i:expr, $f:expr, $submac:ident!( $($args:tt)* )) => (
    flat_map!($i, call!($f), $submac!($($args)*));
  );
);
<|endoftext|>