<commit_before>use tokenizer::new;
use tokenizer::Tokenizer;
use tokenizer::StateFunction;
use token::Token;
use token::Category;

fn initial_state(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            if lexer.token_position == lexer.token_start {
                let remaining_data = lexer.data
                    .slice_from(lexer.token_position).to_string();

                if remaining_data.starts_with("</") {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(2, Category::Text);
                    return Some(StateFunction(inside_tag))
                }
            }
            match c {
                '<' => {
                    lexer.tokenize_next(1, Category::Text);
                    return Some(StateFunction(start_of_tag));
                },
                ' ' | '\n' => {
                    lexer.tokenize(Category::Text);
                    lexer.advance();
                    lexer.states.push(StateFunction(initial_state));
                    return Some(StateFunction(whitespace));
                },
                _ => {
                    lexer.advance();
                }
            }

            Some(StateFunction(initial_state))
        }

        None => {
            lexer.tokenize(Category::Text);
            None
        }
    }
}

fn start_of_tag(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.states.push(StateFunction(inside_tag));
                    return Some(StateFunction(whitespace));
                },
                '>' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                }
                _ => {
                    lexer.advance();
                    Some(StateFunction(start_of_tag))
                }
            }
        }

        None => {
            lexer.tokenize(Category::Identifier);
            None
        }
    }
}

fn inside_tag(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                '"' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.advance();
                    Some(StateFunction(inside_string))
                },
                ' ' | '\n' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.advance();
                    lexer.states.push(StateFunction(inside_tag));
                    return Some(StateFunction(whitespace));
                },
                '=' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::AssignmentOperator);
                    Some(StateFunction(inside_tag))
                }
                '>' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                }
                _ => {
                    if lexer.token_position == lexer.token_start {
                        let remaining_data = lexer.data
                            .slice_from(lexer.token_position).to_string();

                        if remaining_data.starts_with("/>") {
                            lexer.tokenize(Category::Identifier);
                            lexer.tokenize_next(2, Category::Text);
                            return Some(StateFunction(initial_state))
                        }
                    }

                    lexer.advance();
                    Some(StateFunction(inside_tag))
                }
            }
        }

        None => {
            lexer.tokenize(Category::Identifier);
            None
        }
    }
}

fn inside_string(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                '"' => {
                    lexer.advance();
                    lexer.tokenize(Category::String);
                    Some(StateFunction(inside_tag))
                },
                '\\' => {
                    lexer.advance();
                    lexer.advance();
                    Some(StateFunction(inside_string))
                }
                _ => {
                    lexer.advance();
                    Some(StateFunction(inside_string))
                }
            }
        }

        None => {
            lexer.tokenize(Category::String);
            None
        }
    }
}

fn whitespace(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    lexer.advance();
                    Some(StateFunction(whitespace))
                },
                _ => {
                    lexer.tokenize(Category::Whitespace);
                    Some(lexer.states.pop().unwrap())
                }
            }
        }

        None => {
            lexer.tokenize(Category::Whitespace);
            None
        }
    }
}

pub fn lex(data: &str) -> Vec<Token> {
    let mut lexer = new(data);
    let mut state_function = StateFunction(initial_state);
    loop {
        let StateFunction(actual_function) = state_function;
        match actual_function(&mut lexer) {
            Some(f) => state_function = f,
            None => return lexer.tokens(),
        }
    }
}

mod tests {
    use super::lex;
    use token::Token;
    use token::Category;
    use std::old_io::{File, Open, Read};

    #[test]
    fn it_works() {
        let data = File::open_mode(&Path::new("test_data/data.xml"), Open, Read)
            .unwrap().read_to_string().unwrap();
        let tokens = lex(&data);
        let expected_tokens = vec![
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n  ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag_with_attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: "=".to_string(), category: Category::AssignmentOperator },
            Token{ lexeme: "\"value\"".to_string(), category: Category::String },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "</".to_string(), category: Category::Text },
            Token{ lexeme: "tag_with_attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n  ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "self_closing_tag".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "/>".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
            Token{ lexeme: "</".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_garbage() {
        let tokens = lex("} adwyx123&*_ ");
        let expected_tokens = vec![
            Token{ lexeme: "}".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "adwyx123&*_".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_open_strings() {
        let tokens = lex("<tag \"open!>");
        let expected_tokens = vec![
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "\"open!>".to_string(), category: Category::String },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_utf8_data() {
        let tokens = lex("différent");
        let expected_tokens = vec![
            Token{ lexeme: "différent".to_string(), category: Category::Text },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }
}
<commit_msg>Update XML lexer to use include_str! to load test data.<commit_after>use tokenizer::new;
use tokenizer::Tokenizer;
use tokenizer::StateFunction;
use token::Token;
use token::Category;

fn initial_state(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            if lexer.token_position == lexer.token_start {
                let remaining_data = lexer.data
                    .slice_from(lexer.token_position).to_string();

                if remaining_data.starts_with("</") {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(2, Category::Text);
                    return Some(StateFunction(inside_tag))
                }
            }
            match c {
                '<' => {
                    lexer.tokenize_next(1, Category::Text);
                    return Some(StateFunction(start_of_tag));
                },
                ' ' | '\n' => {
                    lexer.tokenize(Category::Text);
                    lexer.advance();
                    lexer.states.push(StateFunction(initial_state));
                    return Some(StateFunction(whitespace));
                },
                _ => {
                    lexer.advance();
                }
            }

            Some(StateFunction(initial_state))
        }

        None => {
            lexer.tokenize(Category::Text);
            None
        }
    }
}

fn start_of_tag(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.states.push(StateFunction(inside_tag));
                    return Some(StateFunction(whitespace));
                },
                '>' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                }
                _ => {
                    lexer.advance();
                    Some(StateFunction(start_of_tag))
                }
            }
        }

        None => {
            lexer.tokenize(Category::Identifier);
            None
        }
    }
}

fn inside_tag(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                '"' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.advance();
                    Some(StateFunction(inside_string))
                },
                ' ' | '\n' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.advance();
                    lexer.states.push(StateFunction(inside_tag));
                    return Some(StateFunction(whitespace));
                },
                '=' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::AssignmentOperator);
                    Some(StateFunction(inside_tag))
                }
                '>' => {
                    lexer.tokenize(Category::Identifier);
                    lexer.tokenize_next(1, Category::Text);
                    Some(StateFunction(initial_state))
                }
                _ => {
                    if lexer.token_position == lexer.token_start {
                        let remaining_data = lexer.data
                            .slice_from(lexer.token_position).to_string();

                        if remaining_data.starts_with("/>") {
                            lexer.tokenize(Category::Identifier);
                            lexer.tokenize_next(2, Category::Text);
                            return Some(StateFunction(initial_state))
                        }
                    }

                    lexer.advance();
                    Some(StateFunction(inside_tag))
                }
            }
        }

        None => {
            lexer.tokenize(Category::Identifier);
            None
        }
    }
}

fn inside_string(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                '"' => {
                    lexer.advance();
                    lexer.tokenize(Category::String);
                    Some(StateFunction(inside_tag))
                },
                '\\' => {
                    lexer.advance();
                    lexer.advance();
                    Some(StateFunction(inside_string))
                }
                _ => {
                    lexer.advance();
                    Some(StateFunction(inside_string))
                }
            }
        }

        None => {
            lexer.tokenize(Category::String);
            None
        }
    }
}

