RUST-1992 Consolidate binary encoding logic (#560)

Author: abr-egn

src/binary.rs

@@ -1,10 +1,7 @@
 #! Module containing functionality related to BSON binary values.
 mod vector;
 
-use std::{
-    convert::TryFrom,
-    fmt::{self, Display},
-};
+use std::fmt::{self, Display};
 
 use crate::{
     base64,
@@ -65,6 +62,8 @@ impl Binary {
 
     #[cfg(feature = "serde")]
     pub(crate) fn from_extended_doc(doc: &crate::Document) -> Option<Self> {
+        use std::convert::TryFrom;
+
         let binary_doc = doc.get_document("$binary").ok()?;
 
         if let Ok(bytes) = binary_doc.get_str("base64") {

src/de.rs

@@ -46,8 +46,6 @@ pub(crate) use self::serde::{convert_unsigned_to_signed_raw, BsonVisitor};
 #[cfg(test)]
 pub(crate) use self::raw::Deserializer as RawDeserializer;
 
-pub(crate) const MAX_BSON_SIZE: i32 = i32::MAX;
-
 /// Hint provided to the deserializer via `deserialize_newtype_struct` as to the type of thing
 /// being deserialized.
 #[derive(Debug, Clone, Copy)]

src/raw/bson_ref.rs

@@ -3,7 +3,7 @@ use std::convert::{TryFrom, TryInto};
 use super::{bson::RawBson, Error, RawArray, RawDocument, Result};
 use crate::{
     oid::{self, ObjectId},
-    raw::RawJavaScriptCodeWithScope,
+    raw::{write_cstring, write_string, RawJavaScriptCodeWithScope},
     spec::{BinarySubtype, ElementType},
     Binary,
     Bson,
@@ -285,6 +285,50 @@ impl<'a> RawBsonRef<'a> {
             }),
         }
     }
+
+    #[inline]
+    pub(crate) fn append_to(self, dest: &mut Vec<u8>) -> Result<()> {
+        match self {
+            Self::Int32(val) => dest.extend(val.to_le_bytes()),
+            Self::Int64(val) => dest.extend(val.to_le_bytes()),
+            Self::Double(val) => dest.extend(val.to_le_bytes()),
+            Self::Binary(b @ RawBinaryRef { subtype, bytes }) => {
+                let len = b.len();
+                dest.extend(len.to_le_bytes());
+                dest.push(subtype.into());
+                if let BinarySubtype::BinaryOld = subtype {
+                    dest.extend((len - 4).to_le_bytes())
+                }
+                dest.extend(bytes);
+            }
+            Self::String(s) => write_string(dest, s),
+            Self::Array(raw_array) => dest.extend(raw_array.as_bytes()),
+            Self::Document(raw_document) => dest.extend(raw_document.as_bytes()),
+            Self::Boolean(b) => dest.push(b as u8),
+            Self::RegularExpression(re) => {
+                write_cstring(dest, re.pattern)?;
+                write_cstring(dest, re.options)?;
+            }
+            Self::JavaScriptCode(js) => write_string(dest, js),
+            Self::JavaScriptCodeWithScope(code_w_scope) => {
+                let len = code_w_scope.len();
+                dest.extend(len.to_le_bytes());
+                write_string(dest, code_w_scope.code);
+                dest.extend(code_w_scope.scope.as_bytes());
+            }
+            Self::Timestamp(ts) => dest.extend(ts.to_le_bytes()),
+            Self::ObjectId(oid) => dest.extend(oid.bytes()),
+            Self::DateTime(dt) => dest.extend(dt.timestamp_millis().to_le_bytes()),
+            Self::Symbol(s) => write_string(dest, s),
+            Self::Decimal128(d) => dest.extend(d.bytes()),
+            Self::DbPointer(dbp) => {
+                write_string(dest, dbp.namespace);
+                dest.extend(dbp.id.bytes());
+            }
+            Self::Null | Self::Undefined | Self::MinKey | Self::MaxKey => {}
+        }
+        Ok(())
+    }
 }
 
 #[cfg(feature = "serde")]

src/raw/document_buf/raw_writer.rs

@@ -1,8 +1,4 @@
-use crate::{
-    raw::{write_cstring, write_string},
-    spec::BinarySubtype,
-    RawBsonRef,
-};
+use crate::{raw::write_cstring, RawBsonRef};
 
 pub(super) struct RawWriter<'a> {
     data: &'a mut Vec<u8>,
@@ -18,72 +14,7 @@ impl<'a> RawWriter<'a> {
         self.data[original_len - 1] = value.element_type() as u8;
 
         write_cstring(self.data, key)?;
-
-        match value {
-            RawBsonRef::Int32(i) => {
-                self.data.extend(i.to_le_bytes());
-            }
-            RawBsonRef::String(s) => {
-                write_string(self.data, s);
-            }
-            RawBsonRef::Document(d) => {
-                self.data.extend(d.as_bytes());
-            }
-            RawBsonRef::Array(a) => {
-                self.data.extend(a.as_bytes());
-            }
-            RawBsonRef::Binary(b) => {
-                let len = b.len();
-                self.data.extend(len.to_le_bytes());
-                self.data.push(b.subtype.into());
-                if let BinarySubtype::BinaryOld = b.subtype {
-                    self.data.extend((len - 4).to_le_bytes())
-                }
-                self.data.extend(b.bytes);
-            }
-            RawBsonRef::Boolean(b) => {
-                self.data.push(b as u8);
-            }
-            RawBsonRef::DateTime(dt) => {
-                self.data.extend(dt.timestamp_millis().to_le_bytes());
-            }
-            RawBsonRef::DbPointer(dbp) => {
-                write_string(self.data, dbp.namespace);
-                self.data.extend(dbp.id.bytes());
-            }
-            RawBsonRef::Decimal128(d) => {
-                self.data.extend(d.bytes());
-            }
-            RawBsonRef::Double(d) => {
-                self.data.extend(d.to_le_bytes());
-            }
-            RawBsonRef::Int64(i) => {
-                self.data.extend(i.to_le_bytes());
-            }
-            RawBsonRef::RegularExpression(re) => {
-                write_cstring(self.data, re.pattern)?;
-                write_cstring(self.data, re.options)?;
-            }
-            RawBsonRef::JavaScriptCode(js) => {
-                write_string(self.data, js);
-            }
-            RawBsonRef::JavaScriptCodeWithScope(code_w_scope) => {
-                let len = code_w_scope.len();
-                self.data.extend(len.to_le_bytes());
-                write_string(self.data, code_w_scope.code);
-                self.data.extend(code_w_scope.scope.as_bytes());
-            }
-            RawBsonRef::Timestamp(ts) => {
-                self.data.extend(ts.to_le_bytes());
-            }
-            RawBsonRef::ObjectId(oid) => {
-                self.data.extend(oid.bytes());
-            }
-            RawBsonRef::Symbol(s) => {
-                write_string(self.data, s);
-            }
-            RawBsonRef::Null | RawBsonRef::Undefined | RawBsonRef::MinKey | RawBsonRef::MaxKey => {}
-        }
+        value.append_to(self.data)?;
 
         // append trailing null byte
         self.data.push(0);

src/raw/serde/seeded_visitor.rs

@@ -271,72 +271,12 @@ impl<'de> Visitor<'de> for SeededVisitor<'_, 'de> {
                         Ok(ElementType::Array)
                     }
                     // Cases that don't
-                    _ => match bson.as_ref() {
-                        RawBsonRef::ObjectId(oid) => {
-                            self.buffer.append_bytes(&oid.bytes());
-                            Ok(ElementType::ObjectId)
-                        }
-                        RawBsonRef::Symbol(s) => {
-                            self.append_string(s);
-                            Ok(ElementType::Symbol)
-                        }
-                        RawBsonRef::Decimal128(d) => {
-                            self.buffer.append_bytes(&d.bytes);
-                            Ok(ElementType::Decimal128)
-                        }
-                        RawBsonRef::RegularExpression(re) => {
-                            self.append_cstring(re.pattern)
-                                .map_err(SerdeError::custom)?;
-                            self.append_cstring(re.options)
-                                .map_err(SerdeError::custom)?;
-                            Ok(ElementType::RegularExpression)
-                        }
-                        RawBsonRef::Undefined => Ok(ElementType::Undefined),
-                        RawBsonRef::DateTime(dt) => {
-                            self.buffer
-                                .append_bytes(&dt.timestamp_millis().to_le_bytes());
-                            Ok(ElementType::DateTime)
-                        }
-                        RawBsonRef::Timestamp(ts) => {
-                            self.buffer.append_bytes(&ts.increment.to_le_bytes());
-                            self.buffer.append_bytes(&ts.time.to_le_bytes());
-                            Ok(ElementType::Timestamp)
-                        }
-                        RawBsonRef::MinKey => Ok(ElementType::MinKey),
-                        RawBsonRef::MaxKey => Ok(ElementType::MaxKey),
-                        RawBsonRef::JavaScriptCode(s) => {
-                            self.append_string(s);
-                            Ok(ElementType::JavaScriptCode)
-                        }
-                        RawBsonRef::JavaScriptCodeWithScope(jsc) => {
-                            let length_index = self.pad_document_length();
-                            self.append_string(jsc.code);
-                            self.buffer.append_bytes(jsc.scope.as_bytes());
-
-                            let length_bytes =
-                                ((self.buffer.len() - length_index) as i32).to_le_bytes();
-                            self.buffer
-                                .copy_from_slice(length_index..length_index + 4, &length_bytes);
-
-                            Ok(ElementType::JavaScriptCodeWithScope)
-                        }
-                        RawBsonRef::DbPointer(dbp) => {
-                            self.append_string(dbp.namespace);
-                            self.buffer.append_bytes(&dbp.id.bytes());
-                            Ok(ElementType::DbPointer)
-                        }
-                        RawBsonRef::Double(d) => self.visit_f64(d),
-                        RawBsonRef::String(s) => self.visit_str(s),
-                        RawBsonRef::Boolean(b) => self.visit_bool(b),
-                        RawBsonRef::Null => self.visit_none(),
-                        RawBsonRef::Int32(i) => self.visit_i32(i),
-                        RawBsonRef::Int64(i) => self.visit_i64(i),
-                        // These are always borrowed and are handled
-                        // at the top of the outer `match`.
-                        RawBsonRef::Array(_) | RawBsonRef::Document(_) | RawBsonRef::Binary(_) => {
-                            unreachable!()
-                        }
-                    },
+                    _ => {
+                        let bson = bson.as_ref();
+                        bson.append_to(self.buffer.get_owned_buffer())
+                            .map_err(A::Error::custom)?;
+                        Ok(bson.element_type())
+                    }
                 }
             }
             MapParse::Aggregate(first_key) => {

src/ser.rs

@@ -30,68 +30,15 @@ pub use self::{
     serde::Serializer,
 };
 
-use std::io::Write;
-
 #[rustfmt::skip]
 use ::serde::{ser::Error as SerdeError, Serialize};
 
 use crate::{
     bson::{Bson, Document},
-    de::MAX_BSON_SIZE,
     ser::serde::SerializerOptions,
-    spec::BinarySubtype,
     RawDocumentBuf,
 };
 
-#[inline]
-pub(crate) fn write_i32<W: Write + ?Sized>(writer: &mut W, val: i32) -> Result<()> {
-    writer
-        .write_all(&val.to_le_bytes())
-        .map(|_| ())
-        .map_err(From::from)
-}
-
-#[inline]
-fn write_i64<W: Write + ?Sized>(writer: &mut W, val: i64) -> Result<()> {
-    writer
-        .write_all(&val.to_le_bytes())
-        .map(|_| ())
-        .map_err(From::from)
-}
-
-#[inline]
-fn write_f64<W: Write + ?Sized>(writer: &mut W, val: f64) -> Result<()> {
-    writer
-        .write_all(&val.to_le_bytes())
-        .map(|_| ())
-        .map_err(From::from)
-}
-
-#[inline]
-fn write_binary<W: Write>(mut writer: W, bytes: &[u8], subtype: BinarySubtype) -> Result<()> {
-    let len = if let BinarySubtype::BinaryOld = subtype {
-        bytes.len() + 4
-    } else {
-        bytes.len()
-    };
-
-    if len > MAX_BSON_SIZE as usize {
-        return Err(Error::custom(format!(
-            "binary length {} exceeded maximum size",
-            bytes.len()
-        )));
-    }
-
-    write_i32(&mut writer, len as i32)?;
-    writer.write_all(&[subtype.into()])?;
-
-    if let BinarySubtype::BinaryOld = subtype {
-        write_i32(&mut writer, len as i32 - 4)?;
-    };
-
-    writer.write_all(bytes).map_err(From::from)
-}
-
 /// Encode a `T` Serializable into a [`Bson`] value.
 ///
 /// The [`Serializer`] used by this function presents itself as human readable, whereas the