From 9df8915b03e03135bf3f8f78fa00435c94aa3ccd Mon Sep 17 00:00:00 2001
From: Dmitry Ivanov <ivadmi5@gmail.com>
Date: Tue, 12 Apr 2022 01:12:07 +0300
Subject: [PATCH] [proxy] `sasl::Mechanism` may return `Output` during exchange

This is needed to forward the `ClientKey` that's required
to connect the proxy to a compute.

Co-authored-by: bojanserafimov <bojan.serafimov7@gmail.com>
---
 proxy/src/sasl.rs           | 13 ++++++++++++-
 proxy/src/sasl/messages.rs  |  1 +
 proxy/src/sasl/stream.rs    | 13 +++++++++----
 proxy/src/scram.rs          |  4 ++--
 proxy/src/scram/exchange.rs | 10 ++++++----
 5 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/proxy/src/sasl.rs b/proxy/src/sasl.rs
index 70a4d9946a..cd9032bfb9 100644
--- a/proxy/src/sasl.rs
+++ b/proxy/src/sasl.rs
@@ -39,9 +39,20 @@ pub enum Error {
 /// A convenient result type for SASL exchange.
 pub type Result<T> = std::result::Result<T, Error>;
 
+/// A result of one SASL exchange.
+pub enum Step<T, R> {
+    /// We should continue exchanging messages.
+    Continue(T),
+    /// The client has been authenticated successfully.
+    Authenticated(R),
+}
+
 /// Every SASL mechanism (e.g. [SCRAM](crate::scram)) is expected to implement this trait.
 pub trait Mechanism: Sized {
+    /// What's produced as a result of successful authentication.
+    type Output;
+
     /// Produce a server challenge to be sent to the client.
     /// This is how this method is called in PostgreSQL (`libpq/sasl.h`).
-    fn exchange(self, input: &str) -> Result<(Option<Self>, String)>;
+    fn exchange(self, input: &str) -> Result<(Step<Self, Self::Output>, String)>;
 }
diff --git a/proxy/src/sasl/messages.rs b/proxy/src/sasl/messages.rs
index 58be6268fe..f48aee4f26 100644
--- a/proxy/src/sasl/messages.rs
+++ b/proxy/src/sasl/messages.rs
@@ -49,6 +49,7 @@ impl<'a> ServerMessage<&'a str> {
         })
     }
 }
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/proxy/src/sasl/stream.rs b/proxy/src/sasl/stream.rs
index 03649b8d11..0e782c5f29 100644
--- a/proxy/src/sasl/stream.rs
+++ b/proxy/src/sasl/stream.rs
@@ -51,18 +51,23 @@ impl<S: AsyncWrite + Unpin> SaslStream<'_, S> {
 impl<S: AsyncRead + AsyncWrite + Unpin> SaslStream<'_, S> {
     /// Perform SASL message exchange according to the underlying algorithm
     /// until user is either authenticated or denied access.
-    pub async fn authenticate(mut self, mut mechanism: impl Mechanism) -> super::Result<()> {
+    pub async fn authenticate<M: Mechanism>(
+        mut self,
+        mut mechanism: M,
+    ) -> super::Result<M::Output> {
         loop {
             let input = self.recv().await?;
             let (moved, reply) = mechanism.exchange(input)?;
+
+            use super::Step::*;
             match moved {
-                Some(moved) => {
+                Continue(moved) => {
                     self.send(&ServerMessage::Continue(&reply)).await?;
                     mechanism = moved;
                 }
-                None => {
+                Authenticated(result) => {
                     self.send(&ServerMessage::Final(&reply)).await?;
-                    return Ok(());
+                    return Ok(result);
                 }
             }
         }
diff --git a/proxy/src/scram.rs b/proxy/src/scram.rs
index 44671084ee..22fce7ac7e 100644
--- a/proxy/src/scram.rs
+++ b/proxy/src/scram.rs
@@ -13,10 +13,10 @@ mod password;
 mod secret;
 mod signature;
 
-pub use secret::*;
-
 pub use exchange::Exchange;
+pub use key::ScramKey;
 pub use secret::ServerSecret;
+pub use secret::*;
 
 use hmac::{Hmac, Mac};
 use sha2::{Digest, Sha256};
diff --git a/proxy/src/scram/exchange.rs b/proxy/src/scram/exchange.rs
index 802fe61db5..cad77e15f5 100644
--- a/proxy/src/scram/exchange.rs
+++ b/proxy/src/scram/exchange.rs
@@ -62,8 +62,10 @@ impl<'a> Exchange<'a> {
 }
 
 impl sasl::Mechanism for Exchange<'_> {
-    fn exchange(mut self, input: &str) -> sasl::Result<(Option<Self>, String)> {
-        use ExchangeState::*;
+    type Output = super::ScramKey;
+
+    fn exchange(mut self, input: &str) -> sasl::Result<(sasl::Step<Self, Self::Output>, String)> {
+        use {sasl::Step::*, ExchangeState::*};
         match &self.state {
             Initial => {
                 let client_first_message =
@@ -82,7 +84,7 @@ impl sasl::Mechanism for Exchange<'_> {
                     server_first_message,
                 };
 
-                Ok((Some(self), msg))
+                Ok((Continue(self), msg))
             }
             SaltSent {
                 cbind_flag,
@@ -124,7 +126,7 @@ impl sasl::Mechanism for Exchange<'_> {
                 let msg = client_final_message
                     .build_server_final_message(signature_builder, &self.secret.server_key);
 
-                Ok((None, msg))
+                Ok((Authenticated(client_key), msg))
             }
         }
     }