<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2014 Jeremy Letang (letang.jeremy@gmail.com)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

#![crate_name = "svg"]
#![desc = "SVG generation in Rust"]
#![license = "MIT"]
#![crate_type = "dylib"]
#![crate_type = "rlib"]
// #![warn(missing_doc)]
#![allow(dead_code)]

use std::io::{Writer, IoResult};
use std::fmt::Show;
use std::vec::Vec;
use std::collections::HashMap;

pub use shapes::{Circle, Rect, RoundedRect, Ellipse, Line, PolyLine, Polygon};
pub use common::{rgb, rgba};
pub use text::Text;
pub use transform::Transform;

mod shapes;
mod text;
mod transform;
mod common;

static DOC_TYPE: &'static str = "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\" \
\"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n";
static XMLNS: &'static str = "version=\"1.1\" xmlns=\"http://www.w3.org/2000/svg\" \
xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n";
static STANDALONE_YES: &'static str = "<?xml version=\"1.0\" standalone=\"yes\"?>\n";
static STANDALONE_NO: &'static str = "<?xml version=\"1.0\" standalone=\"no\"?>\n";

trait SVGEntity {
    fn gen_output(&self) -> String;
}

struct Head {
    pub standalone: bool,
    pub width: i32,
    pub height: i32,
    pub view_box: Option<(i32, i32, i32, i32)>,
    pub desc: Option<String>,
    pub title: Option<String>
}

impl Head {
    pub fn new(width: i32, height: i32) -> Head {
        Head {
            standalone: false,
            width: width,
            height: height,
            view_box: None,
            desc: None,
            title: None
        }
    }
}

pub struct SVG<'a> {
    head: Head,
    content: String
}

fn make_attribs(attribs: &str) -> HashMap<String, String>{
    let mut h = HashMap::new();
    for s in attribs.split(' ') {
        let t: Vec<&str> = s.split('=').collect();
        h.insert(t.get(0).to_string(), t.get(1).to_string());
    }
    h
}

impl<'a> SVG<'a> {
    pub fn new(width: i32, height: i32) -> SVG {
        SVG {
            head: Head::new(width, height),
            content: String::new()
        }
    }

    pub fn standalone(&mut self, standalone: bool) {
        self.head.standalone = standalone;
    }

    pub fn view_box(&mut self,
                    orig_x: i32,
                    orig_y: i32,
                    width: i32,
                    height: i32) {
        self.head.view_box = Some((orig_x, orig_y, width, height));
    }

    pub fn desc(&mut self, text: &str) {
        self.head.desc = Some(text.to_string())
    }

    pub fn title(&mut self, text: &str) {
        self.head.title = Some(text.to_string())
    }

    pub fn add<T: SVGEntity>(&mut self, new_entity: &T) {
        self.content.push_str(new_entity.gen_output().as_slice());
    }