fn whitespace(lexer: &mut Tokenizer) -> Option<StateFunction> {
    match lexer.current_char() {
        Some(c) => {
            match c {
                ' ' | '\n' => {
                    lexer.advance();
                    Some(StateFunction(whitespace))
                },
                _ => {
                    lexer.tokenize(Category::Whitespace);
                    Some(lexer.states.pop().unwrap())
                }
            }
        }

        None => {
            lexer.tokenize(Category::Whitespace);
            None
        }
    }
}

pub fn lex(data: &str) -> Vec<Token> {
    let mut lexer = new(data);
    let mut state_function = StateFunction(initial_state);
    loop {
        let StateFunction(actual_function) = state_function;
        match actual_function(&mut lexer) {
            Some(f) => state_function = f,
            None => return lexer.tokens(),
        }
    }
}

mod tests {
    use super::lex;
    use token::Token;
    use token::Category;
    use std::old_io::{File, Open, Read};

    #[test]
    fn it_works() {
        let data = include_str!("../../test_data/data.xml");
        let tokens = lex(data);
        let expected_tokens = vec![
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n  ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag_with_attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: "=".to_string(), category: Category::AssignmentOperator },
            Token{ lexeme: "\"value\"".to_string(), category: Category::String },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "</".to_string(), category: Category::Text },
            Token{ lexeme: "tag_with_attribute".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n  ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "self_closing_tag".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "/>".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
            Token{ lexeme: "</".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: ">".to_string(), category: Category::Text },
            Token{ lexeme: "\n".to_string(), category: Category::Whitespace },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_garbage() {
        let tokens = lex("} adwyx123&*_ ");
        let expected_tokens = vec![
            Token{ lexeme: "}".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "adwyx123&*_".to_string(), category: Category::Text },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_open_strings() {
        let tokens = lex("<tag \"open!>");
        let expected_tokens = vec![
            Token{ lexeme: "<".to_string(), category: Category::Text },
            Token{ lexeme: "tag".to_string(), category: Category::Identifier },
            Token{ lexeme: " ".to_string(), category: Category::Whitespace },
            Token{ lexeme: "\"open!>".to_string(), category: Category::String },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }

    #[test]
    fn it_can_handle_utf8_data() {
        let tokens = lex("différent");
        let expected_tokens = vec![
            Token{ lexeme: "différent".to_string(), category: Category::Text },
        ];

        for (index, token) in tokens.iter().enumerate() {
            assert_eq!(*token, expected_tokens[index]);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OSRng;

#[cfg(unix)]
mod imp {
    use Rng;
    use reader::ReaderRng;
    use std::io::{IoResult, File};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    ///
    /// This does not block.
    #[cfg(unix)]
    pub struct OSRng {
        inner: ReaderRng<File>
    }

    impl OSRng {
        /// Create a new `OSRng`.
        pub fn new() -> IoResult<OSRng> {
            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OSRng { inner: reader_rng })
        }
    }

    impl Rng for OSRng {
        fn next_u32(&mut self) -> u32 {
            self.inner.next_u32()
        }
        fn next_u64(&mut self) -> u64 {
            self.inner.next_u64()
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            self.inner.fill_bytes(v)
        }
    }
}

#[cfg(windows)]
mod imp {
    use Rng;
    use std::cast;
    use std::io::{IoResult, IoError};
    use std::libc::{c_ulong, DWORD, BYTE, LPCSTR, BOOL};
    use std::os;
    use std::rt::stack;

    type HCRYPTPROV = c_ulong;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    ///
    /// This does not block.
    pub struct OSRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;
    static NTE_BAD_SIGNATURE: DWORD = 0x80090006;

    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OSRng {
        /// Create a new `OSRng`.
        pub fn new() -> IoResult<OSRng> {
            let mut hcp = 0;
            let mut ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            // It turns out that if we can't acquire a context with the
            // NTE_BAD_SIGNATURE error code, the documentation states:
            //
            //     The provider DLL signature could not be verified. Either the
            //     DLL or the digital signature has been tampered with.
            //
            // Sounds fishy, no? As it turns out, our signature can be bad
            // because our Thread Information Block (TIB) isn't exactly what it
            // expects. As to why, I have no idea. The only data we store in the
            // TIB is the stack limit for each thread, but apparently that's
            // enough to make the signature valid.
            //
            // Furthermore, this error only happens the *first* time we call
            // CryptAcquireContext, so we don't have to worry about future
            // calls.
            //
            // Anyway, the fix employed here is that if we see this error, we
            // pray that we're not close to the end of the stack, temporarily
            // set the stack limit to 0 (what the TIB originally was), acquire a
            // context, and then reset the stack limit.
            //
            // Again, I'm not sure why this is the fix, nor why we're getting
            // this error. All I can say is that this seems to allow libnative
            // to progress where it otherwise would be hindered. Who knew?
            if ret == 0 && os::errno() as DWORD == NTE_BAD_SIGNATURE {
                unsafe {
                    let limit = stack::get_sp_limit();
                    stack::record_sp_limit(0);
                    ret = CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                               PROV_RSA_FULL,
                                               CRYPT_VERIFYCONTEXT | CRYPT_SILENT);
                    stack::record_sp_limit(limit);
                }
            }

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OSRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OSRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(v);
            unsafe { cast::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(v);
            unsafe { cast::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                fail!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OSRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                fail!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::OSRng;
    use Rng;
    use std::task;

    #[test]
    fn test_os_rng() {
        let mut r = OSRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(proc() {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OSRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<commit_msg>Fix fallout from std::libc separation<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OSRng;

#[cfg(unix)]
mod imp {
    use Rng;
    use reader::ReaderRng;
    use std::io::{IoResult, File};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    ///
    /// This does not block.
    #[cfg(unix)]
    pub struct OSRng {
        inner: ReaderRng<File>
    }

    impl OSRng {
        /// Create a new `OSRng`.
        pub fn new() -> IoResult<OSRng> {
            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OSRng { inner: reader_rng })
        }
    }

    impl Rng for OSRng {
        fn next_u32(&mut self) -> u32 {
            self.inner.next_u32()
        }
        fn next_u64(&mut self) -> u64 {
            self.inner.next_u64()
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            self.inner.fill_bytes(v)
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use Rng;
    use std::cast;
    use std::io::{IoResult, IoError};
    use std::os;
    use std::rt::stack;
    use self::libc::{c_ulong, DWORD, BYTE, LPCSTR, BOOL};

    type HCRYPTPROV = c_ulong;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    ///
    /// This does not block.
    pub struct OSRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;
    static NTE_BAD_SIGNATURE: DWORD = 0x80090006;

    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OSRng {
        /// Create a new `OSRng`.
        pub fn new() -> IoResult<OSRng> {
            let mut hcp = 0;
            let mut ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            // It turns out that if we can't acquire a context with the
            // NTE_BAD_SIGNATURE error code, the documentation states:
            //
            //     The provider DLL signature could not be verified. Either the
            //     DLL or the digital signature has been tampered with.
            //
            // Sounds fishy, no? As it turns out, our signature can be bad
            // because our Thread Information Block (TIB) isn't exactly what it
            // expects. As to why, I have no idea. The only data we store in the
            // TIB is the stack limit for each thread, but apparently that's
            // enough to make the signature valid.
            //
            // Furthermore, this error only happens the *first* time we call
            // CryptAcquireContext, so we don't have to worry about future
            // calls.
            //
            // Anyway, the fix employed here is that if we see this error, we
            // pray that we're not close to the end of the stack, temporarily
            // set the stack limit to 0 (what the TIB originally was), acquire a
            // context, and then reset the stack limit.
            //
            // Again, I'm not sure why this is the fix, nor why we're getting
            // this error. All I can say is that this seems to allow libnative
            // to progress where it otherwise would be hindered. Who knew?
            if ret == 0 && os::errno() as DWORD == NTE_BAD_SIGNATURE {
                unsafe {
                    let limit = stack::get_sp_limit();
                    stack::record_sp_limit(0);
                    ret = CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                               PROV_RSA_FULL,
                                               CRYPT_VERIFYCONTEXT | CRYPT_SILENT);
                    stack::record_sp_limit(limit);
                }
            }

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OSRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OSRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(v);
            unsafe { cast::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(v);
            unsafe { cast::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                fail!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OSRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                fail!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::OSRng;
    use Rng;
    use std::task;

    #[test]
    fn test_os_rng() {
        let mut r = OSRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(proc() {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OSRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<|endoftext|>


<commit_before>use self::remote_info::RemoteInfo;
use super::shared::{helpers, executor};
use slog::Logger;
use std::{time, thread};
use std::process::{ChildStdout, ChildStdin};
use std::process::Stdio;

pub mod remote_info;
mod rsync;

pub fn run(log: &Logger,
           base_dir: &str,
           remote_dir: &str,
           port: Option<&str>,
           ignore_strings: &mut Vec<String>,
           verbose_mode: bool) {
    let ignores = helpers::process_ignores(ignore_strings);
    let remote_info = RemoteInfo::build(remote_dir, port);

    validate_remote_info(&log, &remote_info);
    rsync::run(&log, &base_dir, &remote_info, &ignores);
    let (remote_reader, remote_writer) =
        start_remote_slave(&log, &remote_info, &ignore_strings, verbose_mode);
    executor::start(&log,
                    base_dir.to_owned(),
                    ignores,
                    remote_reader,
                    remote_writer);
}

fn start_remote_slave(log: &Logger,
                      remote_info: &RemoteInfo,
                      ignores: &Vec<String>,
                      verbose_mode: bool)
                      -> (ChildStdout, ChildStdin) {
    info!(log, "Starting remote slave");
    let ignore_vec: Vec<String> = ignores.iter().map(|i| format!("--ignore \"{}\"", i)).collect();
    let mut cmd = format!("bindrs slave {} {}", remote_info.path, ignore_vec.join(" "));

    if verbose_mode {
        cmd += " -v"
    }

    let mut child = remote_info.generate_command(&mut remote_info.base_command(&cmd), &cmd)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()
        .unwrap();

    thread::sleep(time::Duration::new(1, 0));
    let c_stdout = child.stdout.take().unwrap();
    let c_stdin = child.stdin.take().unwrap();

    (c_stdout, c_stdin)
}

fn validate_remote_info(log: &Logger, remote_info: &RemoteInfo) {
    let cmd = "which bindrs 2>/dev/null";
    if get_cmd_output(remote_info, &cmd) == "" {
        helpers::log_error_and_exit(&log,
                                    "Please install BindRS on the remote machine and add it to \
                                     the path");
    }

    let cmd = "bindrs -V";
    if get_cmd_output(remote_info, &cmd) != format!("BindRS {}", ::VERSION) {
        helpers::log_error_and_exit(&log,
                                    "Please make sure your local and remote versions of BindRS \
                                     match");
    }

    let cmd = format!("test -d {} || echo 'bad'", remote_info.path);
    if get_cmd_output(remote_info, &cmd) == "bad" {
        helpers::log_error_and_exit(&log, "Remote directory does not exist, please create it");
    }
}

fn get_cmd_output(remote_info: &RemoteInfo, cmd: &str) -> String {
    let output =
        remote_info.generate_command(&mut remote_info.base_command(&cmd), &cmd).output().unwrap();
    String::from_utf8_lossy(&output.stdout).trim().to_owned()
}
<commit_msg>Fix which bindrs check<commit_after>use self::remote_info::RemoteInfo;
use super::shared::{helpers, executor};
use slog::Logger;
use std::{time, thread};
use std::process::{ChildStdout, ChildStdin};
use std::process::Stdio;

pub mod remote_info;
mod rsync;

pub fn run(log: &Logger,
           base_dir: &str,
           remote_dir: &str,
           port: Option<&str>,
           ignore_strings: &mut Vec<String>,
           verbose_mode: bool) {
    let ignores = helpers::process_ignores(ignore_strings);
    let remote_info = RemoteInfo::build(remote_dir, port);

    validate_remote_info(&log, &remote_info);
    rsync::run(&log, &base_dir, &remote_info, &ignores);
    let (remote_reader, remote_writer) =
        start_remote_slave(&log, &remote_info, &ignore_strings, verbose_mode);
    executor::start(&log,
                    base_dir.to_owned(),
                    ignores,
                    remote_reader,
                    remote_writer);
}

fn start_remote_slave(log: &Logger,
                      remote_info: &RemoteInfo,
                      ignores: &Vec<String>,
                      verbose_mode: bool)
                      -> (ChildStdout, ChildStdin) {
    info!(log, "Starting remote slave");
    let ignore_vec: Vec<String> = ignores.iter().map(|i| format!("--ignore \"{}\"", i)).collect();
    let mut cmd = format!("bindrs slave {} {}", remote_info.path, ignore_vec.join(" "));

    if verbose_mode {
        cmd += " -v"
    }

    let mut child = remote_info.generate_command(&mut remote_info.base_command(&cmd), &cmd)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()
        .unwrap();

    thread::sleep(time::Duration::new(1, 0));
    let c_stdout = child.stdout.take().unwrap();
    let c_stdin = child.stdin.take().unwrap();

    (c_stdout, c_stdin)
}

fn validate_remote_info(log: &Logger, remote_info: &RemoteInfo) {
    let cmd = "which bindrs";
    if get_cmd_output(remote_info, &cmd) == "bindrs not found" {
        helpers::log_error_and_exit(&log,
                                    "Please install BindRS on the remote machine and add it to \
                                     the path");
    }

    let cmd = "bindrs -V";
    if get_cmd_output(remote_info, &cmd) != format!("BindRS {}", ::VERSION) {
        helpers::log_error_and_exit(&log,
                                    "Please make sure your local and remote versions of BindRS \
                                     match");
    }

    let cmd = format!("test -d {} || echo 'bad'", remote_info.path);
    if get_cmd_output(remote_info, &cmd) == "bad" {
        helpers::log_error_and_exit(&log, "Remote directory does not exist, please create it");
    }
}

fn get_cmd_output(remote_info: &RemoteInfo, cmd: &str) -> String {
    let output =
        remote_info.generate_command(&mut remote_info.base_command(&cmd), &cmd).output().unwrap();
    String::from_utf8_lossy(&output.stdout).trim().to_owned()
}
<|endoftext|>


<commit_before>use std::str::FromStr;
use std::fs::File;
use std::path::Path;
use std::io::prelude::*;
use std::result::Result;
use std::sync::RwLock;

use iron::prelude::*;
use iron::{AfterMiddleware, typemap};
use iron::modifier::Modifier;
use plugin::Plugin as PluginFor;
use iron::headers::ContentType;
use walker::Walker;

use handlebars::Handlebars;
use serialize::json::{ToJson, Json};

#[derive(Clone)]
pub struct Template {
    name: String,
    value: Json
}

impl Template {
    pub fn new<T: ToJson>(name: &str, value: T) -> Template {
        Template {
            name: name.to_string(),
            value: value.to_json()
        }
    }
}

pub struct HandlebarsEngine {
    pub prefix: String,
    pub suffix: String,
    pub registry: RwLock<Box<Handlebars>>
}

impl typemap::Key for HandlebarsEngine {
    type Value = Template;
}

impl Modifier<Response> for Template {
    fn modify(self, resp: &mut Response) {
        resp.extensions.insert::<HandlebarsEngine>(self);
    }
}

impl PluginFor<Response> for HandlebarsEngine {
    type Error = ();

    fn eval(resp: &mut Response) -> Result<Template, ()> {
        match resp.extensions.get::<HandlebarsEngine>(){
            Some(t) => Ok(t.clone()),
            None => Err(())
        }
    }
}

impl HandlebarsEngine {
    pub fn reload(&self) {
        let mut prefix_slash = self.prefix.clone();
        let suffix: &str = self.suffix.as_ref();
        let mut hbs = self.registry.write().unwrap();

        let normalized_prefix = if self.prefix.ends_with("/") {
            prefix_slash
        } else {
            prefix_slash.push('/');
            prefix_slash
        };
        let prefix_path = Path::new(&normalized_prefix);
        let walker = Walker::new(prefix_path);
        if !walker.is_ok() {
            panic!(format!("Failed to list directory: {}", normalized_prefix));
        }

        hbs.clear_templates();
        for p in walker.ok().unwrap().filter_map(Result::ok) {
            let path = p.path();
            let disp = path.to_str().unwrap();
            if disp.ends_with(suffix) {
                if let Ok(mut file) = File::open(&path) {
                    let mut buf = String::new();
                    if let Ok(_) = file.read_to_string(&mut buf) {
                        if let Err(e) = hbs.register_template_string(
                            &disp[normalized_prefix.len() .. disp.len()-suffix.len()], buf){
                            println!("Failed to parse template {}", e);
                        }
                    } else {
                        println!("Failed to read file {}, skipped", disp);
                    }
                } else {
                    println!("Failed to open file {}, skipped.", disp);
                }
            }
        }
    }

    pub fn from(prefix: &str, suffix: &str, custom: Handlebars) -> HandlebarsEngine {
        let eng = HandlebarsEngine {
            prefix: prefix.to_string(),
            suffix: suffix.to_string(),
            registry: RwLock::new(Box::new(custom))
        };
        eng.reload();
        eng
    }

    pub fn new(prefix: &str, suffix: &str) -> HandlebarsEngine {
        HandlebarsEngine::from(prefix,suffix, Handlebars::new())
    }
}

impl AfterMiddleware for HandlebarsEngine {
    fn after(&self, _: &mut Request, r: Response) -> IronResult<Response> {
        let mut resp = r;
        // internally we still extensions.get to avoid clone
        let page = match resp.extensions.get::<HandlebarsEngine>() {
            Some(ref h) => {
                let name = &h.name;
                let value = &h.value;
                let hbs = self.registry.read().unwrap();
                let rendered = hbs.render(name.as_ref(), value);
                match rendered {
                    Ok(r) => Some(r),
                    Err(_) => None
                }
            },
            None => {
                None
            }
        };

        if page.is_some() {
            if !resp.headers.has::<ContentType>() {
                resp.headers.set(ContentType(FromStr::from_str("text/html;charset=utf-8").unwrap()));
            }
            resp.set_mut(page.unwrap());
        }

        Ok(resp)
    }
}

#[cfg(test)]
mod test {
    use serialize::json::ToJson;
    use std::collections::BTreeMap;
    use iron::prelude::*;
    use middleware::*;
    use handlebars::{Handlebars, RenderError, RenderContext, Helper, Context};

    fn hello_world() -> IronResult<Response> {
        let resp = Response::new();

        let mut data = BTreeMap::new();
        data.insert("title".to_string(), "Handlebars on Iron".to_json());

        Ok(resp.set(Template::new("index", data)))
    }

    #[test]
    fn test_resp_set() {
        let mut resp = hello_world().ok().expect("response expected");

        // use response plugin to retrieve a cloned template for testing
        match resp.get::<HandlebarsEngine>() {
            Ok(h) => {
                assert_eq!(h.name, "index".to_string());
                assert_eq!(h.value.as_object().unwrap()
                           .get(&"title".to_string()).unwrap()
                           .as_string().unwrap(),
                           "Handlebars on Iron");
            },
            _ => panic!("template expected")
        }
    }

    #[test]
    fn test_register_helper() {
        let hbs = HandlebarsEngine::new("./examples/templates", ".hbs");
        let mut reg = hbs.registry.write().unwrap();
        reg.register_helper("ignore", Box::new(|_: &Context, _: &Helper, _: &Handlebars, _: &mut RenderContext| -> Result<String, RenderError> {
            Ok("".to_string())
        }));
    }
}
<commit_msg>(fix) fixed test helper as handlebars-rust updated<commit_after>use std::str::FromStr;
use std::fs::File;
use std::path::Path;
use std::io::prelude::*;
use std::result::Result;
use std::sync::RwLock;

use iron::prelude::*;
use iron::{AfterMiddleware, typemap};
use iron::modifier::Modifier;
use plugin::Plugin as PluginFor;
use iron::headers::ContentType;
use walker::Walker;

use handlebars::Handlebars;
use serialize::json::{ToJson, Json};

#[derive(Clone)]
pub struct Template {
    name: String,
    value: Json
}

impl Template {
    pub fn new<T: ToJson>(name: &str, value: T) -> Template {
        Template {
            name: name.to_string(),
            value: value.to_json()
        }
    }
}

pub struct HandlebarsEngine {
    pub prefix: String,
    pub suffix: String,
    pub registry: RwLock<Box<Handlebars>>
}

impl typemap::Key for HandlebarsEngine {
    type Value = Template;
}

impl Modifier<Response> for Template {
    fn modify(self, resp: &mut Response) {
        resp.extensions.insert::<HandlebarsEngine>(self);
    }
}

impl PluginFor<Response> for HandlebarsEngine {
    type Error = ();

    fn eval(resp: &mut Response) -> Result<Template, ()> {
        match resp.extensions.get::<HandlebarsEngine>(){
            Some(t) => Ok(t.clone()),
            None => Err(())
        }
    }
}

impl HandlebarsEngine {
    pub fn reload(&self) {
        let mut prefix_slash = self.prefix.clone();
        let suffix: &str = self.suffix.as_ref();
        let mut hbs = self.registry.write().unwrap();

        let normalized_prefix = if self.prefix.ends_with("/") {
            prefix_slash
        } else {
            prefix_slash.push('/');
            prefix_slash
        };
        let prefix_path = Path::new(&normalized_prefix);
        let walker = Walker::new(prefix_path);
        if !walker.is_ok() {
            panic!(format!("Failed to list directory: {}", normalized_prefix));
        }

        hbs.clear_templates();
        for p in walker.ok().unwrap().filter_map(Result::ok) {
            let path = p.path();
            let disp = path.to_str().unwrap();
            if disp.ends_with(suffix) {
                if let Ok(mut file) = File::open(&path) {
                    let mut buf = String::new();
                    if let Ok(_) = file.read_to_string(&mut buf) {
                        if let Err(e) = hbs.register_template_string(
                            &disp[normalized_prefix.len() .. disp.len()-suffix.len()], buf){
                            println!("Failed to parse template {}", e);
                        }
                    } else {
                        println!("Failed to read file {}, skipped", disp);
                    }
                } else {
                    println!("Failed to open file {}, skipped.", disp);
                }
            }
        }
    }

    pub fn from(prefix: &str, suffix: &str, custom: Handlebars) -> HandlebarsEngine {
        let eng = HandlebarsEngine {
            prefix: prefix.to_string(),
            suffix: suffix.to_string(),
            registry: RwLock::new(Box::new(custom))
        };
        eng.reload();
        eng
    }

    pub fn new(prefix: &str, suffix: &str) -> HandlebarsEngine {
        HandlebarsEngine::from(prefix,suffix, Handlebars::new())
    }
}

impl AfterMiddleware for HandlebarsEngine {
    fn after(&self, _: &mut Request, r: Response) -> IronResult<Response> {
        let mut resp = r;
        // internally we still extensions.get to avoid clone
        let page = match resp.extensions.get::<HandlebarsEngine>() {
            Some(ref h) => {
                let name = &h.name;
                let value = &h.value;
                let hbs = self.registry.read().unwrap();
                let rendered = hbs.render(name.as_ref(), value);
                match rendered {
                    Ok(r) => Some(r),
                    Err(_) => None
                }
            },
            None => {
                None
            }
        };

        if page.is_some() {
            if !resp.headers.has::<ContentType>() {
                resp.headers.set(ContentType(FromStr::from_str("text/html;charset=utf-8").unwrap()));
            }
            resp.set_mut(page.unwrap());
        }

        Ok(resp)
    }
}

#[cfg(test)]
mod test {
    use serialize::json::ToJson;
    use std::collections::BTreeMap;
    use iron::prelude::*;
    use middleware::*;
    use handlebars::{Handlebars, RenderError, RenderContext, Helper, Context};

    fn hello_world() -> IronResult<Response> {
        let resp = Response::new();

        let mut data = BTreeMap::new();
        data.insert("title".to_string(), "Handlebars on Iron".to_json());

        Ok(resp.set(Template::new("index", data)))
    }

    #[test]
    fn test_resp_set() {
        let mut resp = hello_world().ok().expect("response expected");

        // use response plugin to retrieve a cloned template for testing
        match resp.get::<HandlebarsEngine>() {
            Ok(h) => {
                assert_eq!(h.name, "index".to_string());
                assert_eq!(h.value.as_object().unwrap()
                           .get(&"title".to_string()).unwrap()
                           .as_string().unwrap(),
                           "Handlebars on Iron");
            },
            _ => panic!("template expected")
        }
    }

    #[test]
    fn test_register_helper() {
        let hbs = HandlebarsEngine::new("./examples/templates", ".hbs");
        let mut reg = hbs.registry.write().unwrap();
        reg.register_helper("ignore", Box::new(|_: &Context, _: &Helper, _: &Handlebars, _: &mut RenderContext| -> Result<(), RenderError> {
            Ok(())
        }));
    }
}
<|endoftext|>


<commit_before>use std::io::{self, Error, ErrorKind};
use std::net::{UdpSocket, SocketAddr};
use std::fmt;

/// Maximum length for packets received on a `PacketReceiver`.
pub const MAX_PCKT_LEN: usize = 600;

/// A `PacketReceiver` that abstracts over a network socket and reads full packets
/// from the connection. Packets received from this connection are assumed to
/// be no larger than what the typical MTU would be on a standard router.
///
/// See `net::packet::MAX_PCKT_LEN`.
pub struct PacketReceiver(UdpSocket);

impl PacketReceiver {
    /// Create a new PacketReceiver from the given UdpSocket.
    pub fn new(udp: UdpSocket) -> PacketReceiver {
        PacketReceiver(udp)
    }

    /// Receive a packet from the underlying connection.
    pub fn recv_pckt(&self) -> io::Result<(Vec<u8>, SocketAddr)> {
        let mut pckt_buf = vec![0u8; MAX_PCKT_LEN];

        let (size, addr) = try!(self.0.recv_from(&mut pckt_buf));

        // Check For Something That SHOULD NEVER Occur.
        if size > pckt_buf.len() {
            Err(Error::new(ErrorKind::Other, "UdpSocket Reported Receive Length Greater Than Buffer"))
        } else {
            unsafe { pckt_buf.set_len(size) };

            Ok((pckt_buf, addr))
        }
    }
}

impl fmt::Display for PacketReceiver {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.0.local_addr() {
            Ok(addr) => write!(f, "{}", addr),
            Err(err) => write!(f, "{}", err),
        }
    }
}
<commit_msg>Remove an unneeded use of unsafe<commit_after>use std::io::{self, Error, ErrorKind};
use std::net::{UdpSocket, SocketAddr};
use std::fmt;

/// Maximum length for packets received on a `PacketReceiver`.
pub const MAX_PCKT_LEN: usize = 600;

/// A `PacketReceiver` that abstracts over a network socket and reads full packets
/// from the connection. Packets received from this connection are assumed to
/// be no larger than what the typical MTU would be on a standard router.
///
/// See `net::packet::MAX_PCKT_LEN`.
pub struct PacketReceiver(UdpSocket);

impl PacketReceiver {
    /// Create a new PacketReceiver from the given UdpSocket.
    pub fn new(udp: UdpSocket) -> PacketReceiver {
        PacketReceiver(udp)
    }

    /// Receive a packet from the underlying connection.
    pub fn recv_pckt(&self) -> io::Result<(Vec<u8>, SocketAddr)> {
        let mut pckt_buf = vec![0u8; MAX_PCKT_LEN];

        let (size, addr) = try!(self.0.recv_from(&mut pckt_buf));

        // Check For Something That SHOULD NEVER Occur.
        if size > pckt_buf.len() {
            Err(Error::new(ErrorKind::Other, "UdpSocket Reported Receive Length Greater Than Buffer"))
        } else {
            // `truncate` does not reallocate the vec's backing storage
            pckt_buf.truncate(size);

            Ok((pckt_buf, addr))
        }
    }
}

impl fmt::Display for PacketReceiver {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.0.local_addr() {
            Ok(addr) => write!(f, "{}", addr),
            Err(err) => write!(f, "{}", err),
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 The predicates-rs Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/license/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Render `Case` as a tree.

use std::fmt;

use predicates_core::reflection;

/// Render `Self` as a displayable tree.
pub trait CaseTreeExt {
    /// Render `Self` as a displayable tree.
    fn tree(&self) -> CaseTree;
}

impl<'a> CaseTreeExt for reflection::Case<'a> {
    fn tree(&self) -> CaseTree {
        CaseTree(convert(self))
    }
}

type CaseTreeInner = termtree::Tree<Box<dyn fmt::Display + Send + Sync>>;

fn convert(case: &reflection::Case<'_>) -> CaseTreeInner {
    let mut leaves: Vec<CaseTreeInner> = vec![];

    leaves.extend(case.predicate().iter().flat_map(|pred| {
        pred.parameters().map(|item| {
            let root: Box<dyn fmt::Display + Send + Sync> = Box::new(item.to_string());
            termtree::Tree::new(root, vec![])
        })
    }));

    leaves.extend(case.products().map(|item| {
        let root: Box<dyn fmt::Display + Send + Sync> = Box::new(item.to_string());
        termtree::Tree::new(root, vec![])
    }));

    leaves.extend(case.children().map(|item| convert(item)));

    let root = Box::new(case.predicate().map(|p| p.to_string()).unwrap_or_default());
    CaseTreeInner::new(root, leaves)
}

/// A `Case` rendered as a tree for display.
#[allow(missing_debug_implementations)]
pub struct CaseTree(CaseTreeInner);

impl fmt::Display for CaseTree {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}
<commit_msg>fix(tree): Indent items to clarify<commit_after>// Copyright (c) 2018 The predicates-rs Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/license/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Render `Case` as a tree.

use std::fmt;

use predicates_core::reflection;

/// Render `Self` as a displayable tree.
pub trait CaseTreeExt {
    /// Render `Self` as a displayable tree.
    fn tree(&self) -> CaseTree;
}

impl<'a> CaseTreeExt for reflection::Case<'a> {
    fn tree(&self) -> CaseTree {
        CaseTree(convert(self))
    }
}

type CaseTreeInner = termtree::Tree<Box<dyn fmt::Display + Send + Sync>>;

fn convert(case: &reflection::Case<'_>) -> CaseTreeInner {
    let mut leaves: Vec<CaseTreeInner> = vec![];

    leaves.extend(case.predicate().iter().flat_map(|pred| {
        pred.parameters().map(|item| {
            let root: Box<dyn fmt::Display + Send + Sync> = Box::new(item.to_string());
            termtree::Tree::new(root, vec![]).with_multiline(true)
        })
    }));

    leaves.extend(case.products().map(|item| {
        let root: Box<dyn fmt::Display + Send + Sync> = Box::new(item.to_string());
        termtree::Tree::new(root, vec![]).with_multiline(true)
    }));

    leaves.extend(case.children().map(|item| convert(item)));

    let root = Box::new(case.predicate().map(|p| p.to_string()).unwrap_or_default());
    CaseTreeInner::new(root, leaves)
}

/// A `Case` rendered as a tree for display.
#[allow(missing_debug_implementations)]
pub struct CaseTree(CaseTreeInner);

impl fmt::Display for CaseTree {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2019 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

use crate::{
    api_utils::{delete, first, get, post, run_cmd, wait_for_cmd, CmdWrapper},
    display_utils::{display_cmd_state, display_error, generate_table, start_spinner},
    manager_cli_error::{
        DurationParseError, ImlManagerCliError, RunStratagemCommandResult,
        RunStratagemValidationError,
    },
};
use futures::Future;
use iml_wire_types::{ApiList, EndpointName, Filesystem, StratagemConfiguration};
use std::process::exit;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
pub enum StratagemCommand {
    /// Kickoff a Stratagem scan
    #[structopt(name = "scan")]
    Scan(StratagemScanData),
    /// Configure Stratagem scanning
    #[structopt(name = "configure")]
    StratagemConfig(StratagemConfig),
}

#[derive(Debug, StructOpt)]
pub enum StratagemConfig {
    /// List all Stratagem configurations
    #[structopt(name = "list")]
    List,
    /// Add a new Stratagem configuration
    #[structopt(name = "add")]
    Add(StratagemAddData),
    /// Remove a Stratagem configuration
    #[structopt(name = "remove")]
    Remove(StratagemRemoveData),
}

#[derive(Debug, StructOpt, serde::Serialize)]
pub struct StratagemAddData {
    /// Filesystem to configure
    #[structopt(short = "f", long = "filesystem")]
    filesystem: String,
    /// Interval to scan
    #[structopt(short = "i", long = "interval", parse(try_from_str = "parse_duration"))]
    interval: u64,
    /// The report duration
    #[structopt(short = "r", long = "report", parse(try_from_str = "parse_duration"))]
    report_duration: Option<u64>,
    /// The purge duration
    #[structopt(short = "p", long = "purge", parse(try_from_str = "parse_duration"))]
    purge_duration: Option<u64>,
}

#[derive(Debug, StructOpt, serde::Serialize)]
pub struct StratagemRemoveData {
    /// Filesystem to unconfigure
    #[structopt(short = "f", long = "filesystem")]
    name: String,
}

#[derive(serde::Serialize, StructOpt, Debug)]
pub struct StratagemScanData {
    /// The name of the filesystem to scan
    #[structopt(short = "f", long = "filesystem")]
    filesystem: String,
    /// The report duration
    #[structopt(short = "r", long = "report", parse(try_from_str = "parse_duration"))]
    report_duration: Option<u64>,
    /// The purge duration
    #[structopt(short = "p", long = "purge", parse(try_from_str = "parse_duration"))]
    purge_duration: Option<u64>,
}

fn parse_duration(src: &str) -> Result<u64, ImlManagerCliError> {
    if src.len() < 2 {
        return Err(DurationParseError::InvalidValue.into());
    }

    let mut val = String::from(src);
    let unit = val.pop();

    let val = val.parse::<u64>()?;

    match unit {
        Some('h') => Ok(val * 3_600_000),
        Some('d') => Ok(val * 86_400_000),
        Some('m') => Ok(val * 60_000),
        Some('s') => Ok(val * 1_000),
        Some('1'...'9') => Err(DurationParseError::NoUnit.into()),
        _ => Err(DurationParseError::InvalidUnit.into()),
    }
}

fn handle_cmd_resp(
    (resp, body): (iml_manager_client::Response, iml_manager_client::Chunk),
) -> Result<CmdWrapper, ImlManagerCliError> {
    let status = resp.status();
    if status.is_success() {
        Ok(serde_json::from_slice::<CmdWrapper>(&body)?)
    } else if status.is_client_error() {
        log::debug!("body: {:?}", std::str::from_utf8(&body));
        serde_json::from_slice::<RunStratagemValidationError>(&body)
            .map(std::convert::Into::into)
            .map(Err)?
    } else if status.is_server_error() {
        Err(RunStratagemValidationError {
            code: RunStratagemCommandResult::ServerError,
            message: "Internal server error.".into(),
        }
        .into())
    } else {
        Err(RunStratagemValidationError {
            code: RunStratagemCommandResult::UnknownError,
            message: "Unknown error.".into(),
        }
        .into())
    }
}

pub fn stratagem_cli(command: StratagemCommand) {
    match command {
        StratagemCommand::Scan(data) => {
            let fut = post("run_stratagem", data).and_then(handle_cmd_resp);

            let command: Result<CmdWrapper, ImlManagerCliError> = run_cmd(fut);

            match command {
                Ok(CmdWrapper { command }) => {
                    let stop_spinner = start_spinner(&command.message);

                    let command = run_cmd(wait_for_cmd(command)).expect("Could not poll command");

                    stop_spinner(None);

                    display_cmd_state(&command);

                    exit(exitcode::OK)
                }
                Err(validation_error) => {
                    display_error(validation_error);

                    exit(exitcode::CANTCREAT)
                }
            }
        }
        StratagemCommand::StratagemConfig(x) => match x {
            StratagemConfig::List => {
                let stop_spinner = start_spinner("Finding config...");

                let fut = get(
                    StratagemConfiguration::endpoint_name(),
                    serde_json::json!({ "limit": 0 }),
                );
                let result: Result<ApiList<StratagemConfiguration>, _> = run_cmd(fut);

                stop_spinner(None);

                match result {
                    Ok(r) => {
                        if r.objects.is_empty() {
                            return println!("0 Stratagem configuration records found");
                        }

                        let table = generate_table(
                            &["Id", "Filesystem", "State", "Interval", "Purge", "Report"],
                            r.objects.into_iter().map(|x| {
                                vec![
                                    x.id.to_string(),
                                    x.filesystem,
                                    x.state,
                                    x.interval.to_string(),
                                    x.purge_duration.map(|x| x.to_string()).unwrap_or_default(),
                                    x.report_duration.map(|x| x.to_string()).unwrap_or_default(),
                                ]
                            }),
                        );

                        table.printstd();

                        exit(exitcode::OK)
                    }
                    Err(e) => {
                        eprintln!("{}", e);

                        exit(exitcode::SOFTWARE);
                    }
                }
            }
            StratagemConfig::Add(c) => {
                let fut =
                    post(StratagemConfiguration::endpoint_name(), c).and_then(handle_cmd_resp);

                let command: Result<CmdWrapper, ImlManagerCliError> = run_cmd(fut);

                match command {
                    Ok(CmdWrapper { command }) => {
                        let stop_spinner = start_spinner(&command.message);

                        let command =
                            run_cmd(wait_for_cmd(command)).expect("Could not poll command");

                        stop_spinner(None);

                        display_cmd_state(&command);

                        exit(exitcode::OK)
                    }
                    Err(validation_error) => {
                        display_error(validation_error);

                        exit(exitcode::CANTCREAT)
                    }
                }
            }
            StratagemConfig::Remove(StratagemRemoveData { name }) => {
                let fut = get(
                    Filesystem::endpoint_name(),
                    serde_json::json!({
                        "name": name,
                        "dehydrate__mgt": false
                    }),
                )
                .and_then(first)
                .and_then(|fs: Filesystem| {
                    get(
                        StratagemConfiguration::endpoint_name(),
                        serde_json::json!({ "filesystem": fs.id }),
                    )
                })
                .and_then(first)
                .and_then(|x: StratagemConfiguration| {
                    delete(
                        &format!("{}/{}", StratagemConfiguration::endpoint_name(), x.id),
                        Vec::<(String, String)>::new(),
                    )
                })
                .and_then(handle_cmd_resp);

                let x: Result<CmdWrapper, ImlManagerCliError> = run_cmd(fut);

                match x {
                    Ok(CmdWrapper { command }) => {
                        let stop_spinner = start_spinner(&command.message);

                        let command =
                            run_cmd(wait_for_cmd(command)).expect("Could not poll command");

                        stop_spinner(None);

                        display_cmd_state(&command);

                        exit(exitcode::OK)
                    }
                    Err(validation_error) => {
                        display_error(validation_error);

                        exit(exitcode::CANTCREAT)
                    }
                }
            }
        },
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_duration_with_days() {
        match parse_duration("273d") {
            Ok(x) => assert_eq!(x, 23_587_200_000),
            _ => panic!("Duration parser should not have errored!"),
        }
    }

    #[test]
    fn test_parse_duration_with_hours() {
        match parse_duration("273h") {
            Ok(x) => assert_eq!(x, 982_800_000),
            _ => panic!("Duration parser should not have errored!"),
        }
    }

    #[test]
    fn test_parse_duration_with_minutes() {
        match parse_duration("273m") {
            Ok(x) => assert_eq!(x, 16_380_000),
            _ => panic!("Duration parser should not have errored!"),
        }
    }

    #[test]
    fn test_parse_duration_with_seconds() {
        match parse_duration("273s") {
            Ok(x) => assert_eq!(x, 273_000),
            _ => panic!("Duration parser should not have errored!"),
        }
    }

    #[test]
    fn test_parse_duration_with_invalid_unit() {
        match parse_duration("273x") {
            Ok(_) => panic!("Duration should have Errored with InvalidUnit"),
            Err(e) => assert_eq!(
                e.to_string(),
                "Invalid unit. Valid units include 'h' and 'd'."
            ),
        }
    }

    #[test]
    fn test_parse_duration_without_unit() {
        match parse_duration("273") {
            Ok(_) => panic!("Duration should have Errored with NoUnit"),
            Err(e) => assert_eq!(e.to_string(), "No unit specified."),
        }
    }

    #[test]
    fn test_parse_duration_with_insufficient_data() {
        match parse_duration("2") {
            Ok(_) => panic!("Duration should have Errored with InvalidValue"),
            Err(e) => assert_eq!(
                e.to_string(),
                "Invalid value specified. Must be a valid integer."
            ),
        }
    }

    #[test]
    fn test_parse_duration_with_invalid_data() {
        match parse_duration("abch") {
            Ok(_) => panic!("Duration should have Errored with invalid digit"),
            Err(e) => assert_eq!(e.to_string(), "invalid digit found in string"),
        }
    }
}
<commit_msg>Update stratagem cli wording (#1162)<commit_after>// Copyright (c) 2019 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

use crate::{
    api_utils::{delete, first, get, post, run_cmd, wait_for_cmd, CmdWrapper},
    display_utils::{display_cmd_state, display_error, generate_table, start_spinner},
    manager_cli_error::{
        DurationParseError, ImlManagerCliError, RunStratagemCommandResult,
        RunStratagemValidationError,
    },
};
use futures::Future;
use iml_wire_types::{ApiList, EndpointName, Filesystem, StratagemConfiguration};
use std::process::exit;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
pub enum StratagemCommand {
    /// Kickoff a Stratagem scan
    #[structopt(name = "scan")]
    Scan(StratagemScanData),
    /// Configure Stratagem scanning interval
    #[structopt(name = "interval")]
    StratagemInterval(StratagemInterval),
}

#[derive(Debug, StructOpt)]
pub enum StratagemInterval {
    /// List all existing Stratagem intervals
    #[structopt(name = "list")]
    List,
    /// Add a new Stratagem interval
    #[structopt(name = "add")]
    Add(StratagemIntervalConfig),
    /// Remove a Stratagem interval
    #[structopt(name = "remove")]
    Remove(StratagemRemoveData),
}

#[derive(Debug, StructOpt, serde::Serialize)]
pub struct StratagemIntervalConfig {
    /// Filesystem to configure
    #[structopt(short = "f", long = "filesystem")]
    filesystem: String,
    /// Interval to scan
    #[structopt(short = "i", long = "interval", parse(try_from_str = "parse_duration"))]
    interval: u64,
    /// The report duration
    #[structopt(short = "r", long = "report", parse(try_from_str = "parse_duration"))]
    report_duration: Option<u64>,
    /// The purge duration
    #[structopt(short = "p", long = "purge", parse(try_from_str = "parse_duration"))]
    purge_duration: Option<u64>,
}

#[derive(Debug, StructOpt, serde::Serialize)]
pub struct StratagemRemoveData {
    /// Filesystem to unconfigure
    #[structopt(short = "f", long = "filesystem")]
    name: String,
}

#[derive(serde::Serialize, StructOpt, Debug)]
pub struct StratagemScanData {
    /// The name of the filesystem to scan
    #[structopt(short = "f", long = "filesystem")]
    filesystem: String,
    /// The report duration
    #[structopt(short = "r", long = "report", parse(try_from_str = "parse_duration"))]
    report_duration: Option<u64>,
    /// The purge duration
    #[structopt(short = "p", long = "purge", parse(try_from_str = "parse_duration"))]
    purge_duration: Option<u64>,
}

fn parse_duration(src: &str) -> Result<u64, ImlManagerCliError> {
    if src.len() < 2 {
        return Err(DurationParseError::InvalidValue.into());
    }

    let mut val = String::from(src);
    let unit = val.pop();

    let val = val.parse::<u64>()?;

    match unit {
        Some('h') => Ok(val * 3_600_000),
        Some('d') => Ok(val * 86_400_000),
        Some('m') => Ok(val * 60_000),
        Some('s') => Ok(val * 1_000),
        Some('1'..='9') => Err(DurationParseError::NoUnit.into()),
        _ => Err(DurationParseError::InvalidUnit.into()),
    }
}

fn handle_cmd_resp(
    (resp, body): (iml_manager_client::Response, iml_manager_client::Chunk),
) -> Result<CmdWrapper, ImlManagerCliError> {
    let status = resp.status();
    if status.is_success() {
        Ok(serde_json::from_slice::<CmdWrapper>(&body)?)
    } else if status.is_client_error() {
        log::debug!("body: {:?}", std::str::from_utf8(&body));
        serde_json::from_slice::<RunStratagemValidationError>(&body)
            .map(std::convert::Into::into)
            .map(Err)?
    } else if status.is_server_error() {
        Err(RunStratagemValidationError {
            code: RunStratagemCommandResult::ServerError,
            message: "Internal server error.".into(),
        }
        .into())
    } else {
        Err(RunStratagemValidationError {
            code: RunStratagemCommandResult::UnknownError,
            message: "Unknown error.".into(),
        }
        .into())
    }
}

pub fn stratagem_cli(command: StratagemCommand) {
    match command {
        StratagemCommand::Scan(data) => {
            let fut = post("run_stratagem", data).and_then(handle_cmd_resp);

            let command: Result<CmdWrapper, ImlManagerCliError> = run_cmd(fut);

            match command {
                Ok(CmdWrapper { command }) => {
                    let stop_spinner = start_spinner(&command.message);

                    let command = run_cmd(wait_for_cmd(command)).expect("Could not poll command");

                    stop_spinner(None);

                    display_cmd_state(&command);

                    exit(exitcode::OK)
                }
                Err(validation_error) => {
                    display_error(validation_error);

                    exit(exitcode::CANTCREAT)
                }
            }
        }
        StratagemCommand::StratagemInterval(x) => match x {
            StratagemInterval::List => {
                let stop_spinner = start_spinner("Finding existing intervals...");

                let fut = get(
                    StratagemConfiguration::endpoint_name(),
                    serde_json::json!({ "limit": 0 }),
                );
                let result: Result<ApiList<StratagemConfiguration>, _> = run_cmd(fut);

                stop_spinner(None);

                match result {
                    Ok(r) => {
                        if r.objects.is_empty() {
                            return println!("No Stratagem intervals found");
                        }

                        let table = generate_table(
                            &["Id", "Filesystem", "State", "Interval", "Purge", "Report"],
                            r.objects.into_iter().map(|x| {
                                vec![
                                    x.id.to_string(),
                                    x.filesystem,
                                    x.state,
                                    x.interval.to_string(),
                                    x.purge_duration.map(|x| x.to_string()).unwrap_or_default(),
                                    x.report_duration.map(|x| x.to_string()).unwrap_or_default(),
                                ]
                            }),
                        );

                        table.printstd();

                        exit(exitcode::OK)
                    }
                    Err(e) => {
                        eprintln!("{}", e);

                        exit(exitcode::SOFTWARE);
                    }
                }
            }
            StratagemInterval::Add(c) => {
                let fut =
                    post(StratagemConfiguration::endpoint_name(), c).and_then(handle_cmd_resp);

                let command: Result<CmdWrapper, ImlManagerCliError> = run_cmd(fut);

                match command {
                    Ok(CmdWrapper { command }) => {
                        let stop_spinner = start_spinner(&command.message);

                        let command =
                            run_cmd(wait_for_cmd(command)).expect("Could not poll command");

                        stop_spinner(None);

                        display_cmd_state(&command);

                        exit(exitcode::OK)
                    }
                    Err(validation_error) => {
                        display_error(validation_error);

                        exit(exitcode::CANTCREAT)
                    }
                }
            }
            StratagemInterval::Remove(StratagemRemoveData { name }) => {
                let fut = get(
                    Filesystem::endpoint_name(),
                    serde_json::json!({
                        "name": name,
                        "dehydrate__mgt": false
                    }),
                )
                .and_then(first)
                .and_then(|fs: Filesystem| {
                    get(
                        StratagemConfiguration::endpoint_name(),
                        serde_json::json!({ "filesystem": fs.id }),
                    )
                })
                .and_then(first)
                .and_then(|x: StratagemConfiguration| {
                    delete(
                        &format!("{}/{}", StratagemConfiguration::endpoint_name(), x.id),
                        Vec::<(String, String)>::new(),
                    )
                })
                .and_then(handle_cmd_resp);

                let x: Result<CmdWrapper, ImlManagerCliError> = run_cmd(fut);

                match x {
                    Ok(CmdWrapper { command }) => {
                        let stop_spinner = start_spinner(&command.message);

                        let command =
                            run_cmd(wait_for_cmd(command)).expect("Could not poll command");

                        stop_spinner(None);

                        display_cmd_state(&command);

                        exit(exitcode::OK)
                    }
                    Err(validation_error) => {
                        display_error(validation_error);

                        exit(exitcode::CANTCREAT)
                    }
                }
            }
        },
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_duration_with_days() {
        match parse_duration("273d") {
            Ok(x) => assert_eq!(x, 23_587_200_000),
            _ => panic!("Duration parser should not have errored!"),
        }
    }

    #[test]
    fn test_parse_duration_with_hours() {
        match parse_duration("273h") {
            Ok(x) => assert_eq!(x, 982_800_000),
            _ => panic!("Duration parser should not have errored!"),
        }
    }

    #[test]
    fn test_parse_duration_with_minutes() {
        match parse_duration("273m") {
            Ok(x) => assert_eq!(x, 16_380_000),
            _ => panic!("Duration parser should not have errored!"),
        }
    }

    #[test]
    fn test_parse_duration_with_seconds() {
        match parse_duration("273s") {
            Ok(x) => assert_eq!(x, 273_000),
            _ => panic!("Duration parser should not have errored!"),
        }
    }

    #[test]
    fn test_parse_duration_with_invalid_unit() {
        match parse_duration("273x") {
            Ok(_) => panic!("Duration should have Errored with InvalidUnit"),
            Err(e) => assert_eq!(
                e.to_string(),
                "Invalid unit. Valid units include 'h' and 'd'."
            ),
        }
    }

    #[test]
    fn test_parse_duration_without_unit() {
        match parse_duration("273") {
            Ok(_) => panic!("Duration should have Errored with NoUnit"),
            Err(e) => assert_eq!(e.to_string(), "No unit specified."),
        }
    }

    #[test]
    fn test_parse_duration_with_insufficient_data() {
        match parse_duration("2") {
            Ok(_) => panic!("Duration should have Errored with InvalidValue"),
            Err(e) => assert_eq!(
                e.to_string(),
                "Invalid value specified. Must be a valid integer."
            ),
        }
    }

    #[test]
    fn test_parse_duration_with_invalid_data() {
        match parse_duration("abch") {
            Ok(_) => panic!("Duration should have Errored with invalid digit"),
            Err(e) => assert_eq!(e.to_string(), "invalid digit found in string"),
        }
    }
}
<|endoftext|>