tokio_jsoncodec/
lib.rs

1//! This crate integrates [`serde_json`] into a Tokio codec ([`tokio_codec::Decoder`] and
2//! [`Encoder`]).
3
4extern crate bytes;
5#[cfg(test)]
6#[macro_use]
7extern crate maplit;
8extern crate serde;
9extern crate serde_json;
10extern crate tokio_codec;
11
12use bytes::BytesMut;
13use serde::{Deserialize, Serialize};
14use std::fmt;
15use std::io;
16use std::marker::PhantomData;
17use tokio_codec::{Decoder, Encoder};
18
19/// JSON-based codec.
20#[derive(Clone, Debug)]
21pub struct Codec<D, E> {
22    pretty: bool,
23    _priv: (PhantomData<D>, PhantomData<E>),
24}
25
26impl<D, E> Codec<D, E> {
27    /// Creates a new `Codec`.
28    ///
29    /// `pretty` controls whether or not encoded values are pretty-printed.
30    pub fn new(pretty: bool) -> Self {
31        Self {
32            pretty,
33            _priv: (PhantomData, PhantomData),
34        }
35    }
36
37    /// Set whether or not encoded values are pretty-printed.
38    pub fn pretty(&mut self, pretty: bool) {
39        self.pretty = pretty;
40    }
41}
42
43impl<D, E> Default for Codec<D, E> {
44    fn default() -> Self {
45        Self::new(false)
46    }
47}
48
49impl<D, E> Decoder for Codec<D, E>
50where
51    for<'de> D: Deserialize<'de>,
52{
53    type Item = D;
54    type Error = Error;
55
56    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<D>, Error> {
57        let slice = &src.clone();
58        let mut de = serde_json::Deserializer::from_slice(slice).into_iter();
59        match de.next() {
60            Some(Ok(v)) => {
61                src.advance(de.byte_offset());
62                Ok(Some(v))
63            }
64            Some(Err(e)) => {
65                if e.is_eof() {
66                    Ok(None)
67                } else {
68                    Err(e.into())
69                }
70            }
71            None => {
72                // The remaining stream is whitespace; clear the buffer so Decoder::decode_eof
73                // doesn't return an Err
74                src.clear();
75                Ok(None)
76            }
77        }
78    }
79}
80
81impl<D, E> Encoder for Codec<D, E>
82where
83    E: Serialize,
84{
85    type Item = E;
86    type Error = Error;
87
88    fn encode(&mut self, item: E, dst: &mut BytesMut) -> Result<(), Error> {
89        let writer = BytesWriter(dst);
90        if self.pretty {
91            serde_json::to_writer_pretty(writer, &item)?;
92        } else {
93            serde_json::to_writer(writer, &item)?;
94        }
95        Ok(())
96    }
97}
98
99/// The [`Error`][`std::error::Error`] type for this crate.
100///
101/// This is necessary to not lose information about the error. [`Encoder`] requires that the Error
102/// implement `From<std::io::Error>`, and while a [`serde_json::Error`] can possibly be an IO
103/// error, there's no way to combine the two.
104///
105/// If you just want an [`io::Error`], `From<Error>` is implemented for it.
106#[derive(Debug)]
107pub enum Error {
108    /// A [`io::Error`].
109    Io(io::Error),
110    /// A [`serde_json::Error`].
111    Json(serde_json::Error),
112}
113
114impl fmt::Display for Error {
115    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
116        match self {
117            Error::Io(e) => e.fmt(f),
118            Error::Json(e) => e.fmt(f),
119        }
120    }
121}
122
123impl std::error::Error for Error {}
124
125impl From<io::Error> for Error {
126    fn from(err: io::Error) -> Self {
127        Error::Io(err)
128    }
129}
130
131impl From<serde_json::Error> for Error {
132    fn from(err: serde_json::Error) -> Self {
133        Error::Json(err)
134    }
135}
136
137impl From<Error> for io::Error {
138    fn from(err: Error) -> Self {
139        match err {
140            Error::Io(e) => e,
141            Error::Json(e) => e.into(),
142        }
143    }
144}
145
146/// Wrapper for `&mut [BytesMut]` that provides Write.
147///
148/// See also:
149/// * <https://github.com/vorner/tokio-serde-cbor/blob/a347107ad56f2ad8086998eb63ecb70b19f3b71d/src/lib.rs#L167-L181>
150/// * <https://github.com/carllerche/bytes/issues/77>
151struct BytesWriter<'a>(&'a mut BytesMut);
152
153impl<'a> io::Write for BytesWriter<'a> {
154    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
155        self.0.extend(buf);
156        Ok(buf.len())
157    }
158    fn flush(&mut self) -> io::Result<()> {
159        Ok(())
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use bytes::{BufMut, BytesMut};
166    use tokio_codec::{Decoder, Encoder};
167    use Codec;
168
169    #[test]
170    fn decode_empty() {
171        let mut buf = BytesMut::from(&b""[..]);
172        let mut codec: Codec<(), ()> = Codec::default();
173        assert_eq!(codec.decode(&mut buf).unwrap(), None);
174    }
175
176    #[test]
177    fn decode() {
178        let mut buf = BytesMut::from(&b"null null null"[..]);
179        let mut codec: Codec<_, ()> = Codec::default();
180        assert_eq!(codec.decode(&mut buf).unwrap(), Some(()));
181        assert_eq!(codec.decode(&mut buf).unwrap(), Some(()));
182        assert_eq!(codec.decode(&mut buf).unwrap(), Some(()));
183        assert_eq!(codec.decode(&mut buf).unwrap(), None);
184        assert!(buf.is_empty());
185    }
186
187    #[test]
188    fn decode_partial() {
189        let mut buf = BytesMut::from(&b"null null nu"[..]);
190        let mut codec: Codec<_, ()> = Codec::default();
191        assert_eq!(codec.decode(&mut buf).unwrap(), Some(()));
192        assert_eq!(codec.decode(&mut buf).unwrap(), Some(()));
193        assert_eq!(codec.decode(&mut buf).unwrap(), None);
194        assert_eq!(buf, &b" nu"[..]);
195        buf.put(&b"ll"[..]);
196        assert_eq!(codec.decode(&mut buf).unwrap(), Some(()));
197        assert!(buf.is_empty());
198    }
199
200    #[test]
201    fn decode_eof_trailing_whitespae() {
202        let mut buf = BytesMut::from(&b"null\n"[..]);
203        let mut codec: Codec<_, ()> = Codec::default();
204        assert_eq!(codec.decode_eof(&mut buf).unwrap(), Some(()));
205        assert_eq!(codec.decode_eof(&mut buf).unwrap(), None);
206        assert!(buf.is_empty());
207    }
208
209    #[test]
210    fn decode_err() {
211        let mut buf = BytesMut::from(&b"null butts"[..]);
212        let mut codec: Codec<_, ()> = Codec::default();
213        assert_eq!(codec.decode(&mut buf).unwrap(), Some(()));
214        assert!(codec.decode(&mut buf).is_err());
215    }
216
217    #[test]
218    fn encode() {
219        let mut buf = BytesMut::new();
220        let mut codec: Codec<(), _> = Codec::default();
221        codec.encode((), &mut buf).unwrap();
222        assert_eq!(buf, &b"null"[..]);
223    }
224
225    #[test]
226    fn encode_pretty() {
227        let mut buf = BytesMut::new();
228        let mut codec: Codec<(), _> = Codec::default();
229        codec
230            .encode(hashmap! { "butts" => "lol" }, &mut buf)
231            .unwrap();
232        codec.pretty(true);
233        codec
234            .encode(hashmap! { "butts" => "lol" }, &mut buf)
235            .unwrap();
236        assert_eq!(
237            String::from_utf8(buf.to_vec()).unwrap(),
238            r#"{"butts":"lol"}{
239  "butts": "lol"
240}"#
241        );
242    }
243}