1use std::io::{self, Write};
2use std::fmt;
3use std::error;
4use std::i64;
5
6use byteorder::{LittleEndian, WriteBytesExt};
7use chrono::Timelike;
8use serde::ser::{self, Serialize};
9
10use crate::value::Value;
11use crate::serde_impl::encode::Encoder;
12
13#[derive(Debug)]
14pub enum EncodeError {
15 IoError(io::Error),
16 InvalidMapKeyType(Value),
17 Unknown(String),
18 UnsupportedUnsignedType
19}
20
21impl From<io::Error> for EncodeError {
22 fn from(err: io::Error) -> EncodeError {
23 EncodeError::IoError(err)
24 }
25}
26
27impl fmt::Display for EncodeError {
28 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
29 match *self {
30 EncodeError::IoError(ref inner) => inner.fmt(fmt),
31 EncodeError::InvalidMapKeyType(ref bson) => {
32 write!(fmt, "Invalid map key type: {:?}", bson)
33 }
34 EncodeError::Unknown(ref inner) => inner.fmt(fmt),
35 EncodeError::UnsupportedUnsignedType => write!(fmt, "bson does not support unsigned type"),
36 }
37 }
38}
39
40impl error::Error for EncodeError {
41 fn description(&self) -> &str {
42 match *self {
43 EncodeError::IoError(ref inner) => inner.description(),
44 EncodeError::InvalidMapKeyType(_) => "Invalid map key type",
45 EncodeError::Unknown(ref inner) => inner,
46 EncodeError::UnsupportedUnsignedType => "bson does not support unsigned type",
47 }
48 }
49 fn cause(&self) -> Option<&dyn error::Error> {
50 match *self {
51 EncodeError::IoError(ref inner) => Some(inner),
52 _ => None,
53 }
54 }
55}
56
57impl ser::Error for EncodeError {
58 fn custom<T: fmt::Display>(msg: T) -> EncodeError {
59 EncodeError::Unknown(msg.to_string())
60 }
61}
62
63pub type EncodeResult<T> = Result<T, EncodeError>;
64
65pub(crate) fn write_string(writer: &mut impl Write, s: &str) -> EncodeResult<()> {
66 writer.write_i32::<LittleEndian>(s.len() as i32 + 1)?;
67 writer.write_all(s.as_bytes())?;
68 writer.write_u8(0)?;
69 Ok(())
70}
71
72pub(crate) fn write_cstring(writer: &mut impl Write, s: &str) -> EncodeResult<()> {
73 writer.write_all(s.as_bytes())?;
74 writer.write_u8(0)?;
75 Ok(())
76}
77
78#[inline]
79pub(crate) fn write_i32(writer: &mut impl Write, val: i32) -> EncodeResult<()> {
80 writer.write_i32::<LittleEndian>(val).map_err(From::from)
81}
82
83#[inline]
84pub(crate) fn write_i64(writer: &mut impl Write, val: i64) -> EncodeResult<()> {
85 writer.write_i64::<LittleEndian>(val).map_err(From::from)
86}
87
88#[inline]
89pub(crate) fn write_u64(writer: &mut impl Write, val: u64) -> EncodeResult<()> {
90 writer.write_u64::<LittleEndian>(val).map_err(From::from)
91}
92
93#[inline]
94pub(crate) fn write_f64(writer: &mut impl Write, val: f64) -> EncodeResult<()> {
95 writer.write_f64::<LittleEndian>(val).map_err(From::from)
96}
97
98fn encode_array(writer: &mut impl Write, arr: &[Value]) -> EncodeResult<()> {
99 let mut buf = Vec::with_capacity(64);
100 write_i32(&mut buf, 0)?;
101
102 for (key, val) in arr.iter().enumerate() {
103 encode_bson(&mut buf, &key.to_string(), val)?;
104 }
105
106 buf.write_u8(0)?;
107
108 let len_bytes = (buf.len() as i32).to_le_bytes();
109
110 buf[..4].clone_from_slice(&len_bytes);
111
112 writer.write_all(&buf)?;
113 Ok(())
114}
115
116pub fn encode_bson(writer: &mut impl Write, key: &str, val: &Value) -> EncodeResult<()> {
117 writer.write_u8(val.element_type() as u8)?;
118 write_cstring(writer, key)?;
119
120 match *val {
121 Value::Double(v) => write_f64(writer, v),
122 Value::String(ref v) => write_string(writer, &v),
123 Value::Array(ref v) => encode_array(writer, &v),
124 Value::Document(ref v) => encode_document(writer, v),
125 Value::Boolean(v) => writer.write_u8(if v { 0x01 } else { 0x00 }).map_err(From::from),
126 Value::RegExp(ref pat, ref opt) => {
127 write_cstring(writer, pat)?;
128 write_cstring(writer, opt)
129 }
130 Value::JavaScriptCode(ref code) => write_string(writer, &code),
131 Value::ObjectId(ref id) => writer.write_all(&id.bytes()).map_err(From::from),
132 Value::JavaScriptCodeWithScope(ref code, ref scope) => {
133 let mut buf = Vec::new();
134 write_string(&mut buf, code)?;
135 encode_document(&mut buf, scope)?;
136
137 write_i32(writer, buf.len() as i32 + 4)?;
138 writer.write_all(&buf).map_err(From::from)
139 }
140 Value::Int32(v) => write_i32(writer, v),
141 Value::Int64(v) => write_i64(writer, v),
142 Value::TimeStamp(v) => write_u64(writer, v),
143 Value::Binary(subtype, ref data) => {
144 write_i32(writer, data.len() as i32)?;
145 writer.write_u8(From::from(subtype))?;
146 writer.write_all(data).map_err(From::from)
147 }
148 Value::UTCDatetime(ref v) => {
149 write_i64(
150 writer,
151 v.timestamp() * 1000 + i64::from(v.nanosecond() / 1_000_000)
152 )
153 }
154 Value::Null => Ok(()),
155 Value::Symbol(ref v) => write_string(writer, &v)
156 }
157}
158
159pub fn encode_document<'a, S, D> (writer: &mut impl Write, document: D) -> EncodeResult<()>
160 where S: AsRef<str> + 'a, D: IntoIterator<Item = (&'a S, &'a Value)>
161{
162 let mut buf = Vec::with_capacity(64);
163 write_i32(&mut buf, 0)?;
164
165 for (key, val) in document {
166 encode_bson(&mut buf, key.as_ref(), val)?;
167 }
168
169 buf.write_u8(0)?;
170
171 let len_bytes = (buf.len() as i32).to_le_bytes();
172
173 buf[..4].clone_from_slice(&len_bytes);
174
175 writer.write_all(&buf)?;
176 Ok(())
177}
178
179pub fn to_bson<T: ?Sized>(value: &T) -> EncodeResult<Value>
180 where T: Serialize
181{
182 let ser = Encoder::new();
183 value.serialize(ser)
184}
185
186pub fn to_vec<T: ?Sized>(value: &T) -> EncodeResult<Vec<u8>>
187 where T: Serialize
188{
189 let bson = to_bson(value)?;
190
191 if let Value::Document(object) = bson {
192 let mut buf: Vec<u8> = Vec::new();
193 encode_document(&mut buf, &object)?;
194 return Ok(buf)
195 }
196
197 Err(EncodeError::InvalidMapKeyType(bson))
198}
199
200#[cfg(test)]
201mod test {
202 use std::io::Cursor;
203 use crate::encode::encode_document;
204 use crate::decode::decode_document;
205 use crate::doc;
206
207 #[test]
208 fn encode() {
209 let document = doc!{"aa": "bb", "cc": [1, 2, 3, 4]};
210
211 let mut buf: Vec<u8> = Vec::new();
212
213 encode_document(&mut buf, &document).unwrap();
214
215 let mut reader = Cursor::new(buf);
216
217 let document2 = decode_document(&mut reader).unwrap();
218
219 assert_eq!(document, document2);
220 }
221}