RUST-1827 remove generic parameter to Command (#1105)

Author: abr-egn

src/bson_util.rs

@@ -1,8 +1,11 @@
 use std::{
+    collections::HashSet,
     convert::TryFrom,
     io::{Read, Write},
 };
 
+use serde::Serialize;
+
 use crate::{
     bson::{
         oid::ObjectId,
@@ -15,7 +18,7 @@ use crate::{
         RawDocumentBuf,
     },
     checked::Checked,
-    error::{ErrorKind, Result},
+    error::{Error, ErrorKind, Result},
     runtime::SyncLittleEndianRead,
 };
 
@@ -65,6 +68,13 @@ pub(crate) fn to_raw_bson_array(docs: &[Document]) -> Result<RawBson> {
     }
     Ok(RawBson::Array(array))
 }
+pub(crate) fn to_raw_bson_array_ser<T: Serialize>(values: &[T]) -> Result<RawBson> {
+    let mut array = RawArrayBuf::new();
+    for value in values {
+        array.push(bson::to_raw_document_buf(value)?);
+    }
+    Ok(RawBson::Array(array))
+}
 
 pub(crate) fn first_key(document: &Document) -> Option<&str> {
     document.keys().next().map(String::as_str)
@@ -162,13 +172,42 @@ pub(crate) fn extend_raw_document_buf(
     this: &mut RawDocumentBuf,
     other: RawDocumentBuf,
 ) -> Result<()> {
+    let mut keys: HashSet<String> = HashSet::new();
+    for elem in this.iter_elements() {
+        keys.insert(elem?.key().to_owned());
+    }
     for result in other.iter() {
         let (k, v) = result?;
+        if keys.contains(k) {
+            return Err(Error::internal(format!(
+                "duplicate raw document key {:?}",
+                k
+            )));
+        }
         this.append(k, v.to_raw_bson());
     }
     Ok(())
 }
 
+pub(crate) fn append_ser(
+    this: &mut RawDocumentBuf,
+    key: impl AsRef<str>,
+    value: impl Serialize,
+) -> Result<()> {
+    #[derive(Serialize)]
+    struct Helper<T> {
+        value: T,
+    }
+    let raw_doc = bson::to_raw_document_buf(&Helper { value })?;
+    this.append_ref(
+        key,
+        raw_doc
+            .get("value")?
+            .ok_or_else(|| Error::internal("no value"))?,
+    );
+    Ok(())
+}
+
 /// Returns the _id field of this document, prepending the field to the document if one is not
 /// already present.
 pub(crate) fn get_or_prepend_id_field(doc: &mut RawDocumentBuf) -> Result<Bson> {

src/client/auth.rs

@@ -14,6 +14,7 @@ mod x509;
 
 use std::{borrow::Cow, fmt::Debug, str::FromStr};
 
+use bson::RawDocumentBuf;
 use derivative::Derivative;
 use hmac::{digest::KeyInit, Mac};
 use rand::Rng;
@@ -509,9 +510,9 @@ impl Credential {
 
     /// If the mechanism is missing, append the appropriate mechanism negotiation key-value-pair to
     /// the provided hello or legacy hello command document.
-    pub(crate) fn append_needed_mechanism_negotiation(&self, command: &mut Document) {
+    pub(crate) fn append_needed_mechanism_negotiation(&self, command: &mut RawDocumentBuf) {
         if let (Some(username), None) = (self.username.as_ref(), self.mechanism.as_ref()) {
-            command.insert(
+            command.append(
                 "saslSupportedMechs",
                 format!("{}.{}", self.resolved_source(), username),
             );
@@ -613,7 +614,7 @@ pub(crate) enum ClientFirst {
 }
 
 impl ClientFirst {
-    pub(crate) fn to_document(&self) -> Document {
+    pub(crate) fn to_document(&self) -> RawDocumentBuf {
         match self {
             Self::Scram(version, client_first) => client_first.to_command(version).body,
             Self::X509(command) => command.body.clone(),

src/client/auth/oidc.rs

@@ -401,7 +401,7 @@ fn make_spec_auth_command(
     payload: Vec<u8>,
     server_api: Option<&ServerApi>,
 ) -> Command {
-    let body = doc! {
+    let body = rawdoc! {
         "saslStart": 1,
         "mechanism": MONGODB_OIDC_STR,
         "payload": Binary { subtype: BinarySubtype::Generic, bytes: payload },

src/client/auth/sasl.rs

@@ -1,5 +1,7 @@
+use bson::{rawdoc, RawBson};
+
 use crate::{
-    bson::{doc, spec::BinarySubtype, Binary, Bson, Document},
+    bson::{spec::BinarySubtype, Binary, Bson, Document},
     bson_util,
     client::{auth::AuthMechanism, options::ServerApi},
     cmap::Command,
@@ -31,15 +33,15 @@ impl SaslStart {
     }
 
     pub(super) fn into_command(self) -> Command {
-        let mut body = doc! {
+        let mut body = rawdoc! {
             "saslStart": 1,
             "mechanism": self.mechanism.as_str(),
             "payload": Binary { subtype: BinarySubtype::Generic, bytes: self.payload },
         };
         if self.mechanism == AuthMechanism::ScramSha1
             || self.mechanism == AuthMechanism::ScramSha256
         {
-            body.insert("options", doc! { "skipEmptyExchange": true });
+            body.append("options", rawdoc! { "skipEmptyExchange": true });
         }
 
         let mut command = Command::new("saslStart", self.source, body);
@@ -75,9 +77,11 @@ impl SaslContinue {
     }
 
     pub(super) fn into_command(self) -> Command {
-        let body = doc! {
+        // Unwrap safety: the Bson -> RawBson conversion is actually infallible
+        let raw_id: RawBson = self.conversation_id.try_into().unwrap();
+        let body = rawdoc! {
             "saslContinue": 1,
-            "conversationId": self.conversation_id,
+            "conversationId": raw_id,
             "payload": Binary { subtype: BinarySubtype::Generic, bytes: self.payload },
         };
 

src/client/auth/scram.rs

@@ -460,7 +460,7 @@ impl ClientFirst {
         let mut cmd = sasl_start.into_command();
 
         if self.include_db {
-            cmd.body.insert("db", self.source.clone());
+            cmd.body.append("db", self.source.clone());
         }
 
         cmd

src/client/auth/x509.rs

@@ -1,5 +1,7 @@
+use bson::rawdoc;
+
 use crate::{
-    bson::{doc, Document},
+    bson::Document,
     client::options::ServerApi,
     cmap::{Command, Connection, RawCommandResponse},
     error::{Error, Result},
@@ -16,13 +18,13 @@ pub(crate) fn build_client_first(
     credential: &Credential,
     server_api: Option<&ServerApi>,
 ) -> Command {
-    let mut auth_command_doc = doc! {
+    let mut auth_command_doc = rawdoc! {
         "authenticate": 1,
         "mechanism": "MONGODB-X509",
     };
 
     if let Some(ref username) = credential.username {
-        auth_command_doc.insert("username", username);
+        auth_command_doc.append("username", username.as_str());
     }
 
     let mut command = Command::new("authenticate", "$external", auth_command_doc);

src/cmap/conn/command.rs

@@ -17,18 +17,15 @@ use crate::{
 #[serde_with::skip_serializing_none]
 #[derive(Clone, Debug, Serialize, Default)]
 #[serde(rename_all = "camelCase")]
-pub(crate) struct Command<T = Document>
-where
-    T: Serialize,
-{
+pub(crate) struct Command {
     #[serde(skip)]
     pub(crate) name: String,
 
     #[serde(skip)]
     pub(crate) exhaust_allowed: bool,
 
     #[serde(flatten)]
-    pub(crate) body: T,
+    pub(crate) body: RawDocumentBuf,
 
     #[serde(skip)]
     pub(crate) document_sequences: Vec<DocumentSequence>,
@@ -58,11 +55,8 @@ where
     recovery_token: Option<Document>,
 }
 
-impl<T> Command<T>
-where
-    T: Serialize,
-{
-    pub(crate) fn new(name: impl ToString, target_db: impl ToString, body: T) -> Self {
+impl Command {
+    pub(crate) fn new(name: impl ToString, target_db: impl ToString, body: RawDocumentBuf) -> Self {
         Self {
             name: name.to_string(),
             target_db: target_db.to_string(),
@@ -85,7 +79,7 @@ where
         name: String,
         target_db: String,
         read_concern: Option<ReadConcern>,
-        body: T,
+        body: RawDocumentBuf,
     ) -> Self {
         Self {
             name,

src/cmap/conn/wire/message.rs

@@ -2,7 +2,6 @@ use std::io::Read;
 
 use bitflags::bitflags;
 use bson::{doc, Array, Document};
-use serde::Serialize;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 
 #[cfg(any(
@@ -48,10 +47,7 @@ pub(crate) struct DocumentSequence {
 impl Message {
     /// Creates a `Message` from a given `Command`. Note that the `response_to` field must be set
     /// manually.
-    pub(crate) fn from_command<T: Serialize>(
-        command: Command<T>,
-        request_id: Option<i32>,
-    ) -> Result<Self> {
+    pub(crate) fn from_command(command: Command, request_id: Option<i32>) -> Result<Self> {
         let document_payload = bson::to_raw_document_buf(&command)?;
 
         let mut flags = MessageFlags::empty();