    pub fn circle(&mut self,
                  x: i32,
                  y: i32,
                  radius: u32,
                  attribs: &str) {
        self.content.push_str(Circle {
            x: x,
            y: y,
            radius: radius,
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn rect(&mut self,
                x: i32,
                y: i32,
                width: i32,
                height: i32,
                attribs: &str) {
        self.content.push_str(Rect {
            x: x,
            y: y,
            width: width,
            height: height,
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn rounded_rect(&mut self,
                        x: i32,
                        y: i32,
                        width: i32,
                        height: i32,
                        x_round: u32,
                        y_round: u32,
                        attribs: &str) {
        self.content.push_str(RoundedRect {
            x: x,
            y: y,
            width: width,
            height: height,
            x_round: x_round,
            y_round: y_round,
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn ellipse(&mut self,
                   x: i32,
                   y: i32,
                   x_radius: u32,
                   y_radius: u32,
                   attribs: &str) {
        self.content.push_str(Ellipse {
            x: x,
            y: y,
            x_radius: x_radius,
            y_radius: y_radius,
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn line(&mut self,
                x1: i32,
                y1: i32,
                x2: i32,
                y2: i32,
                attribs: &str) {
        self.content.push_str(Line {
            x1: x1,
            y1: y1,
            x2: x2,
            y2: y2,
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn polyline<T: Num + Show + Clone>(&mut self,
                                           points: &Vec<(T, T)>,
                                           attribs: &str) {
        self.content.push_str(PolyLine {
            points: points.clone(),
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn polygon<T: Num + Show + Clone>(&mut self,
                                          points: &Vec<(T, T)>,
                                          attribs: &str) {
        self.content.push_str(Polygon {
            points: points.clone(),
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn text(&mut self,
                x: i32,
                y: i32,
                text: &str,
                attribs: &str) {
        self.content.push_str(Text {
            x: x,
            y: y,
            text: text.to_string(),
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn g_begin(&mut self,
                   id: Option<&str>,
                   transform: Option<&Transform>,
                   attribs: Option<&HashMap<String, String>>) {
        self.content.push_str("<g ");
        match id {
            Some(i) => self.content.push_str(format!("id=\"{}\" ", i).as_slice()),
            None    => {/* nothing to do */}
        }
        match transform {
            Some(t) => self.content.push_str(format!("{} ", t.get()).as_slice()),
            None    => {/* nothing to do */}
        }
        match attribs {
            Some(a) => {
                for (at, value) in a.iter() {
                    self.content.push_str(format!("{}=\"{}\" ", *at, *value).as_slice())
                }
            },
            None    => {/* nothing to do */}
        }
        self.content.push_str(">\n");
    }

    pub fn g_id(&mut self, id: &str) {
        self.g_begin(Some(id.as_slice()), None, None)
    }

    pub fn g_transform(&mut self, transform: &Transform) {
        self.g_begin(None, Some(transform), None)
    }

    pub fn g_attribs(&mut self, attribs: &HashMap<String, String>) {
        self.g_begin(None, None, Some(attribs))
    }

    pub fn g_translate(&mut self, x: i32, y: i32) {
        let mut t = Transform::new();
        t.translate(x, y);
        self.g_begin(None, Some(&t), None)
    }

    pub fn g_rotate(&mut self, angle: i32) {
        let mut t = Transform::new();
        t.rotate(angle);
        self.g_begin(None, Some(&t), None)
    }

    pub fn g_scale(&mut self, x_scale: i32, y_scale: i32) {
        let mut t = Transform::new();
        t.scale(x_scale, y_scale);
        self.g_begin(None, Some(&t), None)
    }

    // FIXME: test if a skew of 0 for y or x don't break
    pub fn g_skew(&mut self, x_factor: i32, y_factor: i32) {
        let mut t = Transform::new();
        t.skew_x(x_factor);
        t.skew_y(y_factor);
        self.g_begin(None, Some(&t), None)
    }

    pub fn g_end(&mut self) {
        self.content.push_str("</g>\n");
    }

    pub fn finalize(&mut self, output: &'a mut Writer) -> IoResult<()>{
        let mut o = String::new();
        // Head
        match self.head.standalone {
            true    => o.push_str(STANDALONE_YES),
            false   => o.push_str(STANDALONE_NO)
        };
        o.push_str(DOC_TYPE);
        o.push_str(format!("<svg width=\"{}cm\" height=\"{}cm\" ",
                           self.head.width, self.head.height).as_slice());
        match self.head.view_box {
            Some((x, y, width, height)) => {
                o.push_str(format!("viewBox=\"{} {} {} {}\" ", x, y, width, height).as_slice())
            },
            None                        => {/* nothing to do */}
        }
        o.push_str(XMLNS);
        match self.head.title {
            Some(ref t) => o.push_str(format!("<title>{}</title>\n", *t).as_slice()),
            None    => {/* nothing to do */}
        }
        match self.head.desc {
            Some(ref d) => o.push_str(format!("<desc>{}</desc>\n", *d).as_slice()),
            None    => {/* nothing to do */}
        }
        // Body
        o.push_str(self.content.clone().as_slice());
        // Close
        o.push_str("</svg>\n");
        output.write_str(o.as_slice())
    }
}
<commit_msg>fix to last Rust - lifetime - index accessing to vector<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2014 Jeremy Letang (letang.jeremy@gmail.com)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

#![crate_name = "svg"]
#![desc = "SVG generation in Rust"]
#![license = "MIT"]
#![crate_type = "dylib"]
#![crate_type = "rlib"]
// #![warn(missing_doc)]
#![allow(dead_code)]

use std::io::{Writer, IoResult};
use std::fmt::Show;
use std::vec::Vec;
use std::collections::HashMap;

pub use shapes::{Circle, Rect, RoundedRect, Ellipse, Line, PolyLine, Polygon};
pub use common::{rgb, rgba};
pub use text::Text;
pub use transform::Transform;

mod shapes;
mod text;
mod transform;
mod common;

static DOC_TYPE: &'static str = "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\" \
\"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n";
static XMLNS: &'static str = "version=\"1.1\" xmlns=\"http://www.w3.org/2000/svg\" \
xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n";
static STANDALONE_YES: &'static str = "<?xml version=\"1.0\" standalone=\"yes\"?>\n";
static STANDALONE_NO: &'static str = "<?xml version=\"1.0\" standalone=\"no\"?>\n";

trait SVGEntity {
    fn gen_output(&self) -> String;
}

struct Head {
    pub standalone: bool,
    pub width: i32,
    pub height: i32,
    pub view_box: Option<(i32, i32, i32, i32)>,
    pub desc: Option<String>,
    pub title: Option<String>
}

impl Head {
    pub fn new(width: i32, height: i32) -> Head {
        Head {
            standalone: false,
            width: width,
            height: height,
            view_box: None,
            desc: None,
            title: None
        }
    }
}

pub struct SVG<'a> {
    head: Head,
    content: String
}

fn make_attribs(attribs: &str) -> HashMap<String, String>{
    let mut h = HashMap::new();
    for s in attribs.split(' ') {
        let t: Vec<&str> = s.split('=').collect();
        h.insert(t[0u].to_string(), t[1u].to_string());
    }
    h
}

impl<'a> SVG<'a> {
    pub fn new(width: i32, height: i32) -> SVG<'a> {
        SVG {
            head: Head::new(width, height),
            content: String::new()
        }
    }

    pub fn standalone(&mut self, standalone: bool) {
        self.head.standalone = standalone;
    }

    pub fn view_box(&mut self,
                    orig_x: i32,
                    orig_y: i32,
                    width: i32,
                    height: i32) {
        self.head.view_box = Some((orig_x, orig_y, width, height));
    }

    pub fn desc(&mut self, text: &str) {
        self.head.desc = Some(text.to_string())
    }

    pub fn title(&mut self, text: &str) {
        self.head.title = Some(text.to_string())
    }

    pub fn add<T: SVGEntity>(&mut self, new_entity: &T) {
        self.content.push_str(new_entity.gen_output().as_slice());
    }

    pub fn circle(&mut self,
                  x: i32,
                  y: i32,
                  radius: u32,
                  attribs: &str) {
        self.content.push_str(Circle {
            x: x,
            y: y,
            radius: radius,
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn rect(&mut self,
                x: i32,
                y: i32,
                width: i32,
                height: i32,
                attribs: &str) {
        self.content.push_str(Rect {
            x: x,
            y: y,
            width: width,
            height: height,
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn rounded_rect(&mut self,
                        x: i32,
                        y: i32,
                        width: i32,
                        height: i32,
                        x_round: u32,
                        y_round: u32,
                        attribs: &str) {
        self.content.push_str(RoundedRect {
            x: x,
            y: y,
            width: width,
            height: height,
            x_round: x_round,
            y_round: y_round,
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn ellipse(&mut self,
                   x: i32,
                   y: i32,
                   x_radius: u32,
                   y_radius: u32,
                   attribs: &str) {
        self.content.push_str(Ellipse {
            x: x,
            y: y,
            x_radius: x_radius,
            y_radius: y_radius,
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn line(&mut self,
                x1: i32,
                y1: i32,
                x2: i32,
                y2: i32,
                attribs: &str) {
        self.content.push_str(Line {
            x1: x1,
            y1: y1,
            x2: x2,
            y2: y2,
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn polyline<T: Num + Show + Clone>(&mut self,
                                           points: &Vec<(T, T)>,
                                           attribs: &str) {
        self.content.push_str(PolyLine {
            points: points.clone(),
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn polygon<T: Num + Show + Clone>(&mut self,
                                          points: &Vec<(T, T)>,
                                          attribs: &str) {
        self.content.push_str(Polygon {
            points: points.clone(),
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn text(&mut self,
                x: i32,
                y: i32,
                text: &str,
                attribs: &str) {
        self.content.push_str(Text {
            x: x,
            y: y,
            text: text.to_string(),
            attribs: make_attribs(attribs),
            transform: None
        }.gen_output().as_slice())
    }

    pub fn g_begin(&mut self,
                   id: Option<&str>,
                   transform: Option<&Transform>,
                   attribs: Option<&HashMap<String, String>>) {
        self.content.push_str("<g ");
        match id {
            Some(i) => self.content.push_str(format!("id=\"{}\" ", i).as_slice()),
            None    => {/* nothing to do */}
        }
        match transform {
            Some(t) => self.content.push_str(format!("{} ", t.get()).as_slice()),
            None    => {/* nothing to do */}
        }
        match attribs {
            Some(a) => {
                for (at, value) in a.iter() {
                    self.content.push_str(format!("{}=\"{}\" ", *at, *value).as_slice())
                }
            },
            None    => {/* nothing to do */}
        }
        self.content.push_str(">\n");
    }

    pub fn g_id(&mut self, id: &str) {
        self.g_begin(Some(id.as_slice()), None, None)
    }

    pub fn g_transform(&mut self, transform: &Transform) {
        self.g_begin(None, Some(transform), None)
    }

    pub fn g_attribs(&mut self, attribs: &HashMap<String, String>) {
        self.g_begin(None, None, Some(attribs))
    }

    pub fn g_translate(&mut self, x: i32, y: i32) {
        let mut t = Transform::new();
        t.translate(x, y);
        self.g_begin(None, Some(&t), None)
    }

    pub fn g_rotate(&mut self, angle: i32) {
        let mut t = Transform::new();
        t.rotate(angle);
        self.g_begin(None, Some(&t), None)
    }

    pub fn g_scale(&mut self, x_scale: i32, y_scale: i32) {
        let mut t = Transform::new();
        t.scale(x_scale, y_scale);
        self.g_begin(None, Some(&t), None)
    }

    // FIXME: test if a skew of 0 for y or x don't break
    pub fn g_skew(&mut self, x_factor: i32, y_factor: i32) {
        let mut t = Transform::new();
        t.skew_x(x_factor);
        t.skew_y(y_factor);
        self.g_begin(None, Some(&t), None)
    }

    pub fn g_end(&mut self) {
        self.content.push_str("</g>\n");
    }

    pub fn finalize(&mut self, output: &'a mut Writer) -> IoResult<()>{
        let mut o = String::new();
        // Head
        match self.head.standalone {
            true    => o.push_str(STANDALONE_YES),
            false   => o.push_str(STANDALONE_NO)
        };
        o.push_str(DOC_TYPE);
        o.push_str(format!("<svg width=\"{}cm\" height=\"{}cm\" ",
                           self.head.width, self.head.height).as_slice());
        match self.head.view_box {
            Some((x, y, width, height)) => {
                o.push_str(format!("viewBox=\"{} {} {} {}\" ", x, y, width, height).as_slice())
            },
            None                        => {/* nothing to do */}
        }
        o.push_str(XMLNS);
        match self.head.title {
            Some(ref t) => o.push_str(format!("<title>{}</title>\n", *t).as_slice()),
            None    => {/* nothing to do */}
        }
        match self.head.desc {
            Some(ref d) => o.push_str(format!("<desc>{}</desc>\n", *d).as_slice()),
            None    => {/* nothing to do */}
        }
        // Body
        o.push_str(self.content.clone().as_slice());
        // Close
        o.push_str("</svg>\n");
        output.write_str(o.as_slice())
    }
}
<|endoftext|>


<commit_before>use std::io::{self, Write, Read, Cursor};
use std::net::{TcpStream, Shutdown};
use std::time::Duration;
use std::borrow::BorrowMut;
use byteorder::{BigEndian, WriteBytesExt};
use bincode::rustc_serialize::{encode, decode};
use bincode::SizeLimit;
use enum_primitive::FromPrimitive;
use {Function, Result, ExceptionCode, Error, Coil, binary, Client};

const MODBUS_PROTOCOL_TCP: u16 = 0x0000;
const MODBUS_TCP_DEFAULT_PORT: u16 = 502;
const MODBUS_HEADER_SIZE: usize = 7;
const MODBUS_MAX_READ_COUNT: usize = 0x7d;
const MODBUS_MAX_WRITE_COUNT: usize = 0x79;

#[derive(RustcEncodable, RustcDecodable)]
#[repr(packed)]
struct Header {
    tid: u16,
    pid: u16,
    len: u16,
    uid: u8,
}

impl Header {
    fn new(transport: &mut Transport, len: u16) -> Header {
        Header {
            tid: transport.new_tid(),
            pid: MODBUS_PROTOCOL_TCP,
            len: len - MODBUS_HEADER_SIZE as u16,
            uid: transport.uid,
        }
    }
}

/// Context object which holds state for all modbus operations.
pub struct Transport {
    tid: u16,
    uid: u8,
    stream: TcpStream,
}

impl Transport {
    /// Create a new context context object and connect it to `addr` on modbus-tcp default
    /// port (502)
    pub fn new(addr: &str) -> io::Result<Transport> {
        Self::new_with_port(addr, MODBUS_TCP_DEFAULT_PORT)
    }

    /// Create a new context object and connect it to `addr` on port `port`
    pub fn new_with_port(addr: &str, port: u16) -> io::Result<Transport> {
        match TcpStream::connect((addr, port)) {
            Ok(s) => {
                // set some sane tcp socket options
                let t = Duration::from_secs(5);
                try!(s.set_read_timeout(Some(t)));
                try!(s.set_write_timeout(Some(t)));
                //                try!(s.set_nodelay(true));
                //                try!(s.set_keepalive(None));
                Ok(Transport {
                    tid: 0,
                    uid: 1,
                    stream: s,
                })
            }
            Err(e) => Err(e),
        }
    }

    // Create a new transaction Id, incrementing the previous one.
    // The Id is wrapping around if the Id reaches `u16::MAX`.
    fn new_tid(&mut self) -> u16 {
        self.tid = self.tid.wrapping_add(1);
        self.tid
    }

    fn read(self: &mut Self, fun: Function) -> Result<Vec<u8>> {
        let packed_size = |v: u16| {
            v / 8 +
            if v % 8 > 0 {
                1
            } else {
                0
            }
        };
        let (addr, count, expected_bytes) = match fun {
            Function::ReadCoils(a, c) => (a, c, packed_size(c) as usize),
            Function::ReadDiscreteInputs(a, c) => (a, c, packed_size(c) as usize),
            Function::ReadHoldingRegisters(a, c) => (a, c, 2 * c as usize),
            Function::ReadInputRegisters(a, c) => (a, c, 2 * c as usize),
            _ => panic!("Unexpected modbus function"),
        };

        if count < 1 || count as usize > MODBUS_MAX_READ_COUNT {
            return Err(Error::InvalidData);
        }

        let header = Header::new(self, MODBUS_HEADER_SIZE as u16 + 6u16);
        let mut buff = try!(encode(&header, SizeLimit::Infinite));
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(count));

        match self.stream.write_all(&buff[..]) {
            Ok(_s) => {
                let mut reply = vec![0; MODBUS_HEADER_SIZE + expected_bytes + 2];
                match self.stream.read(&mut reply) {
                    Ok(_s) => {
                        let resp_hd = try!(decode(&reply[..MODBUS_HEADER_SIZE]));
                        try!(Transport::validate_response_header(&header, &resp_hd));
                        try!(Transport::validate_response_code(&buff, &reply[..]));
                        Transport::get_reply_data(&reply, expected_bytes)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

    fn validate_response_header(req: &Header, resp: &Header) -> Result<()> {
        if req.tid != resp.tid || resp.pid != MODBUS_PROTOCOL_TCP {
            Err(Error::InvalidResponse)
        } else {
            Ok(())
        }
    }

    fn validate_response_code(req: &[u8], resp: &[u8]) -> Result<()> {
        if req[7] + 0x80 == resp[7] {
            match ExceptionCode::from_u8(resp[8]) {
                Some(code) => Err(Error::Exception(code)),
                None => Err(Error::InvalidResponse)
            }
        } else if req[7] != resp[7] {
            Err(Error::InvalidResponse)
        } else {
            Ok(())
        }
    }

    fn get_reply_data(reply: &[u8], expected_bytes: usize) -> Result<Vec<u8>> {
        if reply[8] as usize != expected_bytes ||
           reply.len() != MODBUS_HEADER_SIZE + expected_bytes + 2 {
            Err(Error::InvalidData)
        } else {
            let mut d = Vec::new();
            d.extend(reply[MODBUS_HEADER_SIZE + 2..].iter());
            Ok(d)
        }
    }

    fn write_single(self: &mut Self, fun: Function) -> Result<()> {
        let (addr, value) = match fun {
            Function::WriteSingleCoil(a, v) => (a, v),
            Function::WriteSingleRegister(a, v) => (a, v),
            _ => panic!("Unexpected modbus function"),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(value));
        self.write(&mut buff)
    }

    fn write_multiple(self: &mut Self, fun: Function) -> Result<()> {
        let (addr, quantity, values) = match fun {
            Function::WriteMultipleCoils(a, q, v) => (a, q, v),
            Function::WriteMultipleRegisters(a, q, v) => (a, q, v),
            _ => panic!("Unexpected modbus function"),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(quantity));
        try!(buff.write_u8(values.len() as u8));
        for v in values {
            try!(buff.write_u8(*v));
        }
        self.write(&mut buff)
    }

    fn write(self: &mut Self, buff: &mut [u8]) -> Result<()> {
        if buff.len() < 1 || buff.len() > MODBUS_MAX_WRITE_COUNT {
            return Err(Error::InvalidData);
        }
        let header = Header::new(self, buff.len() as u16 + 1u16);
        let head_buff = try!(encode(&header, SizeLimit::Infinite));
        {
            let mut start = Cursor::new(buff.borrow_mut());
            try!(start.write(&head_buff[..]));
        }
        match self.stream.write_all(&buff[..]) {
            Ok(_s) => {
                let reply = &mut [0; 12];
                match self.stream.read(reply) {
                    Ok(_s) => {
                        let resp_hd = try!(decode(&reply[..MODBUS_HEADER_SIZE]));
                        try!(Transport::validate_response_header(&header, &resp_hd));
                        Transport::validate_response_code(&buff, reply)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

}

impl Drop for Transport {
    fn drop(&mut self) {
        self.stream.shutdown(Shutdown::Both).unwrap();
    }
}

impl Client for Transport {
    /// Read `count` bits starting at address `addr`.
    fn read_coils(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = try!(self.read(Function::ReadCoils(addr, count)));
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` input bits starting at address `addr`.
    fn read_discrete_inputs(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = try!(self.read(Function::ReadDiscreteInputs(addr, count)));
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` 16bit registers starting at address `addr`.
    fn read_holding_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = try!(self.read(Function::ReadHoldingRegisters(addr, count)));
        binary::pack_bytes(&bytes[..])
    }

    /// Read `count` 16bit input registers starting at address `addr`.
    fn read_input_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = try!(self.read(Function::ReadInputRegisters(addr, count)));
        binary::pack_bytes(&bytes[..])
    }

    /// Write a single coil (bit) to address `addr`.
    fn write_single_coil(self: &mut Self, addr: u16, value: Coil) -> Result<()> {
        self.write_single(Function::WriteSingleCoil(addr, value.code()))
    }

    /// Write a single 16bit register to address `addr`.
    fn write_single_register(self: &mut Self, addr: u16, value: u16) -> Result<()> {
        self.write_single(Function::WriteSingleRegister(addr, value))
    }

    /// Write a multiple coils (bits) starting at address `addr`.
    fn write_multiple_coils(self: &mut Self, addr: u16, values: &[Coil]) -> Result<()> {
        let bytes = binary::pack_bits(values);
        self.write_multiple(Function::WriteMultipleCoils(addr, values.len() as u16, &bytes[..]))
    }

    /// Write a multiple 16bit registers starting at address `addr`.
    fn write_multiple_registers(self: &mut Self, addr: u16, values: &[u16]) -> Result<()> {
        let bytes = binary::unpack_bytes(values);
        self.write_multiple(Function::WriteMultipleRegisters(addr, values.len() as u16, &bytes[..]))
    }
}
<commit_msg>manually close a the TCP connection<commit_after>use std::io::{self, Write, Read, Cursor};
use std::net::{TcpStream, Shutdown};
use std::time::Duration;
use std::borrow::BorrowMut;
use byteorder::{BigEndian, WriteBytesExt};
use bincode::rustc_serialize::{encode, decode};
use bincode::SizeLimit;
use enum_primitive::FromPrimitive;
use {Function, Result, ExceptionCode, Error, Coil, binary, Client};

const MODBUS_PROTOCOL_TCP: u16 = 0x0000;
const MODBUS_TCP_DEFAULT_PORT: u16 = 502;
const MODBUS_HEADER_SIZE: usize = 7;
const MODBUS_MAX_READ_COUNT: usize = 0x7d;
const MODBUS_MAX_WRITE_COUNT: usize = 0x79;

#[derive(RustcEncodable, RustcDecodable)]
#[repr(packed)]
struct Header {
    tid: u16,
    pid: u16,
    len: u16,
    uid: u8,
}

impl Header {
    fn new(transport: &mut Transport, len: u16) -> Header {
        Header {
            tid: transport.new_tid(),
            pid: MODBUS_PROTOCOL_TCP,
            len: len - MODBUS_HEADER_SIZE as u16,
            uid: transport.uid,
        }
    }
}

/// Context object which holds state for all modbus operations.
pub struct Transport {
    tid: u16,
    uid: u8,
    stream: TcpStream,
}

impl Transport {
    /// Create a new context context object and connect it to `addr` on modbus-tcp default
    /// port (502)
    pub fn new(addr: &str) -> io::Result<Transport> {
        Self::new_with_port(addr, MODBUS_TCP_DEFAULT_PORT)
    }

    /// Create a new context object and connect it to `addr` on port `port`
    pub fn new_with_port(addr: &str, port: u16) -> io::Result<Transport> {
        match TcpStream::connect((addr, port)) {
            Ok(s) => {
                // set some sane tcp socket options
                let t = Duration::from_secs(5);
                try!(s.set_read_timeout(Some(t)));
                try!(s.set_write_timeout(Some(t)));
                //                try!(s.set_nodelay(true));
                //                try!(s.set_keepalive(None));
                Ok(Transport {
                    tid: 0,
                    uid: 1,
                    stream: s,
                })
            }
            Err(e) => Err(e),
        }
    }

    // Create a new transaction Id, incrementing the previous one.
    // The Id is wrapping around if the Id reaches `u16::MAX`.
    fn new_tid(&mut self) -> u16 {
        self.tid = self.tid.wrapping_add(1);
        self.tid
    }

    fn read(self: &mut Self, fun: Function) -> Result<Vec<u8>> {
        let packed_size = |v: u16| {
            v / 8 +
            if v % 8 > 0 {
                1
            } else {
                0
            }
        };
        let (addr, count, expected_bytes) = match fun {
            Function::ReadCoils(a, c) => (a, c, packed_size(c) as usize),
            Function::ReadDiscreteInputs(a, c) => (a, c, packed_size(c) as usize),
            Function::ReadHoldingRegisters(a, c) => (a, c, 2 * c as usize),
            Function::ReadInputRegisters(a, c) => (a, c, 2 * c as usize),
            _ => panic!("Unexpected modbus function"),
        };

        if count < 1 || count as usize > MODBUS_MAX_READ_COUNT {
            return Err(Error::InvalidData);
        }

        let header = Header::new(self, MODBUS_HEADER_SIZE as u16 + 6u16);
        let mut buff = try!(encode(&header, SizeLimit::Infinite));
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(count));

        match self.stream.write_all(&buff[..]) {
            Ok(_s) => {
                let mut reply = vec![0; MODBUS_HEADER_SIZE + expected_bytes + 2];
                match self.stream.read(&mut reply) {
                    Ok(_s) => {
                        let resp_hd = try!(decode(&reply[..MODBUS_HEADER_SIZE]));
                        try!(Transport::validate_response_header(&header, &resp_hd));
                        try!(Transport::validate_response_code(&buff, &reply[..]));
                        Transport::get_reply_data(&reply, expected_bytes)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

    fn validate_response_header(req: &Header, resp: &Header) -> Result<()> {
        if req.tid != resp.tid || resp.pid != MODBUS_PROTOCOL_TCP {
            Err(Error::InvalidResponse)
        } else {
            Ok(())
        }
    }

    fn validate_response_code(req: &[u8], resp: &[u8]) -> Result<()> {
        if req[7] + 0x80 == resp[7] {
            match ExceptionCode::from_u8(resp[8]) {
                Some(code) => Err(Error::Exception(code)),
                None => Err(Error::InvalidResponse)
            }
        } else if req[7] != resp[7] {
            Err(Error::InvalidResponse)
        } else {
            Ok(())
        }
    }

    fn get_reply_data(reply: &[u8], expected_bytes: usize) -> Result<Vec<u8>> {
        if reply[8] as usize != expected_bytes ||
           reply.len() != MODBUS_HEADER_SIZE + expected_bytes + 2 {
            Err(Error::InvalidData)
        } else {
            let mut d = Vec::new();
            d.extend(reply[MODBUS_HEADER_SIZE + 2..].iter());
            Ok(d)
        }
    }

    fn write_single(self: &mut Self, fun: Function) -> Result<()> {
        let (addr, value) = match fun {
            Function::WriteSingleCoil(a, v) => (a, v),
            Function::WriteSingleRegister(a, v) => (a, v),
            _ => panic!("Unexpected modbus function"),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(value));
        self.write(&mut buff)
    }

    fn write_multiple(self: &mut Self, fun: Function) -> Result<()> {
        let (addr, quantity, values) = match fun {
            Function::WriteMultipleCoils(a, q, v) => (a, q, v),
            Function::WriteMultipleRegisters(a, q, v) => (a, q, v),
            _ => panic!("Unexpected modbus function"),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        try!(buff.write_u8(fun.code()));
        try!(buff.write_u16::<BigEndian>(addr));
        try!(buff.write_u16::<BigEndian>(quantity));
        try!(buff.write_u8(values.len() as u8));
        for v in values {
            try!(buff.write_u8(*v));
        }
        self.write(&mut buff)
    }

    fn write(self: &mut Self, buff: &mut [u8]) -> Result<()> {
        if buff.len() < 1 || buff.len() > MODBUS_MAX_WRITE_COUNT {
            return Err(Error::InvalidData);
        }
        let header = Header::new(self, buff.len() as u16 + 1u16);
        let head_buff = try!(encode(&header, SizeLimit::Infinite));
        {
            let mut start = Cursor::new(buff.borrow_mut());
            try!(start.write(&head_buff[..]));
        }
        match self.stream.write_all(&buff[..]) {
            Ok(_s) => {
                let reply = &mut [0; 12];
                match self.stream.read(reply) {
                    Ok(_s) => {
                        let resp_hd = try!(decode(&reply[..MODBUS_HEADER_SIZE]));
                        try!(Transport::validate_response_header(&header, &resp_hd));
                        Transport::validate_response_code(&buff, reply)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

    pub fn close(self: &mut Self) -> Result<()> {
        self.stream.shutdown(Shutdown::Both).map_err(Error::Io)
    }

}

impl Client for Transport {
    /// Read `count` bits starting at address `addr`.
    fn read_coils(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = try!(self.read(Function::ReadCoils(addr, count)));
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` input bits starting at address `addr`.
    fn read_discrete_inputs(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = try!(self.read(Function::ReadDiscreteInputs(addr, count)));
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` 16bit registers starting at address `addr`.
    fn read_holding_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = try!(self.read(Function::ReadHoldingRegisters(addr, count)));
        binary::pack_bytes(&bytes[..])
    }

    /// Read `count` 16bit input registers starting at address `addr`.
    fn read_input_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = try!(self.read(Function::ReadInputRegisters(addr, count)));
        binary::pack_bytes(&bytes[..])
    }

    /// Write a single coil (bit) to address `addr`.
    fn write_single_coil(self: &mut Self, addr: u16, value: Coil) -> Result<()> {
        self.write_single(Function::WriteSingleCoil(addr, value.code()))
    }

    /// Write a single 16bit register to address `addr`.
    fn write_single_register(self: &mut Self, addr: u16, value: u16) -> Result<()> {
        self.write_single(Function::WriteSingleRegister(addr, value))
    }

    /// Write a multiple coils (bits) starting at address `addr`.
    fn write_multiple_coils(self: &mut Self, addr: u16, values: &[Coil]) -> Result<()> {
        let bytes = binary::pack_bits(values);
        self.write_multiple(Function::WriteMultipleCoils(addr, values.len() as u16, &bytes[..]))
    }

    /// Write a multiple 16bit registers starting at address `addr`.
    fn write_multiple_registers(self: &mut Self, addr: u16, values: &[u16]) -> Result<()> {
        let bytes = binary::unpack_bytes(values);
        self.write_multiple(Function::WriteMultipleRegisters(addr, values.len() as u16, &bytes[..]))
    }
}
<|endoftext|>


<commit_before>use super::RW;
use super::evented::{Evented, EventedImpl, MioAdapter};
use std::io;
use std::net::SocketAddr;
use mio_orig;
use std;

pub use mio_orig::tcp::Shutdown;

/// TCP Listener
pub type TcpListener = MioAdapter<mio_orig::tcp::TcpListener>;

impl TcpListener {
    /// Local address
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.shared().io_ref().local_addr()
    }

    /// TODO: document
    pub fn take_socket_error(&self) -> io::Result<()> {
        self.shared().io_ref().take_socket_error()
    }

    /// Try cloning the listener descriptor.
    pub fn try_clone(&self) -> io::Result<TcpListener> {
        self.shared().io_ref().try_clone().map(MioAdapter::new)
    }
}

impl TcpListener {
    /// Bind to a port
    pub fn bind(addr: &SocketAddr) -> io::Result<Self> {
        mio_orig::tcp::TcpListener::bind(addr).map(MioAdapter::new)
    }

    /// Creates a new TcpListener from an instance of a `std::net::TcpListener` type.
    pub fn from_listener(listener: std::net::TcpListener, addr: &SocketAddr) -> io::Result<Self> {
        mio_orig::tcp::TcpListener::from_listener(listener, addr).map(MioAdapter::new)
    }
}

/// TCP Stream
pub type TcpStream = MioAdapter<mio_orig::tcp::TcpStream>;

impl TcpStream {
    /// Create a new TCP stream an issue a non-blocking connect to the specified address.
    pub fn connect(addr: &SocketAddr) -> io::Result<Self> {
        mio_orig::tcp::TcpStream::connect(addr).map(|t| {
            let stream = MioAdapter::new(t);
            stream.block_on(RW::write());
            stream
        })
    }

    /// Creates a new TcpStream from the pending socket inside the given
    /// `std::net::TcpBuilder`, connecting it to the address specified.
    pub fn connect_stream(stream: std::net::TcpStream, addr: &SocketAddr) -> io::Result<Self> {
        mio_orig::tcp::TcpStream::connect_stream(stream, addr).map(|t| {
            let stream = MioAdapter::new(t);
            stream.block_on(RW::write());
            stream
        })
    }

    /// Local address of connection.
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.shared().io_ref().local_addr()
    }

    /// Peer address of connection.
    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
        self.shared().io_ref().peer_addr()
    }

    /// Shutdown the connection.
    pub fn shutdown(&self, how: mio_orig::tcp::Shutdown) -> io::Result<()> {
        self.shared().io_ref().shutdown(how)
    }

    /// Set `no_delay`.
    pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
        self.shared().io_ref().set_nodelay(nodelay)
    }

    /// Set keepalive.
    pub fn set_keepalive(&self, seconds: Option<u32>) -> io::Result<()> {
        self.shared().io_ref().set_keepalive(seconds)
    }

    /// TODO: document
    pub fn take_socket_error(&self) -> io::Result<()> {
        self.shared().io_ref().take_socket_error()
    }

    /// Try cloning the socket descriptor.
    pub fn try_clone(&self) -> io::Result<TcpStream> {
        self.shared().io_ref().try_clone().map(MioAdapter::new)
    }
}
<commit_msg>Block on TcpStream::connect and check error.<commit_after>use super::RW;
use super::evented::{Evented, EventedImpl, MioAdapter};
use std::io;
use std::net::SocketAddr;
use mio_orig;
use std;

pub use mio_orig::tcp::Shutdown;

/// TCP Listener
pub type TcpListener = MioAdapter<mio_orig::tcp::TcpListener>;

impl TcpListener {
    /// Local address
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.shared().io_ref().local_addr()
    }

    /// TODO: document
    pub fn take_socket_error(&self) -> io::Result<()> {
        self.shared().io_ref().take_socket_error()
    }

    /// Try cloning the listener descriptor.
    pub fn try_clone(&self) -> io::Result<TcpListener> {
        self.shared().io_ref().try_clone().map(MioAdapter::new)
    }
}

impl TcpListener {
    /// Bind to a port
    pub fn bind(addr: &SocketAddr) -> io::Result<Self> {
        mio_orig::tcp::TcpListener::bind(addr).map(MioAdapter::new)
    }

    /// Creates a new TcpListener from an instance of a `std::net::TcpListener` type.
    pub fn from_listener(listener: std::net::TcpListener, addr: &SocketAddr) -> io::Result<Self> {
        mio_orig::tcp::TcpListener::from_listener(listener, addr).map(MioAdapter::new)
    }
}

/// TCP Stream
pub type TcpStream = MioAdapter<mio_orig::tcp::TcpStream>;

impl TcpStream {
    /// Create a new TCP stream an issue a non-blocking connect to the specified address.
    pub fn connect(addr: &SocketAddr) -> io::Result<Self> {
        let stream = mio_orig::tcp::TcpStream::connect(addr).map(|t| {
            let stream = MioAdapter::new(t);
            stream.block_on(RW::write());
            stream
        });

        if let Ok(ref stream) = stream {
            if let Err(err) = stream.shared().io_ref().take_socket_error() {
                return Err(err)
            }
        }

        stream

    }

    /// Creates a new TcpStream from the pending socket inside the given
    /// `std::net::TcpBuilder`, connecting it to the address specified.
    pub fn connect_stream(stream: std::net::TcpStream, addr: &SocketAddr) -> io::Result<Self> {
        let stream = mio_orig::tcp::TcpStream::connect_stream(stream, addr).map(|t| {
            let stream = MioAdapter::new(t);
            stream.block_on(RW::write());

            stream
        });

        if let Ok(ref stream) = stream {
            if let Err(err) = stream.shared().io_ref().take_socket_error() {
                return Err(err)
            }
        }

        stream
    }

    /// Local address of connection.
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.shared().io_ref().local_addr()
    }

    /// Peer address of connection.
    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
        self.shared().io_ref().peer_addr()
    }

    /// Shutdown the connection.
    pub fn shutdown(&self, how: mio_orig::tcp::Shutdown) -> io::Result<()> {
        self.shared().io_ref().shutdown(how)
    }

    /// Set `no_delay`.
    pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
        self.shared().io_ref().set_nodelay(nodelay)
    }

    /// Set keepalive.
    pub fn set_keepalive(&self, seconds: Option<u32>) -> io::Result<()> {
        self.shared().io_ref().set_keepalive(seconds)
    }

    /// TODO: document
    pub fn take_socket_error(&self) -> io::Result<()> {
        self.shared().io_ref().take_socket_error()
    }

    /// Try cloning the socket descriptor.
    pub fn try_clone(&self) -> io::Result<TcpStream> {
        self.shared().io_ref().try_clone().map(MioAdapter::new)
    }
}
<|endoftext|>


<commit_before>extern crate libc;

use std;
use std::mem;
use libc::c_char;
use std::ffi::CStr;

use nom::*;

use rparser::*;

use tls_parser::tls::{TlsMessage,TlsMessageHandshake,tls_parser_many};
use tls_parser::tls_ciphers::TlsCipherSuite;
use tls_parser::tls_extensions::parse_tls_extensions;
use tls_parser::tls_states::{TlsState,tls_state_transition};

// --------------------------------------------
// example implementation: TLS parser

#[repr(u32)]
pub enum TlsParserEvents {
    HeartbeatOverflow = 1,
    InvalidState = 2,
}

pub struct TlsParserState<'a> {
    pub o: Option<&'a[u8]>,

    events: Vec<u32>,

    cipher: u16,
    state: TlsState,
}

impl<'a> TlsParserState<'a> {
    pub fn new(i: &'a[u8]) -> TlsParserState<'a> {
        TlsParserState{
            o:Some(i),
            events:Vec::new(),
            cipher:0,
            state:TlsState::None
        }
    }

}

r_declare_state_new!(r_tls_state_new,TlsParserState,b"blah");
r_declare_state_free!(r_tls_state_free,TlsParserState,{ () });

struct TlsParser;

impl<'a> RParser<TlsParserState<'a>> for TlsParser {
    fn new_state() -> TlsParserState<'a> {
        TlsParserState::new(b"blah")
    }

    fn probe(i: &[u8]) -> bool {
        if i.len() <= 2 { return false; }
        // first byte is record type (between 0x14 and 0x17, 0x16 is handhake)
        // second is TLS version major (0x3)
        // third is TLS version minor (0x0 for SSLv3, 0x1 for TLSv1.0, etc.)
        match (i[0],i[1],i[2]) {
            (0x14...0x17,0x03,0...3) => true,
            _ => false,
        }
    }

    fn parse(this: &mut TlsParserState, i: &[u8], direction: u8) -> u32 {
        debug!("[TLS->parse: direction={}, len={}]",direction,i.len());

        let d = tls_parser_many(i);

        let mut status = R_STATUS_OK;

        match d {
            IResult::Done(rem,p) => {
                for ref record in &p {
                    debug!("{:?}", record);
                    for msg in &record.msg {
                        if this.state == TlsState::ClientChangeCipherSpec {
                            // Ignore records from now on, they are encrypted
                            break
                        };
                        // update state machine
                        match tls_state_transition(this.state, msg) {
                            Ok(s)  => this.state = s,
                            Err(_) => {
                                this.state = TlsState::Invalid;
                                this.events.push(TlsParserEvents::InvalidState as u32);
                                status |= R_STATUS_EVENTS;
                            },
                        };

                        // extract variables
                        match *msg {
                            TlsMessage::Handshake(ref m) => {
                                match *m {
                                    TlsMessageHandshake::ClientHello(ref content) => {
                                        let blah = parse_tls_extensions(content.ext.unwrap_or(b""));
                                        debug!("ext {:?}", blah);
                                    },
                                    TlsMessageHandshake::ServerHello(ref content) => {
                                        this.cipher = content.cipher;
                                        match TlsCipherSuite::from_id(content.cipher) {
                                            Some(c) => debug!("Selected cipher: {:?}", c),
                                            _ => warn!("Unknown cipher 0x{:x}", content.cipher),
                                        };
                                        let blah = parse_tls_extensions(content.ext.unwrap_or(b""));
                                        debug!("extensions: {:?}", blah);
                                    },
                                    _ => (),
                                }
                            },
                            TlsMessage::Heartbeat(ref d) => {
                                if d.payload_len as usize > d.payload.len() {
                                    warn!("Heartbeat message with incorrect length {}. Heartbleed attempt ?",d.payload.len());
                                    this.events.push(TlsParserEvents::HeartbeatOverflow as u32);
                                    status |= R_STATUS_EVENTS;
                                }
                            },
                            _ => (),
                        }
                    }
                }
                if rem.len() > 0 { warn!("** unparsed ** {:?}",rem); };
            },
            IResult::Error(e) => error!("TLS parser reported an error: {:?}", e),
            IResult::Incomplete(e) => error!("TLS parser reported incomplete input: {:?}", e),
        };

        status
    }
}

r_implement_probe!(r_tls_probe,TlsParser);
r_implement_parse!(r_tls_parse,TlsParser);

// --------------------------------------------

#[no_mangle]
pub extern fn r_tls_get_next_event(ptr: *mut libc::c_char) -> u32
// pub extern fn r_tls_get_next_event<'a>(this: &mut TlsParserState<'a>) -> u32
{
    assert!(!ptr.is_null());
    let this: &mut TlsParserState = unsafe { mem::transmute(ptr) };
    match this.events.pop() {
        None     => 0xffffffff,
        Some(ev) => ev,
    }
}

#[no_mangle]
pub extern fn rusticata_tls_get_cipher(ptr: *mut libc::c_char) -> u32
// or
// pub extern fn rusticata_tls_get_cipher<'a>(this: &TlsParserState<'a>) -> u32
// but this gives a warning:
// warning: generic functions must be mangled, #[warn(no_mangle_generic_items)] on by default
{
    let this: &Box<TlsParserState> = unsafe { mem::transmute(ptr) };
    this.cipher as u32
}

#[no_mangle]
pub extern fn rusticata_tls_cipher_of_string(value: *const c_char) -> u32
{
    let c_str = unsafe { CStr::from_ptr(value) };
    let s = c_str.to_str().unwrap();
    match TlsCipherSuite::from_name(s) {
        Some(c) => c.id as u32,
        None    => 0,
    }
}

<commit_msg>Do not report parsing errors if input is empty (connection closed)<commit_after>extern crate libc;

use std;
use std::mem;
use libc::c_char;
use std::ffi::CStr;

use nom::*;

use rparser::*;

use tls_parser::tls::{TlsMessage,TlsMessageHandshake,tls_parser_many};
use tls_parser::tls_ciphers::TlsCipherSuite;
use tls_parser::tls_extensions::parse_tls_extensions;
use tls_parser::tls_states::{TlsState,tls_state_transition};

// --------------------------------------------
// example implementation: TLS parser

#[repr(u32)]
pub enum TlsParserEvents {
    HeartbeatOverflow = 1,
    InvalidState = 2,
}

pub struct TlsParserState<'a> {
    pub o: Option<&'a[u8]>,

    events: Vec<u32>,

    cipher: u16,
    state: TlsState,
}

impl<'a> TlsParserState<'a> {
    pub fn new(i: &'a[u8]) -> TlsParserState<'a> {
        TlsParserState{
            o:Some(i),
            events:Vec::new(),
            cipher:0,
            state:TlsState::None
        }
    }

}

r_declare_state_new!(r_tls_state_new,TlsParserState,b"blah");
r_declare_state_free!(r_tls_state_free,TlsParserState,{ () });

struct TlsParser;

impl<'a> RParser<TlsParserState<'a>> for TlsParser {
    fn new_state() -> TlsParserState<'a> {
        TlsParserState::new(b"blah")
    }

    fn probe(i: &[u8]) -> bool {
        if i.len() <= 2 { return false; }
        // first byte is record type (between 0x14 and 0x17, 0x16 is handhake)
        // second is TLS version major (0x3)
        // third is TLS version minor (0x0 for SSLv3, 0x1 for TLSv1.0, etc.)
        match (i[0],i[1],i[2]) {
            (0x14...0x17,0x03,0...3) => true,
            _ => false,
        }
    }

    fn parse(this: &mut TlsParserState, i: &[u8], direction: u8) -> u32 {
        debug!("[TLS->parse: direction={}, len={}]",direction,i.len());

        if i.len() == 0 {
            // Connection closed ?
            return R_STATUS_OK;
        };

        let d = tls_parser_many(i);

        let mut status = R_STATUS_OK;

        match d {
            IResult::Done(rem,p) => {
                for ref record in &p {
                    debug!("{:?}", record);
                    for msg in &record.msg {
                        if this.state == TlsState::ClientChangeCipherSpec {
                            // Ignore records from now on, they are encrypted
                            break
                        };
                        // update state machine
                        match tls_state_transition(this.state, msg) {
                            Ok(s)  => this.state = s,
                            Err(_) => {
                                this.state = TlsState::Invalid;
                                this.events.push(TlsParserEvents::InvalidState as u32);
                                status |= R_STATUS_EVENTS;
                            },
                        };

                        // extract variables
                        match *msg {
                            TlsMessage::Handshake(ref m) => {
                                match *m {
                                    TlsMessageHandshake::ClientHello(ref content) => {
                                        let blah = parse_tls_extensions(content.ext.unwrap_or(b""));
                                        debug!("ext {:?}", blah);
                                    },
                                    TlsMessageHandshake::ServerHello(ref content) => {
                                        this.cipher = content.cipher;
                                        match TlsCipherSuite::from_id(content.cipher) {
                                            Some(c) => debug!("Selected cipher: {:?}", c),
                                            _ => warn!("Unknown cipher 0x{:x}", content.cipher),
                                        };
                                        let blah = parse_tls_extensions(content.ext.unwrap_or(b""));
                                        debug!("extensions: {:?}", blah);
                                    },
                                    _ => (),
                                }
                            },
                            TlsMessage::Heartbeat(ref d) => {
                                if d.payload_len as usize > d.payload.len() {
                                    warn!("Heartbeat message with incorrect length {}. Heartbleed attempt ?",d.payload.len());
                                    this.events.push(TlsParserEvents::HeartbeatOverflow as u32);
                                    status |= R_STATUS_EVENTS;
                                }
                            },
                            _ => (),
                        }
                    }
                }
                if rem.len() > 0 { warn!("** unparsed ** {:?}",rem); };
            },
            IResult::Error(e) => error!("TLS parser reported an error: {:?}", e),
            IResult::Incomplete(e) => error!("TLS parser reported incomplete input: {:?}", e),
        };

        status
    }
}

r_implement_probe!(r_tls_probe,TlsParser);
r_implement_parse!(r_tls_parse,TlsParser);

// --------------------------------------------

#[no_mangle]
pub extern fn r_tls_get_next_event(ptr: *mut libc::c_char) -> u32
// pub extern fn r_tls_get_next_event<'a>(this: &mut TlsParserState<'a>) -> u32
{
    assert!(!ptr.is_null());
    let this: &mut TlsParserState = unsafe { mem::transmute(ptr) };
    match this.events.pop() {
        None     => 0xffffffff,
        Some(ev) => ev,
    }
}

#[no_mangle]
pub extern fn rusticata_tls_get_cipher(ptr: *mut libc::c_char) -> u32
// or
// pub extern fn rusticata_tls_get_cipher<'a>(this: &TlsParserState<'a>) -> u32
// but this gives a warning:
// warning: generic functions must be mangled, #[warn(no_mangle_generic_items)] on by default
{
    let this: &Box<TlsParserState> = unsafe { mem::transmute(ptr) };
    this.cipher as u32
}

#[no_mangle]
pub extern fn rusticata_tls_cipher_of_string(value: *const c_char) -> u32
{
    let c_str = unsafe { CStr::from_ptr(value) };
    let s = c_str.to_str().unwrap();
    match TlsCipherSuite::from_name(s) {
        Some(c) => c.id as u32,
        None    => 0,
    }
}

<|endoftext|>


<commit_before>use std::sync::mpsc::Sender;
use std::sync::mpsc::Receiver;
use std::sync::mpsc;
use std::thread;

use platform::time::sleep_secs;

fn transport(id: u64, sender: Sender<u64>) {
    for i in 0..10 {
        println!("[transport {}] Sending: {:?}", id, i);
        sender.send(i).unwrap();
        println!("[transport {}] Sent: {:?}", id, i);
        sleep_secs(1.0);
    }
    println!("[transport {}] Stopping", id);
}

fn driver(receiver: Receiver<u64>) {
    loop {
        println!("[driver]      Receiving...");
        let val = receiver.recv().unwrap();
        println!("[driver]      Received: {:?}", val);
    }
}

pub fn run_it() {
    let (tx, rx) = mpsc::channel();

    println!("[listen]      Launching driver");
    thread::spawn(move || {
        driver(rx);
    });

    for i in 1..5 {
        println!("[listen]      Launching transport {}", i);

        let tx_clone = tx.clone();
        thread::spawn(move || {
            transport(i, tx_clone);
        });
    }

    loop {
        sleep_secs(1.0);
    }
}
<commit_msg>working single transport + driver demo<commit_after>use std::sync::mpsc::Sender;
use std::sync::mpsc::Receiver;
use std::sync::mpsc;
use std::thread;

use platform::time::sleep_secs;
use protocol::cmd::Cmd;
use protocol::cmd::Get;
use protocol::cmd::Resp;
use protocol::cmd::Value;

type CmdSender = Sender<Cmd>;
type CmdReceiver = Receiver<Cmd>;

type RespSender = Sender<Resp>;
type RespReceiver = Receiver<Resp>;

fn transport(id: u64, sender: CmdSender, receiver: RespReceiver) {
    for i in 0..5 {
        let val = Cmd::Get(Get::new("x"));
        println!("[transport {}] Sending: {:?}", id, val.clone());
        sender.send(val.clone()).unwrap();
        println!("[transport {}] Sent: {:?}", id, val.clone());

        println!("[transport {}] Receiving...", id);
        let val = receiver.recv().unwrap();
        println!("[transport {}] Received: {:?}", id, val);

        sleep_secs(1.0);
    }

    println!("[transport {}] Stopping", id);
}

fn driver(sender: RespSender, receiver: CmdReceiver) {
    loop {
        println!("[driver]      Receiving...");
        let val = receiver.recv().unwrap();
        println!("[driver]      Received: {:?}", val);

        let val = Resp::Value(Value::new("x", vec![1, 2, 3]));
        println!("[driver]      Sending: {:?}", val.clone());
        sender.send(val.clone()).unwrap();
        println!("[driver]      Sent: {:?}", val.clone());
    }

    println!("[driver]      Stopping");
}

pub fn run_it() {
    let (c_tx, c_rx) = mpsc::channel();
    let (r_tx, r_rx) = mpsc::channel();

    println!("[listen]      Launching driver");
    thread::spawn(move || {
        driver(r_tx, c_rx);
    });

    println!("[listen]      Launching transport {}", 1);
    thread::spawn(move || {
        transport(1, c_tx, r_rx);
    });

    loop {
        sleep_secs(1.0);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;

extern crate abstract_ns;
extern crate futures;
extern crate getopts;
extern crate ns_dns_tokio;
extern crate rand;
extern crate tokio_core;
extern crate tokio_io;
extern crate net2;

// the device under test
extern crate net_interceptor;

use futures::{Future, Stream, IntoFuture};
use futures::sync::oneshot;
use futures::sync::oneshot::{Sender};
use tokio_core::net::{TcpListener};
use tokio_core::reactor::{Core, Timeout};

use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::sync::{Arc, Barrier};
use std::thread;

use net_interceptor::logging::unittestlogger;

// Timeout for reactor based tests in seconds
const TIMEOUT: u64 = 5;

fn stopper(oneshot_sender_stream: Sender<i64>) -> Arc<Barrier> {
    // barrier awaits the releaser thread and the signalling one
    let barrier = Arc::new(Barrier::new(2));
    let copy_for_thread = barrier.clone();
    // the releaser thread
    let builder = thread::Builder::new()
        .name("barrier to signal thread".into());
    builder.spawn(move || {
        debug!("waiting for end");
        copy_for_thread.wait();
        debug!("signalling end");
        // release the future blockig core.run()
        oneshot_sender_stream.send(1).unwrap();
    }).unwrap();
    barrier
}

#[test]
fn test_infra_connect() {
    unittestlogger();

    // this is the main event loop, powered by tokio core
    let mut core = Core::new().unwrap();
    let handle = core.handle();

    let bind_socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 0);

    let listener = net2::TcpBuilder::new_v4().unwrap()
        // .reuse_address(true).unwrap()
        .bind(bind_socket_addr).unwrap()
        // 2 = allow two connection requests in the backlog
        .listen(2).unwrap();

    let listen_port = listener.local_addr().unwrap();
    let contextid = listen_port.port();
    warn!("Listening on {} in context {}", listen_port, contextid);

    let sock = TcpListener::from_listener(listener, &bind_socket_addr, &handle).unwrap();

    let (tx, signal) = oneshot::channel::<i64>();
    let barrier = stopper(tx);

    let result = sock.incoming().for_each(move |(socket, addr)| {
        //protocol.bind_connection(&handle, socket, addr, AlphaBravo::new(&path));
        warn!("connection from {} in context {}", addr, contextid);
        barrier.wait();
        Ok(())
    });

    // register the acceptor in the reactor core
    // the map_err is a hack to satisfy connections.for_each above
    handle.spawn(result.map_err(|_|()));

    let fut = Timeout::new(std::time::Duration::from_secs(TIMEOUT), &handle).into_future();
    let timeout = fut.flatten();

    let timeout_expired = timeout.and_then(|_| -> Result<i64, std::io::Error> {
        panic!("test timeout expired: test failed!")
    });

    match core.run(timeout_expired.select2(signal)) {
        Err(_) => { panic!("oh no!"); },
        _ => { },
    };

    //drop(result);
}

#[test]
fn test_two_times() {
    test_infra_connect();
    test_infra_connect();
}
<commit_msg>simple integration test accepts connections and terminates without timeout<commit_after>#[macro_use]
extern crate log;

extern crate abstract_ns;
extern crate futures;
extern crate getopts;
extern crate ns_dns_tokio;
extern crate rand;
extern crate tokio_core;
extern crate tokio_io;
extern crate net2;

// the device under test
extern crate net_interceptor;

use futures::{Future, Stream, IntoFuture};
use futures::sync::oneshot;
use futures::sync::oneshot::{Sender};
use tokio_core::net::{TcpListener};
use tokio_core::reactor::{Core, Timeout};

use std::net::{IpAddr, Ipv4Addr, SocketAddr, TcpStream};
use std::sync::{Arc, Barrier};
use std::thread;

use net_interceptor::logging::unittestlogger;

// Timeout for reactor based tests in seconds
const TIMEOUT: u64 = 5;

fn stopper(oneshot_sender_stream: Sender<i64>) -> Arc<Barrier> {
    // barrier awaits the releaser thread and the signalling one
    let barrier = Arc::new(Barrier::new(2));
    let copy_for_thread = barrier.clone();
    // the releaser thread
    let builder = thread::Builder::new()
        .name("barrier to signal thread".into());
    builder.spawn(move || {
        debug!("waiting for end");
        copy_for_thread.wait();
        debug!("signalling end");
        // release the future blockig core.run()
        oneshot_sender_stream.send(1).unwrap();
    }).unwrap();
    barrier
}

#[test]
fn test_infra_connect() {
    unittestlogger();

    // this is the main event loop, powered by tokio core
    let mut core = Core::new().unwrap();
    let handle = core.handle();

    let bind_socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 0);

    let listener = net2::TcpBuilder::new_v4().unwrap()
        // .reuse_address(true).unwrap()
        .bind(bind_socket_addr).unwrap()
        // 2 = allow two connection requests in the backlog
        .listen(2).unwrap();

    let listen_port = listener.local_addr().unwrap();
    let contextid = listen_port.port();
    warn!("Listening on {} in context {}", listen_port, contextid);

    let sock = TcpListener::from_listener(listener, &bind_socket_addr, &handle).unwrap();

    let (tx, signal) = oneshot::channel::<i64>();
    let barrier = stopper(tx);

    let result = sock.incoming().for_each(move |(socket, addr)| {
        //protocol.bind_connection(&handle, socket, addr, AlphaBravo::new(&path));
        warn!("connection from {} in context {}", addr, contextid);
        barrier.wait();
        Ok(())
    });

    // register the acceptor in the reactor core
    // the map_err is a hack to satisfy connections.for_each above
    handle.spawn(result.map_err(|_|()));

    let fut = Timeout::new(std::time::Duration::from_secs(TIMEOUT), &handle).into_future();
    let timeout = fut.flatten();

    let timeout_expired = timeout.and_then(|_| -> Result<i64, std::io::Error> {
        panic!("test timeout expired: test failed!")
    });

    // issue a connect
    let stream = TcpStream::connect(listen_port).unwrap();
    drop(stream);
//    let _ = stream.write(&[1]).unwrap(); // ignore the Result
//    let _ = stream.read(&mut [0; 128]); // ignore this too

    match core.run(timeout_expired.select2(signal)) {
        Err(_) => { panic!("oh no!"); },
        _ => { },
    };

    //drop(result);
}

#[test]
fn test_two_times() {
    test_infra_connect();
    test_infra_connect();
}
<|endoftext|>


<commit_before>use std::collections::HashSet;
use std::hash::Hash;
use std::hash::Hasher;
use std::sync::Arc;

use crate::descriptor::DescriptorProto;
use crate::descriptor::FileDescriptorProto;
use crate::reflect::file::dynamic::DynamicFileDescriptor;
use crate::reflect::file::fds::build_fds;
use crate::reflect::file::index::EnumIndices;
use crate::reflect::file::index::FileDescriptorCommon;
use crate::reflect::file::index::MessageIndices;
use crate::reflect::name::protobuf_name_starts_with_package;
use crate::reflect::service::ServiceDescriptor;
use crate::reflect::EnumDescriptor;
use crate::reflect::FieldDescriptor;
use crate::reflect::GeneratedFileDescriptor;
use crate::reflect::MessageDescriptor;
use crate::reflect::Syntax;

pub(crate) mod building;
pub(crate) mod dynamic;
pub(crate) mod fds;
pub(crate) mod generated;
pub(crate) mod index;
pub(crate) mod syntax;

#[derive(Clone, Debug)]
pub(crate) enum FileDescriptorImpl {
    Generated(&'static GeneratedFileDescriptor),
    Dynamic(Arc<DynamicFileDescriptor>),
}

impl PartialEq for FileDescriptorImpl {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (FileDescriptorImpl::Generated(a), FileDescriptorImpl::Generated(b)) => {
                *a as *const GeneratedFileDescriptor == *b as *const GeneratedFileDescriptor
            }
            (FileDescriptorImpl::Dynamic(a), FileDescriptorImpl::Dynamic(b)) => Arc::ptr_eq(a, b),
            _ => false,
        }
    }
}

impl Hash for FileDescriptorImpl {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match self {
            FileDescriptorImpl::Generated(g) => {
                Hash::hash(&(*g as *const GeneratedFileDescriptor), state)
            }
            FileDescriptorImpl::Dynamic(a) => {
                Hash::hash(&(&**a as *const DynamicFileDescriptor), state)
            }
        }
    }
}

impl Eq for FileDescriptorImpl {}

/// Reflection for objects defined in `.proto` file (messages, enums, etc).
///
/// The object is refcounted: clone is shallow.
///
/// The equality performs pointer comparison: two clones of the same `FileDescriptor`
/// objects are equal, but two `FileDescriptor` objects created from the same `FileDescriptorProto`
/// objects are **not** equal.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct FileDescriptor {
    pub(crate) imp: FileDescriptorImpl,
}

impl FileDescriptor {
    pub(crate) fn common(&self) -> &FileDescriptorCommon {
        match &self.imp {
            FileDescriptorImpl::Generated(g) => &g.common,
            FileDescriptorImpl::Dynamic(d) => &d.common,
        }
    }

    /// Same as `common`, but returns `&'static`.
    pub(crate) fn common_for_generated_descriptor(&self) -> &'static FileDescriptorCommon {
        match &self.imp {
            FileDescriptorImpl::Generated(g) => &g.common,
            FileDescriptorImpl::Dynamic(..) => panic!("not generated"),
        }
    }

    pub(crate) fn message_indices(&self, index: usize) -> &MessageIndices {
        &self.common().messages[index]
    }

    pub(crate) fn message_by_index(&self, index: usize) -> MessageDescriptor {
        MessageDescriptor {
            file_descriptor: self.clone(),
            index,
        }
    }

    pub(crate) fn message_proto_by_index(&self, index: usize) -> &DescriptorProto {
        &self.common().messages[index].proto
    }

    pub(crate) fn enum_indices(&self, index: usize) -> &EnumIndices {
        &self.common().enums[index]
    }

    /// Syntax of current file.
    pub fn syntax(&self) -> Syntax {
        Syntax::parse(self.proto().syntax()).unwrap_or(Syntax::Proto2)
    }

    /// Top-level messages.
    pub fn messages(&self) -> impl Iterator<Item = MessageDescriptor> + '_ {
        self.common()
            .top_level_messages
            .iter()
            .map(|i| MessageDescriptor::new(self.clone(), *i))
    }

    /// Get top-level enums.
    pub fn enums(&self) -> impl Iterator<Item = EnumDescriptor> + '_ {
        self.proto()
            .enum_type
            .iter()
            .enumerate()
            .map(|(i, _)| EnumDescriptor::new(self.clone(), i))
    }

    /// Get services defined in `.proto` file.
    pub fn services(&self) -> impl Iterator<Item = ServiceDescriptor> + '_ {
        self.proto()
            .service
            .iter()
            .enumerate()
            .map(|(i, _)| ServiceDescriptor::new(self.clone(), i))
    }

    /// Extension fields.
    pub fn extensions(&self) -> impl Iterator<Item = FieldDescriptor> + '_ {
        self.common()
            .extension_field_range()
            .map(move |index| FieldDescriptor {
                file_descriptor: self.clone(),
                index,
            })
    }

    /// Find message by name relative to the package.
    ///
    /// Only search in the current file, not in any dependencies.
    pub fn message_by_package_relative_name(&self, name: &str) -> Option<MessageDescriptor> {
        self.common()
            .message_by_name_to_package
            .get(name)
            .map(|&index| MessageDescriptor::new(self.clone(), index))
    }

    /// Find message by name relative to the package.
    ///
    /// Only search in the current file, not in any dependencies.
    pub fn enum_by_package_relative_name(&self, name: &str) -> Option<EnumDescriptor> {
        self.common()
            .enums_by_name_to_package
            .get(name)
            .map(|&index| EnumDescriptor::new(self.clone(), index))
    }

    /// Find message by fully-qualified name.
    ///
    /// Only search in the current file, not in any dependencies.
    pub fn message_by_full_name(&self, name: &str) -> Option<MessageDescriptor> {
        if let Some(name_to_package) =
            protobuf_name_starts_with_package(name, self.proto().package())
        {
            self.message_by_package_relative_name(name_to_package)
        } else {
            None
        }
    }

    /// Find enum by name fully-qualified name.
    ///
    /// Only search in the current file, not in any dependencies.
    pub fn enum_by_full_name(&self, name: &str) -> Option<EnumDescriptor> {
        if let Some(name_to_package) =
            protobuf_name_starts_with_package(name, self.proto().package())
        {
            self.enum_by_package_relative_name(name_to_package)
        } else {
            None
        }
    }

    /// This function is called from generated code, it is not stable, and should not be called.
    #[doc(hidden)]
    pub fn new_generated_2(generated: &'static GeneratedFileDescriptor) -> FileDescriptor {
        FileDescriptor {
            imp: FileDescriptorImpl::Generated(generated),
        }
    }

    /// Dynamic message created from [`FileDescriptorProto`] without generated files.
    pub fn new_dynamic(
        proto: FileDescriptorProto,
        dependencies: &[FileDescriptor],
    ) -> crate::Result<FileDescriptor> {
        Ok(FileDescriptor {
            imp: FileDescriptorImpl::Dynamic(Arc::new(DynamicFileDescriptor::new(
                proto,
                dependencies,
            )?)),
        })
    }

    /// Create a set of file descriptors from individual file descriptors.
    pub fn new_dynamic_fds(
        protos: Vec<FileDescriptorProto>,
        dependencies: &[FileDescriptor],
    ) -> crate::Result<Vec<FileDescriptor>> {
        build_fds(protos, dependencies)
    }

    /// `.proto` data for this file.
    pub fn proto(&self) -> &FileDescriptorProto {
        match &self.imp {
            FileDescriptorImpl::Generated(g) => &g.proto,
            FileDescriptorImpl::Dynamic(d) => &d.proto,
        }
    }

    /// Direct dependencies of this file.
    pub fn deps(&self) -> &[FileDescriptor] {
        &self.common().dependencies
    }

    /// Subset of dependencies which are public
    pub fn public_deps(&self) -> impl Iterator<Item = FileDescriptor> + '_ {
        self.proto()
            .public_dependency
            .iter()
            .map(|&i| self.deps()[i as usize].clone())
    }

    fn _all_files(&self) -> Vec<&FileDescriptor> {
        let mut r = Vec::new();
        let mut visited = HashSet::new();

        let mut stack = Vec::new();
        stack.push(self);
        while let Some(file) = stack.pop() {
            if !visited.insert(file) {
                continue;
            }

            r.push(file);
            stack.extend(file.deps());
        }

        r
    }
}

#[cfg(test)]
mod test {
    use crate::descriptor;

    #[test]
    #[cfg_attr(miri, ignore)]
    fn eq() {
        assert!(descriptor::file_descriptor() == descriptor::file_descriptor().clone());
    }
}
<commit_msg>FileDescriptor::name, package<commit_after>use std::collections::HashSet;
use std::hash::Hash;
use std::hash::Hasher;
use std::sync::Arc;

use crate::descriptor::DescriptorProto;
use crate::descriptor::FileDescriptorProto;
use crate::reflect::file::dynamic::DynamicFileDescriptor;
use crate::reflect::file::fds::build_fds;
use crate::reflect::file::index::EnumIndices;
use crate::reflect::file::index::FileDescriptorCommon;
use crate::reflect::file::index::MessageIndices;
use crate::reflect::name::protobuf_name_starts_with_package;
use crate::reflect::service::ServiceDescriptor;
use crate::reflect::EnumDescriptor;
use crate::reflect::FieldDescriptor;
use crate::reflect::GeneratedFileDescriptor;
use crate::reflect::MessageDescriptor;
use crate::reflect::Syntax;

pub(crate) mod building;
pub(crate) mod dynamic;
pub(crate) mod fds;
pub(crate) mod generated;
pub(crate) mod index;
pub(crate) mod syntax;

#[derive(Clone, Debug)]
pub(crate) enum FileDescriptorImpl {
    Generated(&'static GeneratedFileDescriptor),
    Dynamic(Arc<DynamicFileDescriptor>),
}

impl PartialEq for FileDescriptorImpl {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (FileDescriptorImpl::Generated(a), FileDescriptorImpl::Generated(b)) => {
                *a as *const GeneratedFileDescriptor == *b as *const GeneratedFileDescriptor
            }
            (FileDescriptorImpl::Dynamic(a), FileDescriptorImpl::Dynamic(b)) => Arc::ptr_eq(a, b),
            _ => false,
        }
    }
}

impl Hash for FileDescriptorImpl {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match self {
            FileDescriptorImpl::Generated(g) => {
                Hash::hash(&(*g as *const GeneratedFileDescriptor), state)
            }
            FileDescriptorImpl::Dynamic(a) => {
                Hash::hash(&(&**a as *const DynamicFileDescriptor), state)
            }
        }
    }
}

impl Eq for FileDescriptorImpl {}

/// Reflection for objects defined in `.proto` file (messages, enums, etc).
///
/// The object is refcounted: clone is shallow.
///
/// The equality performs pointer comparison: two clones of the same `FileDescriptor`
/// objects are equal, but two `FileDescriptor` objects created from the same `FileDescriptorProto`
/// objects are **not** equal.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct FileDescriptor {
    pub(crate) imp: FileDescriptorImpl,
}

impl FileDescriptor {
    pub(crate) fn common(&self) -> &FileDescriptorCommon {
        match &self.imp {
            FileDescriptorImpl::Generated(g) => &g.common,
            FileDescriptorImpl::Dynamic(d) => &d.common,
        }
    }

    /// Same as `common`, but returns `&'static`.
    pub(crate) fn common_for_generated_descriptor(&self) -> &'static FileDescriptorCommon {
        match &self.imp {
            FileDescriptorImpl::Generated(g) => &g.common,
            FileDescriptorImpl::Dynamic(..) => panic!("not generated"),
        }
    }

    pub(crate) fn message_indices(&self, index: usize) -> &MessageIndices {
        &self.common().messages[index]
    }

    pub(crate) fn message_by_index(&self, index: usize) -> MessageDescriptor {
        MessageDescriptor {
            file_descriptor: self.clone(),
            index,
        }
    }

    pub(crate) fn message_proto_by_index(&self, index: usize) -> &DescriptorProto {
        &self.common().messages[index].proto
    }

    pub(crate) fn enum_indices(&self, index: usize) -> &EnumIndices {
        &self.common().enums[index]
    }

    /// The file name.
    pub fn name(&self) -> &str {
        self.proto().name()
    }

    /// Protobuf package.
    pub fn package(&self) -> &str {
        self.proto().package()
    }

    /// Syntax of current file.
    pub fn syntax(&self) -> Syntax {
        Syntax::parse(self.proto().syntax()).unwrap_or(Syntax::Proto2)
    }

    /// Top-level messages.
    pub fn messages(&self) -> impl Iterator<Item = MessageDescriptor> + '_ {
        self.common()
            .top_level_messages
            .iter()
            .map(|i| MessageDescriptor::new(self.clone(), *i))
    }

    /// Get top-level enums.
    pub fn enums(&self) -> impl Iterator<Item = EnumDescriptor> + '_ {
        self.proto()
            .enum_type
            .iter()
            .enumerate()
            .map(|(i, _)| EnumDescriptor::new(self.clone(), i))
    }

    /// Get services defined in `.proto` file.
    pub fn services(&self) -> impl Iterator<Item = ServiceDescriptor> + '_ {
        self.proto()
            .service
            .iter()
            .enumerate()
            .map(|(i, _)| ServiceDescriptor::new(self.clone(), i))
    }

    /// Extension fields.
    pub fn extensions(&self) -> impl Iterator<Item = FieldDescriptor> + '_ {
        self.common()
            .extension_field_range()
            .map(move |index| FieldDescriptor {
                file_descriptor: self.clone(),
                index,
            })
    }

    /// Find message by name relative to the package.
    ///
    /// Only search in the current file, not in any dependencies.
    pub fn message_by_package_relative_name(&self, name: &str) -> Option<MessageDescriptor> {
        self.common()
            .message_by_name_to_package
            .get(name)
            .map(|&index| MessageDescriptor::new(self.clone(), index))
    }

    /// Find message by name relative to the package.
    ///
    /// Only search in the current file, not in any dependencies.
    pub fn enum_by_package_relative_name(&self, name: &str) -> Option<EnumDescriptor> {
        self.common()
            .enums_by_name_to_package
            .get(name)
            .map(|&index| EnumDescriptor::new(self.clone(), index))
    }

    /// Find message by fully-qualified name.
    ///
    /// Only search in the current file, not in any dependencies.
    pub fn message_by_full_name(&self, name: &str) -> Option<MessageDescriptor> {
        if let Some(name_to_package) =
            protobuf_name_starts_with_package(name, self.proto().package())
        {
            self.message_by_package_relative_name(name_to_package)
        } else {
            None
        }
    }

    /// Find enum by name fully-qualified name.
    ///
    /// Only search in the current file, not in any dependencies.
    pub fn enum_by_full_name(&self, name: &str) -> Option<EnumDescriptor> {
        if let Some(name_to_package) =
            protobuf_name_starts_with_package(name, self.proto().package())
        {
            self.enum_by_package_relative_name(name_to_package)
        } else {
            None
        }
    }

    /// This function is called from generated code, it is not stable, and should not be called.
    #[doc(hidden)]
    pub fn new_generated_2(generated: &'static GeneratedFileDescriptor) -> FileDescriptor {
        FileDescriptor {
            imp: FileDescriptorImpl::Generated(generated),
        }
    }

    /// Dynamic message created from [`FileDescriptorProto`] without generated files.
    pub fn new_dynamic(
        proto: FileDescriptorProto,
        dependencies: &[FileDescriptor],
    ) -> crate::Result<FileDescriptor> {
        Ok(FileDescriptor {
            imp: FileDescriptorImpl::Dynamic(Arc::new(DynamicFileDescriptor::new(
                proto,
                dependencies,
            )?)),
        })
    }

    /// Create a set of file descriptors from individual file descriptors.
    pub fn new_dynamic_fds(
        protos: Vec<FileDescriptorProto>,
        dependencies: &[FileDescriptor],
    ) -> crate::Result<Vec<FileDescriptor>> {
        build_fds(protos, dependencies)
    }

    /// `.proto` data for this file.
    pub fn proto(&self) -> &FileDescriptorProto {
        match &self.imp {
            FileDescriptorImpl::Generated(g) => &g.proto,
            FileDescriptorImpl::Dynamic(d) => &d.proto,
        }
    }

    /// Direct dependencies of this file.
    pub fn deps(&self) -> &[FileDescriptor] {
        &self.common().dependencies
    }

    /// Subset of dependencies which are public
    pub fn public_deps(&self) -> impl Iterator<Item = FileDescriptor> + '_ {
        self.proto()
            .public_dependency
            .iter()
            .map(|&i| self.deps()[i as usize].clone())
    }

    fn _all_files(&self) -> Vec<&FileDescriptor> {
        let mut r = Vec::new();
        let mut visited = HashSet::new();

        let mut stack = Vec::new();
        stack.push(self);
        while let Some(file) = stack.pop() {
            if !visited.insert(file) {
                continue;
            }

            r.push(file);
            stack.extend(file.deps());
        }

        r
    }
}

#[cfg(test)]
mod test {
    use crate::descriptor;

    #[test]
    #[cfg_attr(miri, ignore)]
    fn eq() {
        assert!(descriptor::file_descriptor() == descriptor::file_descriptor().clone());
    }
}
<|endoftext|>