<commit_before>use std::ascii::OwnedAsciiExt;
use std::collections::HashMap;
use std::io::File;

use time;
use time::{
    Timespec,
    Tm,
};

use command::command::{
    Attribute,
    Envelope,
    Flags,
    InternalDate,
    RFC822,
    Body,
    BodyPeek,
    BodyStructure,
    UID
};
use command::command::{
    AllRFC822,
    HeaderRFC822,
    SizeRFC822,
    TextRFC822
};
use error::{
    Error, ImapResult, InternalIoError, MessageDecodeError
};

#[deriving(Show)]
enum Flag {
    Answered,
    Draft,
    Flagged,
    Seen
}

#[deriving(Show)]
pub struct Message {
    uid: u32,
    pub path: String,
    headers: HashMap<String, String>,
    body: Vec<MIMEPart>,
    flags: Vec<Flag>,
    pub deleted: bool,
    size: uint,
    raw_contents: String,
    raw_header: String
}

#[deriving(Show)]
pub struct MIMEPart {
    mime_header: String,
    mime_body: String
}

impl Message {
    pub fn parse(arg_path: &Path) -> ImapResult<Message> {
        // Load the file contents.
        let file = match File::open(arg_path) {
            Ok(mut file) => match file.read_to_end() {
                Ok(contents) => contents,
                Err(e) => return Err(Error::simple(InternalIoError(e), "Failed to read mail file."))
            },
            Err(e) => return Err(Error::simple(InternalIoError(e), "Failed to open mail file."))
        };
        let raw_contents = String::from_utf8_lossy(file.as_slice()).to_string();
        let size = raw_contents.as_slice().len();

        let mut path = arg_path.filename_str().unwrap().splitn(1, ':');
        let filename = path.next().unwrap();
        let path_flags = path.next();

        // Retrieve the UID from the provided filename.
        let uid = match from_str::<u32>(filename) {
            Some(uid) => uid,
            None => return Err(Error::simple(MessageDecodeError, "Failed to retrieve UID from filename."))
        };
        // Parse the flags from the filename.
        let flags = match path_flags {
            None => Vec::new(),
            Some(flags) => {
                let unparsed_flags = flags.splitn(1, ',').skip(1).next().unwrap();
                let mut vec_flags: Vec<Flag> = Vec::new();
                for flag in unparsed_flags.chars() {
                    let parsed_flag = match flag {
                        'D' => Some(Draft),
                        'F' => Some(Flagged),
                        'R' => Some(Answered),
                        'S' => Some(Seen),
                        _ => None
                    };
                    match parsed_flag {
                        Some(enum_flag) => vec_flags.push(enum_flag),
                        None => { }
                    }
                }
                vec_flags
            }
        };

        let header_boundary = raw_contents.as_slice().find_str("\n\n").unwrap();
        let raw_header = raw_contents.as_slice().slice_to(header_boundary);
        let raw_body = raw_contents.as_slice().slice_from(header_boundary + 2);

        // Iterate over the lines of the header in reverse.
        // If a line with leading whitespace is detected, it is merged to the line before it.
        // This "unwraps" the header as indicated in RFC 2822 2.2.3
        let mut parsed_header: Vec<String> = Vec::new();
        let mut iterator = raw_header.lines().rev();
        loop {
            let line = match iterator.next() {
                Some(line) => line,
                None => break
            };
            let mut field = line.to_string();
            if line.as_slice().starts_with(" ") || line.as_slice().starts_with("\t") {
                field = field.as_slice().trim_left_chars(' ').trim_left_chars('\t').to_string();
                loop {
                    let next = iterator.next().unwrap().to_string();
                    let mut trimmed_next = next.as_slice().trim_left_chars(' ').trim_left_chars('\t').to_string();
                    // Add a space between the merged lines.
                    trimmed_next.push_str(" ".as_slice());
                    trimmed_next.push_str(field.as_slice());
                    field = trimmed_next;
                    if !next.as_slice().starts_with(" ") && !next.as_slice().starts_with("\t") {
                        break;
                    }
                }
            }

            parsed_header.push(field.to_string());
        }
        let parsed_header: Vec<&String> = parsed_header.iter().rev().collect();

        let mut headers = HashMap::new();
        for line in parsed_header.iter() {
            let split: Vec<&str> = line.as_slice().splitn(1, ':').collect();
            headers.insert(split[0].to_string().into_ascii_upper().to_string(), split[1].slice_from(1).clone().to_string());
        }
        // Remove the "Received" key from the HashMap.
        headers.remove(&"RECEIVED".to_string());

        // Determine whether the message is MULTIPART or not.
        let content_type = headers["CONTENT-TYPE".to_string()].clone();
        let mut body = Vec::new();
        if content_type.as_slice().contains("MULTIPART") {
            let mime_boundary = {
                let value: Vec<&str> = content_type.as_slice().split_str("BOUNDARY=\"").collect();
                if value.len() < 2 { return Err(Error::simple(MessageDecodeError, "Failed to determine MULTIPART boundary.")) }
                let value: Vec<&str> = value[1].splitn(1, '"').collect();
                if value.len() < 1 { return Err(Error::simple(MessageDecodeError, "Failed to determine MULTIPART boundary.")) }
                format!("--{}--\n", value[0])
            };

            let first_content_type_index = match raw_body.find_str("Content-Type") {
                Some(val) => val,
                None => return Err(Error::simple(MessageDecodeError, "Missing Content-Type for body part"))
            };
            let raw_body = raw_body.slice_from(first_content_type_index);
            let raw_body: Vec<&str> = raw_body.split_str(mime_boundary.as_slice()).collect();
            let raw_body = raw_body.slice_to(raw_body.len() - 1);

            for part in raw_body.iter() {
                let header_boundary = part.as_slice().find_str("\n\n").unwrap();
                let header = part.as_slice().slice_to(header_boundary);
                let mut content_type = String::new();
                for line in header.lines() {
                    let split_line: Vec<&str> = line.splitn(1, ':').collect();
                    if split_line[0] == "Content-Type" {
                        let content_type_values: Vec<&str> = split_line[1].splitn(1, ';').collect();
                        content_type = content_type_values[0].slice_from(1).to_string();
                        break;
                    }
                }
                let body_part = MIMEPart {
                    mime_header: content_type.to_string(),
                    mime_body: raw_body.to_string()
                };
                body.push(body_part);
            }
        } else {
            let body_part = MIMEPart {
                mime_header: content_type.to_string(),
                mime_body: raw_body.to_string()
            };
            body.push(body_part);
        }

        let message = Message {
            uid: uid,
            path: arg_path.display().to_string(),
            headers: headers,
            body: body,
            flags: flags,
            deleted: false,
            size: size,
            raw_contents: raw_contents.clone(),
            raw_header: raw_header.to_string()
        };

        Ok(message)
    }

    // TODO: Make sure that returning a pointer is fine.
    pub fn envelope_from(&self) -> &String {
        self.headers.find(&"FROM".to_string()).unwrap()
    }

    // TODO: Make sure that returning a pointer is fine.
    pub fn envelope_to(&self) -> &String {
        self.headers.find(&"TO".to_string()).unwrap()
    }

    pub fn fetch(&self, attributes: &Vec<Attribute>) -> String {
        let mut res = String::new();
        for attr in attributes.iter() {
            let attr_str = match attr {
                &Envelope => { format!("ENVELOPE {}", self.get_envelope()) }, // TODO: Finish implementing this.
                &Flags => { "".to_string() },
                &InternalDate => { "".to_string() },
                &RFC822(ref attr) => {
                    let rfc_attr = match attr {
                        &AllRFC822 => { "".to_string() },
                        &HeaderRFC822 => { format!(".HEADER {{{}}}\n{}", self.raw_header.len(), self.raw_header) },
                        &SizeRFC822 => { format!(".SIZE {}", self.size) },
                        &TextRFC822 => { "".to_string() }
                    };
                    format!("RFC822{}", rfc_attr)
                },
                &Body(ref section, ref octets) => { "".to_string() },
                &BodyPeek(ref section, ref octets) => { "".to_string() },
                &BodyStructure => { "".to_string() },
                &UID => { format!("UID {}", self.uid) }
            };
            res = format!("{} {}", res, attr_str);
        }
        res
    }

    /**
     * RFC3501 - 7.4.2 - P.76-77
     *
     * Returns a parenthesized list that described the envelope structure of a
     * message.
     * Computed by parsing the [RFC-2822] header into the component parts,
     * defaulting various fields as necessary.
     *
     * Requires (in the following order): date, subject, from, sender, reply-to,
     * to, cc, bcc, in-reply-to, and message-id.
     * The date, subject, in-reply-to, and message-id fields are strings.
     * The from, sender, reply-to, to, cc, and bcc fields are parenthesized
     * lists of address structures.
     */
    // TODO: Finish implementing this.
    fn get_envelope(&self) -> String {
        let date = self.get_quoted_field_or_nil("DATE".to_string());
        let subject = self.get_quoted_field_or_nil("SUBJECT".to_string());
        let from = self.get_parenthesized_addresses("FROM".to_string());
        let sender = self.get_parenthesized_addresses("SENDER".to_string());
        let reply_to = self.get_parenthesized_addresses("REPLY-TO".to_string());
        let to = self.get_parenthesized_addresses("TO".to_string());
        let cc = self.get_parenthesized_addresses("CC".to_string());
        let bcc = self.get_parenthesized_addresses("BCC".to_string());
        let in_reply_to = self.get_quoted_field_or_nil("IN-REPLY-TO".to_string());
        let message_id = self.get_quoted_field_or_nil("MESSAGE-ID".to_string());

        format!(
            "({} {} {} {} {} {} {} {} {} {})",
            date,
            subject,
            from,
            sender,
            reply_to,
            to,
            cc,
            bcc,
            in_reply_to,
            message_id)
    }

