From 92787159f76e2c18419c7c189234ea2cc59fd103 Mon Sep 17 00:00:00 2001
From: Bojan Serafimov <bojan.serafimov7@gmail.com>
Date: Fri, 18 Feb 2022 21:37:12 -0500
Subject: [PATCH] Add client auth method option

---
 proxy/src/auth.rs   | 16 +++++++++++-----
 proxy/src/config.rs | 27 +++++++++++++++++++++++++++
 proxy/src/main.rs   |  9 +++++++++
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/proxy/src/auth.rs b/proxy/src/auth.rs
index 20beb6ac79..17d4b0860e 100644
--- a/proxy/src/auth.rs
+++ b/proxy/src/auth.rs
@@ -1,5 +1,5 @@
 use crate::compute::DatabaseInfo;
-use crate::config::ProxyConfig;
+use crate::config::{ClientAuthMethod, ProxyConfig};
 use crate::cplane_api::{self, CPlaneApi};
 use crate::stream::PqStream;
 use anyhow::{anyhow, bail, Context};
@@ -38,10 +38,16 @@ impl ClientCredentials {
         config: &ProxyConfig,
         client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
     ) -> anyhow::Result<DatabaseInfo> {
-        let db_info = if self.user.ends_with("@zenith") {
-            handle_existing_user(config, client, self).await
-        } else {
-            handle_new_user(config, client).await
+        let db_info = match config.client_auth_method {
+            ClientAuthMethod::Mixed => {
+                if self.user.ends_with("@zenith") {
+                    handle_existing_user(config, client, self).await
+                } else {
+                    handle_new_user(config, client).await
+                }
+            }
+            ClientAuthMethod::Password => handle_existing_user(config, client, self).await,
+            ClientAuthMethod::Link => handle_new_user(config, client).await,
         };
 
         db_info.context("failed to authenticate client")
diff --git a/proxy/src/config.rs b/proxy/src/config.rs
index a39980321b..eeea7743a3 100644
--- a/proxy/src/config.rs
+++ b/proxy/src/config.rs
@@ -1,14 +1,41 @@
 use anyhow::{anyhow, ensure, Context};
 use rustls::{internal::pemfile, NoClientAuth, ProtocolVersion, ServerConfig};
 use std::net::SocketAddr;
+use std::str::FromStr;
 use std::sync::Arc;
 
 pub type TlsConfig = Arc<ServerConfig>;
 
+#[non_exhaustive]
+pub enum ClientAuthMethod {
+    Password,
+    Link,
+
+    /// Use password auth only if username ends with "@zenith"
+    Mixed,
+}
+
+impl FromStr for ClientAuthMethod {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> anyhow::Result<Self> {
+        use ClientAuthMethod::*;
+        match s {
+            "password" => Ok(Password),
+            "link" => Ok(Link),
+            "mixed" => Ok(Mixed),
+            _ => Err(anyhow::anyhow!("Invlid option for router")),
+        }
+    }
+}
+
 pub struct ProxyConfig {
     /// main entrypoint for users to connect to
     pub proxy_address: SocketAddr,
 
+    /// method of assigning compute nodes
+    pub client_auth_method: ClientAuthMethod,
+
     /// internally used for status and prometheus metrics
     pub http_address: SocketAddr,
 
diff --git a/proxy/src/main.rs b/proxy/src/main.rs
index a72c3b0b1e..84230de65c 100644
--- a/proxy/src/main.rs
+++ b/proxy/src/main.rs
@@ -44,6 +44,14 @@ async fn main() -> anyhow::Result<()> {
                 .help("listen for incoming client connections on ip:port")
                 .default_value("127.0.0.1:4432"),
         )
+        .arg(
+            Arg::new("auth-method")
+                .short('a')
+                .long("router")
+                .takes_value(true)
+                .help("Possible values: password | link | mixed")
+                .default_value("mixed"),
+        )
         .arg(
             Arg::new("mgmt")
                 .short('m')
@@ -103,6 +111,7 @@ async fn main() -> anyhow::Result<()> {
 
     let config: &ProxyConfig = Box::leak(Box::new(ProxyConfig {
         proxy_address: arg_matches.value_of("proxy").unwrap().parse()?,
+        client_auth_method: arg_matches.value_of("auth-method").unwrap().parse()?,
         mgmt_address: arg_matches.value_of("mgmt").unwrap().parse()?,
         http_address: arg_matches.value_of("http").unwrap().parse()?,
         redirect_uri: arg_matches.value_of("uri").unwrap().parse()?,