RUST-598 Perform handshake when establishing monitoring connections (#271)

Author: patrickfreed

src/cmap/establish/handshake/mod.rs

@@ -10,9 +10,9 @@ use crate::{
     bson::{doc, Bson, Document},
     client::auth::{ClientFirst, FirstRound},
     cmap::{options::ConnectionPoolOptions, Command, Connection, StreamDescription},
-    error::Result,
+    error::{ErrorKind, Result},
     is_master::{IsMasterCommandResponse, IsMasterReply},
-    options::{AuthMechanism, Credential},
+    options::{AuthMechanism, ClientOptions, Credential, DriverInfo, ServerApi},
 };
 
 #[cfg(feature = "tokio-runtime")]
@@ -136,95 +136,167 @@ lazy_static! {
 
 /// Contains the logic needed to handshake a connection.
 #[derive(Debug, Clone)]
-pub(super) struct Handshaker {
+pub(crate) struct Handshaker {
     /// The `isMaster` command to send when handshaking. This will always be identical
     /// given the same pool options, so it can be created at the time the Handshaker is created.
     command: Command,
+    credential: Option<Credential>,
 }
 
 impl Handshaker {
     /// Creates a new Handshaker.
-    pub(super) fn new(options: Option<&ConnectionPoolOptions>) -> Self {
+    pub(crate) fn new(options: Option<HandshakerOptions>) -> Self {
         let mut metadata = BASE_CLIENT_METADATA.clone();
+        let mut credential = None;
+        let mut db = None;
+        let mut server_api = None;
 
-        if let Some(app_name) = options.as_ref().and_then(|opts| opts.app_name.as_ref()) {
-            metadata.application = Some(AppMetadata {
-                name: app_name.to_string(),
-            });
-        }
-
-        if let Some(driver_info) = options.as_ref().and_then(|opts| opts.driver_info.as_ref()) {
-            metadata.driver.name.push('|');
-            metadata.driver.name.push_str(&driver_info.name);
+        let mut body = doc! {
+            "isMaster": 1,
+        };
 
-            if let Some(ref version) = driver_info.version {
-                metadata.driver.version.push('|');
-                metadata.driver.version.push_str(version);
+        if let Some(options) = options {
+            if let Some(app_name) = options.app_name {
+                metadata.application = Some(AppMetadata { name: app_name });
             }
 
-            if let Some(ref mut platform) = metadata.platform {
-                if let Some(ref driver_info_platform) = driver_info.platform {
-                    platform.push('|');
-                    platform.push_str(driver_info_platform);
+            if let Some(driver_info) = options.driver_info {
+                metadata.driver.name.push('|');
+                metadata.driver.name.push_str(&driver_info.name);
+
+                if let Some(ref version) = driver_info.version {
+                    metadata.driver.version.push('|');
+                    metadata.driver.version.push_str(version);
+                }
+
+                if let Some(ref mut platform) = metadata.platform {
+                    if let Some(ref driver_info_platform) = driver_info.platform {
+                        platform.push('|');
+                        platform.push_str(driver_info_platform);
+                    }
                 }
             }
+
+            if let Some(cred) = options.credential {
+                cred.append_needed_mechanism_negotiation(&mut body);
+                db = Some(cred.resolved_source().to_string());
+                credential = Some(cred);
+            }
+
+            server_api = options.server_api;
         }
 
-        let mut db = "admin";
+        body.insert("client", metadata);
 
-        let mut body = doc! {
-            "isMaster": 1,
-            "client": metadata,
-        };
+        let mut command = Command::new_read(
+            "isMaster".to_string(),
+            db.unwrap_or_else(|| "admin".to_string()),
+            None,
+            body,
+        );
 
-        if let Some(credential) = options.as_ref().and_then(|opts| opts.credential.as_ref()) {
-            credential.append_needed_mechanism_negotiation(&mut body);
-            db = credential.resolved_source();
+        if let Some(server_api) = server_api {
+            command.set_server_api(&server_api)
         }
 
-        let mut command = Command::new_read("isMaster".to_string(), db.to_string(), None, body);
-        if let Some(server_api) = options.as_ref().and_then(|opts| opts.server_api.as_ref()) {
-            command.set_server_api(server_api)
+        Self {
+            command,
+            credential,
         }
-
-        Self { command }
     }
 
     /// Handshakes a connection.
-    pub(super) async fn handshake(
-        &self,
-        conn: &mut Connection,
-        credential: Option<&Credential>,
-    ) -> Result<Option<FirstRound>> {
+    pub(crate) async fn handshake(&self, conn: &mut Connection) -> Result<HandshakeResult> {
         let mut command = self.command.clone();
 
-        let client_first = set_speculative_auth_info(&mut command.body, credential)?;
-
-        let start_time = PreciseTime::now();
-        let response = conn.send_command(command, None).await?;
-        let end_time = PreciseTime::now();
+        let client_first = set_speculative_auth_info(&mut command.body, self.credential.as_ref())?;
 
-        response.validate()?;
-        let mut command_response: IsMasterCommandResponse = response.body()?;
+        let mut is_master_reply = is_master(command, conn).await?;
+        conn.stream_description = Some(StreamDescription::from_is_master(is_master_reply.clone()));
 
         // Record the client's message and the server's response from speculative authentication if
         // the server did send a response.
         let first_round = client_first.and_then(|client_first| {
-            command_response
+            is_master_reply
+                .command_response
                 .speculative_authenticate
                 .take()
                 .map(|server_first| client_first.into_first_round(server_first))
         });
 
-        let is_master_reply = IsMasterReply {
-            command_response,
-            round_trip_time: Some(start_time.to(end_time).to_std().unwrap()),
-            cluster_time: None,
-        };
+        Ok(HandshakeResult {
+            first_round,
+            is_master_reply,
+        })
+    }
+}
+
+/// The information returned from the server as part of the handshake.
+///
+/// Also optionally includes the first round of speculative authentication
+/// if applicable.
+#[derive(Debug)]
+pub(crate) struct HandshakeResult {
+    /// The response from the server.
+    pub(crate) is_master_reply: IsMasterReply,
+
+    /// The first round of speculative authentication, if applicable.
+    pub(crate) first_round: Option<FirstRound>,
+}
+
+#[derive(Debug)]
+pub(crate) struct HandshakerOptions {
+    app_name: Option<String>,
+    credential: Option<Credential>,
+    driver_info: Option<DriverInfo>,
+    server_api: Option<ServerApi>,
+}
+
+impl From<ConnectionPoolOptions> for HandshakerOptions {
+    fn from(options: ConnectionPoolOptions) -> Self {
+        Self {
+            app_name: options.app_name,
+            credential: options.credential,
+            driver_info: options.driver_info,
+            server_api: options.server_api,
+        }
+    }
+}
+
+impl From<ClientOptions> for HandshakerOptions {
+    fn from(options: ClientOptions) -> Self {
+        Self {
+            app_name: options.app_name,
+            credential: options.credential,
+            driver_info: options.driver_info,
+            server_api: options.server_api,
+        }
+    }
+}
 
-        conn.stream_description = Some(StreamDescription::from_is_master(is_master_reply));
-        Ok(first_round)
+/// Run the given isMaster command.
+///
+/// If the given command is not an isMaster, this function will return an error.
+pub(crate) async fn is_master(command: Command, conn: &mut Connection) -> Result<IsMasterReply> {
+    if !command.name.eq_ignore_ascii_case("ismaster") {
+        return Err(ErrorKind::OperationError {
+            message: format!("invalid ismaster command: {}", command.name),
+        }
+        .into());
     }
+    let start_time = PreciseTime::now();
+    let response = conn.send_command(command, None).await?;
+    let end_time = PreciseTime::now();
+
+    response.validate()?;
+    let cluster_time = response.cluster_time().cloned();
+    let command_response: IsMasterCommandResponse = response.body()?;
+
+    Ok(IsMasterReply {
+        command_response,
+        round_trip_time: Some(start_time.to(end_time).to_std().unwrap()),
+        cluster_time,
+    })
 }
 
 /// Updates the handshake command document with the speculative authenitication info.

src/cmap/establish/handshake/test.rs

@@ -34,7 +34,7 @@ fn metadata_with_options() {
         )
         .build();
 
-    let handshaker = Handshaker::new(Some(&options));
+    let handshaker = Handshaker::new(Some(options.into()));
 
     let metadata = handshaker.command.body.get_document("client").unwrap();
     assert_eq!(

src/cmap/establish/mod.rs

@@ -1,4 +1,4 @@
-mod handshake;
+pub(super) mod handshake;
 #[cfg(test)]
 mod test;
 
@@ -24,7 +24,7 @@ pub(super) struct ConnectionEstablisher {
 impl ConnectionEstablisher {
     /// Creates a new ConnectionEstablisher from the given options.
     pub(super) fn new(http_client: HttpClient, options: Option<&ConnectionPoolOptions>) -> Self {
-        let handshaker = Handshaker::new(options);
+        let handshaker = Handshaker::new(options.cloned().map(Into::into));
 
         Self {
             handshaker,
@@ -43,8 +43,9 @@ impl ConnectionEstablisher {
 
         let first_round = self
             .handshaker
-            .handshake(&mut connection, self.credential.as_ref())
-            .await?;
+            .handshake(&mut connection)
+            .await?
+            .first_round;
 
         if let Some(ref credential) = self.credential {
             credential

src/cmap/establish/test.rs

@@ -39,16 +39,13 @@ async fn speculative_auth_test(
         .tls_options(CLIENT_OPTIONS.tls_options())
         .build();
 
-    let handshaker = Handshaker::new(Some(&pool_options));
+    let handshaker = Handshaker::new(Some(pool_options.clone().into()));
 
     let mut conn = Connection::new_testing(1, Default::default(), 1, Some(pool_options.into()))
         .await
         .unwrap();
 
-    let first_round = handshaker
-        .handshake(&mut conn, Some(&credential))
-        .await
-        .unwrap();
+    let first_round = handshaker.handshake(&mut conn).await.unwrap().first_round;
 
     // We expect that the server will return a response with the `speculativeAuthenticate` field if
     // and only if it's new enough to support it.

src/cmap/mod.rs

@@ -13,8 +13,11 @@ use std::{sync::Arc, time::Duration};
 use derivative::Derivative;
 
 pub use self::conn::ConnectionInfo;
-pub(crate) use self::conn::{Command, CommandResponse, Connection, StreamDescription};
 use self::options::ConnectionPoolOptions;
+pub(crate) use self::{
+    conn::{Command, CommandResponse, Connection, StreamDescription},
+    establish::handshake::{is_master, Handshaker},
+};
 use crate::{
     error::{ErrorKind, Result},
     event::cmap::{