    fn get_quoted_field_or_nil(&self, key: String) -> String {
        match self.headers.find(&key) {
            Some(v) => format!("\"{}\"", v),
            None => "NIL".to_string()
        }
    }

    /**
     * RFC3501 - 7.4.2 - P.76
     */
    // TODO: Finish implementing this.
    fn get_parenthesized_addresses(&self, key: String) -> String {
        let addresses = match self.headers.find(&key) {
            Some(v) => v, // TODO: this is not parenthesized.
            None => return "NIL".to_string()
        };

        addresses.clone()
    }

    fn date_received(&self) -> String {
        // Retrieve the date received from the UID.
        let date_received = Timespec { sec: self.uid as i64, nsec: 0i32 };
        let date_received_tm = time::at_utc(date_received);

        let month = match date_received_tm.tm_mon {
            0 => "Jan".to_string(),
            1 => "Feb".to_string(),
            2 => "Mar".to_string(),
            3 => "Apr".to_string(),
            4 => "May".to_string(),
            5 => "Jun".to_string(),
            6 => "Jul".to_string(),
            7 => "Aug".to_string(),
            8 => "Sep".to_string(),
            9 => "Oct".to_string(),
            10 => "Nov".to_string(),
            11 => "Dec".to_string(),
            _ => fail!("Unable to determine month!") // NOTE: this should never happen.
        };

        format!(
            "{:0>2}-{}-{:0>2} {:0>2}:{:0>2}:{:0>2} -0000",
            date_received_tm.tm_mday,
            month,
            date_received_tm.tm_year + 1900i32,
            date_received_tm.tm_hour,
            date_received_tm.tm_min,
            date_received_tm.tm_sec)
    }
}
<commit_msg>Implement INTERNALDATE<commit_after>use std::ascii::OwnedAsciiExt;
use std::collections::HashMap;
use std::io::File;

use time;
use time::{
    Timespec,
    Tm,
};

use command::command::{
    Attribute,
    Envelope,
    Flags,
    InternalDate,
    RFC822,
    Body,
    BodyPeek,
    BodyStructure,
    UID
};
use command::command::{
    AllRFC822,
    HeaderRFC822,
    SizeRFC822,
    TextRFC822
};
use error::{
    Error, ImapResult, InternalIoError, MessageDecodeError
};

#[deriving(Show)]
enum Flag {
    Answered,
    Draft,
    Flagged,
    Seen
}

#[deriving(Show)]
pub struct Message {
    uid: u32,
    pub path: String,
    headers: HashMap<String, String>,
    body: Vec<MIMEPart>,
    flags: Vec<Flag>,
    pub deleted: bool,
    size: uint,
    raw_contents: String,
    raw_header: String
}

#[deriving(Show)]
pub struct MIMEPart {
    mime_header: String,
    mime_body: String
}

impl Message {
    pub fn parse(arg_path: &Path) -> ImapResult<Message> {
        // Load the file contents.
        let file = match File::open(arg_path) {
            Ok(mut file) => match file.read_to_end() {
                Ok(contents) => contents,
                Err(e) => return Err(Error::simple(InternalIoError(e), "Failed to read mail file."))
            },
            Err(e) => return Err(Error::simple(InternalIoError(e), "Failed to open mail file."))
        };
        let raw_contents = String::from_utf8_lossy(file.as_slice()).to_string();
        let size = raw_contents.as_slice().len();

        let mut path = arg_path.filename_str().unwrap().splitn(1, ':');
        let filename = path.next().unwrap();
        let path_flags = path.next();

        // Retrieve the UID from the provided filename.
        let uid = match from_str::<u32>(filename) {
            Some(uid) => uid,
            None => return Err(Error::simple(MessageDecodeError, "Failed to retrieve UID from filename."))
        };
        // Parse the flags from the filename.
        let flags = match path_flags {
            None => Vec::new(),
            Some(flags) => {
                let unparsed_flags = flags.splitn(1, ',').skip(1).next().unwrap();
                let mut vec_flags: Vec<Flag> = Vec::new();
                for flag in unparsed_flags.chars() {
                    let parsed_flag = match flag {
                        'D' => Some(Draft),
                        'F' => Some(Flagged),
                        'R' => Some(Answered),
                        'S' => Some(Seen),
                        _ => None
                    };
                    match parsed_flag {
                        Some(enum_flag) => vec_flags.push(enum_flag),
                        None => { }
                    }
                }
                vec_flags
            }
        };

        let header_boundary = raw_contents.as_slice().find_str("\n\n").unwrap();
        let raw_header = raw_contents.as_slice().slice_to(header_boundary);
        let raw_body = raw_contents.as_slice().slice_from(header_boundary + 2);

        // Iterate over the lines of the header in reverse.
        // If a line with leading whitespace is detected, it is merged to the line before it.
        // This "unwraps" the header as indicated in RFC 2822 2.2.3
        let mut parsed_header: Vec<String> = Vec::new();
        let mut iterator = raw_header.lines().rev();
        loop {
            let line = match iterator.next() {
                Some(line) => line,
                None => break
            };
            let mut field = line.to_string();
            if line.as_slice().starts_with(" ") || line.as_slice().starts_with("\t") {
                field = field.as_slice().trim_left_chars(' ').trim_left_chars('\t').to_string();
                loop {
                    let next = iterator.next().unwrap().to_string();
                    let mut trimmed_next = next.as_slice().trim_left_chars(' ').trim_left_chars('\t').to_string();
                    // Add a space between the merged lines.
                    trimmed_next.push_str(" ".as_slice());
                    trimmed_next.push_str(field.as_slice());
                    field = trimmed_next;
                    if !next.as_slice().starts_with(" ") && !next.as_slice().starts_with("\t") {
                        break;
                    }
                }
            }

            parsed_header.push(field.to_string());
        }
        let parsed_header: Vec<&String> = parsed_header.iter().rev().collect();

        let mut headers = HashMap::new();
        for line in parsed_header.iter() {
            let split: Vec<&str> = line.as_slice().splitn(1, ':').collect();
            headers.insert(split[0].to_string().into_ascii_upper().to_string(), split[1].slice_from(1).clone().to_string());
        }
        // Remove the "Received" key from the HashMap.
        headers.remove(&"RECEIVED".to_string());

        // Determine whether the message is MULTIPART or not.
        let content_type = headers["CONTENT-TYPE".to_string()].clone();
        let mut body = Vec::new();
        if content_type.as_slice().contains("MULTIPART") {
            let mime_boundary = {
                let value: Vec<&str> = content_type.as_slice().split_str("BOUNDARY=\"").collect();
                if value.len() < 2 { return Err(Error::simple(MessageDecodeError, "Failed to determine MULTIPART boundary.")) }
                let value: Vec<&str> = value[1].splitn(1, '"').collect();
                if value.len() < 1 { return Err(Error::simple(MessageDecodeError, "Failed to determine MULTIPART boundary.")) }
                format!("--{}--\n", value[0])
            };

            let first_content_type_index = match raw_body.find_str("Content-Type") {
                Some(val) => val,
                None => return Err(Error::simple(MessageDecodeError, "Missing Content-Type for body part"))
            };
            let raw_body = raw_body.slice_from(first_content_type_index);
            let raw_body: Vec<&str> = raw_body.split_str(mime_boundary.as_slice()).collect();
            let raw_body = raw_body.slice_to(raw_body.len() - 1);

            for part in raw_body.iter() {
                let header_boundary = part.as_slice().find_str("\n\n").unwrap();
                let header = part.as_slice().slice_to(header_boundary);
                let mut content_type = String::new();
                for line in header.lines() {
                    let split_line: Vec<&str> = line.splitn(1, ':').collect();
                    if split_line[0] == "Content-Type" {
                        let content_type_values: Vec<&str> = split_line[1].splitn(1, ';').collect();
                        content_type = content_type_values[0].slice_from(1).to_string();
                        break;
                    }
                }
                let body_part = MIMEPart {
                    mime_header: content_type.to_string(),
                    mime_body: raw_body.to_string()
                };
                body.push(body_part);
            }
        } else {
            let body_part = MIMEPart {
                mime_header: content_type.to_string(),
                mime_body: raw_body.to_string()
            };
            body.push(body_part);
        }

        let message = Message {
            uid: uid,
            path: arg_path.display().to_string(),
            headers: headers,
            body: body,
            flags: flags,
            deleted: false,
            size: size,
            raw_contents: raw_contents.clone(),
            raw_header: raw_header.to_string()
        };

        Ok(message)
    }

    // TODO: Make sure that returning a pointer is fine.
    pub fn envelope_from(&self) -> &String {
        self.headers.find(&"FROM".to_string()).unwrap()
    }

    // TODO: Make sure that returning a pointer is fine.
    pub fn envelope_to(&self) -> &String {
        self.headers.find(&"TO".to_string()).unwrap()
    }

