From 9a159869f2ec281d11c0a131d47a37a5bd0eba51 Mon Sep 17 00:00:00 2001
From: Andrew Hauck <ahauck@cloudflare.com>
Date: Tue, 22 Oct 2024 08:45:24 -0700
Subject: Use OkOrErr on connect() + accept() when stream doesn't exist

---
 .bleep                                          |  2 +-
 pingora-core/src/protocols/tls/rustls/stream.rs | 46 ++++++++++---------------
 2 files changed, 20 insertions(+), 28 deletions(-)

diff --git a/.bleep b/.bleep
index 6be531e..bd43cdc 100644
--- a/.bleep
+++ b/.bleep
@@ -1 +1 @@
-93a8806fd5a1e6df065f20bc40e9594fde0a21db
\ No newline at end of file
+d715e2d7a34ba00f8872339d4596b2e1e68de304
\ No newline at end of file
diff --git a/pingora-core/src/protocols/tls/rustls/stream.rs b/pingora-core/src/protocols/tls/rustls/stream.rs
index 9cdd2a7..a2c6811 100644
--- a/pingora-core/src/protocols/tls/rustls/stream.rs
+++ b/pingora-core/src/protocols/tls/rustls/stream.rs
@@ -27,7 +27,7 @@ use crate::protocols::{
 };
 use crate::utils::tls::get_organization_serial_bytes;
 use pingora_error::ErrorType::{AcceptError, ConnectError, InternalError, TLSHandshakeFailure};
-use pingora_error::{Error, ImmutStr, OrErr, Result};
+use pingora_error::{OkOrErr, OrErr, Result};
 use pingora_rustls::TlsStream as RusTlsStream;
 use pingora_rustls::{hash_certificate, NoDebug};
 use pingora_rustls::{Accept, Connect, ServerName, TlsConnector};
@@ -273,40 +273,32 @@ impl<T: AsyncRead + AsyncWrite + Unpin + Send> InnerStream<T> {
     /// Connect to the remote TLS server as a client
     pub(crate) async fn connect(&mut self) -> Result<()> {
         let connect = &mut (*self.connect);
+        let connect = connect.take().or_err(
+            ConnectError,
+            "TLS connect not available to perform handshake.",
+        )?;
 
-        if let Some(connect) = connect.take() {
-            let stream = connect
-                .await
-                .or_err(TLSHandshakeFailure, "tls connect error")?;
-            self.stream = Some(RusTlsStream::Client(stream));
-
-            Ok(())
-        } else {
-            Error::e_explain(
-                ConnectError,
-                ImmutStr::from("TLS connect not available to perform handshake."),
-            )
-        }
+        let stream = connect
+            .await
+            .or_err(TLSHandshakeFailure, "tls connect error")?;
+        self.stream = Some(RusTlsStream::Client(stream));
+        Ok(())
     }
 
     /// Finish the TLS handshake from client as a server
     /// no-op implementation within Rustls, handshake is performed during creation of stream.
     pub(crate) async fn accept(&mut self) -> Result<()> {
         let accept = &mut (*self.accept);
+        let accept = accept.take().or_err(
+            AcceptError,
+            "TLS accept not available to perform handshake.",
+        )?;
 
-        if let Some(ref mut accept) = accept.take() {
-            let stream = accept
-                .await
-                .explain_err(TLSHandshakeFailure, |e| format!("tls connect error: {e}"))?;
-            self.stream = Some(RusTlsStream::Server(stream));
-
-            Ok(())
-        } else {
-            Err(Error::explain(
-                AcceptError,
-                ImmutStr::from("TLS accept not available to perform handshake."),
-            ))
-        }
+        let stream = accept
+            .await
+            .explain_err(TLSHandshakeFailure, |e| format!("tls connect error: {e}"))?;
+        self.stream = Some(RusTlsStream::Server(stream));
+        Ok(())
     }
 
     pub(crate) fn digest(&mut self) -> Option<Arc<SslDigest>> {
-- 
cgit v1.2.3