<commit_before>use hir::db::HirDatabase;
use ra_syntax::{ast::AstNode, ast::Literal, SyntaxText, TextRange, TextUnit};

use crate::{assist_ctx::AssistBuilder, Assist, AssistCtx, AssistId};

pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() == ra_syntax::SyntaxKind::STRING {
        ctx.add_action(AssistId("make_raw_string"), "make raw string", |edit| {
            edit.target(literal.syntax().text_range());
            edit.insert(literal.syntax().text_range().start(), "r");
        });
        ctx.build()
    } else {
        None
    }
}

pub(crate) fn make_usual_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() == ra_syntax::SyntaxKind::RAW_STRING {
        ctx.add_action(AssistId("make_usual_string"), "make usual string", |edit| {
            let text = literal.syntax().text();
            let usual_start_pos = text.find_char('"').unwrap(); // we have a RAW_STRING
            let end = literal.syntax().text_range().end();
            let mut i = 0;
            let mut pos = 0;
            let mut c = text.char_at(end - TextUnit::from(i));
            while c != Some('"') {
                if c != None {
                    pos += 1;
                }
                i += 1;
                c = text.char_at(end - TextUnit::from(i));
            }

            edit.target(literal.syntax().text_range());
            edit.delete(TextRange::from_to(
                literal.syntax().text_range().start(),
                literal.syntax().text_range().start() + usual_start_pos,
            ));
            edit.delete(TextRange::from_to(
                literal.syntax().text_range().end() - TextUnit::from(pos),
                literal.syntax().text_range().end(),
            ));
            // parse inside string to escape `"`
            let start_of_inside = usual_start_pos + TextUnit::from(1);
            let end_of_inside = text.len() - usual_start_pos - TextUnit::from(1);
            let inside_str = text.slice(TextRange::from_to(start_of_inside, end_of_inside));
            escape_quote(
                edit,
                &inside_str,
                literal.syntax().text_range().start() + start_of_inside,
            );
        });
        ctx.build()
    } else {
        None
    }
}

pub(crate) fn add_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() == ra_syntax::SyntaxKind::RAW_STRING {
        ctx.add_action(AssistId("add_hash"), "add hash to raw string", |edit| {
            edit.target(literal.syntax().text_range());
            edit.insert(literal.syntax().text_range().start() + TextUnit::from(1), "#");
            edit.insert(literal.syntax().text_range().end(), "#");
        });
        ctx.build()
    } else {
        None
    }
}

fn escape_quote(edit: &mut AssistBuilder, inside_str: &SyntaxText, offset: TextUnit) {
    let mut start = TextUnit::from(0);
    inside_str.for_each_chunk(|chunk| {
        let end = start + TextUnit::of_str(chunk);
        let mut i = 0;
        for c in chunk.to_string().chars() {
            if c == '"' {
                edit.insert(offset + start + TextUnit::from(i), "\\");
            }
            i += 1;
        }
        start = end;
    });
}

pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() == ra_syntax::SyntaxKind::RAW_STRING {
        if !literal.syntax().text().contains_char('#') {
            return None;
        }
        ctx.add_action(AssistId("remove_hash"), "remove hash from raw string", |edit| {
            edit.target(literal.syntax().text_range());
            edit.delete(TextRange::from_to(
                literal.syntax().text_range().start() + TextUnit::from(1),
                literal.syntax().text_range().start() + TextUnit::from(2),
            ));
            edit.delete(TextRange::from_to(
                literal.syntax().text_range().end() - TextUnit::from(1),
                literal.syntax().text_range().end(),
            ));
            let text = literal.syntax().text();
            if text.char_at(TextUnit::from(2)) == Some('"') {
                // no more hash after assist, need to escape any `"` in the string
                let inside_str = text
                    .slice(TextRange::from_to(TextUnit::from(3), text.len() - TextUnit::from(2)));
                escape_quote(
                    edit,
                    &inside_str,
                    literal.syntax().text_range().start() + TextUnit::from(3),
                );
            }
        });
        ctx.build()
    } else {
        None
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn make_raw_string_target() {
        check_assist_target(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r#""random string""#,
        );
    }

    #[test]
    fn make_raw_string_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_raw_string_with_escaped_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r#"
            fn f() {
                let s = <|>r"random\nstring";
            }
            "#,
        )
    }

    #[test]
    fn make_raw_string_not_works() {
        check_assist_not_applicable(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn add_hash_target() {
        check_assist_target(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r#"r"random string""#,
        );
    }

    #[test]
    fn add_hash_works() {
        check_assist(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_more_hash_works() {
        check_assist(
            add_hash,
            r##"
            fn f() {
                let s = <|>r#"random"string"#;
            }
            "##,
            r###"
            fn f() {
                let s = <|>r##"random"string"##;
            }
            "###,
        )
    }

    #[test]
    fn add_hash_not_works() {
        check_assist_not_applicable(
            add_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_target() {
        check_assist_target(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn remove_hash_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn remove_hash_with_quote_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn remove_more_hash_works() {
        check_assist(
            remove_hash,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn remove_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_no_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn make_usual_string_target() {
        check_assist_target(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn make_usual_string_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_with_quote_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_more_hash_works() {
        check_assist(
            make_usual_string,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>"random string";
            }
            "##,
        )
    }

    #[test]
    fn make_usual_string_not_works() {
        check_assist_not_applicable(
            make_usual_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }
}
<commit_msg>feat(assists): Rename `escape_quote` to `escape_double_quote`<commit_after>use hir::db::HirDatabase;
use ra_syntax::{ast::AstNode, ast::Literal, SyntaxText, TextRange, TextUnit};

use crate::{assist_ctx::AssistBuilder, Assist, AssistCtx, AssistId};

pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() == ra_syntax::SyntaxKind::STRING {
        ctx.add_action(AssistId("make_raw_string"), "make raw string", |edit| {
            edit.target(literal.syntax().text_range());
            edit.insert(literal.syntax().text_range().start(), "r");
        });
        ctx.build()
    } else {
        None
    }
}

pub(crate) fn make_usual_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() == ra_syntax::SyntaxKind::RAW_STRING {
        ctx.add_action(AssistId("make_usual_string"), "make usual string", |edit| {
            let text = literal.syntax().text();
            let usual_start_pos = text.find_char('"').unwrap(); // we have a RAW_STRING
            let end = literal.syntax().text_range().end();
            let mut i = 0;
            let mut pos = 0;
            let mut c = text.char_at(end - TextUnit::from(i));
            while c != Some('"') {
                if c != None {
                    pos += 1;
                }
                i += 1;
                c = text.char_at(end - TextUnit::from(i));
            }

            edit.target(literal.syntax().text_range());
            edit.delete(TextRange::from_to(
                literal.syntax().text_range().start(),
                literal.syntax().text_range().start() + usual_start_pos,
            ));
            edit.delete(TextRange::from_to(
                literal.syntax().text_range().end() - TextUnit::from(pos),
                literal.syntax().text_range().end(),
            ));
            // parse inside string to escape `"`
            let start_of_inside = usual_start_pos + TextUnit::from(1);
            let end_of_inside = text.len() - usual_start_pos - TextUnit::from(1);
            let inside_str = text.slice(TextRange::from_to(start_of_inside, end_of_inside));
            escape_double_quote(
                edit,
                &inside_str,
                literal.syntax().text_range().start() + start_of_inside,
            );
        });
        ctx.build()
    } else {
        None
    }
}

pub(crate) fn add_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() == ra_syntax::SyntaxKind::RAW_STRING {
        ctx.add_action(AssistId("add_hash"), "add hash to raw string", |edit| {
            edit.target(literal.syntax().text_range());
            edit.insert(literal.syntax().text_range().start() + TextUnit::from(1), "#");
            edit.insert(literal.syntax().text_range().end(), "#");
        });
        ctx.build()
    } else {
        None
    }
}

fn escape_double_quote(edit: &mut AssistBuilder, inside_str: &SyntaxText, offset: TextUnit) {
    let mut start = TextUnit::from(0);
    inside_str.for_each_chunk(|chunk| {
        let end = start + TextUnit::of_str(chunk);
        let mut i = 0;
        for c in chunk.to_string().chars() {
            if c == '"' {
                edit.insert(offset + start + TextUnit::from(i), "\\");
            }
            i += 1;
        }
        start = end;
    });
}

pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let literal = ctx.node_at_offset::<Literal>()?;
    if literal.token().kind() == ra_syntax::SyntaxKind::RAW_STRING {
        if !literal.syntax().text().contains_char('#') {
            return None;
        }
        ctx.add_action(AssistId("remove_hash"), "remove hash from raw string", |edit| {
            edit.target(literal.syntax().text_range());
            edit.delete(TextRange::from_to(
                literal.syntax().text_range().start() + TextUnit::from(1),
                literal.syntax().text_range().start() + TextUnit::from(2),
            ));
            edit.delete(TextRange::from_to(
                literal.syntax().text_range().end() - TextUnit::from(1),
                literal.syntax().text_range().end(),
            ));
            let text = literal.syntax().text();
            if text.char_at(TextUnit::from(2)) == Some('"') {
                // no more hash after assist, need to escape any `"` in the string
                let inside_str = text
                    .slice(TextRange::from_to(TextUnit::from(3), text.len() - TextUnit::from(2)));
                escape_double_quote(
                    edit,
                    &inside_str,
                    literal.syntax().text_range().start() + TextUnit::from(3),
                );
            }
        });
        ctx.build()
    } else {
        None
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn make_raw_string_target() {
        check_assist_target(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r#""random string""#,
        );
    }

    #[test]
    fn make_raw_string_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_raw_string_with_escaped_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r#"
            fn f() {
                let s = <|>r"random\nstring";
            }
            "#,
        )
    }

    #[test]
    fn make_raw_string_not_works() {
        check_assist_not_applicable(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn add_hash_target() {
        check_assist_target(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r#"r"random string""#,
        );
    }

    #[test]
    fn add_hash_works() {
        check_assist(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_more_hash_works() {
        check_assist(
            add_hash,
            r##"
            fn f() {
                let s = <|>r#"random"string"#;
            }
            "##,
            r###"
            fn f() {
                let s = <|>r##"random"string"##;
            }
            "###,
        )
    }

    #[test]
    fn add_hash_not_works() {
        check_assist_not_applicable(
            add_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_target() {
        check_assist_target(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn remove_hash_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn remove_hash_with_quote_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn remove_more_hash_works() {
        check_assist(
            remove_hash,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn remove_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_no_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn make_usual_string_target() {
        check_assist_target(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn make_usual_string_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_with_quote_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_more_hash_works() {
        check_assist(
            make_usual_string,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>"random string";
            }
            "##,
        )
    }

    #[test]
    fn make_usual_string_not_works() {
        check_assist_not_applicable(
            make_usual_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }
}
<|endoftext|>


<commit_before>use std::io::net::ip::{Ipv4Addr, SocketAddr};

mod libtorresmo {
    extern crate time;

    use std::io::net::udp::UdpSocket;
    use std::io::net::ip::SocketAddr;
    use std::io::IoResult;
    use std::mem::transmute;
    use std::rand::random;
    use std::fmt;

    #[allow(dead_code,non_camel_case_types)]
    enum UtpPacketType {
        ST_DATA  = 0,
        ST_FIN   = 1,
        ST_STATE = 2,
        ST_RESET = 3,
        ST_SYN   = 4,
    }

    #[allow(dead_code)]
    #[packed]
    struct UtpPacketHeader {
        type_ver: u8, // type: u4, ver: u4
        extension: u8,
        connection_id: u16,
        timestamp_microseconds: u32,
        timestamp_difference_microseconds: u32,
        wnd_size: u32,
        seq_nr: u16,
        ack_nr: u16,
    }

    impl UtpPacketHeader {
        fn bytes(&self) -> &[u8] {
            unsafe {
                let buf: &[u8, ..20] = transmute(self);
                return buf.as_slice();
            }
        }

        fn len(&self) -> uint {
            static HEADER_SIZE: uint = 20;
            return HEADER_SIZE;
        }
    }

    impl fmt::Show for UtpPacketHeader {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            let t: UtpPacketType = unsafe { transmute(self.type_ver >> 4) };
            let ptype = match t {
                ST_DATA  => "ST_DATA",
                ST_FIN   => "ST_FIN",
                ST_STATE => "ST_STATE",
                ST_RESET => "ST_RESET",
                ST_SYN   => "ST_SYN",
            };
            write!(f, "(type: {}, version: {}, extension: {}, connection_id: {}, timestamp_microseconds: {}, timestamp_difference_microseconds: {}, wnd_size: {}, seq_nr: {}, ack_nr: {})", ptype, self.type_ver & 0x0F, self.extension, self.connection_id, self.timestamp_microseconds, self.timestamp_difference_microseconds, self.wnd_size, self.seq_nr, self.ack_nr)
        }
    }

    #[allow(dead_code)]
    struct UtpPacket {
        header: UtpPacketHeader,
        payload: Vec<u8>,
    }

    impl UtpPacket {
        /// Constructs a new, empty UtpPacket.
        fn new() -> UtpPacket {
            UtpPacket {
                header: UtpPacketHeader {
                    type_ver: ST_DATA as u8 << 4 | 1,
                    extension: 0,
                    connection_id: 0,
                    timestamp_microseconds: 0,
                    timestamp_difference_microseconds: 0,
                    wnd_size: 0,
                    seq_nr: 0,
                    ack_nr: 0,
                },
                payload: Vec::new(),
            }
        }

        fn set_type(&mut self, t: UtpPacketType) {
            let version = 0x0F & self.header.type_ver;
            self.header.type_ver = t as u8 << 4 | version;
        }

        fn bytes(&self) -> Vec<u8> {
            let mut buf: Vec<u8> = Vec::with_capacity(self.len());
            buf.push_all(self.header.bytes());
            buf.push_all(self.payload.as_slice());
            return buf;
        }

        fn len(&self) -> uint {
            self.header.len() + self.payload.len()
        }

        fn decode(buf: &[u8]) -> UtpPacket {
            let header: UtpPacketHeader = UtpPacketHeader {
                type_ver: buf[0],
                extension: buf[1],
                connection_id: buf[2] as u16 << 8 | buf[3] as u16,
                timestamp_microseconds: buf[4] as u32 << 24 | buf[5] as u32 << 16 | buf[6] as u32 << 8 | buf[7] as u32,
                timestamp_difference_microseconds: buf[8] as u32 << 24 | buf[9] as u32 << 16 | buf[10] as u32 << 8 | buf[11] as u32,
                wnd_size: buf[12] as u32 << 24 | buf[13] as u32 << 16 | buf[14] as u32 << 8 | buf[15] as u32,
                seq_nr: buf[16] as u16 << 8 | buf[17] as u16,
                ack_nr: buf[18] as u16 << 8 | buf[19] as u16,
            };
            UtpPacket { header: header, payload: Vec::from_slice(buf.slice(20, buf.len())) }
        }
    }

    pub struct UtpSocket {
        socket: UdpSocket,
    }

    impl UtpSocket {
        pub fn bind(addr: SocketAddr) -> IoResult<UtpSocket> {
            let skt = UdpSocket::bind(addr);
            match skt {
                Ok(x)  => Ok(UtpSocket { socket: x }),
                Err(e) => Err(e)
            }
        }

        pub fn connect(self, other: SocketAddr) -> UtpStream {
            let stream = UtpStream::new(self.socket, other);
            stream.connect()
        }

        pub fn recvfrom(&mut self, buf: &mut[u8]) -> IoResult<(uint,SocketAddr)> {
            self.socket.recvfrom(buf)
        }
    }

    impl Clone for UtpSocket {
        fn clone(&self) -> UtpSocket {
            UtpSocket {
                socket: self.socket.clone(),
            }
        }
    }

    #[allow(dead_code)]
    pub struct UtpStream {
        socket: UdpSocket,
        connected_to: SocketAddr,
        connection_id: u16,
        seq_nr: u16,
        ack_nr: u16,
    }

    #[allow(dead_code)]
    impl UtpStream {
        pub fn new(socket: UdpSocket, conn: SocketAddr) -> UtpStream {
            let r: u16 = random();
            UtpStream {
                socket: socket,
                connected_to: conn,
                connection_id: r.to_be(),
                seq_nr: 1,
                ack_nr: 0,
            }
        }

        pub fn connect(mut self) -> UtpStream {
            let mut packet = UtpPacket::new();
            packet.set_type(ST_SYN);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = self.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let _result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;

            self
        }

        pub fn terminate(mut self) -> UtpStream {
            let mut packet = UtpPacket::new();
            packet.set_type(ST_FIN);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = self.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let _result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;

            self
        }

        pub fn disconnect(self) -> UdpSocket {
            self.socket
        }
    }

    impl Reader for UtpStream {
        fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
            match self.socket.recvfrom(buf) {
                Ok((_nread, src)) if src != self.connected_to => Ok(0),
                Ok((nread, _src)) => Ok(nread),
                Err(e) => Err(e),
            }
        }
    }
    impl Writer for UtpStream {
        fn write(&mut self, buf: &[u8]) -> IoResult<()> {
            let mut packet = UtpPacket::new();
            packet.payload = Vec::from_slice(buf);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = self.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;
            result
        }
    }
}

fn usage() {
    println!("Usage: torresmo [-s|-c] <address> <port>");
}

fn main() {
    use libtorresmo::{UtpSocket, UtpStream};
    use std::from_str::FromStr;

    // Defaults
    let mut addr = SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 };

    let args = std::os::args();

    if args.len() == 4 {
        let ip = FromStr::from_str(args.get(2).as_slice()).unwrap();
        let port = FromStr::from_str(args.get(3).as_slice()).unwrap();
        addr = SocketAddr { ip: ip, port: port };
    }

    match args.get(1).as_slice() {
        "-s" => {
            let mut buf = [0, ..512];
            let sock = UtpSocket::bind(addr).unwrap();
            println!("Serving on {}", addr);

            loop {
                let mut sock = sock.clone();
                match sock.recvfrom(buf) {
                    Ok((_, src)) => {
                        spawn(proc() {
                            let mut stream = sock.connect(src);
                            let payload = String::from_str("Hello\n").into_bytes();

                            // Send uTP packet
                            let _ = stream.write(payload.as_slice());
                        })
                    }
                    Err(_) => {}
                }
            }
        }
        "-c" => {
            let sock = UtpSocket::bind(SocketAddr { ip: Ipv4Addr(127,0,0,1), port: std::rand::random() }).unwrap();
            let mut stream = sock.connect(addr);
            let _ = stream.write([0]);
            let stream = stream.terminate();
            drop(stream);
        }
        _ => usage(),
    }
}
<commit_msg>Implement Show for UtpPacketHeader.<commit_after>use std::io::net::ip::{Ipv4Addr, SocketAddr};

mod libtorresmo {
    extern crate time;

    use std::io::net::udp::UdpSocket;
    use std::io::net::ip::SocketAddr;
    use std::io::IoResult;
    use std::mem::transmute;
    use std::rand::random;
    use std::fmt;

    #[allow(dead_code,non_camel_case_types)]
    enum UtpPacketType {
        ST_DATA  = 0,
        ST_FIN   = 1,
        ST_STATE = 2,
        ST_RESET = 3,
        ST_SYN   = 4,
    }

    impl fmt::Show for UtpPacketType {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            let s = match *self {
                ST_DATA  => "ST_DATA",
                ST_FIN   => "ST_FIN",
                ST_STATE => "ST_STATE",
                ST_RESET => "ST_RESET",
                ST_SYN   => "ST_SYN",
            };
            write!(f, "{}", s)
        }
    }

    #[allow(dead_code)]
    #[packed]
    struct UtpPacketHeader {
        type_ver: u8, // type: u4, ver: u4
        extension: u8,
        connection_id: u16,
        timestamp_microseconds: u32,
        timestamp_difference_microseconds: u32,
        wnd_size: u32,
        seq_nr: u16,
        ack_nr: u16,
    }

    impl UtpPacketHeader {
        fn bytes(&self) -> &[u8] {
            unsafe {
                let buf: &[u8, ..20] = transmute(self);
                return buf.as_slice();
            }
        }

        fn len(&self) -> uint {
            static HEADER_SIZE: uint = 20;
            return HEADER_SIZE;
        }
    }

    impl fmt::Show for UtpPacketHeader {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            let t: UtpPacketType = unsafe { transmute(self.type_ver >> 4) };
            let ptype = match t {
                ST_DATA  => "ST_DATA",
                ST_FIN   => "ST_FIN",
                ST_STATE => "ST_STATE",
                ST_RESET => "ST_RESET",
                ST_SYN   => "ST_SYN",
            };
            write!(f, "(type: {}, version: {}, extension: {}, connection_id: {}, timestamp_microseconds: {}, timestamp_difference_microseconds: {}, wnd_size: {}, seq_nr: {}, ack_nr: {})", ptype, self.type_ver & 0x0F, self.extension, self.connection_id, self.timestamp_microseconds, self.timestamp_difference_microseconds, self.wnd_size, self.seq_nr, self.ack_nr)
        }
    }

    #[allow(dead_code)]
    struct UtpPacket {
        header: UtpPacketHeader,
        payload: Vec<u8>,
    }

    impl UtpPacket {
        /// Constructs a new, empty UtpPacket.
        fn new() -> UtpPacket {
            UtpPacket {
                header: UtpPacketHeader {
                    type_ver: ST_DATA as u8 << 4 | 1,
                    extension: 0,
                    connection_id: 0,
                    timestamp_microseconds: 0,
                    timestamp_difference_microseconds: 0,
                    wnd_size: 0,
                    seq_nr: 0,
                    ack_nr: 0,
                },
                payload: Vec::new(),
            }
        }

        fn set_type(&mut self, t: UtpPacketType) {
            let version = 0x0F & self.header.type_ver;
            self.header.type_ver = t as u8 << 4 | version;
        }

        fn bytes(&self) -> Vec<u8> {
            let mut buf: Vec<u8> = Vec::with_capacity(self.len());
            buf.push_all(self.header.bytes());
            buf.push_all(self.payload.as_slice());
            return buf;
        }

        fn len(&self) -> uint {
            self.header.len() + self.payload.len()
        }

        fn decode(buf: &[u8]) -> UtpPacket {
            let header: UtpPacketHeader = UtpPacketHeader {
                type_ver: buf[0],
                extension: buf[1],
                connection_id: buf[2] as u16 << 8 | buf[3] as u16,
                timestamp_microseconds: buf[4] as u32 << 24 | buf[5] as u32 << 16 | buf[6] as u32 << 8 | buf[7] as u32,
                timestamp_difference_microseconds: buf[8] as u32 << 24 | buf[9] as u32 << 16 | buf[10] as u32 << 8 | buf[11] as u32,
                wnd_size: buf[12] as u32 << 24 | buf[13] as u32 << 16 | buf[14] as u32 << 8 | buf[15] as u32,
                seq_nr: buf[16] as u16 << 8 | buf[17] as u16,
                ack_nr: buf[18] as u16 << 8 | buf[19] as u16,
            };
            UtpPacket { header: header, payload: Vec::from_slice(buf.slice(20, buf.len())) }
        }
    }

    pub struct UtpSocket {
        socket: UdpSocket,
    }

    impl UtpSocket {
        pub fn bind(addr: SocketAddr) -> IoResult<UtpSocket> {
            let skt = UdpSocket::bind(addr);
            match skt {
                Ok(x)  => Ok(UtpSocket { socket: x }),
                Err(e) => Err(e)
            }
        }

        pub fn connect(self, other: SocketAddr) -> UtpStream {
            let stream = UtpStream::new(self.socket, other);
            stream.connect()
        }

        pub fn recvfrom(&mut self, buf: &mut[u8]) -> IoResult<(uint,SocketAddr)> {
            self.socket.recvfrom(buf)
        }
    }

    impl Clone for UtpSocket {
        fn clone(&self) -> UtpSocket {
            UtpSocket {
                socket: self.socket.clone(),
            }
        }
    }

    #[allow(dead_code)]
    pub struct UtpStream {
        socket: UdpSocket,
        connected_to: SocketAddr,
        connection_id: u16,
        seq_nr: u16,
        ack_nr: u16,
    }

    #[allow(dead_code)]
    impl UtpStream {
        pub fn new(socket: UdpSocket, conn: SocketAddr) -> UtpStream {
            let r: u16 = random();
            UtpStream {
                socket: socket,
                connected_to: conn,
                connection_id: r.to_be(),
                seq_nr: 1,
                ack_nr: 0,
            }
        }

        pub fn connect(mut self) -> UtpStream {
            let mut packet = UtpPacket::new();
            packet.set_type(ST_SYN);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = self.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let _result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;

            self
        }

        pub fn terminate(mut self) -> UtpStream {
            let mut packet = UtpPacket::new();
            packet.set_type(ST_FIN);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = self.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let _result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;

            self
        }

        pub fn disconnect(self) -> UdpSocket {
            self.socket
        }
    }

    impl Reader for UtpStream {
        fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
            match self.socket.recvfrom(buf) {
                Ok((_nread, src)) if src != self.connected_to => Ok(0),
                Ok((nread, _src)) => Ok(nread),
                Err(e) => Err(e),
            }
        }
    }
    impl Writer for UtpStream {
        fn write(&mut self, buf: &[u8]) -> IoResult<()> {
            let mut packet = UtpPacket::new();
            packet.payload = Vec::from_slice(buf);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = self.seq_nr.to_be();
            let t = time::get_time();
            packet.header.timestamp_microseconds = ((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32).to_be();

            // Send packet
            let result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;
            result
        }
    }
}

fn usage() {
    println!("Usage: torresmo [-s|-c] <address> <port>");
}

fn main() {
    use libtorresmo::{UtpSocket, UtpStream};
    use std::from_str::FromStr;

    // Defaults
    let mut addr = SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 };

    let args = std::os::args();

    if args.len() == 4 {
        let ip = FromStr::from_str(args.get(2).as_slice()).unwrap();
        let port = FromStr::from_str(args.get(3).as_slice()).unwrap();
        addr = SocketAddr { ip: ip, port: port };
    }

    match args.get(1).as_slice() {
        "-s" => {
            let mut buf = [0, ..512];
            let sock = UtpSocket::bind(addr).unwrap();
            println!("Serving on {}", addr);

            loop {
                let mut sock = sock.clone();
                match sock.recvfrom(buf) {
                    Ok((_, src)) => {
                        spawn(proc() {
                            let mut stream = sock.connect(src);
                            let payload = String::from_str("Hello\n").into_bytes();

                            // Send uTP packet
                            let _ = stream.write(payload.as_slice());
                        })
                    }
                    Err(_) => {}
                }
            }
        }
        "-c" => {
            let sock = UtpSocket::bind(SocketAddr { ip: Ipv4Addr(127,0,0,1), port: std::rand::random() }).unwrap();
            let mut stream = sock.connect(addr);
            let _ = stream.write([0]);
            let stream = stream.terminate();
            drop(stream);
        }
        _ => usage(),
    }
}
<|endoftext|>


<commit_before>use std::io::net::ip::{Ipv4Addr, SocketAddr};

mod libtorresmo {
    extern crate time;

    use std::io::net::udp::UdpSocket;
    use std::io::net::ip::SocketAddr;
    use std::io::IoResult;
    use std::mem::{to_be32, to_be16, transmute};
    use std::rand::random;

    #[allow(dead_code,non_camel_case_types)]
    enum UtpPacketType {
        ST_DATA  = 0,
        ST_FIN   = 1,
        ST_STATE = 2,
        ST_RESET = 3,
        ST_SYN   = 4,
    }

    #[allow(dead_code)]
    #[packed]
    struct UtpPacketHeader {
        type_ver: u8, // type: u4, ver: u4
        extension: u8,
        connection_id: u16,
        timestamp_microseconds: u32,
        timestamp_difference_microseconds: u32,
        wnd_size: u32,
        seq_nr: u16,
        ack_nr: u16,
    }

    impl UtpPacketHeader {
        fn bytes(&self) -> &[u8] {
            unsafe {
                let buf: &[u8, ..20] = transmute(self);
                return buf.as_slice();
            }
        }

        fn len(&self) -> uint {
            static HEADER_SIZE: uint = 20;
            return HEADER_SIZE;
        }
    }

    #[allow(dead_code)]
    struct UtpPacket {
        header: UtpPacketHeader,
        payload: Vec<u8>,
    }

    impl UtpPacket {
        /// Constructs a new, empty UtpPacket.
        fn new() -> UtpPacket {
            UtpPacket {
                header: UtpPacketHeader {
                    type_ver: ST_DATA as u8 << 4 | 1,
                    extension: 0,
                    connection_id: 0,
                    timestamp_microseconds: 0,
                    timestamp_difference_microseconds: 0,
                    wnd_size: 0,
                    seq_nr: 0,
                    ack_nr: 0,
                },
                payload: Vec::new(),
            }
        }

        fn set_type(&mut self, t: UtpPacketType) {
            let version = 0x0F & self.header.type_ver;
            self.header.type_ver = t as u8 << 4 | version;
        }

        fn bytes(&self) -> Vec<u8> {
            let mut buf: Vec<u8> = Vec::with_capacity(self.len());
            buf.push_all(self.header.bytes());
            buf.push_all(self.payload.as_slice());
            return buf;
        }

        fn len(&self) -> uint {
            self.header.len() + self.payload.len()
        }
    }

    pub struct UtpSocket {
        socket: UdpSocket,
    }

    impl UtpSocket {
        pub fn bind(addr: SocketAddr) -> IoResult<UtpSocket> {
            let skt = UdpSocket::bind(addr);
            match skt {
                Ok(x)  => Ok(UtpSocket { socket: x }),
                Err(e) => Err(e)
            }
        }

        pub fn connect(self, other: SocketAddr) -> UtpStream {
            UtpStream::new(self.socket, other)
        }

        pub fn recvfrom(&mut self, buf: &mut[u8]) -> IoResult<(uint,SocketAddr)> {
            self.socket.recvfrom(buf)
        }
    }

    impl Clone for UtpSocket {
        fn clone(&self) -> UtpSocket {
            UtpSocket {
                socket: self.socket.clone(),
            }
        }
    }

    #[allow(dead_code)]
    pub struct UtpStream {
        socket: UdpSocket,
        connected_to: SocketAddr,
        connection_id: u16,
        seq_nr: u16,
        ack_nr: u16,
    }

    #[allow(dead_code)]
    impl UtpStream {
        pub fn new(socket: UdpSocket, conn: SocketAddr) -> UtpStream {
            UtpStream {
                socket: socket,
                connected_to: conn,
                connection_id: to_be16(random()),
                seq_nr: 1,
                ack_nr: 0,
            }
        }

        pub fn disconnect(self) -> UdpSocket {
            self.socket
        }
    }

    impl Reader for UtpStream {
        fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
            match self.socket.recvfrom(buf) {
                Ok((_nread, src)) if src != self.connected_to => Ok(0),
                Ok((nread, _src)) => Ok(nread),
                Err(e) => Err(e),
            }
        }
    }
    impl Writer for UtpStream {
        fn write(&mut self, buf: &[u8]) -> IoResult<()> {
            let mut packet = UtpPacket::new();
            packet.payload = Vec::from_slice(buf);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = to_be16(self.seq_nr);
            let t = time::get_time();
            packet.header.timestamp_microseconds = to_be32((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32);

            // Send packet
            let result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;
            result
        }
    }
}

fn usage() {
    println!("Usage: torresmo [-s|-c] <address> <port>");
}

fn main() {
    use libtorresmo::{UtpSocket, UtpStream};
    use std::from_str::FromStr;

    // Defaults
    let mut addr = SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 };

    let args = std::os::args();

    if args.len() == 4 {
        let ip = FromStr::from_str(args.get(2).as_slice()).unwrap();
        let port = FromStr::from_str(args.get(3).as_slice()).unwrap();
        addr = SocketAddr { ip: ip, port: port };
    }

    match args.get(1).as_slice() {
        "-s" => {
            let mut buf = [0, ..512];
            let sock = UtpSocket::bind(addr).unwrap();
            println!("Serving on {}", addr);

            loop {
                let mut sock = sock.clone();
                match sock.recvfrom(buf) {
                    Ok((_, src)) => {
                        spawn(proc() {
                            let mut stream = sock.connect(src);
                            let payload = String::from_str("Hello\n").into_bytes();

                            // Send uTP packet
                            let _ = stream.write(payload.as_slice());
                        })
                    }
                    Err(_) => {}
                }
            }
        }
        "-c" => {
            let sock = UtpSocket::bind(SocketAddr { ip: Ipv4Addr(127,0,0,1), port: std::rand::random() }).unwrap();
            let mut stream = sock.connect(addr);
            let _ = stream.write([0]);
            drop(stream);
        }
        _ => usage(),
    }
}
<commit_msg>Send proper SYN packet on connection start.<commit_after>use std::io::net::ip::{Ipv4Addr, SocketAddr};

mod libtorresmo {
    extern crate time;

    use std::io::net::udp::UdpSocket;
    use std::io::net::ip::SocketAddr;
    use std::io::IoResult;
    use std::mem::{to_be32, to_be16, transmute};
    use std::rand::random;

    #[allow(dead_code,non_camel_case_types)]
    enum UtpPacketType {
        ST_DATA  = 0,
        ST_FIN   = 1,
        ST_STATE = 2,
        ST_RESET = 3,
        ST_SYN   = 4,
    }

    #[allow(dead_code)]
    #[packed]
    struct UtpPacketHeader {
        type_ver: u8, // type: u4, ver: u4
        extension: u8,
        connection_id: u16,
        timestamp_microseconds: u32,
        timestamp_difference_microseconds: u32,
        wnd_size: u32,
        seq_nr: u16,
        ack_nr: u16,
    }

    impl UtpPacketHeader {
        fn bytes(&self) -> &[u8] {
            unsafe {
                let buf: &[u8, ..20] = transmute(self);
                return buf.as_slice();
            }
        }

        fn len(&self) -> uint {
            static HEADER_SIZE: uint = 20;
            return HEADER_SIZE;
        }
    }

    #[allow(dead_code)]
    struct UtpPacket {
        header: UtpPacketHeader,
        payload: Vec<u8>,
    }

    impl UtpPacket {
        /// Constructs a new, empty UtpPacket.
        fn new() -> UtpPacket {
            UtpPacket {
                header: UtpPacketHeader {
                    type_ver: ST_DATA as u8 << 4 | 1,
                    extension: 0,
                    connection_id: 0,
                    timestamp_microseconds: 0,
                    timestamp_difference_microseconds: 0,
                    wnd_size: 0,
                    seq_nr: 0,
                    ack_nr: 0,
                },
                payload: Vec::new(),
            }
        }

        fn set_type(&mut self, t: UtpPacketType) {
            let version = 0x0F & self.header.type_ver;
            self.header.type_ver = t as u8 << 4 | version;
        }

        fn bytes(&self) -> Vec<u8> {
            let mut buf: Vec<u8> = Vec::with_capacity(self.len());
            buf.push_all(self.header.bytes());
            buf.push_all(self.payload.as_slice());
            return buf;
        }

        fn len(&self) -> uint {
            self.header.len() + self.payload.len()
        }
    }

    pub struct UtpSocket {
        socket: UdpSocket,
    }

    impl UtpSocket {
        pub fn bind(addr: SocketAddr) -> IoResult<UtpSocket> {
            let skt = UdpSocket::bind(addr);
            match skt {
                Ok(x)  => Ok(UtpSocket { socket: x }),
                Err(e) => Err(e)
            }
        }

        pub fn connect(self, other: SocketAddr) -> UtpStream {
            let stream = UtpStream::new(self.socket, other);
            stream.connect()
        }

        pub fn recvfrom(&mut self, buf: &mut[u8]) -> IoResult<(uint,SocketAddr)> {
            self.socket.recvfrom(buf)
        }
    }

    impl Clone for UtpSocket {
        fn clone(&self) -> UtpSocket {
            UtpSocket {
                socket: self.socket.clone(),
            }
        }
    }

    #[allow(dead_code)]
    pub struct UtpStream {
        socket: UdpSocket,
        connected_to: SocketAddr,
        connection_id: u16,
        seq_nr: u16,
        ack_nr: u16,
    }

    #[allow(dead_code)]
    impl UtpStream {
        pub fn new(socket: UdpSocket, conn: SocketAddr) -> UtpStream {
            UtpStream {
                socket: socket,
                connected_to: conn,
                connection_id: to_be16(random()),
                seq_nr: 1,
                ack_nr: 0,
            }
        }

        pub fn connect(mut self) -> UtpStream {
            let mut packet = UtpPacket::new();
            packet.set_type(ST_SYN);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = to_be16(self.seq_nr);
            let t = time::get_time();
            packet.header.timestamp_microseconds = to_be32((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32);

            // Send packet
            let _result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;

            self
        }

        pub fn disconnect(self) -> UdpSocket {
            self.socket
        }
    }

    impl Reader for UtpStream {
        fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
            match self.socket.recvfrom(buf) {
                Ok((_nread, src)) if src != self.connected_to => Ok(0),
                Ok((nread, _src)) => Ok(nread),
                Err(e) => Err(e),
            }
        }
    }
    impl Writer for UtpStream {
        fn write(&mut self, buf: &[u8]) -> IoResult<()> {
            let mut packet = UtpPacket::new();
            packet.payload = Vec::from_slice(buf);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = to_be16(self.seq_nr);
            let t = time::get_time();
            packet.header.timestamp_microseconds = to_be32((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32);

            // Send packet
            let result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;
            result
        }
    }
}

fn usage() {
    println!("Usage: torresmo [-s|-c] <address> <port>");
}

fn main() {
    use libtorresmo::{UtpSocket, UtpStream};
    use std::from_str::FromStr;

    // Defaults
    let mut addr = SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 };

    let args = std::os::args();

    if args.len() == 4 {
        let ip = FromStr::from_str(args.get(2).as_slice()).unwrap();
        let port = FromStr::from_str(args.get(3).as_slice()).unwrap();
        addr = SocketAddr { ip: ip, port: port };
    }

    match args.get(1).as_slice() {
        "-s" => {
            let mut buf = [0, ..512];
            let sock = UtpSocket::bind(addr).unwrap();
            println!("Serving on {}", addr);

            loop {
                let mut sock = sock.clone();
                match sock.recvfrom(buf) {
                    Ok((_, src)) => {
                        spawn(proc() {
                            let mut stream = sock.connect(src);
                            let payload = String::from_str("Hello\n").into_bytes();

                            // Send uTP packet
                            let _ = stream.write(payload.as_slice());
                        })
                    }
                    Err(_) => {}
                }
            }
        }
        "-c" => {
            let sock = UtpSocket::bind(SocketAddr { ip: Ipv4Addr(127,0,0,1), port: std::rand::random() }).unwrap();
            let mut stream = sock.connect(addr);
            let _ = stream.write([0]);
            drop(stream);
        }
        _ => usage(),
    }
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! ELF file header
use super::{ElfResult, ElfWord, Section, section};
use super::ValidatesWord;
use core::{fmt, mem, convert};

/// Trait representing an ELF File Header.
///
/// This trait allows [HeaderRepr] to provide a consistent API regardless
/// of whether the header section uses 32- or 64-bit [ELF word]s. A number of
/// field values in the header of various sizes are converted to `usize` by
/// this API so that they can be used as indices, etc.
///
/// For more information on ELF File Headers, refer to:
/// + the ELF [specification]
/// + the [OS Dev Wiki]
///
/// [ELF word]: ../type.ElfWord.html
/// [HeaderRepr]: struct.HeaderRepr.html
/// [specification](http://www.sco.com/developers/gabi/latest/ch4.eheader.html)
/// [OS Dev Wiki](http://wiki.osdev.org/ELF#Header)
pub trait Header: Sized {
    type Word: ElfWord;

    /// Attempt to extract an ELF file header from a slice of bytes.
    fn from_slice<'a>(input: &'a [u8]) -> ElfResult<&'a Self>;

    /// Attempt to extract a section header from a slice of bytes.
    /// TODO: can/should the index be `usize`?
    //          - eliza, 03/08/2017
    fn parse_section<'a>(&'a self, input: &'a [u8], idx: u16)
                            -> ElfResult<&'a Section>;

    // Field accessors -------------------------------------------
    fn ident(&self) -> Ident;
    fn get_type(&self) -> Type;
    fn machine(&self) -> Machine;
    /// Offset of the program entry point
    fn entry_point(&self) -> usize;
    /// Offset of the start of program headers
    fn ph_offset(&self) -> usize;
    /// Number of program headers.
    fn ph_count(&self) -> usize;
    /// Size of a program header.
    /// TODO: add function for using this and `ph_count` to make a memrange
    ///       representing the program header entries?
    //          - eliza, 03/08/2017
    fn ph_entry_size(&self) -> usize;
    /// Offset of the start of [section header]s.
    ///
    /// [section header]: ../section/struct.Header.html
    fn sh_offset(&self) -> usize;
    /// Number of [section header]s.
    ///
    /// [section header]: ../section/struct.Header.html
    fn sh_count(&self) -> usize;
    /// Size of a [section header].
    ///
    /// [section header]: ../section/struct.Header.html
    /// TODO: add function for using this and `sh_count` to make a memrange
    ///       representing the section header entries?
    //          - eliza, 03/08/2017
    fn sh_entry_size(&self) -> usize;
    /// TODO: can this return the flags type?
    //          - eliza, 03/08/2017
    fn flags(&self) -> u32;
    /// Index of the section header [string table].
    ///
    /// [string table]: ../section/struct.StrTable.html"]
    fn sh_str_idx(&self) -> usize;

}

macro_rules! impl_getters {
    ($(#[$attr:meta])* pub fn $name:ident(&self) -> $ty:ident; $($rest:tt)*) => {
        $(#[$attr])* #[inline] pub fn $name(&self) -> $ty { self.$name as $ty }
        impl_getters!{ $( $rest )* }
    };
    ($(#[$attr:meta])* fn $name:ident(&self) -> $ty:ident; $($rest:tt)*) => {
        $(#[$attr])* #[inline] fn $name(&self) -> $ty { self.$name as $ty }
        impl_getters!{ $( $rest )* }
    };
    ( $(#[$attr:meta])* pub fn $name: ident (&self)-> $ty: ident; ) => {
        $(#[$attr])* #[inline] pub fn $name(&self) -> $ty { self.$name as $ty }
    };
    ( $(#[$attr:meta])* fn $name: ident (&self)-> $ty: ident; ) => {
        $(#[$attr])* #[inline] fn $name(&self) -> $ty { self.$name as $ty }
    };
    () => {};
}

macro_rules! Header {
    (($($size:ty),+) $(pub)* enum $name:ident $($tail:tt)* ) => {
        Header! { @impl $name, $($size)+ }
    };
    (($($size:ty),+) $(pub)* struct $name:ident $($tail:tt)*) => {
        Header! { @impl $name, $($size)+ }
    };
    (@impl $name:ident, $($size:ty)+) => {
        $(impl Header for $name<$size> {
            type Word = $size;
            /// Attempt to extract an ELF file header from a slice of bytes.
            fn from_slice<'a>(input: &'a [u8]) -> ElfResult<&'a Self> {
                if input.len() < mem::size_of::<Self>() {
                    Err("Input too short to extract ELF header")
                } else {
                    unsafe { Ok(&super::extract_from_slice::<Self>(input, 0, 1)[0]) }
                }
            }

            /// Attempt to extract a section header from a slice of bytes.
            fn parse_section<'a>(&'a self, input: &'a [u8], idx: u16)
                                    -> ElfResult<&'a Section>
            {
                if idx < section::SHN_LORESERVE {
                    Err("Cannot parse reserved section.")
                } else {
                    // use ValidatesWord to check if this section's Class field
                    // will let us interpret the section with the requested
                    // word length. this is a bit of a hack around the type
                    // system not letting me do this the way I wanted to to....
                    let validator: &ValidatesWord<Self::Word>
                        = &self.ident.class;
                    validator.check()?;

                    // start offset for section
                    let start = self.sh_offset() + idx as usize *
                                self.sh_entry_size();
                    // end offset for section
                    let end = start + self.sh_entry_size();
                    let raw = &input[start .. end];

                    unsafe {
                        Ok(&*(raw as *const [u8]
                                  as *const _
                                  as *const Section))
                    }
                }
            }

            #[inline] fn get_type(&self) -> Type { self.elftype.as_type() }

            impl_getters! {
                #[doc = "Index for the start of [section header]s. \
                         [section header]: ../section/struct.Header.html"]
                fn sh_offset(&self) -> usize;
                fn sh_entry_size(&self) -> usize;
                fn sh_count(&self) -> usize;

                #[doc = "Index for the start of program headers"]
                fn ph_offset(&self) -> usize;
                fn ph_entry_size(&self) -> usize;
                fn ph_count(&self) -> usize;

                #[doc = "Index for the program entry point"]
                fn entry_point(&self) -> usize;
                #[doc = "Index of the section header [string table] \
                         [string table]: ../section/struct.StrTable.html"]
                fn sh_str_idx(&self) -> usize;
                fn flags(&self) -> u32;
                fn ident(&self) -> Ident;
                fn machine(&self) -> Machine;
            }
        }

        impl<'a> convert::TryFrom<&'a [u8]> for $name<$size>
        where $name<$size>: Header {
            type Err = &'static str;
            #[inline]
            fn try_from(slice: &'a [u8]) -> ElfResult<Self> {
                <Self as Header>::from_slice(slice).map(|x| *x)
            }
        }
        )+
    }
}

macro_attr! {
    /// Raw representation of an ELF file header.
    #[derive(Copy, Clone, Debug, Header!(u32, u64))]
    #[repr(C, packed)]
    pub struct HeaderRepr<W: ElfWord> {
        /// the ELF [file identifier](struct.Ident.html)
        pub ident: Ident
      , elftype: TypeRepr
      , pub machine: Machine
      , /// Program entry point
        entry_point: W
      , /// Offset for start of program headers
        ph_offset: W
      , /// Offset for start of [section header]s.
        /// [section header]: ../section/struct.Header.html
        sh_offset: W
      , pub flags: u32
      , pub header_size: u16
      , pub ph_entry_size: u16
      , pub ph_count: u16
      , pub sh_entry_size: u16
      , pub sh_count: u16
      , /// Index of the section header string table
        sh_str_idx: u16
    }
}

/// ELF header magic
pub const MAGIC: Magic = [0x7f, b'E', b'L', b'F'];

/// Type of header magic
pub type Magic = [u8; 4];

/// ELF identifier (`e_ident` in the ELF standard)
#[derive(Copy, Clone, Debug)]
#[repr(C, packed)]
pub struct Ident {
    /// ELF magic numbers. Must be equal to the [ELF magic], `[0x7, E, L, F]`.
    ///
    /// [ELF magic]: constant.MAGIC.html
    pub magic: Magic
  , /// ELF [file class](enum.Class.html) (32- or 64-bit)
    pub class: Class
  , /// ELF [data encoding](enum.DataEncoding.html) (big- or little-endian)
    pub encoding: DataEncoding
  , /// ELF file [version](enum.Version.html)
    pub version: Version
  , /// What [operating system ABI] this file was compiled for.
    ///
    /// [operating system ABI]: enum.OsAbi.html
    pub abi: OsAbi
  , /// ABI version (often this is just padding)
    pub abi_version: u8
  , _padding: [u8; 7]
}

impl Ident {
    #[inline] pub fn check_magic(&self) -> bool { self.magic == MAGIC }

    /// Returns true if the identifier section identifies a valid ELF file.
    #[inline] pub fn is_valid(&self) -> bool {
        // the ELF magic number is correct
        self.check_magic() &&
        // the file class is either 32- or 64-bits
        self.class.is_valid() &&
        // the data encoding is either big- or little-endian
        self.encoding.is_valid()
    }
}

/// Identifies the class of the ELF file
#[derive(Copy, Clone, PartialEq, Debug)]
#[repr(u8)]
pub enum Class {
    /// Invalid ELF class file (`ELFCLASSNONE` in the standard)
    None  = 0
  , /// 32-bit ELF file (`ELFCLASS32` in the standard)
    Elf32 = 1
  , /// 64-bit ELF file (`ELFCLASS64` in the standard)
    Elf64 = 2
}

impl Class {
    /// Returns true if the class field for this file is valid.
    #[inline]
    pub fn is_valid(&self) -> bool {
        match *self { Class::None => false
                    , _ => true
                    }
    }

}

impl ValidatesWord<u64> for Class {
    #[inline]
    fn check(&self) -> ElfResult<()> {
        use self::Class::*;
        match *self {
            None => Err("Invalid ELF type ELFCLASSNONE!")
          , Elf32 => Err("Cannot extract 64-bit section from 32-bit ELF")
          , Elf64 => Ok(())
        }
    }
}

impl ValidatesWord<u32> for Class {
    #[inline]
    fn check(&self) -> ElfResult<()> {
        use self::Class::*;
        match *self {
            None => Err("Invalid ELF type ELFCLASSNONE!")
          , Elf64 =>
                // TODO: is this actually true?
                //          - eliza, 03/08/2017
                Err("Cannot extract 32-bit section from 64-bit ELF")
          , Elf32 => Ok(())
        }
    }
}

/// Identifies the data encoding of the ELF file
#[derive(Copy, Clone, PartialEq, Debug)]
#[repr(u8)]
pub enum DataEncoding {
    /// Invalid data encoding (`ELFDATANONE` in the standard)
    None  = 0
  , /// Twos-complement little-endian data encoding
    /// (`ELFDATA2LSB` in the standard)
    LittleEndian = 1
  , /// Twos-complement big-endian data encoding
    /// (`ELFDATA2MSB` in the standard)
    BigEndian = 2
}

impl DataEncoding {
    /// Returns true if the data encoding field for this file is valid.
    #[inline]
    pub fn is_valid(&self) -> bool {
        match *self { DataEncoding::None => false
                    , _ => true
                    }
    }
}

/// Operating system ABI
#[derive(Copy, Clone, Debug)]
#[repr(u8)]
pub enum OsAbi { /// Ox00 also represents "none"
                 SystemV = 0x00
               , HpUx    = 0x01
               , NetBsd  = 0x02
               , Linux   = 0x03
               , Solaris = 0x06
               , Aix     = 0x07
               , Irix    = 0x08
               , FreeBsd = 0x09
               , OpenBsd = 0x0C
               , OpenVms = 0x0D
               }

/// Identifies the version of the ELF file
#[derive(Copy, Clone, PartialEq, Debug)]
#[repr(u8)]
pub enum Version { None = 0, Current = 1 }

#[derive(Clone, Copy, PartialEq)]
struct TypeRepr(u16);

impl TypeRepr {
    /// TODO: rewrite this as a `convert::Into` implementation
    ///         - eliza, 03/09/2017
    pub fn as_type(&self) -> Type {
        match self.0 {
            0 => Type::None
          , 1 => Type::Relocatable
          , 2 => Type::Executable
          , 3 => Type::SharedObject
          , 4 => Type::Core
          , anything => Type::Other(anything)
        }
    }
}

impl fmt::Debug for TypeRepr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.as_type().fmt(f)
    }
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum Type { None
              , Relocatable
              , Executable
              , SharedObject
              , Core
              , Other(u16)
              }

#[allow(non_camel_case_types)]
#[derive(Clone, Copy, PartialEq, Debug)]
#[repr(u16)]
pub enum Machine { None    = 0
                 , Sparc   = 0x02
                 , X86     = 0x03
                 , Mips    = 0x08
                 , PowerPc = 0x14
                 , Arm     = 0x28
                 , SuperH  = 0x2A
                 , Ia64    = 0x32
                 , X86_64  = 0x3E
                 , AArch64 = 0xB7
                 }
<commit_msg>refac(elf): use header info to find section/program header indices<commit_after>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! ELF file header
use super::{ElfResult, ElfWord, Section, section};
use super::ValidatesWord;

use core::{fmt, mem, convert};
use core::ops::Range;

/// Trait representing an ELF File Header.
///
/// This trait allows [HeaderRepr] to provide a consistent API regardless
/// of whether the header section uses 32- or 64-bit [ELF word]s. A number of
/// field values in the header of various sizes are converted to `usize` by
/// this API so that they can be used as indices, etc.
///
/// For more information on ELF File Headers, refer to:
/// + the ELF [specification]
/// + the [OS Dev Wiki]
///
/// [ELF word]: ../type.ElfWord.html
/// [HeaderRepr]: struct.HeaderRepr.html
/// [specification](http://www.sco.com/developers/gabi/latest/ch4.eheader.html)
/// [OS Dev Wiki](http://wiki.osdev.org/ELF#Header)
pub trait Header: Sized {
    type Word: ElfWord;

    /// Attempt to extract an ELF file header from a slice of bytes.
    fn from_slice<'a>(input: &'a [u8]) -> ElfResult<&'a Self>;

    /// Attempt to extract a section header from a slice of bytes.
    /// TODO: can/should the index be `usize`?
    //          - eliza, 03/08/2017
    fn parse_section<'a>(&'a self, input: &'a [u8], idx: u16)
                            -> ElfResult<&'a Section>;

    /// Calculate the index for a [section header]
    ///
    /// TODO: should this check the index is reasonable & return a `Result`
    //          - eliza, 03/10/2017
    /// [section header]: ../section/struct.Header.html
    fn section_index(&self, idx: usize) -> Range<usize> {
        let size = self.sh_entry_size();
        let start = self.sh_offset() + (idx * size);
        start .. start + size
    }

    /// Calculate the index for a program header
    ///
    /// TODO: should this check the index is reasonable & return a `Result`
    //          - eliza, 03/10/2017
    fn program_header_index(&self, idx: usize) -> Range<usize> {
        let size = self.ph_entry_size();
        let start = self.ph_offset() + (idx * size);
        start .. start + size
    }

    // Field accessors -------------------------------------------
    fn ident(&self) -> Ident;
    fn get_type(&self) -> Type;
    fn machine(&self) -> Machine;
    /// Offset of the program entry point
    fn entry_point(&self) -> usize;
    /// Offset of the start of program headers
    fn ph_offset(&self) -> usize;
    /// Number of program headers.
    fn ph_count(&self) -> usize;
    /// Size of a program header.
    fn ph_entry_size(&self) -> usize;
    /// Offset of the start of [section header]s.
    ///
    /// [section header]: ../section/struct.Header.html
    fn sh_offset(&self) -> usize;
    /// Number of [section header]s.
    ///
    /// [section header]: ../section/struct.Header.html
    fn sh_count(&self) -> usize;
    /// Size of a [section header].
    ///
    /// [section header]: ../section/struct.Header.html
    fn sh_entry_size(&self) -> usize;
    /// TODO: can this return the flags type?
    //          - eliza, 03/08/2017
    fn flags(&self) -> u32;
    /// Index of the section header [string table].
    ///
    /// [string table]: ../section/struct.StrTable.html"]
    fn sh_str_idx(&self) -> usize;
}

macro_rules! impl_getters {
    ($(#[$attr:meta])* pub fn $name:ident(&self) -> $ty:ident; $($rest:tt)*) => {
        $(#[$attr])* #[inline] pub fn $name(&self) -> $ty { self.$name as $ty }
        impl_getters!{ $( $rest )* }
    };
    ($(#[$attr:meta])* fn $name:ident(&self) -> $ty:ident; $($rest:tt)*) => {
        $(#[$attr])* #[inline] fn $name(&self) -> $ty { self.$name as $ty }
        impl_getters!{ $( $rest )* }
    };
    ( $(#[$attr:meta])* pub fn $name: ident (&self)-> $ty: ident; ) => {
        $(#[$attr])* #[inline] pub fn $name(&self) -> $ty { self.$name as $ty }
    };
    ( $(#[$attr:meta])* fn $name: ident (&self)-> $ty: ident; ) => {
        $(#[$attr])* #[inline] fn $name(&self) -> $ty { self.$name as $ty }
    };
    () => {};
}

macro_rules! Header {
    (($($size:ty),+) $(pub)* enum $name:ident $($tail:tt)* ) => {
        Header! { @impl $name, $($size)+ }
    };
    (($($size:ty),+) $(pub)* struct $name:ident $($tail:tt)*) => {
        Header! { @impl $name, $($size)+ }
    };
    (@impl $name:ident, $($size:ty)+) => {
        $(impl Header for $name<$size> {
            type Word = $size;
            /// Attempt to extract an ELF file header from a slice of bytes.
            fn from_slice<'a>(input: &'a [u8]) -> ElfResult<&'a Self> {
                if input.len() < mem::size_of::<Self>() {
                    Err("Input too short to extract ELF header")
                } else {
                    unsafe { Ok(&super::extract_from_slice::<Self>(input, 0, 1)[0]) }
                }
            }

            /// Attempt to extract a [section header] from a slice of bytes.
            ///
            /// TODO: should this move to the `File` type since it owns the
            ///       byte slice (which then wouldn't have to be passed as an
            ///       argument)?
            //          - eliza, 03/10/2017
            ///
            /// [section header]: ../section/struct.Header.html
            fn parse_section<'a>(&'a self, input: &'a [u8], idx: u16)
                                -> ElfResult<&'a Section>
            {
                if idx < section::SHN_LORESERVE {
                    Err("Cannot parse reserved section.")
                } else {
                    // use ValidatesWord to check if this section's Class field
                    // will let us interpret the section with the requested
                    // word length. this is a bit of a hack around the type
                    // system not letting me do this the way I wanted to to....
                    let validator: &ValidatesWord<Self::Word>
                        = &self.ident.class;
                    validator.check()?;

                    let raw = &input[self.section_index(idx as usize)];

                    unsafe {
                        Ok(&*(raw as *const [u8]
                                  as *const _
                                  as *const Section))
                    }
                }
            }

            #[inline] fn get_type(&self) -> Type { self.elftype.as_type() }

            impl_getters! {
                #[doc = "Index for the start of [section header]s. \
                         [section header]: ../section/struct.Header.html"]
                fn sh_offset(&self) -> usize;
                fn sh_entry_size(&self) -> usize;
                fn sh_count(&self) -> usize;

                #[doc = "Index for the start of program headers"]
                fn ph_offset(&self) -> usize;
                fn ph_entry_size(&self) -> usize;
                fn ph_count(&self) -> usize;

                #[doc = "Index for the program entry point"]
                fn entry_point(&self) -> usize;
                #[doc = "Index of the section header [string table] \
                         [string table]: ../section/struct.StrTable.html"]
                fn sh_str_idx(&self) -> usize;
                fn flags(&self) -> u32;
                fn ident(&self) -> Ident;
                fn machine(&self) -> Machine;
            }
        }

        impl<'a> convert::TryFrom<&'a [u8]> for $name<$size>
        where $name<$size>: Header {
            type Err = &'static str;
            #[inline]
            fn try_from(slice: &'a [u8]) -> ElfResult<Self> {
                <Self as Header>::from_slice(slice).map(|x| *x)
            }
        }
        )+
    }
}

macro_attr! {
    /// Raw representation of an ELF file header.
    #[derive(Copy, Clone, Debug, Header!(u32, u64))]
    #[repr(C, packed)]
    pub struct HeaderRepr<W: ElfWord> {
        /// the ELF [file identifier](struct.Ident.html)
        pub ident: Ident
      , elftype: TypeRepr
      , pub machine: Machine
      , /// Program entry point
        entry_point: W
      , /// Offset for start of program headers
        ph_offset: W
      , /// Offset for start of [section header]s.
        /// [section header]: ../section/struct.Header.html
        sh_offset: W
      , pub flags: u32
      , pub header_size: u16
      , pub ph_entry_size: u16
      , pub ph_count: u16
      , pub sh_entry_size: u16
      , pub sh_count: u16
      , /// Index of the section header string table
        sh_str_idx: u16
    }
}

/// ELF header magic
pub const MAGIC: Magic = [0x7f, b'E', b'L', b'F'];

/// Type of header magic
pub type Magic = [u8; 4];

/// ELF identifier (`e_ident` in the ELF standard)
#[derive(Copy, Clone, Debug)]
#[repr(C, packed)]
pub struct Ident {
    /// ELF magic numbers. Must be equal to the [ELF magic], `[0x7, E, L, F]`.
    ///
    /// [ELF magic]: constant.MAGIC.html
    pub magic: Magic
  , /// ELF [file class](enum.Class.html) (32- or 64-bit)
    pub class: Class
  , /// ELF [data encoding](enum.DataEncoding.html) (big- or little-endian)
    pub encoding: DataEncoding
  , /// ELF file [version](enum.Version.html)
    pub version: Version
  , /// What [operating system ABI] this file was compiled for.
    ///
    /// [operating system ABI]: enum.OsAbi.html
    pub abi: OsAbi
  , /// ABI version (often this is just padding)
    pub abi_version: u8
  , _padding: [u8; 7]
}

impl Ident {
    #[inline] pub fn check_magic(&self) -> bool { self.magic == MAGIC }

    /// Returns true if the identifier section identifies a valid ELF file.
    #[inline] pub fn is_valid(&self) -> bool {
        // the ELF magic number is correct
        self.check_magic() &&
        // the file class is either 32- or 64-bits
        self.class.is_valid() &&
        // the data encoding is either big- or little-endian
        self.encoding.is_valid()
    }
}

/// Identifies the class of the ELF file
#[derive(Copy, Clone, PartialEq, Debug)]
#[repr(u8)]
pub enum Class {
    /// Invalid ELF class file (`ELFCLASSNONE` in the standard)
    None  = 0
  , /// 32-bit ELF file (`ELFCLASS32` in the standard)
    Elf32 = 1
  , /// 64-bit ELF file (`ELFCLASS64` in the standard)
    Elf64 = 2
}

impl Class {
    /// Returns true if the class field for this file is valid.
    #[inline]
    pub fn is_valid(&self) -> bool {
        match *self { Class::None => false
                    , _ => true
                    }
    }

}

impl ValidatesWord<u64> for Class {
    #[inline]
    fn check(&self) -> ElfResult<()> {
        use self::Class::*;
        match *self {
            None => Err("Invalid ELF type ELFCLASSNONE!")
          , Elf32 => Err("Cannot extract 64-bit section from 32-bit ELF")
          , Elf64 => Ok(())
        }
    }
}

impl ValidatesWord<u32> for Class {
    #[inline]
    fn check(&self) -> ElfResult<()> {
        use self::Class::*;
        match *self {
            None => Err("Invalid ELF type ELFCLASSNONE!")
          , Elf64 =>
                // TODO: is this actually true?
                //          - eliza, 03/08/2017
                Err("Cannot extract 32-bit section from 64-bit ELF")
          , Elf32 => Ok(())
        }
    }
}

/// Identifies the data encoding of the ELF file
#[derive(Copy, Clone, PartialEq, Debug)]
#[repr(u8)]
pub enum DataEncoding {
    /// Invalid data encoding (`ELFDATANONE` in the standard)
    None  = 0
  , /// Twos-complement little-endian data encoding
    /// (`ELFDATA2LSB` in the standard)
    LittleEndian = 1
  , /// Twos-complement big-endian data encoding
    /// (`ELFDATA2MSB` in the standard)
    BigEndian = 2
}

impl DataEncoding {
    /// Returns true if the data encoding field for this file is valid.
    #[inline]
    pub fn is_valid(&self) -> bool {
        match *self { DataEncoding::None => false
                    , _ => true
                    }
    }
}

/// Operating system ABI
#[derive(Copy, Clone, Debug)]
#[repr(u8)]
pub enum OsAbi { /// Ox00 also represents "none"
                 SystemV = 0x00
               , HpUx    = 0x01
               , NetBsd  = 0x02
               , Linux   = 0x03
               , Solaris = 0x06
               , Aix     = 0x07
               , Irix    = 0x08
               , FreeBsd = 0x09
               , OpenBsd = 0x0C
               , OpenVms = 0x0D
               }

/// Identifies the version of the ELF file
#[derive(Copy, Clone, PartialEq, Debug)]
#[repr(u8)]
pub enum Version { None = 0, Current = 1 }

#[derive(Clone, Copy, PartialEq)]
struct TypeRepr(u16);

impl TypeRepr {
    /// TODO: rewrite this as a `convert::Into` implementation
    ///         - eliza, 03/09/2017
    pub fn as_type(&self) -> Type {
        match self.0 {
            0 => Type::None
          , 1 => Type::Relocatable
          , 2 => Type::Executable
          , 3 => Type::SharedObject
          , 4 => Type::Core
          , anything => Type::Other(anything)
        }
    }
}

impl fmt::Debug for TypeRepr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.as_type().fmt(f)
    }
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum Type { None
              , Relocatable
              , Executable
              , SharedObject
              , Core
              , Other(u16)
              }

#[allow(non_camel_case_types)]
#[derive(Clone, Copy, PartialEq, Debug)]
#[repr(u16)]
pub enum Machine { None    = 0
                 , Sparc   = 0x02
                 , X86     = 0x03
                 , Mips    = 0x08
                 , PowerPc = 0x14
                 , Arm     = 0x28
                 , SuperH  = 0x2A
                 , Ia64    = 0x32
                 , X86_64  = 0x3E
                 , AArch64 = 0xB7
                 }
<|endoftext|>


<commit_before>use alloc::arc::Arc;

use collections::string::{String, ToString};
use collections::vec::Vec;

use core::cell::UnsafeCell;
use core::ops::DerefMut;
use core::{mem, ptr};

use common::elf::Elf;
use common::memory;

use scheduler::context::{CONTEXT_STACK_SIZE, CONTEXT_STACK_ADDR, context_switch,
context_userspace, Context, ContextMemory};

use schemes::Url;

use sync::Intex;

/// Execute an executable
pub fn execute(url: Url, mut args: Vec<String>) {
    let mut context_ptr: *mut Context = 0 as *mut Context;
    let mut entry: usize = 0;

    if let Ok(mut resource) = url.open() {
        let mut vec: Vec<u8> = Vec::new();
        resource.read_to_end(&mut vec);

        let executable = Elf::from_data(vec.as_ptr() as usize);
        entry = unsafe { executable.entry() };
        let mut memory = Vec::new();
        unsafe {
            for segment in executable.load_segment().iter() {
                let virtual_address = segment.vaddr as usize;
                let virtual_size = segment.mem_len as usize;

                //TODO: Warning: Investigate this hack!
                let hack = virtual_address % 4096;

                let physical_address = memory::alloc(virtual_size + hack);

                if physical_address > 0 {
                    debugln!("VADDR: {:X} OFF: {:X} FLG: {:X} HACK: {:X}", segment.vaddr, segment.off, segment.flags, hack);

                    // Copy progbits
                    ::memcpy((physical_address + hack) as *mut u8,
                             (executable.data + segment.off as usize) as *const u8,
                             segment.file_len as usize);
                    // Zero bss
                    if segment.mem_len > segment.file_len {
                        debugln!("BSS: {:X} {}", segment.vaddr + segment.file_len, segment.mem_len - segment.file_len);
                        ::memset((physical_address + hack + segment.file_len as usize) as *mut u8,
                                0,
                                segment.mem_len as usize - segment.file_len as usize);
                    }

                    memory.push(ContextMemory {
                        physical_address: physical_address,
                        virtual_address: virtual_address - hack,
                        virtual_size: virtual_size + hack,
                        writeable: segment.flags & 2 == 2
                    });
                }
            }
        }

        if entry > 0 && ! memory.is_empty() {
            args.insert(0, url.to_string());

            let mut contexts = ::env().contexts.lock();
            if let Some(mut context) = contexts.current_mut() {
                context.name = url.string;
                context.args = Arc::new(UnsafeCell::new(args));
                context.cwd = Arc::new(UnsafeCell::new(unsafe { (*context.cwd.get()).clone() }));

                unsafe { context.unmap() };
                context.memory = Arc::new(UnsafeCell::new(memory));
                unsafe { context.map() };

                context_ptr = context.deref_mut();
            }
        } else {
            debug!("{}: Invalid memory or entry\n", url.string);
        }
    } else {
        debug!("{}: Failed to open\n", url.string);
    }

    if context_ptr as usize > 0 {
        Context::spawn("kexec".to_string(), box move || {
            unsafe {
                let _intex = Intex::static_lock();

                let context = &mut *context_ptr;

                let mut context_args: Vec<usize> = Vec::new();
                context_args.push(0); // ENVP
                context_args.push(0); // ARGV NULL
                let mut argc = 0;
                for i in 0..(*context.args.get()).len() {
                    let reverse_i = (*context.args.get()).len() - i - 1;
                    if let Some(ref mut arg) = (*context.args.get()).get_mut(reverse_i) {
                        if ! arg.ends_with('\0') {
                            arg.push('\0');
                        }
                        context_args.push(arg.as_ptr() as usize);
                        argc += 1;
                    }
                }
                context_args.push(argc);

                context.sp = context.kernel_stack + CONTEXT_STACK_SIZE - 128;

                context.stack = Some(ContextMemory {
                    physical_address: memory::alloc(CONTEXT_STACK_SIZE),
                    virtual_address: CONTEXT_STACK_ADDR,
                    virtual_size: CONTEXT_STACK_SIZE,
                    writeable: true
                });

                let user_sp = if let Some(ref stack) = context.stack {
                    let mut sp = stack.physical_address + stack.virtual_size - 128;
                    for arg in context_args.iter() {
                        sp -= mem::size_of::<usize>();
                        ptr::write(sp as *mut usize, *arg);
                    }
                    sp - stack.physical_address + stack.virtual_address
                } else {
                    0
                };

                context.push(0x20 | 3);
                context.push(user_sp);
                context.push(/*1 << 9*/ 0);
                context.push(0x18 | 3);
                context.push(entry);
                context.push(context_userspace as usize);
            }
        });

        loop {
            unsafe { context_switch(false) };
        }
    }
}
<commit_msg>Segment unsafe code from safe code in executor.rs<commit_after>use alloc::arc::Arc;

use collections::string::{String, ToString};
use collections::vec::Vec;

use core::cell::UnsafeCell;
use core::ops::DerefMut;
use core::{mem, ptr};

use common::elf::Elf;
use common::memory;

use scheduler::context::{CONTEXT_STACK_SIZE, CONTEXT_STACK_ADDR, context_switch,
context_userspace, Context, ContextMemory};

use schemes::Url;

use sync::Intex;

/// Execute an executable
pub fn execute(url: Url, mut args: Vec<String>) {
    let mut context_ptr: *mut Context = 0 as *mut Context;
    let mut entry: usize = 0;

    if let Ok(mut resource) = url.open() {
        let mut vec: Vec<u8> = Vec::new();
        resource.read_to_end(&mut vec);

        let executable = Elf::from_data(vec.as_ptr() as usize);
        entry = unsafe { executable.entry() };
        let mut memory = Vec::new();
        unsafe {
            for segment in executable.load_segment().iter() {
                let virtual_address = segment.vaddr as usize;
                let virtual_size = segment.mem_len as usize;

                //TODO: Warning: Investigate this hack!
                let hack = virtual_address % 4096;

                let physical_address = memory::alloc(virtual_size + hack);

                if physical_address > 0 {
                    debugln!("VADDR: {:X} OFF: {:X} FLG: {:X} HACK: {:X}", segment.vaddr, segment.off, segment.flags, hack);

                    // Copy progbits
                    ::memcpy((physical_address + hack) as *mut u8,
                             (executable.data + segment.off as usize) as *const u8,
                             segment.file_len as usize);
                    // Zero bss
                    if segment.mem_len > segment.file_len {
                        debugln!("BSS: {:X} {}", segment.vaddr + segment.file_len, segment.mem_len - segment.file_len);
                        ::memset((physical_address + hack + segment.file_len as usize) as *mut u8,
                                0,
                                segment.mem_len as usize - segment.file_len as usize);
                    }

                    memory.push(ContextMemory {
                        physical_address: physical_address,
                        virtual_address: virtual_address - hack,
                        virtual_size: virtual_size + hack,
                        writeable: segment.flags & 2 == 2
                    });
                }
            }
        }

        if entry > 0 && ! memory.is_empty() {
            args.insert(0, url.to_string());

            let mut contexts = ::env().contexts.lock();
            if let Some(mut context) = contexts.current_mut() {
                context.name = url.string;
                context.args = Arc::new(UnsafeCell::new(args));
                context.cwd = Arc::new(UnsafeCell::new(unsafe { (*context.cwd.get()).clone() }));

                unsafe { context.unmap() };
                context.memory = Arc::new(UnsafeCell::new(memory));
                unsafe { context.map() };

                context_ptr = context.deref_mut();
            }
        } else {
            debug!("{}: Invalid memory or entry\n", url.string);
        }
    } else {
        debug!("{}: Failed to open\n", url.string);
    }

    if context_ptr as usize > 0 {
        Context::spawn("kexec".to_string(), box move || {
            let _intex = Intex::static_lock();

            let context = unsafe { &mut *context_ptr };

            let mut context_args: Vec<usize> = Vec::new();
            context_args.push(0); // ENVP
            context_args.push(0); // ARGV NULL
            let mut argc = 0;
            for i in 0..unsafe { (*context.args.get()).len() } {
                let reverse_i = unsafe { (*context.args.get()).len() } - i - 1;
                if let Some(ref mut arg) = unsafe { (*context.args.get()).get_mut(reverse_i) } {
                    if ! arg.ends_with('\0') {
                        arg.push('\0');
                    }
                    context_args.push(arg.as_ptr() as usize);
                    argc += 1;
                }
            }
            context_args.push(argc);

            context.sp = context.kernel_stack + CONTEXT_STACK_SIZE - 128;

            context.stack = Some(ContextMemory {
                physical_address: unsafe { memory::alloc(CONTEXT_STACK_SIZE) },
                virtual_address: CONTEXT_STACK_ADDR,
                virtual_size: CONTEXT_STACK_SIZE,
                writeable: true
            });

            let user_sp = if let Some(ref stack) = context.stack {
                let mut sp = stack.physical_address + stack.virtual_size - 128;
                for arg in context_args.iter() {
                    sp -= mem::size_of::<usize>();
                    unsafe { ptr::write(sp as *mut usize, *arg) };
                }
                sp - stack.physical_address + stack.virtual_address
            } else {
                0
            };

            unsafe {
                context.push(0x20 | 3);
                context.push(user_sp);
                context.push(/*1 << 9*/ 0);
                context.push(0x18 | 3);
                context.push(entry);
                context.push(context_userspace as usize);
            }
        });

        loop {
            unsafe { context_switch(false) };
        }
    }
}
<|endoftext|>


<commit_before>extern crate pcap;

fn main() {
    // list all of the devices pcap tells us are available
    for device in pcap::Devices::list_all().unwrap() {
        println!("Found device! {:?}", device);

        // now you can .open() on this device to get a Capture if you want
        let mut cap = pcap::Capture::new(device).unwrap();

        // get a packet from this device
        {
        	let packet = cap.next();

        	println!("got a packet! {:?}", packet);
        }
    }
}<commit_msg>Fixed comments<commit_after>extern crate pcap;

fn main() {
    // list all of the devices pcap tells us are available
    for device in pcap::Devices::list_all().unwrap() {
        println!("Found device! {:?}", device);

        // now you can create a Capture with this Device if you want.
        let mut cap = pcap::Capture::new(device).unwrap();

        // get a packet from this capture
        {
        	let packet = cap.next();

        	println!("got a packet! {:?}", packet);
        }
    }
}<|endoftext|>


<commit_before>use nom::*;
use std::str;

#[derive(Debug)]
pub struct Paragraph<'a> {
    parts: Vec<Text<'a>>,
}

#[derive(Debug, PartialEq)]
pub enum Text<'a> {
    Plain(&'a str),
    Bold(&'a str),
}

impl<'a> Paragraph<'a> {
    pub fn to_html(&self) -> String {
        self.parts.iter()
            .map(|text| match *text {
                Text::Plain(ref s) => s.to_string(),
                Text::Bold(ref s) => ["", &s[..], ""].concat(),
            })
            .collect::<Vec<_>>()
            .concat()
    }
}

named!(paragraph<Paragraph>,
    chain!(
        vec: many1!(alt!(plain | bold)),
        || { Paragraph { parts: vec } }
    )
);

named!(plain<Text>,
    chain!(
        bytes: is_not!("*") ~
        s:     expr_res!(str::from_utf8(bytes)),
        || { Text::Plain(s) }
    )
);

named!(bold<Text>,
    delimited!(
        char!('*'),
        chain!(
            bytes: is_not!("*") ~
            s:     expr_res!(str::from_utf8(bytes)),
            || { Text::Bold(s) }
        ),
        char!('*')
    )
);

#[test]
fn test_paragraph_parse() {
    if let IResult::Done(input, output) = paragraph(b"Hello *World*") {
        assert_eq!(input, b"");
        assert_eq!(output.parts, vec![Text::Plain("Hello "),
                                      Text::Bold("World")]);
    } else {
        panic!(r#" Paragraph parse failed with "Hello *World*" "#);
    }

    if let IResult::Done(input, output) = paragraph(b"*This* is *a test*") {
        assert_eq!(input, b"");
        assert_eq!(output.parts, vec![Text::Bold("This"),
                                      Text::Plain(" is "),
                                      Text::Bold("a test")]);
    } else {
        panic!(r#" Paragraph parse failed with "Hello *World*" "#);
    }
}

#[test]
fn test_plain_parse() {
    if let IResult::Done(input, output) = plain(b"Hello *World*") {
        assert_eq!(input, b"*World*");
        assert_eq!(output, Text::Plain("Hello "));
    } else {
        panic!(r#" Plain parse failed with "Hello *World*" "#);
    }

    match plain(b"*Hello World*") {
        IResult::Error(_) => (),
        _ => panic!(r#" Plain failed with "*Hello World*" "#),
    }
}

#[test]
fn test_html1() {
    let p = Paragraph {
        parts: vec![Text::Plain("plaintext "),
                    Text::Bold("bold text")],
    };

    assert!(p.to_html() == "plaintext bold text".to_owned());
}

#[test]
fn test_bold_parse() {
    if let IResult::Done(input, output) = bold(b"*Hello* World") {
        assert_eq!(input, b" World");
        assert_eq!(output, Text::Bold("Hello"));
    } else {
        panic!(r#" Bold failed with "*Hello* World" "#);
    }

    match bold(b"Hello World") {
        IResult::Error(_) => (),
        _ => panic!(r#" Bold failed with "Hello World" "#),
    };

    match bold(b"*Hello World") {
        IResult::Incomplete(_) => (),
        _ => panic!(r#" Bold failed with "*Hello World" "#),
    }
}

#[test]
#[should_panic]                 // TODO: Fix
fn test_bold_escaped() {
    if let IResult::Done(input, output) = bold(br"*Hell\*o* World") {
        assert_eq!(output, Text::Bold("Hell*o"));
        assert_eq!(input, b" World");
    } else {
        panic!("Bold failed with \"*Hell\\*o* World");
    }
}
<commit_msg>Move Text to use a tree-structure<commit_after>use nom::*;
use std::str;

#[derive(Debug, PartialEq)]
pub struct Paragraph<'a> {
    parts: Vec<Text<'a>>,
}

#[derive(Debug, PartialEq)]
pub enum Text<'a> {
    Plain(&'a str),
    Bold(Paragraph<'a>),
}

impl<'a> Paragraph<'a> {
    pub fn to_html(&self) -> String {
        self.parts.iter()
            .map(|text| match *text {
                Text::Plain(ref s) => s.to_string(),
                Text::Bold(ref p) => ["", &p.to_html(), ""].concat(),
            })
            .collect::<Vec<_>>()
            .concat()
    }
}

named!(paragraph<Paragraph>,
    chain!(
        vec: many1!(alt!(plain | bold)),
        || { Paragraph { parts: vec } }
    )
);

named!(only_paragraph<Paragraph>, complete!(paragraph));

named!(plain<Text>,
    chain!(
        s: map_res!(is_not!("*"), str::from_utf8),
        || { Text::Plain(s) }
    )
);

named!(bold<Text>,
    map_res!(
        delimited!(char!('*'), is_not!("*"), char!('*')),
        |bytes| {
            match only_paragraph(bytes) {
                IResult::Done(input, output) => Ok(Text::Bold(output)),
                IResult::Error(err) => Err(err),
                _ => panic!("Nom complete! macro returned an IResult::Incomplete")
            }
        }
    )
);

#[test]
fn test_plain_parse() {
    if let IResult::Done(input, output) = plain(b"Hello *World*") {
        assert_eq!(input, b"*World*");
        assert_eq!(output, Text::Plain("Hello "));
    } else {
        panic!(r#" Plain parse failed with "Hello *World*" "#);
    }

    match plain(b"*Hello World*") {
        IResult::Error(_) => (),
        _ => panic!(r#" Plain failed with "*Hello World*" "#),
    }
}

#[test]
fn test_bold_parse() {
    if let IResult::Done(input, output) = bold(b"*Hello* World") {
        assert_eq!(input, b" World");
        assert_eq!(output, Text::Bold(Paragraph {
            parts: vec![Text::Plain("Hello")]
        }));
    } else {
        panic!(r#" Bold failed with "*Hello* World" "#);
    }

    match bold(b"Hello World") {
        IResult::Error(_) => (),
        _ => panic!(r#" Bold failed with "Hello World" "#),
    };

    match bold(b"*Hello World") {
        IResult::Incomplete(_) => (),
        _ => panic!(r#" Bold failed with "*Hello World" "#),
    }
}

#[test]
fn test_paragraph_parse() {
    if let IResult::Done(input, output) = paragraph(b"Hello *World*") {
        assert_eq!(input, b"");
        assert_eq!(output.parts, vec![
            Text::Plain("Hello "),
            Text::Bold(Paragraph {parts: vec![Text::Plain("World")]}),
        ]);
    } else {
        panic!(r#" Paragraph parse failed with "Hello *World*" "#);
    }

    if let IResult::Done(input, output) = paragraph(b"*This* *is* *a* *test*") {
        assert_eq!(input, b"");
        assert_eq!(output.parts, vec![
            Text::Bold(Paragraph {parts: vec![Text::Plain("This")]}),
            Text::Plain(" "),
            Text::Bold(Paragraph {parts: vec![Text::Plain("is")]}),
            Text::Plain(" "),
            Text::Bold(Paragraph {parts: vec![Text::Plain("a")]}),
            Text::Plain(" "),
            Text::Bold(Paragraph {parts: vec![Text::Plain("test")]}),
        ]);
    } else {
        panic!(r#" Paragraph parse failed with "Hello *World*" "#);
    }
}

#[test]
fn test_html1() {
    let p = Paragraph { parts: vec![
            Text::Plain("plaintext "),
            Text::Bold(Paragraph { parts: vec![Text::Plain("bold text")]})
    ]};

    assert!(p.to_html() == "plaintext bold text".to_owned());
}
<|endoftext|>