    pub fn fetch(&self, attributes: &Vec<Attribute>) -> String {
        let mut res = String::new();
        for attr in attributes.iter() {
            let attr_str = match attr {
                &Envelope => { format!("ENVELOPE {}", self.get_envelope()) }, // TODO: Finish implementing this.
                &Flags => { "".to_string() },
                &InternalDate => { format!(" INTERNALDATE \"{}\"", self.date_received()) }
                &RFC822(ref attr) => {
                    let rfc_attr = match attr {
                        &AllRFC822 => { "".to_string() },
                        &HeaderRFC822 => { format!(".HEADER {{{}}}\n{}", self.raw_header.len(), self.raw_header) },
                        &SizeRFC822 => { format!(".SIZE {}", self.size) },
                        &TextRFC822 => { "".to_string() }
                    };
                    format!("RFC822{}", rfc_attr)
                },
                &Body(ref section, ref octets) => { "".to_string() },
                &BodyPeek(ref section, ref octets) => { "".to_string() },
                &BodyStructure => { "".to_string() },
                &UID => { format!("UID {}", self.uid) }
            };
            res = format!("{} {}", res, attr_str);
        }
        res
    }

    /**
     * RFC3501 - 7.4.2 - P.76-77
     *
     * Returns a parenthesized list that described the envelope structure of a
     * message.
     * Computed by parsing the [RFC-2822] header into the component parts,
     * defaulting various fields as necessary.
     *
     * Requires (in the following order): date, subject, from, sender, reply-to,
     * to, cc, bcc, in-reply-to, and message-id.
     * The date, subject, in-reply-to, and message-id fields are strings.
     * The from, sender, reply-to, to, cc, and bcc fields are parenthesized
     * lists of address structures.
     */
    // TODO: Finish implementing this.
    fn get_envelope(&self) -> String {
        let date = self.get_quoted_field_or_nil("DATE".to_string());
        let subject = self.get_quoted_field_or_nil("SUBJECT".to_string());
        let from = self.get_parenthesized_addresses("FROM".to_string());
        let sender = self.get_parenthesized_addresses("SENDER".to_string());
        let reply_to = self.get_parenthesized_addresses("REPLY-TO".to_string());
        let to = self.get_parenthesized_addresses("TO".to_string());
        let cc = self.get_parenthesized_addresses("CC".to_string());
        let bcc = self.get_parenthesized_addresses("BCC".to_string());
        let in_reply_to = self.get_quoted_field_or_nil("IN-REPLY-TO".to_string());
        let message_id = self.get_quoted_field_or_nil("MESSAGE-ID".to_string());

        format!(
            "({} {} {} {} {} {} {} {} {} {})",
            date,
            subject,
            from,
            sender,
            reply_to,
            to,
            cc,
            bcc,
            in_reply_to,
            message_id)
    }

    fn get_quoted_field_or_nil(&self, key: String) -> String {
        match self.headers.find(&key) {
            Some(v) => format!("\"{}\"", v),
            None => "NIL".to_string()
        }
    }

    /**
     * RFC3501 - 7.4.2 - P.76
     */
    // TODO: Finish implementing this.
    fn get_parenthesized_addresses(&self, key: String) -> String {
        let addresses = match self.headers.find(&key) {
            Some(v) => v, // TODO: this is not parenthesized.
            None => return "NIL".to_string()
        };

        addresses.clone()
    }

    fn date_received(&self) -> String {
        // Retrieve the date received from the UID.
        let date_received = Timespec { sec: self.uid as i64, nsec: 0i32 };
        let date_received_tm = time::at_utc(date_received);

        let month = match date_received_tm.tm_mon {
            0 => "Jan".to_string(),
            1 => "Feb".to_string(),
            2 => "Mar".to_string(),
            3 => "Apr".to_string(),
            4 => "May".to_string(),
            5 => "Jun".to_string(),
            6 => "Jul".to_string(),
            7 => "Aug".to_string(),
            8 => "Sep".to_string(),
            9 => "Oct".to_string(),
            10 => "Nov".to_string(),
            11 => "Dec".to_string(),
            _ => fail!("Unable to determine month!") // NOTE: this should never happen.
        };

        format!(
            "{:0>2}-{}-{:0>2} {:0>2}:{:0>2}:{:0>2} -0000",
            date_received_tm.tm_mday,
            month,
            date_received_tm.tm_year + 1900i32,
            date_received_tm.tm_hour,
            date_received_tm.tm_min,
            date_received_tm.tm_sec)
    }
}
<|endoftext|>


<commit_before>use std::any::{Any, TypeId};
use std::cell::RefCell;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::mem::transmute;
use std::rc::Rc;

/// Trait for setting up mocks
///
/// The trait is implemented for all functions, so its methods can be called on any function.
///
/// Note: methods have any effect only if called on functions [annotated as mockable](https://docs.rs/mocktopus_macros).
pub trait Mockable<T, O> {

    /// Core function for setting up mocks
    ///
    /// The passed closure is called whenever the mocked function is called. Depending on variant of returned
    /// [MockResult](enum.MockResult.html) the mocked function continues to run or returns immediately.
    /// In case of continuation the function arguments can be modified or replaced.
    ///
    /// The mock closure is saved in a
    /// [thread local static storage](https://doc.rust-lang.org/std/macro.thread_local.html),
    /// so it has effect only in thread, where it was set.
    /// Each Rust test is executed in separate thread, so mocks do not leak between them.
    /// # Safety
    /// It is up to the user to make sure, that the closure is valid long enough to serve all calls to mocked function.
    /// If the mock closure uses any non-static values or references, it will silently become invalid at some point of
    /// host thread lifetime.
    ///
    /// ```
    /// #[mockable]
    /// fn get_string(context: &Context) -> &String {
    ///     context.get_string()
    /// }
    ///
    /// #[test]
    /// fn get_string_test() {
    ///     let mocked = "mocked".to_string();
    ///     unsafe {
    ///         get_string.mock_raw(|_| MockResult::Return(&mocked));
    ///     }
    ///
    ///     assert_eq!("mocked", get_string(&Context::default()));
    /// }
    /// ```
    unsafe fn mock_raw<M: FnMut<T, Output=MockResult<T, O>>>(&self, mock: M);

    /// A safe variant of [mock_raw](#tymethod.mock_raw) for static closures
    ///
    /// The safety is guaranteed by forcing passed closure to be static.
    /// This eliminates the problem of using non-static values, which may not live long enough.
    ///
    /// ```
    /// #[mockable]
    /// fn get_string() -> String {
    ///     "not mocked".to_string()
    /// }
    ///
    /// #[test]
    /// fn get_string_test() {
    ///     get_string.mock_safe(|| MockResult::Return("mocked".to_string()));
    ///
    ///     assert_eq!("mocked", get_string());
    /// }
    /// ```
    fn mock_safe<M: FnMut<T, Output=MockResult<T, O>> + 'static>(&self, mock: M);

    /// Stop mocking this function.
    ///
    /// All future invocations will be forwarded to the real implementation.
    fn clear_mock(&self);

    #[doc(hidden)]
    /// Called before every execution of a mockable function. Checks if mock is set and if it is, calls it.
    fn call_mock(&self, input: T) -> MockResult<T, O>;

    #[doc(hidden)]
    /// Returns a unique ID of the function, which is used to set and get its mock.
    unsafe fn get_mock_id(&self) -> TypeId;
}

/// Controls mocked function behavior when returned from [mock closure](trait.Mockable.html)
pub enum MockResult<T, O> {
    /// Function runs normally as if it was called with given arguments.
    /// The arguments are passed inside enum variant as a tuple.
    Continue(T),

    /// Function returns immediately with a given value. The returned value is passed inside enum variant.
    Return(O),
}

thread_local!{
    static MOCK_STORE: RefCell<HashMap<TypeId, Rc<RefCell<Box<FnMut<(), Output=()>>>>>> = RefCell::new(HashMap::new())
}

/// Clear all mocks in the ThreadLocal; only necessary if tests share threads
pub fn clear_mocks() {
    MOCK_STORE.with(|mock_ref_cell| {
        mock_ref_cell.borrow_mut().clear();
    });
}

struct ScopedMock<'a> {
    phantom: PhantomData<&'a ()>,
    id: TypeId,
}

impl<'a> ScopedMock<'a> {
    unsafe fn new<T, O, M: Mockable<T, O> + 'a, F: FnMut<T, Output=MockResult<T, O>>>(
        mockable: &M,
        mock: F,
    ) -> Self {
        mockable.mock_raw(mock);
        ScopedMock {
            phantom: PhantomData,
            id: mockable.get_mock_id(),
        }
    }
}

impl<'a> Drop for ScopedMock<'a> {
    fn drop(&mut self) {
        clear_id(self.id);
    }
}

fn clear_id(id: TypeId) {
    MOCK_STORE.with(|mock_ref_cell| {
        mock_ref_cell.borrow_mut().remove(&id);
    });
}

impl<T, O, F: FnOnce<T, Output=O>> Mockable<T, O> for F {
    unsafe fn mock_raw<M: FnMut<T, Output=MockResult<T, O>>>(&self, mock: M) {
        let id = self.get_mock_id();
        MOCK_STORE.with(|mock_ref_cell| {
            let real = Rc::new(RefCell::new(Box::new(mock) as Box<FnMut<_, Output=_>>));
            let stored = transmute(real);
            mock_ref_cell.borrow_mut()
                .insert(id, stored);
        })
    }

    fn mock_safe<M: FnMut<T, Output=MockResult<T, O>> + 'static>(&self, mock: M) {
        unsafe {
            self.mock_raw(mock)
        }
    }

    fn clear_mock(&self) {
        let id = unsafe { self.get_mock_id() };
        clear_id(id);
    }

    fn call_mock(&self, input: T) -> MockResult<T, O> {
        unsafe {
            let id = self.get_mock_id();
            let rc_opt = MOCK_STORE.with(|mock_ref_cell|
                mock_ref_cell.borrow()
                    .get(&id)
                    .cloned()
            );
            let stored_opt = rc_opt.as_ref()
                .and_then(|rc| rc.try_borrow_mut().ok());
            match stored_opt {
                Some(mut stored) => {
                    let real: &mut Box<FnMut<_, Output=_>> = transmute(&mut*stored);
                    real.call_mut(input)
                }
                None => MockResult::Continue(input),
            }
        }
    }

    unsafe fn get_mock_id(&self) -> TypeId {
        (||()).type_id()
    }
}

