Add TLS first support

NATS server used to send INFO when client connected, before establishing
TLS.
As server added support to establish TLS before getting any INFO,
this commit adds `tls_first` option to be able to used that feature.

Signed-off-by: Tomasz Pietrek <[email protected]>

async-nats/src/connector.rs

@@ -51,6 +51,7 @@ pub(crate) struct ConnectorOptions {
     pub(crate) client_cert: Option<PathBuf>,
     pub(crate) client_key: Option<PathBuf>,
     pub(crate) tls_client_config: Option<rustls::ClientConfig>,
+    pub(crate) tls_first: bool,
     pub(crate) auth: Auth,
     pub(crate) no_echo: bool,
     pub(crate) connection_timeout: Duration,
@@ -299,24 +300,7 @@ impl Connector {
             self.options.read_buffer_capacity.into(),
         );
 
-        let op = connection.read_op().await?;
-        let info = match op {
-            Some(ServerOp::Info(info)) => info,
-            Some(op) => {
-                return Err(ConnectError::with_source(
-                    crate::ConnectErrorKind::Io,
-                    format!("expected INFO, got {:?}", op),
-                ))
-            }
-            None => {
-                return Err(ConnectError::with_source(
-                    crate::ConnectErrorKind::Io,
-                    "expected INFO, got nothing",
-                ))
-            }
-        };
-
-        if self.options.tls_required || info.tls_required || tls_required {
+        let tls_connection = |connection: Connection| async {
             let tls_config = Arc::new(
                 tls::config_tls(&self.options)
                     .await
@@ -334,10 +318,40 @@ impl Connector {
             let domain = rustls::ServerName::try_from(tls_host)
                 .map_err(|err| ConnectError::with_source(crate::ConnectErrorKind::Tls, err))?;
 
-            connection = Connection::new(
-                Box::new(tls_connector.connect(domain, connection.stream).await?),
-                0,
-            );
+            let tls_stream = tls_connector.connect(domain, connection.stream).await?;
+
+            Ok::<Connection, ConnectError>(Connection::new(Box::new(tls_stream), 0))
+        };
+
+        // If `tls_first` was set, establish TLS connection before getting INFO.
+        // There is no point in  checking if tls is required, because
+        // the connection has to be be upgraded to TLS anyway as it's different flow.
+        if self.options.tls_first {
+            connection = tls_connection(connection).await?;
+        }
+
+        let op = connection.read_op().await?;
+        let info = match op {
+            Some(ServerOp::Info(info)) => info,
+            Some(op) => {
+                return Err(ConnectError::with_source(
+                    crate::ConnectErrorKind::Io,
+                    format!("expected INFO, got {:?}", op),
+                ))
+            }
+            None => {
+                return Err(ConnectError::with_source(
+                    crate::ConnectErrorKind::Io,
+                    "expected INFO, got nothing",
+                ))
+            }
+        };
+
+        // If `tls_first` was not set, establish TLS connection if it is required.
+        if !self.options.tls_first
+            && (self.options.tls_required || info.tls_required || tls_required)
+        {
+            connection = tls_connection(connection).await?;
         };
 
         Ok((*info, connection))

async-nats/src/lib.rs

@@ -794,6 +794,7 @@ pub async fn connect_with_options<A: ToServerAddrs>(
             client_key: options.client_key,
             client_cert: options.client_cert,
             tls_client_config: options.tls_client_config,
+            tls_first: options.tls_first,
             auth: options.auth,
             no_echo: options.no_echo,
             connection_timeout: options.connection_timeout,

async-nats/src/options.rs

@@ -51,6 +51,7 @@ pub struct ConnectOptions {
     pub(crate) connection_timeout: Duration,
     pub(crate) auth: Auth,
     pub(crate) tls_required: bool,
+    pub(crate) tls_first: bool,
     pub(crate) certificates: Vec<PathBuf>,
     pub(crate) client_cert: Option<PathBuf>,
     pub(crate) client_key: Option<PathBuf>,
@@ -83,6 +84,7 @@ impl fmt::Debug for ConnectOptions {
             .entry(&"client_cert", &self.client_cert)
             .entry(&"client_key", &self.client_key)
             .entry(&"tls_client_config", &"XXXXXXXX")
+            .entry(&"tls_first", &self.tls_first)
             .entry(&"ping_interval", &self.ping_interval)
             .entry(&"sender_capacity", &self.sender_capacity)
             .entry(&"inbox_prefix", &self.inbox_prefix)
@@ -102,6 +104,7 @@ impl Default for ConnectOptions {
             max_reconnects: Some(60),
             connection_timeout: Duration::from_secs(5),
             tls_required: false,
+            tls_first: false,
             certificates: Vec::new(),
             client_cert: None,
             client_key: None,
@@ -565,6 +568,15 @@ impl ConnectOptions {
         self
     }
 
+    /// Changes how tls connection is established. If `tls_first` is set,
+    /// client will try to establish tls before getting info from the server.
+    /// That requires the server to enable `handshake_first` option in the config.
+    pub fn tls_first(mut self) -> ConnectOptions {
+        self.tls_first = true;
+        self.tls_required = true;
+        self
+    }
+
     /// Sets how often Client sends PING message to the server.
     ///
     /// # Examples

async-nats/tests/configs/tls_first.conf

@@ -0,0 +1,17 @@
+# this needs to be here for testing localhost tls.
+listen: localhost:4545
+
+tls {
+  cert_file:  "./tests/configs/certs/server-cert.pem"
+  key_file:   "./tests/configs/certs/server-key.pem"
+  ca_file:    "./tests/configs/certs/rootCA.pem"
+  verify :    true
+  timeout:    2
+  handshake_first: true
+}
+
+authorization {
+  user:     derek
+  password: porkchop
+  timeout:  1
+}

async-nats/tests/configs/tls_first_auto.conf

@@ -0,0 +1,17 @@
+# this needs to be here for testing localhost tls.
+listen: localhost:4545
+
+tls {
+  cert_file:  "./tests/configs/certs/server-cert.pem"
+  key_file:   "./tests/configs/certs/server-key.pem"
+  ca_file:    "./tests/configs/certs/rootCA.pem"
+  verify :    true
+  timeout:    2
+  handshake_first: 300ms
+}
+
+authorization {
+  user:     derek
+  password: porkchop
+  timeout:  1
+}

async-nats/tests/tls_tests.rs

@@ -12,7 +12,7 @@
 // limitations under the License.
 
 mod client {
-    use std::path::PathBuf;
+    use std::{path::PathBuf, time::Duration};
 
     use futures::StreamExt;
 
@@ -141,4 +141,78 @@ mod client {
         client.flush().await.unwrap();
         assert!(subscription.next().await.is_some());
     }
+
+    #[tokio::test]
+    async fn tls_first() {
+        let _server =
+            nats_server::run_server_with_port("tests/configs/tls_first.conf", Some("9090"));
+
+        // For some reason tls-first makes server starup longer.
+        tokio::time::sleep(Duration::from_secs(2)).await;
+
+        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+
+        let client =
+            async_nats::ConnectOptions::with_user_and_password("derek".into(), "porkchop".into())
+                .add_root_certificates(path.join("tests/configs/certs/rootCA.pem"))
+                .add_client_certificate(
+                    path.join("tests/configs/certs/client-cert.pem"),
+                    path.join("tests/configs/certs/client-key.pem"),
+                )
+                .require_tls(true)
+                .tls_first()
+                .connect("tls://localhost:9090")
+                .await
+                .unwrap();
+
+        assert_eq!(client.server_info().tls_required, true);
+
+        async_nats::ConnectOptions::with_user_and_password("derek".into(), "porkchop".into())
+            .add_root_certificates(path.join("tests/configs/certs/rootCA.pem"))
+            .add_client_certificate(
+                path.join("tests/configs/certs/client-cert.pem"),
+                path.join("tests/configs/certs/client-key.pem"),
+            )
+            .require_tls(true)
+            .connect("tls://localhost:9090")
+            .await
+            .unwrap_err();
+    }
+
+    #[tokio::test]
+    async fn tls_auto() {
+        let _server =
+            nats_server::run_server_with_port("tests/configs/tls_first_auto.conf", Some("9898"));
+
+        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+
+        // For some reason tls-first makes server starup longer.
+        tokio::time::sleep(Duration::from_secs(2)).await;
+
+        async_nats::ConnectOptions::with_user_and_password("derek".into(), "porkchop".into())
+            .add_root_certificates(path.join("tests/configs/certs/rootCA.pem"))
+            .add_client_certificate(
+                path.join("tests/configs/certs/client-cert.pem"),
+                path.join("tests/configs/certs/client-key.pem"),
+            )
+            .require_tls(true)
+            .connect("tls://localhost:9898")
+            .await
+            .unwrap();
+
+        let client =
+            async_nats::ConnectOptions::with_user_and_password("derek".into(), "porkchop".into())
+                .add_root_certificates(path.join("tests/configs/certs/rootCA.pem"))
+                .add_client_certificate(
+                    path.join("tests/configs/certs/client-cert.pem"),
+                    path.join("tests/configs/certs/client-key.pem"),
+                )
+                .require_tls(true)
+                .tls_first()
+                .connect("tls://localhost:9898")
+                .await
+                .unwrap();
+
+        assert_eq!(client.server_info().tls_required, true);
+    }
 }