From 9a396e1feb9f35e4f2d57d38a2ac07070ecc1b4b Mon Sep 17 00:00:00 2001
From: Stas Kelvich <stas.kelvich@gmail.com>
Date: Mon, 2 May 2022 00:35:15 +0300
Subject: [PATCH] Support SNI-based routing in proxy

---
 proxy/src/auth.rs                 |  2 ++
 proxy/src/auth/credentials.rs     |  6 +++---
 proxy/src/auth_backend/console.rs | 15 +++++++++++----
 proxy/src/proxy.rs                |  7 +++++--
 4 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/proxy/src/auth.rs b/proxy/src/auth.rs
index d4e21d78a0..2463f31645 100644
--- a/proxy/src/auth.rs
+++ b/proxy/src/auth.rs
@@ -6,6 +6,7 @@ use crate::config::{AuthBackendType, ProxyConfig};
 use crate::error::UserFacingError;
 use crate::stream::PqStream;
 use crate::{auth_backend, compute, waiters};
+use console::ConsoleAuthError::SniMissing;
 use std::io;
 use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -72,6 +73,7 @@ impl UserFacingError for AuthError {
         match self.0.as_ref() {
             Console(e) => e.to_string_client(),
             MalformedPassword => self.to_string(),
+            GetAuthInfo(e) if matches!(e, SniMissing) => e.to_string(),
             _ => "Internal error".to_string(),
         }
     }
diff --git a/proxy/src/auth/credentials.rs b/proxy/src/auth/credentials.rs
index 88677de511..9d2272b5ad 100644
--- a/proxy/src/auth/credentials.rs
+++ b/proxy/src/auth/credentials.rs
@@ -24,9 +24,9 @@ pub struct ClientCredentials {
     pub user: String,
     pub dbname: String,
 
-    // New console API requires SNI info to determine cluster name.
+    // New console API requires SNI info to determine the cluster name.
     // Other Auth backends don't need it.
-    pub sni_cluster: Option<String>,
+    pub sni_data: Option<String>,
 }
 
 impl ClientCredentials {
@@ -52,7 +52,7 @@ impl TryFrom<HashMap<String, String>> for ClientCredentials {
         Ok(Self {
             user,
             dbname: db,
-            sni_cluster: None,
+            sni_data: None,
         })
     }
 }
diff --git a/proxy/src/auth_backend/console.rs b/proxy/src/auth_backend/console.rs
index 863e929489..55a0889af4 100644
--- a/proxy/src/auth_backend/console.rs
+++ b/proxy/src/auth_backend/console.rs
@@ -22,10 +22,12 @@ pub enum ConsoleAuthError {
     #[error("Bad client credentials: {0:?}")]
     BadCredentials(crate::auth::ClientCredentials),
 
-    /// For passwords that couldn't be processed by [`parse_password`].
-    #[error("Absend SNI information")]
+    #[error("SNI info is missing, please upgrade the postgres client library")]
     SniMissing,
 
+    #[error("Unexpected SNI content")]
+    SniWrong,
+
     #[error(transparent)]
     BadUrl(#[from] url::ParseError),
 
@@ -166,10 +168,15 @@ pub async fn handle_user(
     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
     creds: &ClientCredentials,
 ) -> Result<compute::NodeInfo, crate::auth::AuthError> {
+    // Determine cluster name from SNI.
     let cluster = creds
-        .sni_cluster
+        .sni_data
         .as_ref()
-        .ok_or(ConsoleAuthError::SniMissing)?;
+        .ok_or(ConsoleAuthError::SniMissing)?
+        .split_once('.')
+        .ok_or(ConsoleAuthError::SniWrong)?
+        .0;
+
     let user = creds.user.as_str();
 
     // Step 1: get the auth secret
diff --git a/proxy/src/proxy.rs b/proxy/src/proxy.rs
index 4bdbac8510..821ce377f5 100644
--- a/proxy/src/proxy.rs
+++ b/proxy/src/proxy.rs
@@ -144,11 +144,14 @@ async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
                 }
 
                 // Here and forth: `or_else` demands that we use a future here
-                let creds = async { params.try_into() }
+                let mut creds: auth::ClientCredentials = async { params.try_into() }
                     .or_else(|e| stream.throw_error(e))
                     .await?;
 
-                // TODO: set creds.cluster here when SNI info is available
+                // Set SNI info when available
+                if let Stream::Tls { tls } = stream.get_ref() {
+                    creds.sni_data = tls.get_ref().1.sni_hostname().map(|s| s.to_owned());
+                }
 
                 break Ok(Some((stream, creds)));
             }