RUST-1405 Expose full server response in Error (#1474)

Author: isabelatkinson

src/bson_compat.rs

@@ -36,6 +36,8 @@ pub(crate) trait RawDocumentBufExt: Sized {
         key: impl AsRef<CStr>,
         value: impl Into<crate::bson::raw::RawBsonRef<'a>> + 'a,
     );
+
+    fn to_document(&self) -> crate::error::Result<crate::bson::Document>;
 }
 
 #[cfg(feature = "bson-3")]
@@ -47,6 +49,10 @@ impl RawDocumentBufExt for crate::bson::RawDocumentBuf {
     ) {
         self.append(key, value);
     }
+
+    fn to_document(&self) -> crate::error::Result<crate::bson::Document> {
+        self.try_into().map_err(Into::into)
+    }
 }
 
 #[cfg(feature = "bson-3")]

src/client.rs

@@ -496,6 +496,7 @@ impl Client {
             operation_name,
             start_time,
             timeout,
+            self.options().tracing_max_document_length_bytes,
         );
         #[cfg(feature = "tracing-unstable")]
         event_emitter.emit_started_event(self.inner.topology.latest().description.clone());

src/client/csfle/state_machine.rs

@@ -17,7 +17,7 @@ use tokio::{
 use crate::{
     client::{csfle::options::KmsProvidersTlsOptions, options::ServerAddress, WeakClient},
     error::{Error, Result},
-    operation::{run_command::RunCommand, RawOutput},
+    operation::{raw_output::RawOutput, run_command::RunCommand},
     options::ReadConcern,
     runtime::{process::Process, AsyncStream, TlsConfig},
     Client,

src/client/executor.rs

@@ -532,9 +532,13 @@ impl Client {
             let start_time = Instant::now();
             let command_result = match connection.send_message(message).await {
                 Ok(response) => {
-                    let is_sharded =
-                        connection.stream_description()?.initial_server_type == ServerType::Mongos;
-                    self.parse_response(op, session, is_sharded, response).await
+                    match self
+                        .parse_response(op, session, connection.is_sharded()?, &response)
+                        .await
+                    {
+                        Ok(()) => Ok(response),
+                        Err(error) => Err(error.with_server_response(&response)),
+                    }
                 }
                 Err(err) => Err(err),
             };
@@ -612,15 +616,15 @@ impl Client {
                         effective_criteria: effective_criteria.clone(),
                     };
 
-                    match op.handle_response(response, context).await {
+                    match op.handle_response(&response, context).await {
                         Ok(response) => Ok(response),
                         Err(mut err) => {
                             err.add_labels_and_update_pin(
                                 Some(connection.stream_description()?),
                                 session,
                                 Some(retryability),
                             );
-                            Err(err)
+                            Err(err.with_server_response(&response))
                         }
                     }
                 }
@@ -820,8 +824,8 @@ impl Client {
         op: &T,
         session: &mut Option<&mut ClientSession>,
         is_sharded: bool,
-        response: RawCommandResponse,
-    ) -> Result<RawCommandResponse> {
+        response: &RawCommandResponse,
+    ) -> Result<()> {
         let raw_doc = RawDocument::from_bytes(response.as_bytes())?;
 
         let ok = match raw_doc.get("ok")? {
@@ -870,7 +874,7 @@ impl Client {
                 }
             }
 
-            Ok(response)
+            Ok(())
         } else {
             Err(response
                 .body::<CommandErrorBody>()

src/cmap.rs

@@ -69,14 +69,23 @@ impl ConnectionPool {
         address: ServerAddress,
         connection_establisher: ConnectionEstablisher,
         server_updater: TopologyUpdater,
-        topology_id: ObjectId,
         options: Option<ConnectionPoolOptions>,
+        #[cfg(feature = "tracing-unstable")] topology_id: ObjectId,
     ) -> Self {
         let event_handler = options
             .as_ref()
             .and_then(|opts| opts.cmap_event_handler.clone());
 
-        let event_emitter = CmapEventEmitter::new(event_handler, topology_id);
+        #[cfg(feature = "tracing-unstable")]
+        let event_emitter = CmapEventEmitter::new(
+            event_handler,
+            topology_id,
+            options
+                .as_ref()
+                .and_then(|options| options.max_document_length_bytes),
+        );
+        #[cfg(not(feature = "tracing-unstable"))]
+        let event_emitter = CmapEventEmitter::new(event_handler);
 
         let (manager, connection_requester, generation_subscriber) = ConnectionPoolWorker::start(
             address.clone(),
@@ -114,7 +123,10 @@ impl ConnectionPool {
             manager,
             connection_requester,
             generation_subscriber,
-            event_emitter: CmapEventEmitter::new(None, ObjectId::new()),
+            #[cfg(feature = "tracing-unstable")]
+            event_emitter: CmapEventEmitter::new(None, ObjectId::new(), None),
+            #[cfg(not(feature = "tracing-unstable"))]
+            event_emitter: CmapEventEmitter::new(None),
         }
     }
 

src/cmap/conn/pooled.rs

@@ -30,6 +30,7 @@ use crate::{
         ConnectionReadyEvent,
     },
     runtime::AsyncStream,
+    ServerType,
 };
 
 /// A wrapper around the [`Connection`] type that represents a connection within a connection pool.
@@ -205,6 +206,11 @@ impl PooledConnection {
         Instant::now().duration_since(available_time) >= max_idle_time
     }
 
+    /// Whether this connection is to a mongos.
+    pub(crate) fn is_sharded(&self) -> Result<bool> {
+        Ok(self.stream_description()?.initial_server_type == ServerType::Mongos)
+    }
+
     /// Nullifies the internal state of this connection and returns it in a new [PooledConnection]
     /// with the given state.
     fn take(&mut self, new_state: PooledConnectionState) -> Self {

src/cmap/options.rs

@@ -70,6 +70,10 @@ pub(crate) struct ConnectionPoolOptions {
     ///
     /// The default is 2.
     pub(crate) max_connecting: Option<u32>,
+
+    /// The maximum length for documents in tracing events.
+    #[cfg(feature = "tracing-unstable")]
+    pub(crate) max_document_length_bytes: Option<usize>,
 }
 
 impl ConnectionPoolOptions {
@@ -86,6 +90,8 @@ impl ConnectionPoolOptions {
             load_balanced: options.load_balanced,
             credential: options.credential.clone(),
             max_connecting: options.max_connecting,
+            #[cfg(feature = "tracing-unstable")]
+            max_document_length_bytes: options.tracing_max_document_length_bytes,
         }
     }
 

src/cmap/test.rs

@@ -165,8 +165,9 @@ impl Executor {
             ConnectionEstablisher::new(EstablisherOptions::from(get_client_options().await))
                 .unwrap(),
             updater,
-            crate::bson::oid::ObjectId::new(),
             Some(self.pool_options),
+            #[cfg(feature = "tracing-unstable")]
+            crate::bson::oid::ObjectId::new(),
         );
 
         // Mock a monitoring task responding to errors reported by the pool.

src/cmap/test/integration.rs

@@ -49,8 +49,9 @@ async fn acquire_connection_and_send_command() {
         client_options.hosts[0].clone(),
         ConnectionEstablisher::new(EstablisherOptions::from(&client_options)).unwrap(),
         TopologyUpdater::channel().0,
-        crate::bson::oid::ObjectId::new(),
         Some(pool_options),
+        #[cfg(feature = "tracing-unstable")]
+        crate::bson::oid::ObjectId::new(),
     );
     let mut connection = pool.check_out().await.unwrap();
 
@@ -116,8 +117,9 @@ async fn concurrent_connections() {
         get_client_options().await.hosts[0].clone(),
         ConnectionEstablisher::new(EstablisherOptions::from(&client_options)).unwrap(),
         TopologyUpdater::channel().0,
-        crate::bson::oid::ObjectId::new(),
         Some(options),
+        #[cfg(feature = "tracing-unstable")]
+        crate::bson::oid::ObjectId::new(),
     );
 
     let tasks = (0..2).map(|_| {
@@ -191,8 +193,9 @@ async fn connection_error_during_establishment() {
         client_options.hosts[0].clone(),
         ConnectionEstablisher::new(EstablisherOptions::from(&client_options)).unwrap(),
         TopologyUpdater::channel().0,
-        crate::bson::oid::ObjectId::new(),
         Some(options),
+        #[cfg(feature = "tracing-unstable")]
+        crate::bson::oid::ObjectId::new(),
     );
 
     pool.check_out().await.expect_err("check out should fail");

src/error.rs

@@ -13,7 +13,8 @@ use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
 use crate::{
-    bson::{Bson, Document},
+    bson::{doc, rawdoc, Bson, Document, RawDocumentBuf},
+    cmap::RawCommandResponse,
     options::ServerAddress,
     sdam::{ServerType, TopologyVersion},
 };
@@ -52,20 +53,32 @@ pub type Result<T> = std::result::Result<T, Error>;
 #[derive(Clone, Debug, Error)]
 #[cfg_attr(
     test,
-    error("Kind: {kind}, labels: {labels:?}, source: {source:?}, backtrace: {bt}")
+    error(
+        "Kind: {kind}, labels: {labels:?}, source: {source:?}, backtrace: {bt}, server response: \
+         {server_response:?}"
+    )
 )]
 #[cfg_attr(
     not(test),
-    error("Kind: {kind}, labels: {labels:?}, source: {source:?}")
+    error(
+        "Kind: {kind}, labels: {labels:?}, source: {source:?}, server response: \
+         {server_response:?}"
+    )
 )]
 #[non_exhaustive]
 pub struct Error {
     /// The type of error that occurred.
     pub kind: Box<ErrorKind>,
+
     labels: HashSet<String>,
+
     pub(crate) wire_version: Option<i32>,
+
     #[source]
     pub(crate) source: Option<Box<Error>>,
+
+    pub(crate) server_response: Option<Box<RawDocumentBuf>>,
+
     #[cfg(test)]
     bt: Arc<std::backtrace::Backtrace>,
 }
@@ -99,6 +112,7 @@ impl Error {
             labels,
             wire_version: None,
             source: None,
+            server_response: None,
             #[cfg(test)]
             bt: Arc::new(std::backtrace::Backtrace::capture()),
         }
@@ -288,6 +302,20 @@ impl Error {
         self.labels.insert(label);
     }
 
+    /// The full response returned from the server. This can be used to inspect error fields that
+    /// are not represented in the `Error` type.
+    pub fn server_response(&self) -> Option<&RawDocumentBuf> {
+        self.server_response.as_deref()
+    }
+
+    /// Adds the server's response to this error if it is not already present.
+    pub(crate) fn with_server_response(mut self, response: &RawCommandResponse) -> Self {
+        if self.server_response.is_none() {
+            self.server_response = Some(Box::new(response.raw_body().to_owned()));
+        }
+        self
+    }
+
     #[cfg(feature = "dns-resolver")]
     pub(crate) fn from_resolve_error(error: hickory_resolver::error::ResolveError) -> Self {
         ErrorKind::DnsResolve {
@@ -496,6 +524,10 @@ impl Error {
             source.redact();
         }
 
+        if self.server_response.is_some() {
+            self.server_response = Some(Box::new(rawdoc! { "redacted": true }));
+        }
+
         // This is intentionally written without a catch-all branch so that if new error
         // kinds are added we remember to reason about whether they need to be redacted.
         match *self.kind {