pub struct MockContext<'a> {
    planned_mocks: Vec<Box<FnOnce() -> ScopedMock<'a> + 'a>>,
}

impl<'a> MockContext<'a> {
    pub fn new() -> Self {
        MockContext {
            planned_mocks: Vec::new(),
        }
    }

    pub fn mock_safe<
        Args,
        Output,
        M: Mockable<Args, Output> + 'a,
        F: FnMut<Args, Output = MockResult<Args, Output>> + 'a,
    >(
        mut self,
        mock: M,
        body: F,
    ) -> Self {
        self.planned_mocks
            .push(Box::new(move || unsafe { ScopedMock::new(&mock, body) }));
        self
    }

    pub fn run<T, F: FnOnce() -> T>(self, f: F) -> T {
        let _scoped_mocks = self
            .planned_mocks
            .into_iter()
            .map(|f| f())
            .collect::<Vec<_>>();
        f()
    }
}
<commit_msg>Derive Default and convert to HashMap<commit_after>use std::any::{Any, TypeId};
use std::cell::RefCell;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::mem::transmute;
use std::rc::Rc;

/// Trait for setting up mocks
///
/// The trait is implemented for all functions, so its methods can be called on any function.
///
/// Note: methods have any effect only if called on functions [annotated as mockable](https://docs.rs/mocktopus_macros).
pub trait Mockable<T, O> {

    /// Core function for setting up mocks
    ///
    /// The passed closure is called whenever the mocked function is called. Depending on variant of returned
    /// [MockResult](enum.MockResult.html) the mocked function continues to run or returns immediately.
    /// In case of continuation the function arguments can be modified or replaced.
    ///
    /// The mock closure is saved in a
    /// [thread local static storage](https://doc.rust-lang.org/std/macro.thread_local.html),
    /// so it has effect only in thread, where it was set.
    /// Each Rust test is executed in separate thread, so mocks do not leak between them.
    /// # Safety
    /// It is up to the user to make sure, that the closure is valid long enough to serve all calls to mocked function.
    /// If the mock closure uses any non-static values or references, it will silently become invalid at some point of
    /// host thread lifetime.
    ///
    /// ```
    /// #[mockable]
    /// fn get_string(context: &Context) -> &String {
    ///     context.get_string()
    /// }
    ///
    /// #[test]
    /// fn get_string_test() {
    ///     let mocked = "mocked".to_string();
    ///     unsafe {
    ///         get_string.mock_raw(|_| MockResult::Return(&mocked));
    ///     }
    ///
    ///     assert_eq!("mocked", get_string(&Context::default()));
    /// }
    /// ```
    unsafe fn mock_raw<M: FnMut<T, Output=MockResult<T, O>>>(&self, mock: M);

    /// A safe variant of [mock_raw](#tymethod.mock_raw) for static closures
    ///
    /// The safety is guaranteed by forcing passed closure to be static.
    /// This eliminates the problem of using non-static values, which may not live long enough.
    ///
    /// ```
    /// #[mockable]
    /// fn get_string() -> String {
    ///     "not mocked".to_string()
    /// }
    ///
    /// #[test]
    /// fn get_string_test() {
    ///     get_string.mock_safe(|| MockResult::Return("mocked".to_string()));
    ///
    ///     assert_eq!("mocked", get_string());
    /// }
    /// ```
    fn mock_safe<M: FnMut<T, Output=MockResult<T, O>> + 'static>(&self, mock: M);

    /// Stop mocking this function.
    ///
    /// All future invocations will be forwarded to the real implementation.
    fn clear_mock(&self);

    #[doc(hidden)]
    /// Called before every execution of a mockable function. Checks if mock is set and if it is, calls it.
    fn call_mock(&self, input: T) -> MockResult<T, O>;

    #[doc(hidden)]
    /// Returns a unique ID of the function, which is used to set and get its mock.
    unsafe fn get_mock_id(&self) -> TypeId;
}

/// Controls mocked function behavior when returned from [mock closure](trait.Mockable.html)
pub enum MockResult<T, O> {
    /// Function runs normally as if it was called with given arguments.
    /// The arguments are passed inside enum variant as a tuple.
    Continue(T),

    /// Function returns immediately with a given value. The returned value is passed inside enum variant.
    Return(O),
}

thread_local!{
    static MOCK_STORE: RefCell<HashMap<TypeId, Rc<RefCell<Box<FnMut<(), Output=()>>>>>> = RefCell::new(HashMap::new())
}

/// Clear all mocks in the ThreadLocal; only necessary if tests share threads
pub fn clear_mocks() {
    MOCK_STORE.with(|mock_ref_cell| {
        mock_ref_cell.borrow_mut().clear();
    });
}

struct ScopedMock<'a> {
    phantom: PhantomData<&'a ()>,
    id: TypeId,
}

impl<'a> ScopedMock<'a> {
    unsafe fn new<T, O, M: Mockable<T, O> + 'a, F: FnMut<T, Output=MockResult<T, O>>>(
        mockable: &M,
        mock: F,
    ) -> Self {
        mockable.mock_raw(mock);
        ScopedMock {
            phantom: PhantomData,
            id: mockable.get_mock_id(),
        }
    }
}

impl<'a> Drop for ScopedMock<'a> {
    fn drop(&mut self) {
        clear_id(self.id);
    }
}

fn clear_id(id: TypeId) {
    MOCK_STORE.with(|mock_ref_cell| {
        mock_ref_cell.borrow_mut().remove(&id);
    });
}

impl<T, O, F: FnOnce<T, Output=O>> Mockable<T, O> for F {
    unsafe fn mock_raw<M: FnMut<T, Output=MockResult<T, O>>>(&self, mock: M) {
        let id = self.get_mock_id();
        MOCK_STORE.with(|mock_ref_cell| {
            let real = Rc::new(RefCell::new(Box::new(mock) as Box<FnMut<_, Output=_>>));
            let stored = transmute(real);
            mock_ref_cell.borrow_mut()
                .insert(id, stored);
        })
    }

    fn mock_safe<M: FnMut<T, Output=MockResult<T, O>> + 'static>(&self, mock: M) {
        unsafe {
            self.mock_raw(mock)
        }
    }

    fn clear_mock(&self) {
        let id = unsafe { self.get_mock_id() };
        clear_id(id);
    }

    fn call_mock(&self, input: T) -> MockResult<T, O> {
        unsafe {
            let id = self.get_mock_id();
            let rc_opt = MOCK_STORE.with(|mock_ref_cell|
                mock_ref_cell.borrow()
                    .get(&id)
                    .cloned()
            );
            let stored_opt = rc_opt.as_ref()
                .and_then(|rc| rc.try_borrow_mut().ok());
            match stored_opt {
                Some(mut stored) => {
                    let real: &mut Box<FnMut<_, Output=_>> = transmute(&mut*stored);
                    real.call_mut(input)
                }
                None => MockResult::Continue(input),
            }
        }
    }

    unsafe fn get_mock_id(&self) -> TypeId {
        (||()).type_id()
    }
}

#[derive(Default)]
pub struct MockContext<'a> {
    planned_mocks: HashMap<TypeId, Box<FnOnce() -> ScopedMock<'a> + 'a>>,
}

impl<'a> MockContext<'a> {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn mock_safe<
        Args,
        Output,
        M: Mockable<Args, Output> + 'a,
        F: FnMut<Args, Output = MockResult<Args, Output>> + 'a,
    >(
        mut self,
        mock: M,
        body: F,
    ) -> Self {
        self.planned_mocks.insert(
            unsafe { mock.get_mock_id() },
            Box::new(move || unsafe { ScopedMock::new(&mock, body) }),
        );
        self
    }

    pub fn run<T, F: FnOnce() -> T>(self, f: F) -> T {
        let _scoped_mocks = self
            .planned_mocks
            .into_iter()
            .map(|entry| entry.1())
            .collect::<Vec<_>>();
        f()
    }
}
<|endoftext|>


<commit_before>#![allow(unused_imports)]

use std::io;
use std::process;

use libc;
use tokio_core::reactor::Handle;

use futures::{self, Future, Stream};

use plugin::Plugin;

