zmtp/
writer.rs

1extern crate byteorder;
2
3use types::*;
4use self::byteorder::{ByteOrder, BigEndian};
5use std::borrow::{Borrow};
6use std::io;
7use std::io::{Write};
8use std::marker::PhantomData;
9use std::{u8, u64, usize};
10
11
12#[derive(Debug)]
13pub struct FrameWriter<W: Write>(PhantomData<W>);
14impl<W: Write> FrameWriter<W> {
15    pub fn write(frame: &Frame, writer: &mut W) -> io::Result<()> {
16        let mut fw = FrameWriter(PhantomData);
17        let mut result = Ok(frame);
18        result = fw.write_flags(result, writer);
19        result = fw.write_size(result, writer);
20        result = fw.write_body(result, writer);
21        result.map(|_| ())
22    }
23
24    fn write_flags<'l>(&mut self,
25                       result: io::Result<&'l Frame>,
26                       writer: &mut W) -> io::Result<&'l Frame> {
27        let frame = try!{ result };
28        let mut flags: u8 = 0b0000_0000;
29        flags |= match *frame.get_kind() {
30            FrameKind::Cmd => 0b100,
31            FrameKind::Msg => 0b000,
32        };
33        flags |= match *frame.get_size() {
34            Size::Short(_) => 0b000,
35            Size::Long(_)  => 0b010,
36        };
37        flags |= !frame.get_is_last() as u8;
38        assert_eq!(try!{ writer.write(&[flags]) },  1);
39        Ok(frame)
40    }
41
42    fn write_size<'l>(&mut self,
43                      result: io::Result<&'l Frame>,
44                      writer: &mut W) -> io::Result<&'l Frame> {
45        let frame = try!{ result };
46        match *frame.get_size() {
47            Size::Short(size) =>
48                assert_eq!(try!{ writer.write(&[size]) },  1),
49            Size::Long(size) => {
50                let mut bytes = [0; BYTES_PER_U64 as usize];
51                BigEndian::write_u64(&mut bytes, size);
52                assert_eq!(try!{ writer.write(&bytes) }, BYTES_PER_U64);
53            },
54        };
55        Ok(frame)
56    }
57
58    fn write_body<'l>(&mut self,
59                      result: io::Result<&'l Frame>,
60                      writer: &mut W) -> io::Result<&'l Frame> {
61        let frame = try!{ result };
62        writer.write(frame.view_body()).map(|size| {
63            assert_eq!(size,  frame.body_len());
64            frame
65        })
66    }
67}
68
69#[derive(Debug)]
70pub struct CommandWriter<W: Write>(PhantomData<W>);
71impl<W: Write> CommandWriter<W> {
72    pub fn write(cmd: &Command, writer: &mut W) -> io::Result<()> {
73        let mut frame = Frame::new()
74            .kind(FrameKind::Cmd)
75            .is_last(true);
76        let mut cmd_writer = CommandWriter::<W>(PhantomData);
77        match cmd.get_name_str().borrow() {
78            "READY" => {
79                frame = cmd_writer.write_size(cmd, frame);
80                frame = cmd_writer.write_name(cmd, frame);
81                frame = cmd_writer.write_metadata(cmd, frame);
82            },
83            _ => {
84                frame = cmd_writer.write_size(cmd, frame);
85                frame = cmd_writer.write_name(cmd, frame);
86                frame = cmd_writer.write_data(cmd, frame);
87            },
88        };
89        FrameWriter::write(&frame, writer)
90    }
91
92    fn write_size(&mut self, cmd: &Command, frame: Frame) -> Frame {
93        frame.set_size(match cmd.size().as_usize() {
94            size @ 0...0xFF => Size::Short(size as u8),
95            size            => Size::Long(size as u64),
96        })
97    }
98
99    fn write_name(&mut self, cmd: &Command, mut frame: Frame) -> Frame {
100        let cmd_name = cmd.get_name();
101        frame.edit_body().push(cmd_name.len() as u8);
102        frame.edit_body().extend_from_slice(cmd_name);
103        frame
104    }
105
106    fn write_data(&mut self, cmd: &Command, mut frame: Frame) -> Frame {
107        frame.edit_body().extend_from_slice(&cmd.get_data());
108        frame
109    }
110
111    fn write_metadata(&mut self, cmd: &Command, mut frame: Frame) -> Frame {
112        for (name, val) in cmd.get_metadata().iter() {
113            let name_bytes = name.as_bytes();
114            let name_len = name_bytes.len();
115            assert!(1 <= name_len && name_len <= 255);
116            let val_len = val.len() as u32;
117            let mut val_len_bytes = [0 as u8; 4];
118            BigEndian::write_u32(&mut val_len_bytes, val_len);
119            frame.edit_body().extend_from_slice(&[name_len as u8]);
120            frame.edit_body().extend_from_slice(name_bytes);
121            frame.edit_body().extend_from_slice(&val_len_bytes);
122            frame.edit_body().extend_from_slice(val);
123        }
124        frame
125    }
126}
127
128#[derive(Debug)]
129pub struct MessageWriter<W: Write>(PhantomData<W>);
130impl<W: Write> MessageWriter<W> {
131    pub fn write(msg: &Message, writer: &mut W) -> io::Result<()> {
132        let mut result = Ok(());
133        for frame in msg.frame_iter() {
134            result = FrameWriter::write(&frame, writer);
135            if result.is_err() { break }
136        }
137        result
138    }
139}
140
141#[derive(Debug)]
142pub struct GreetingWriter<W: Write>(PhantomData<W>);
143impl<W: Write> GreetingWriter<W> {
144    pub fn write(g: &Greeting, writer: &mut W) -> io::Result<()> {
145        try!{ Self::write_signature(g, writer) }
146        try!{ Self::write_version(  g, writer) }
147        try!{ Self::write_mechanism(g, writer) }
148        try!{ Self::write_as_server(g, writer) }
149        try!{ Self::write_filler(   g, writer) }
150        Ok(assert_eq!(try!{ writer.flush() }, ()))
151    }
152
153    pub fn write_signature(g: &Greeting, writer: &mut W) -> io::Result<()> {
154        Ok(assert_eq!(try!{ writer.write(&g.signature) },  g.signature.len()))
155    }
156
157    pub fn write_version(g: &Greeting, writer: &mut W) -> io::Result<()> {
158        let (major, minor) = g.version;
159        Ok(assert_eq!(try!{ writer.write(&[major, minor]) },  2))
160    }
161
162    pub fn write_mechanism(g: &Greeting, writer: &mut W) -> io::Result<()> {
163        Ok(assert_eq!(try!{ writer.write(&g.mechanism) },   g.mechanism.len()))
164    }
165
166    pub fn write_as_server(g: &Greeting, writer: &mut W) -> io::Result<()> {
167        Ok(assert_eq!(try!{ writer.write(&[g.as_server as u8]) },  1))
168    }
169
170    pub fn write_filler(g: &Greeting, writer: &mut W) -> io::Result<()> {
171        Ok(assert_eq!(try!{ writer.write(&g.filler) },  g.filler.len()))
172    }
173}
174
175#[derive(Debug)]
176pub struct TrafficWriter<W: Write>(PhantomData<W>);
177impl<W: Write> TrafficWriter<W> {
178    pub fn write(traffic: &Traffic, writer: &mut W) -> io::Result<()> {
179        let mut result = Ok(());
180        for item in traffic.items() {
181            result = match *item {
182                TrafficItem::Cmd(ref cmd) => CommandWriter::write(cmd, writer),
183                TrafficItem::Msg(ref msg) => MessageWriter::write(msg, writer),
184            };
185            if result.is_err() { break }
186        }
187        result
188    }
189}