#[cfg(unix)]
pub fn monitor_signal(handle: &Handle, plugins: Vec<Plugin>) {
    use tokio_signal::unix::Signal;

    // Monitor SIGCHLD, triggered if subprocess (plugin) is exited.
    let fut1 = Signal::new(libc::SIGCHLD, handle)
        .and_then(|signal| {
                      signal.take(1)
                            .for_each(|_| -> Result<(), io::Error> {
                                          error!("Plugin exited unexpectly (SIGCHLD)");
                                          Ok(())
                                      })
                            .map(|_| libc::SIGCHLD)
                  })
        .map_err(|err| {
                     error!("Failed to monitor SIGCHLD, err: {:?}", err);
                 });

    // Monitor SIGTERM, triggered if shadowsocks is exited gracefully. (Kill by user).
    let fut2 = Signal::new(libc::SIGTERM, handle)
        .and_then(|sigterm| {
                      sigterm.take(1)
                             .for_each(|_| -> Result<(), io::Error> {
                                           error!("Received SIGTERM, aborting process");
                                           Ok(())
                                       })
                             .map(|_| libc::SIGTERM)
                  })
        .map_err(|err| {
                     error!("Failed to monitor SIGTERM, err: {:?}", err);
                 });

    // Monitor SIGINT, triggered by CTRL-C
    let fut3 = Signal::new(libc::SIGINT, handle)
        .and_then(|sigint| {
                      sigint.take(1)
                            .for_each(|_| -> Result<(), io::Error> {
                                          error!("Received SIGINT, aborting process");
                                          Ok(())
                                      })
                            .map(|_| libc::SIGINT)
                  })
        .map_err(|err| {
                     error!("Failed to monitor SIGINT, err: {:?}", err);
                 });

    // Join them all, if any of them is triggered, kill all subprocesses and exit.
    let fut = fut1.select(fut2)
                  .map(|(sig, _)| sig)
                  .map_err(|(e, _)| e)
                  .select(fut3)
                  .map(|(sig, _)| sig)
                  .map_err(|(e, _)| e)
                  .then(|r| {
        // Something happened ... killing all subprocesses
        info!("Killing {} plugins and then exit", plugins.len());
        drop(plugins);

        match r {
            Ok(exit_code) => {
                process::exit(128 + exit_code);
            }
            Err(..) => Err(()),
        }
    });

    handle.spawn(fut);
}

#[cfg(windows)]
pub fn monitor_signal(handle: &Handle, plugins: Vec<Plugin>) {
    // FIXME: How to handle SIGTERM equavalent in Windows?

    use tokio_signal::windows::Event;

    let fut1 = Event::ctrl_c(handle)
        .and_then(|ev| {
                      ev.take(1).for_each(|_| -> Result<(), io::Error> {
                                              error!("Received Ctrl-C event");
                                              Ok(())
                                          })
                  })
        .map_err(|err| {
                     error!("Failed to monitor Ctrl-C event: {:?}", err);
                 });

    let fut2 = Event::ctrl_break(handle)
        .and_then(|ev| {
                      ev.take(1).for_each(|_| -> Result<(), io::Error> {
                                              error!("Received Ctrl-Break event");
                                              Ok(())
                                          })
                  })
        .map_err(|err| {
                     error!("Failed to monitor Ctrl-Break event: {:?}", err);
                 });

    let fut = fut1.select(fut2).then(|_| -> Result<(), ()> {
                                         drop(plugins);
                                         process::exit(libc::EXIT_FAILURE);
                                     });
    handle.spawn(fut);
}

#[cfg(not(any(windows, unix)))]
pub fn monitor_signal(handle: &Handle, plugins: Vec<Plugin>) {
    // FIXME: What can I do ...
    // Blocks forever
    let fut = futures::empty::<(), ()>().and_then(|_| { drop(plugins); });
    handle.spawn(fut);
}
<commit_msg>Lower log level<commit_after>#![allow(unused_imports)]

use std::io;
use std::process;

use libc;
use tokio_core::reactor::Handle;

use futures::{self, Future, Stream};

use plugin::Plugin;

#[cfg(unix)]
pub fn monitor_signal(handle: &Handle, plugins: Vec<Plugin>) {
    use tokio_signal::unix::Signal;

    // Monitor SIGCHLD, triggered if subprocess (plugin) is exited.
    let fut1 = Signal::new(libc::SIGCHLD, handle)
        .and_then(|signal| {
                      signal.take(1)
                            .for_each(|_| -> Result<(), io::Error> {
                                          error!("Plugin exited unexpectly (SIGCHLD)");
                                          Ok(())
                                      })
                            .map(|_| libc::SIGCHLD)
                  })
        .map_err(|err| {
                     error!("Failed to monitor SIGCHLD, err: {:?}", err);
                 });

    // Monitor SIGTERM, triggered if shadowsocks is exited gracefully. (Kill by user).
    let fut2 = Signal::new(libc::SIGTERM, handle)
        .and_then(|sigterm| {
                      sigterm.take(1)
                             .for_each(|_| -> Result<(), io::Error> {
                                           info!("Received SIGTERM, aborting process");
                                           Ok(())
                                       })
                             .map(|_| libc::SIGTERM)
                  })
        .map_err(|err| {
                     error!("Failed to monitor SIGTERM, err: {:?}", err);
                 });

    // Monitor SIGINT, triggered by CTRL-C
    let fut3 = Signal::new(libc::SIGINT, handle)
        .and_then(|sigint| {
                      sigint.take(1)
                            .for_each(|_| -> Result<(), io::Error> {
                                          error!("Received SIGINT, aborting process");
                                          Ok(())
                                      })
                            .map(|_| libc::SIGINT)
                  })
        .map_err(|err| {
                     error!("Failed to monitor SIGINT, err: {:?}", err);
                 });

    // Join them all, if any of them is triggered, kill all subprocesses and exit.
    let fut = fut1.select(fut2)
                  .map(|(sig, _)| sig)
                  .map_err(|(e, _)| e)
                  .select(fut3)
                  .map(|(sig, _)| sig)
                  .map_err(|(e, _)| e)
                  .then(|r| {
        // Something happened ... killing all subprocesses
        info!("Killing {} plugins and then exit", plugins.len());
        drop(plugins);

        match r {
            Ok(exit_code) => {
                process::exit(128 + exit_code);
            }
            Err(..) => Err(()),
        }
    });

    handle.spawn(fut);
}

#[cfg(windows)]
pub fn monitor_signal(handle: &Handle, plugins: Vec<Plugin>) {
    // FIXME: How to handle SIGTERM equavalent in Windows?

    use tokio_signal::windows::Event;

    let fut1 = Event::ctrl_c(handle)
        .and_then(|ev| {
                      ev.take(1).for_each(|_| -> Result<(), io::Error> {
                                              error!("Received Ctrl-C event");
                                              Ok(())
                                          })
                  })
        .map_err(|err| {
                     error!("Failed to monitor Ctrl-C event: {:?}", err);
                 });

    let fut2 = Event::ctrl_break(handle)
        .and_then(|ev| {
                      ev.take(1).for_each(|_| -> Result<(), io::Error> {
                                              error!("Received Ctrl-Break event");
                                              Ok(())
                                          })
                  })
        .map_err(|err| {
                     error!("Failed to monitor Ctrl-Break event: {:?}", err);
                 });

    let fut = fut1.select(fut2).then(|_| -> Result<(), ()> {
                                         // Something happened ... killing all subprocesses
                                         info!("Killing {} plugins and then exit", plugins.len());
                                         drop(plugins);
                                         process::exit(libc::EXIT_FAILURE);
                                     });
    handle.spawn(fut);
}

#[cfg(not(any(windows, unix)))]
pub fn monitor_signal(handle: &Handle, plugins: Vec<Plugin>) {
    // FIXME: What can I do ...
    // Blocks forever
    let fut = futures::empty::<(), ()>().and_then(|_| { drop(plugins); });
    handle.spawn(fut);
}
<|endoftext|>


<commit_before>use std::convert::From;
use std::io::{Read, Write};
use std::sync::mpsc::{Sender, channel};
use std::fmt;
use std::fmt::{Display, Formatter};
use std::error::Error;

use VERSION;

use bincode::serde::{DeserializeError, deserialize, serialize};
use bincode::SizeLimit;
use byteorder::{ByteOrder, LittleEndian, WriteBytesExt};
use mio::{EventLoop as MioEventLoop, EventSet, Handler as MioHandler, PollOpt, Token};
use mio::tcp::{Shutdown, TcpListener, TcpStream};
use mio::util::Slab;
use slab::Index;

pub mod client;
pub mod server;
#[cfg(Test)]
pub mod test;

/// The default port for a server to listen on.
pub const STANDARD_PORT: u16 = 25566;
/// Standard number to ensure network connections are syncronized and the same protocol is being used.
/// Reexported incase it is of use for something not-networking.
pub const NET_MAGIC_NUMBER: u32 = 0xCB011043; //0xcafebade + 0x25565, because programming references.
const MAX_CONNECTIONS: usize = 1024;

pub enum HandlerMessage {
    Send(NetworkPacket, Token),
    AddStream(TcpStream, Sender<Token>),
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum NetworkPacket {
    /// Sent on connection to verify everything is in sync.
    Init {
        version: String,
    },
    /// An error that should crash the game and show an error to the user, but only on a client.
    Error(NetworkError),
    /// Used to unit test networking. When recived, increments a value. TODO: Locate that value.
    #[cfg(test)]
    Test,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum NetworkError {
    VersionMismatch(String, String),
}

impl Display for NetworkError {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        match *self {
            NetworkError::VersionMismatch(ref ver1, ref ver2) => {
                write!(fmt,
                       "VersionMismatch: The versions of the client and server attempting to connect mismatch. ver1: {}, ver2: {}",
                       ver1,
                       ver2)
            }
        }
    }
}

impl Error for NetworkError {
    fn description(&self) -> &str {
        match *self {
            NetworkError::VersionMismatch(_, _) => "VersionMismatch: The versions of the client and server attempting to connect mismatch.",
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            NetworkError::VersionMismatch(_, _) => None,
        }
    }
}

pub type EventLoop = MioEventLoop<Handler>;
pub struct Handler {
    connections: Slab<Connection>,
    listener: Option<TcpListener>,
}

impl Handler {
    /// Creates a new instance of a Handler. Does not listen for connections.
    pub fn new() -> Self {
        Handler {
            connections: Slab::new_starting_at(Token::from_usize(1), MAX_CONNECTIONS),
            listener: None,
        }
    }

    /// Creates a new instance of a Handler. Does listen for connections.
    pub fn new_listener(addr: &SocketAddr) -> Self {
        Handler {
            connections: Slab::new_starting_at(Token::from_usize(1), MAX_CONNECTIONS),
            listener: TcpListener::bind(addr),
        }
    }
}

#[derive(Debug)]
struct Connection {
    message_queue: Vec<NetworkPacket>,
    stream: TcpStream,
}

impl Connection {
    fn new(stream: TcpStream) -> Connection {
        Connection {
            stream: stream,
            message_queue: Vec::new(),
        }
    }
}

impl MioHandler for Handler {
    type Timeout = ();
    type Message = HandlerMessage;
    fn ready(&mut self, event_loop: &mut MioEventLoop<Handler>, token: Token, events: EventSet) {
        if events.is_readable() {
            // Read the 6 byte header of each packet, throw it into get_packet_length, then read that number of bytes.
            // Then throw those bytes into deserialize_packet. Afterwards, throw it to handle_packet.

            let mut header = [0; 6];
            self.connections[token]
                .stream
                .read(&mut header)
                .expect(&format!("An error occured reading from socket {:?}", token));    // TODO: Figure out possible errors and take care of them.
            let length = get_packet_length(header).unwrap_or(0);
            if length == 0 {
                event_loop.deregister(self.connections[token].stream);
                self.connections[token].stream.shutdown(Shutdown::Both);
                self.connections.remove(token);
                // I directly kill the connection, becasue if the magic number doesn't match,
                // the peer probably doesn't share the same protocol. It wouldn't understand a normal error packet.
                return;
            }
            let mut packet = Vec::new();
            (&mut self.connections[token].stream).take(length as u64).read_to_end(&mut packet).unwrap();
            // I do this because Read.take takes a self, instead of a reasonable alternitive.
            // However &mut Read is it's self a Reader. So I use that instead.
            let dese = deserialize_packet(&packet).unwrap();
            handle_packet(dese, token, event_loop);
        }

        if events.is_writable() {
            // Get all the messages that should be send to that token, seralize all of them, then send all of them.
            // Afterwards, flush the buffer.
            let to_send = self.connections[token].message_queue.clone();
            self.connections[token].message_queue = Vec::new();
            if !to_send.is_empty() {
                for send in to_send {
                    let bytes = seralize_packet(&send);
                    self.connections[token]
                        .stream
                        .write(&bytes)
                        .expect(&format!("Error writing {:?} to connection {:?}/{:?}",
                                         send,
                                         token,
                                         self.connections[token]));
                }
                self.connections[token]
                    .stream
                    .flush()
                    .expect(&format!("Error flushing connection {:?}/{:?}",
                                     token,
                                     self.connections[token]));;
            }
        }

        if self.listener.is_some() {
            match self.listener.as_ref().unwrap().accept() {
                Ok(Some((socket, address))) => {
                    let token = self.connections.insert(Connection::new(socket)).unwrap();
                    send(event_loop,
                         NetworkPacket::Init { version: VERSION.to_owned() },
                         token);
                    info!("Accepted connection {}.", address);
                }
                Ok(None) => debug!("Recived Ok(None) in listener.accept()."),
                Err(error) => panic!("listener.accept() errored with {:?}!", error),
            }
        }
    }

    fn notify(&mut self, event_loop: &mut MioEventLoop<Handler>, msg: HandlerMessage) {
        match msg {
            HandlerMessage::AddStream(stream, tx) => {
                let token = self.connections.insert(Connection::new(stream)).unwrap();
                event_loop.register(&self.connections[token].stream,
                                    token,
                                    EventSet::all(),
                                    PollOpt::level())
                          .unwrap();
                tx.send(token).unwrap();
            }
            HandlerMessage::Send(packet, token) => {
                self.connections[token].message_queue.push(packet);
            }
        }
    }
}

fn add_socket(event_loop: &EventLoop, socket: TcpStream) -> Token {
    let (tx, rx) = channel();   // How expensive is this? Should I be creating a whole new channel for just 1 message?
    event_loop.channel().send(HandlerMessage::AddStream(socket, tx)).unwrap();
    rx.recv().unwrap() // I should feel bad.
}

/// Send a NetworkPacket over the network.
///
/// This is a static function instead of an impl function because you can't impl on external structs. In this case, mio::EventLoop.
pub fn send(event_loop: &EventLoop, to_send: NetworkPacket, token: Token) {
    event_loop.channel().send(HandlerMessage::Send(to_send, token)).unwrap();
}

fn seralize_packet(to_ser: &NetworkPacket) -> Vec<u8> {
    let mut result: Vec<u8> = Vec::new();
    result.write_u32::<LittleEndian>(NET_MAGIC_NUMBER).unwrap();   // No possible errors here.
    // The NET_MAGIC_NUMBER is used before every packet, so incase the stream is desynced for whatever reason, the game doesn't just read arbratrary data and crash badly.
    // Instead, it can either recover somehow, by disconnecting and reconnecting, or just erroring gracefully.
    let mut encoded = serialize(to_ser, SizeLimit::Infinite).unwrap();
    // Since the size limit is infinite and i'm not encoding to a stream, there is no error and I can safely unwrap();
    result.write_u16::<LittleEndian>(encoded.len() as u16).unwrap();
    result.append(&mut encoded);
    result
}

#[derive(Debug)]
enum PacketDeseError {
    MagicNumberMismatch,
    InvalidEncoding(DeserializeError),
}

impl From<DeserializeError> for PacketDeseError {
    fn from(err: DeserializeError) -> PacketDeseError {
        PacketDeseError::InvalidEncoding(err)
    }
}

/// Returns the length of a given packet, or a None if the first four bytes do not match NET_MAGIC_NUMBER.
fn get_packet_length(to_ln: [u8; 6]) -> Option<u16> {
    let (first_four, next_two) = to_ln.split_at(4);
    let should_be_magic_num = LittleEndian::read_u32(&first_four);
    if should_be_magic_num != NET_MAGIC_NUMBER {
        return None;
    }
    let length = LittleEndian::read_u16(&next_two);
    Some(length)
}

fn deserialize_packet(to_de: &[u8]) -> Result<NetworkPacket, PacketDeseError> {
    Ok(try!(deserialize::<NetworkPacket>(to_de)))
}

fn handle_packet(to_handle: NetworkPacket, sender: Token, event_loop: &EventLoop) {
    match to_handle {
        NetworkPacket::Init{version} => {
            if version != VERSION {
                send(event_loop,
                     NetworkPacket::Error(NetworkError::VersionMismatch(version.to_owned(), VERSION.to_owned())),
                     sender)
            }
        }

        #[cfg(Test)]
        NetworkPacket::Test => unimplemented!(),

        NetworkPacket::Error(error) => if ::check_should_crash() { panic!(error) } else { unimplemented!() },
    }
}
<commit_msg>Add TODO for error handling that requires testing<commit_after>use std::convert::From;
use std::io::{Read, Write};
use std::sync::mpsc::{Sender, channel};
use std::fmt;
use std::fmt::{Display, Formatter};
use std::error::Error;

use VERSION;

use bincode::serde::{DeserializeError, deserialize, serialize};
use bincode::SizeLimit;
use byteorder::{ByteOrder, LittleEndian, WriteBytesExt};
use mio::{EventLoop as MioEventLoop, EventSet, Handler as MioHandler, PollOpt, Token};
use mio::tcp::{Shutdown, TcpListener, TcpStream};
use mio::util::Slab;
use slab::Index;

pub mod client;
pub mod server;
#[cfg(Test)]
pub mod test;

/// The default port for a server to listen on.
pub const STANDARD_PORT: u16 = 25566;
/// Standard number to ensure network connections are syncronized and the same protocol is being used.
/// Reexported incase it is of use for something not-networking.
pub const NET_MAGIC_NUMBER: u32 = 0xCB011043; //0xcafebade + 0x25565, because programming references.
const MAX_CONNECTIONS: usize = 1024;

pub enum HandlerMessage {
    Send(NetworkPacket, Token),
    AddStream(TcpStream, Sender<Token>),
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum NetworkPacket {
    /// Sent on connection to verify everything is in sync.
    Init {
        version: String,
    },
    /// An error that should crash the game and show an error to the user, but only on a client.
    Error(NetworkError),
    /// Used to unit test networking. When recived, increments a value. TODO: Locate that value.
    #[cfg(test)]
    Test,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum NetworkError {
    VersionMismatch(String, String),
}

impl Display for NetworkError {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        match *self {
            NetworkError::VersionMismatch(ref ver1, ref ver2) => {
                write!(fmt,
                       "VersionMismatch: The versions of the client and server attempting to connect mismatch. ver1: {}, ver2: {}",
                       ver1,
                       ver2)
            }
        }
    }
}

impl Error for NetworkError {
    fn description(&self) -> &str {
        match *self {
            NetworkError::VersionMismatch(_, _) => "VersionMismatch: The versions of the client and server attempting to connect mismatch.",
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            NetworkError::VersionMismatch(_, _) => None,
        }
    }
}

pub type EventLoop = MioEventLoop<Handler>;
pub struct Handler {
    connections: Slab<Connection>,
    listener: Option<TcpListener>,
}

impl Handler {
    /// Creates a new instance of a Handler. Does not listen for connections.
    pub fn new() -> Self {
        Handler {
            connections: Slab::new_starting_at(Token::from_usize(1), MAX_CONNECTIONS),
            listener: None,
        }
    }

    /// Creates a new instance of a Handler. Does listen for connections.
    pub fn new_listener(addr: &SocketAddr) -> Self {
        Handler {
            connections: Slab::new_starting_at(Token::from_usize(1), MAX_CONNECTIONS),
            listener: TcpListener::bind(addr),
        }
    }
}

#[derive(Debug)]
struct Connection {
    message_queue: Vec<NetworkPacket>,
    stream: TcpStream,
}

impl Connection {
    fn new(stream: TcpStream) -> Connection {
        Connection {
            stream: stream,
            message_queue: Vec::new(),
        }
    }
}

impl MioHandler for Handler {
    type Timeout = ();
    type Message = HandlerMessage;
    fn ready(&mut self, event_loop: &mut MioEventLoop<Handler>, token: Token, events: EventSet) {
        if events.is_readable() {
            // Read the 6 byte header of each packet, throw it into get_packet_length, then read that number of bytes.
            // Then throw those bytes into deserialize_packet. Afterwards, throw it to handle_packet.

            let mut header = [0; 6];
            self.connections[token]
                .stream
                .read(&mut header)
                .expect(&format!("An error occured reading from socket {:?}", token));    // TODO: Figure out possible errors and take care of them.
            let length = get_packet_length(header).unwrap_or(0);
            if length == 0 {
                event_loop.deregister(self.connections[token].stream);
                self.connections[token].stream.shutdown(Shutdown::Both);
                self.connections.remove(token);
                // I directly kill the connection, becasue if the magic number doesn't match,
                // the peer probably doesn't share the same protocol. It wouldn't understand a normal error packet.
                return;
            }
            let mut packet = Vec::new();
            (&mut self.connections[token].stream).take(length as u64).read_to_end(&mut packet).unwrap();
            // I do this because Read.take takes a self, instead of a reasonable alternitive.
            // However &mut Read is it's self a Reader. So I use that instead.
            let dese = deserialize_packet(&packet).unwrap();
            handle_packet(dese, token, event_loop);
        }

        if events.is_writable() {
            // Get all the messages that should be send to that token, seralize all of them, then send all of them.
            // Afterwards, flush the buffer.
            let to_send = self.connections[token].message_queue.clone();
            self.connections[token].message_queue = Vec::new();
            if !to_send.is_empty() {
                for send in to_send {
                    let bytes = seralize_packet(&send);
                    self.connections[token]
                        .stream
                        .write(&bytes)
                        .expect(&format!("Error writing {:?} to connection {:?}/{:?}",
                                         send,
                                         token,
                                         self.connections[token]));
                }
                self.connections[token]
                    .stream
                    .flush()
                    .expect(&format!("Error flushing connection {:?}/{:?}",
                                     token,
                                     self.connections[token]));;
            }
        }

        if self.listener.is_some() {
            match self.listener.as_ref().unwrap().accept() {
                Ok(Some((socket, address))) => {
                    let token = self.connections.insert(Connection::new(socket)).unwrap();
                    send(event_loop,
                         NetworkPacket::Init { version: VERSION.to_owned() },
                         token);
                    info!("Accepted connection {}.", address);
                }
                Ok(None) => debug!("Recived Ok(None) in listener.accept()."),
                Err(error) => panic!("listener.accept() errored with {:?}!", error),
                // TODO: See if any of these errors can be handled better.
            }
        }
    }

    fn notify(&mut self, event_loop: &mut MioEventLoop<Handler>, msg: HandlerMessage) {
        match msg {
            HandlerMessage::AddStream(stream, tx) => {
                let token = self.connections.insert(Connection::new(stream)).unwrap();
                event_loop.register(&self.connections[token].stream,
                                    token,
                                    EventSet::all(),
                                    PollOpt::level())
                          .unwrap();
                tx.send(token).unwrap();
            }
            HandlerMessage::Send(packet, token) => {
                self.connections[token].message_queue.push(packet);
            }
        }
    }
}

fn add_socket(event_loop: &EventLoop, socket: TcpStream) -> Token {
    let (tx, rx) = channel();   // How expensive is this? Should I be creating a whole new channel for just 1 message?
    event_loop.channel().send(HandlerMessage::AddStream(socket, tx)).unwrap();
    rx.recv().unwrap() // I should feel bad.
}

/// Send a NetworkPacket over the network.
///
/// This is a static function instead of an impl function because you can't impl on external structs. In this case, mio::EventLoop.
pub fn send(event_loop: &EventLoop, to_send: NetworkPacket, token: Token) {
    event_loop.channel().send(HandlerMessage::Send(to_send, token)).unwrap();
}

fn seralize_packet(to_ser: &NetworkPacket) -> Vec<u8> {
    let mut result: Vec<u8> = Vec::new();
    result.write_u32::<LittleEndian>(NET_MAGIC_NUMBER).unwrap();   // No possible errors here.
    // The NET_MAGIC_NUMBER is used before every packet, so incase the stream is desynced for whatever reason, the game doesn't just read arbratrary data and crash badly.
    // Instead, it can either recover somehow, by disconnecting and reconnecting, or just erroring gracefully.
    let mut encoded = serialize(to_ser, SizeLimit::Infinite).unwrap();
    // Since the size limit is infinite and i'm not encoding to a stream, there is no error and I can safely unwrap();
    result.write_u16::<LittleEndian>(encoded.len() as u16).unwrap();
    result.append(&mut encoded);
    result
}

#[derive(Debug)]
enum PacketDeseError {
    MagicNumberMismatch,
    InvalidEncoding(DeserializeError),
}

impl From<DeserializeError> for PacketDeseError {
    fn from(err: DeserializeError) -> PacketDeseError {
        PacketDeseError::InvalidEncoding(err)
    }
}

/// Returns the length of a given packet, or a None if the first four bytes do not match NET_MAGIC_NUMBER.
fn get_packet_length(to_ln: [u8; 6]) -> Option<u16> {
    let (first_four, next_two) = to_ln.split_at(4);
    let should_be_magic_num = LittleEndian::read_u32(&first_four);
    if should_be_magic_num != NET_MAGIC_NUMBER {
        return None;
    }
    let length = LittleEndian::read_u16(&next_two);
    Some(length)
}

fn deserialize_packet(to_de: &[u8]) -> Result<NetworkPacket, PacketDeseError> {
    Ok(try!(deserialize::<NetworkPacket>(to_de)))
}

fn handle_packet(to_handle: NetworkPacket, sender: Token, event_loop: &EventLoop) {
    match to_handle {
        NetworkPacket::Init{version} => {
            if version != VERSION {
                send(event_loop,
                     NetworkPacket::Error(NetworkError::VersionMismatch(version.to_owned(), VERSION.to_owned())),
                     sender)
            }
        }

        #[cfg(Test)]
        NetworkPacket::Test => unimplemented!(),

        NetworkPacket::Error(error) => if ::check_should_crash() { panic!(error) } else { unimplemented!() },
    }
}
<|endoftext|>


<commit_before>use address::Address;
use messages::outgoing;
use mio;
use node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use transaction::{TransactionId, TransactionIdGenerator};

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

#[derive(Debug)]
pub enum Event {
    ReceivedPacket(Address, Vec<u8>)
}

pub enum OneshotTask {
    Incoming(Vec<u8>),
    StartBootstrap,
    SendPacket(Address, Vec<u8>)
}

enum TableAction {
    Bootstrap(mio::Timeout),
    HealthCheck(mio::Timeout),
    RefreshBucket(mio::Timeout)
}

pub type TaskSender = mio::Sender<OneshotTask>;

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Box<node::Node + 'static>,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>
}

impl Network {
    pub fn new<N: node::Node + 'static, T: ToSocketAddrs>(self_node: N, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let self_address = self_node.address().clone();
        let routing_table = RoutingTable::new(8, self_address, routers);

        Network {
            host: host,
            routing_table: routing_table,
            self_node: Box::new(self_node),
            transaction_ids: TransactionIdGenerator::new(),
            status: Status::Idle,
            pending_actions: HashMap::new(),
            event_listeners: vec![]
        }
    }

    pub fn run(self) -> TaskSender {
        let mut event_loop_config = mio::EventLoopConfig::new();
        event_loop_config.notify_capacity(16384);
        let mut event_loop = mio::EventLoop::configured(event_loop_config).unwrap();

        create_incoming_udp_channel(self.host, event_loop.channel());
        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

    pub fn register_event_listener(&mut self, event_listener: mpsc::Sender<Event>) {
        self.event_listeners.push(event_listener);
    }

    fn handle_incoming(&mut self, data: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        use messages::incoming::{Message, Query, Response, self};
        let mut data = Cursor::new(data);
        let message = incoming::parse_from_reader(&mut data).unwrap();

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.address();
                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_to(&target, false));
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        {
                            if let Some(origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_query(transaction_id);
                            }
                        }
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(origin_address, payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        if let Some(origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_query(transaction_id);
                        }
                    }
                }
            }
            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                match response {
                    Response::FindNode(mut nodes) => {
                        let mut encounted_new_node = false;
                        let mut tail = vec![origin];
                        for node in nodes.drain(..).chain(tail.drain(..)) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

                        {
                            if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_response(transaction_id);
                            }
                        }

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

                        match self.status {
                            Status::Bootstrapping => {
                                self.status = Status::Idle;
                            }
                            _ => { }
                        }
                    }
                    Response::Packet => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                    Response::Ping => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                }
            }
        }
    }

    fn start_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        self.status = Status::Bootstrapping;
        self.continue_bootstrap(event_loop);
    }

    fn continue_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let address = &self.self_node.address();
        let transaction_id = self.find_node(address);
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueBootstrap, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::Bootstrap(timeout));
    }

    fn continue_health_check(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let transaction_id = self.health_check();
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueHealthCheck, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::HealthCheck(timeout));
    }

    fn continue_refresh(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueRefresh, 1000).unwrap();

        match self.refresh_bucket() {
            Some(transaction_id) => {
                self.pending_actions.insert(transaction_id, TableAction::RefreshBucket(timeout));
            }
            None => { }
        }
    }

    fn find_node(&mut self, address: &Address) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        let query = outgoing::create_find_node_query(
            transaction_id,
            &self.self_node,
            address);
        for node in self.routing_table.nearest() {
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn health_check(&mut self) -> TransactionId {
        self.routing_table.remove_bad_nodes();
        let transaction_id = self.transaction_ids.generate();
        if let Some(node) = self.routing_table.questionable_nodes().get_mut(0) {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.send(query);
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn refresh_bucket(&mut self) -> Option<TransactionId> {
        if let Some(address) = self.address_to_find_for_refresh() {
            Some(self.find_node(&address))
        } else {
            None
        }
    }

    fn address_to_find_for_refresh(&self) -> Option<Address> {
        match self.routing_table.bucket_needing_refresh() {
            Some(bucket) => {
                Some(bucket.random_address_in_space())
            }
            None => None
        }
    }

    fn insert_node(&mut self, node: Box<node::Node>) -> InsertionResult {
        self.routing_table.insert(node, &self.self_node, &mut self.transaction_ids)
    }

    fn send_packet(&mut self, recipient: Address, payload: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        for node in self.routing_table.nearest_to(&recipient, false) {
            let transaction_id = self.transaction_ids.generate();
            let query = outgoing::create_packet_query(
                transaction_id, &self.self_node, payload.clone());
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
    }
}

struct Handler {
    network: Network
}

impl Handler {
    fn new(network: Network) -> Handler {
        Handler {
            network: network
        }
    }
}

impl mio::Handler for Handler {
    type Timeout = ScheduledTask;
    type Message = OneshotTask;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Handler>, task: OneshotTask) {
        match task {
            OneshotTask::Incoming(data) => self.network.handle_incoming(data, event_loop),
            OneshotTask::StartBootstrap => self.network.start_bootstrap(event_loop),
            OneshotTask::SendPacket(recipient, payload) => {
                self.network.send_packet(recipient, payload, event_loop);
            }
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Handler>, timeout: ScheduledTask) {
        match timeout {
            ScheduledTask::ContinueBootstrap => self.network.continue_bootstrap(event_loop),
            ScheduledTask::ContinueHealthCheck => self.network.continue_health_check(event_loop),
            ScheduledTask::ContinueRefresh => self.network.continue_refresh(event_loop)
        }
    }
}

fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) {
    thread::spawn(move || {
        let host = ("0.0.0.0", host.port());
        let socket = UdpSocket::bind(host).unwrap();
        loop {
            let mut buf = [0; 4096];
            match socket.recv_from(&mut buf) {
                Ok((size, _src)) => {
                    sender
                        .send(OneshotTask::Incoming(buf[..size].iter().cloned().collect()))
                        .unwrap_or_else(|err| info!("Couldn't handling incoming: {:?}", err));
                }
                Err(e) => panic!("Error receiving from server: {}", e)
            }
        }
    });
}
<commit_msg>Make match branch syntax consistent<commit_after>use address::Address;
use messages::outgoing;
use mio;
use node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use transaction::{TransactionId, TransactionIdGenerator};

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

#[derive(Debug)]
pub enum Event {
    ReceivedPacket(Address, Vec<u8>)
}

pub enum OneshotTask {
    Incoming(Vec<u8>),
    StartBootstrap,
    SendPacket(Address, Vec<u8>)
}

enum TableAction {
    Bootstrap(mio::Timeout),
    HealthCheck(mio::Timeout),
    RefreshBucket(mio::Timeout)
}

pub type TaskSender = mio::Sender<OneshotTask>;

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Box<node::Node + 'static>,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>
}

impl Network {
    pub fn new<N: node::Node + 'static, T: ToSocketAddrs>(self_node: N, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let self_address = self_node.address().clone();
        let routing_table = RoutingTable::new(8, self_address, routers);

        Network {
            host: host,
            routing_table: routing_table,
            self_node: Box::new(self_node),
            transaction_ids: TransactionIdGenerator::new(),
            status: Status::Idle,
            pending_actions: HashMap::new(),
            event_listeners: vec![]
        }
    }

    pub fn run(self) -> TaskSender {
        let mut event_loop_config = mio::EventLoopConfig::new();
        event_loop_config.notify_capacity(16384);
        let mut event_loop = mio::EventLoop::configured(event_loop_config).unwrap();

        create_incoming_udp_channel(self.host, event_loop.channel());
        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

    pub fn register_event_listener(&mut self, event_listener: mpsc::Sender<Event>) {
        self.event_listeners.push(event_listener);
    }

    fn handle_incoming(&mut self, data: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        use messages::incoming::{Message, Query, Response, self};
        let mut data = Cursor::new(data);
        let message = incoming::parse_from_reader(&mut data).unwrap();

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.address();
                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_to(&target, false));
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        {
                            if let Some(origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_query(transaction_id);
                            }
                        }
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(origin_address, payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        if let Some(origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_query(transaction_id);
                        }
                    }
                }
            }
            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                match response {
                    Response::FindNode(mut nodes) => {
                        let mut encounted_new_node = false;
                        let mut tail = vec![origin];
                        for node in nodes.drain(..).chain(tail.drain(..)) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

                        {
                            if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_response(transaction_id);
                            }
                        }

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

                        match self.status {
                            Status::Bootstrapping => {
                                self.status = Status::Idle;
                            }
                            _ => { }
                        }
                    }
                    Response::Packet => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                    Response::Ping => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                }
            }
        }
    }

    fn start_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        self.status = Status::Bootstrapping;
        self.continue_bootstrap(event_loop);
    }

    fn continue_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let address = &self.self_node.address();
        let transaction_id = self.find_node(address);
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueBootstrap, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::Bootstrap(timeout));
    }

    fn continue_health_check(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let transaction_id = self.health_check();
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueHealthCheck, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::HealthCheck(timeout));
    }

    fn continue_refresh(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueRefresh, 1000).unwrap();

        match self.refresh_bucket() {
            Some(transaction_id) => {
                self.pending_actions.insert(transaction_id, TableAction::RefreshBucket(timeout));
            }
            None => { }
        }
    }

    fn find_node(&mut self, address: &Address) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        let query = outgoing::create_find_node_query(
            transaction_id,
            &self.self_node,
            address);
        for node in self.routing_table.nearest() {
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn health_check(&mut self) -> TransactionId {
        self.routing_table.remove_bad_nodes();
        let transaction_id = self.transaction_ids.generate();
        if let Some(node) = self.routing_table.questionable_nodes().get_mut(0) {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.send(query);
            node.sent_query(transaction_id);
        }
        transaction_id
    }

    fn refresh_bucket(&mut self) -> Option<TransactionId> {
        if let Some(address) = self.address_to_find_for_refresh() {
            Some(self.find_node(&address))
        } else {
            None
        }
    }

    fn address_to_find_for_refresh(&self) -> Option<Address> {
        match self.routing_table.bucket_needing_refresh() {
            Some(bucket) => {
                Some(bucket.random_address_in_space())
            }
            None => None
        }
    }

    fn insert_node(&mut self, node: Box<node::Node>) -> InsertionResult {
        self.routing_table.insert(node, &self.self_node, &mut self.transaction_ids)
    }

    fn send_packet(&mut self, recipient: Address, payload: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        for node in self.routing_table.nearest_to(&recipient, false) {
            let transaction_id = self.transaction_ids.generate();
            let query = outgoing::create_packet_query(
                transaction_id, &self.self_node, payload.clone());
            node.send(query.clone());
            node.sent_query(transaction_id);
        }
    }
}

struct Handler {
    network: Network
}

impl Handler {
    fn new(network: Network) -> Handler {
        Handler {
            network: network
        }
    }
}

impl mio::Handler for Handler {
    type Timeout = ScheduledTask;
    type Message = OneshotTask;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Handler>, task: OneshotTask) {
        match task {
            OneshotTask::Incoming(data) => self.network.handle_incoming(data, event_loop),
            OneshotTask::StartBootstrap => self.network.start_bootstrap(event_loop),
            OneshotTask::SendPacket(recipient, payload) =>
                self.network.send_packet(recipient, payload, event_loop)
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Handler>, timeout: ScheduledTask) {
        match timeout {
            ScheduledTask::ContinueBootstrap => self.network.continue_bootstrap(event_loop),
            ScheduledTask::ContinueHealthCheck => self.network.continue_health_check(event_loop),
            ScheduledTask::ContinueRefresh => self.network.continue_refresh(event_loop)
        }
    }
}

fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) {
    thread::spawn(move || {
        let host = ("0.0.0.0", host.port());
        let socket = UdpSocket::bind(host).unwrap();
        loop {
            let mut buf = [0; 4096];
            match socket.recv_from(&mut buf) {
                Ok((size, _src)) => {
                    sender
                        .send(OneshotTask::Incoming(buf[..size].iter().cloned().collect()))
                        .unwrap_or_else(|err| info!("Couldn't handling incoming: {:?}", err));
                }
                Err(e) => panic!("Error receiving from server: {}", e)
            }
        }
    });
}
<|endoftext|>