feat(frontend): implement OAuth authentication

proto/meta.proto

@@ -557,6 +557,7 @@ message SystemParams {
   optional bool pause_on_next_bootstrap = 13;
   optional string wasm_storage_url = 14;
   optional bool enable_tracing = 15;
+  optional string oauth_jwks_url = 16;
 }
 
 message GetSystemParamsRequest {}

proto/user.proto

@@ -14,6 +14,7 @@ message AuthInfo {
     PLAINTEXT = 1;
     SHA256 = 2;
     MD5 = 3;
+    OAUTH = 4;
   }
   EncryptionType encryption_type = 1;
   bytes encrypted_value = 2;

src/common/src/system_param/mod.rs

@@ -88,6 +88,7 @@ macro_rules! for_all_params {
             { pause_on_next_bootstrap,                  bool,   Some(false),                                true,   "Whether to pause all data sources on next bootstrap.", },
             { wasm_storage_url,                         String, Some("fs://.risingwave/data".to_string()),  false,  "", },
             { enable_tracing,                           bool,   Some(false),                                true,   "Whether to enable distributed tracing.", },
+            { oauth_jwks_url,                           String, None,                                       true,   "Url to get JSON Web Key Set(JWKS) for oauth authentication.", },
         }
     };
 }
@@ -376,6 +377,7 @@ macro_rules! impl_system_params_for_test {
             ret.state_store = Some("hummock+memory".to_string());
             ret.backup_storage_url = Some("memory".into());
             ret.backup_storage_directory = Some("backup".into());
+            ret.oauth_jwks_url = Some("https://auth-static.confluent.io/jwks".into());
             ret
         }
     };
@@ -442,6 +444,7 @@ mod tests {
             (PAUSE_ON_NEXT_BOOTSTRAP_KEY, "false"),
             (WASM_STORAGE_URL_KEY, "a"),
             (ENABLE_TRACING_KEY, "true"),
+            (OAUTH_JWKS_URL_KEY, "a"),
             ("a_deprecated_param", "foo"),
         ];
 

src/common/src/system_param/reader.rs

@@ -167,4 +167,8 @@ where
             .as_ref()
             .unwrap_or(&default::WASM_STORAGE_URL)
     }
+
+    fn oauth_jwks_url(&self) -> &str {
+        self.inner().oauth_jwks_url.as_ref().unwrap()
+    }
 }

src/config/docs.md

@@ -146,6 +146,7 @@ This page is automatically generated by `./risedev generate-example-config`
 | data_directory | Remote directory for storing data and metadata objects. |  |
 | enable_tracing | Whether to enable distributed tracing. | false |
 | max_concurrent_creating_streaming_jobs | Max number of concurrent creating streaming jobs. | 1 |
+| oauth_jwks_url | Url to get JSON Web Key Set(JWKS) for oauth authentication. |  |
 | parallel_compact_size_mb |  | 512 |
 | pause_on_next_bootstrap | Whether to pause all data sources on next bootstrap. | false |
 | sstable_size_mb | Target size of the Sstable. | 256 |

src/frontend/src/handler/alter_user.rs

@@ -23,7 +23,7 @@ use crate::catalog::CatalogError;
 use crate::error::ErrorCode::{InternalError, PermissionDenied};
 use crate::error::Result;
 use crate::handler::HandlerArgs;
-use crate::user::user_authentication::encrypted_password;
+use crate::user::user_authentication::{build_oauth_info, encrypted_password};
 use crate::user::user_catalog::UserCatalog;
 
 fn alter_prost_user_info(
@@ -111,6 +111,10 @@ fn alter_prost_user_info(
                 }
                 update_fields.push(UpdateField::AuthInfo);
             }
+            UserOption::OAuth => {
+                user_info.auth_info = build_oauth_info();
+                update_fields.push(UpdateField::AuthInfo)
+            }
         }
     }
     Ok((user_info, update_fields))

src/frontend/src/handler/create_user.rs

@@ -23,7 +23,7 @@ use crate::catalog::{CatalogError, DatabaseId};
 use crate::error::ErrorCode::PermissionDenied;
 use crate::error::Result;
 use crate::handler::HandlerArgs;
-use crate::user::user_authentication::encrypted_password;
+use crate::user::user_authentication::{build_oauth_info, encrypted_password};
 use crate::user::user_catalog::UserCatalog;
 
 fn make_prost_user_info(
@@ -91,6 +91,7 @@ fn make_prost_user_info(
                     user_info.auth_info = encrypted_password(&user_info.name, &password.0);
                 }
             }
+            UserOption::OAuth => user_info.auth_info = build_oauth_info(),
         }
     }
 

src/frontend/src/session.rs

@@ -46,6 +46,7 @@ use risingwave_common::session_config::{ConfigMap, ConfigReporter, VisibilityMod
 use risingwave_common::system_param::local_manager::{
     LocalSystemParamsManager, LocalSystemParamsManagerRef,
 };
+use risingwave_common::system_param::reader::{SystemParamsRead, SystemParamsReader};
 use risingwave_common::telemetry::manager::TelemetryManager;
 use risingwave_common::telemetry::telemetry_env_enabled;
 use risingwave_common::types::DataType;
@@ -924,26 +925,29 @@ pub struct SessionManagerImpl {
 impl SessionManager for SessionManagerImpl {
     type Session = SessionImpl;
 
-    fn connect(
+    async fn connect(
         &self,
-        database: &str,
-        user_name: &str,
+        database: String,
+        user_name: String,
         peer_addr: AddressRef,
     ) -> std::result::Result<Arc<Self::Session>, BoxedError> {
-        let catalog_reader = self.env.catalog_reader();
-        let reader = catalog_reader.read_guard();
-        let database_id = reader
-            .get_database_by_name(database)
-            .map_err(|_| {
-                Box::new(Error::new(
-                    ErrorKind::InvalidInput,
-                    format!("database \"{}\" does not exist", database),
-                ))
-            })?
-            .id();
-        let user_reader = self.env.user_info_reader();
-        let reader = user_reader.read_guard();
-        if let Some(user) = reader.get_user_by_name(user_name) {
+        let database_id = {
+            let catalog_reader = self.env.catalog_reader().read_guard();
+            catalog_reader
+                .get_database_by_name(&database)
+                .map_err(|_| {
+                    Box::new(Error::new(
+                        ErrorKind::InvalidInput,
+                        format!("database \"{}\" does not exist", database),
+                    ))
+                })?
+                .id()
+        };
+        let user = {
+            let user_reader = self.env.user_info_reader().read_guard();
+            user_reader.get_user_by_name(&user_name).cloned()
+        };
+        if let Some(user) = user {
             if !user.can_login {
                 return Err(Box::new(Error::new(
                     ErrorKind::InvalidInput,
@@ -974,6 +978,15 @@ impl SessionManager for SessionManagerImpl {
                             ),
                             salt,
                         }
+                    } else if auth_info.encryption_type == EncryptionType::Oauth as i32 {
+                        let reader = self
+                            .env
+                            .meta_client()
+                            .get_system_params()
+                            .await
+                            .map_err(|e| PsqlError::StartupError(e.into()))?;
+                        let oauth_jwks_url = reader.oauth_jwks_url().to_string();
+                        UserAuthenticator::OAuth(oauth_jwks_url)
                     } else {
                         return Err(Box::new(Error::new(
                             ErrorKind::Unsupported,
@@ -1087,6 +1100,10 @@ impl Session for SessionImpl {
         &self.user_authenticator
     }
 
+    async fn get_system_params(&self) -> std::result::Result<SystemParamsReader, BoxedError> {
+        Ok(self.env.meta_client.get_system_params().await?)
+    }
+
     fn id(&self) -> SessionId {
         self.id
     }

src/frontend/src/test_utils.rs

@@ -80,10 +80,10 @@ pub struct LocalFrontend {
 impl SessionManager for LocalFrontend {
     type Session = SessionImpl;
 
-    fn connect(
+    async fn connect(
         &self,
-        _database: &str,
-        _user_name: &str,
+        _database: String,
+        _user_name: String,
         _peer_addr: AddressRef,
     ) -> std::result::Result<Arc<Self::Session>, BoxedError> {
         Ok(self.session_ref())

src/frontend/src/user/user_authentication.rs

@@ -24,6 +24,15 @@ const MD5_ENCRYPTED_PREFIX: &str = "md5";
 const VALID_SHA256_ENCRYPTED_LEN: usize = SHA256_ENCRYPTED_PREFIX.len() + 64;
 const VALID_MD5_ENCRYPTED_LEN: usize = MD5_ENCRYPTED_PREFIX.len() + 32;
 
+/// Build `AuthInfo` for `OAuth`.
+#[inline(always)]
+pub fn build_oauth_info() -> Option<AuthInfo> {
+    Some(AuthInfo {
+        encryption_type: EncryptionType::Oauth as i32,
+        encrypted_value: Vec::new(),
+    })
+}
+
 /// Try to extract the encryption password from given password. The password is always stored
 /// encrypted in the system catalogs. The ENCRYPTED keyword has no effect, but is accepted for
 /// backwards compatibility. The method of encryption is by default SHA-256-encrypted. If the
@@ -81,6 +90,7 @@ pub fn encrypted_raw_password(info: &AuthInfo) -> String {
         EncryptionType::Plaintext => "",
         EncryptionType::Sha256 => SHA256_ENCRYPTED_PREFIX,
         EncryptionType::Md5 => MD5_ENCRYPTED_PREFIX,
+        EncryptionType::Oauth => "",
     };
     format!("{}{}", prefix, encrypted_pwd)
 }

src/sqlparser/src/ast/statement.rs

@@ -715,6 +715,7 @@ pub enum UserOption {
     NoLogin,
     EncryptedPassword(AstString),
     Password(Option<AstString>),
+    OAuth,
 }
 
 impl fmt::Display for UserOption {
@@ -731,6 +732,7 @@ impl fmt::Display for UserOption {
             UserOption::EncryptedPassword(p) => write!(f, "ENCRYPTED PASSWORD {}", p),
             UserOption::Password(None) => write!(f, "PASSWORD NULL"),
             UserOption::Password(Some(p)) => write!(f, "PASSWORD {}", p),
+            UserOption::OAuth => write!(f, "OAUTH"),
         }
     }
 }
@@ -818,10 +820,11 @@ impl ParseTo for UserOptions {
                             UserOption::EncryptedPassword(AstString::parse_to(parser)?),
                         )
                     }
+                    Keyword::OAUTH => (&mut builder.password, UserOption::OAuth),
                     _ => {
                         parser.expected(
                             "SUPERUSER | NOSUPERUSER | CREATEDB | NOCREATEDB | LOGIN \
-                            | NOLOGIN | CREATEUSER | NOCREATEUSER | [ENCRYPTED] PASSWORD | NULL",
+                            | NOLOGIN | CREATEUSER | NOCREATEUSER | [ENCRYPTED] PASSWORD | NULL | OAUTH",
                             token,
                         )?;
                         unreachable!()
@@ -831,7 +834,7 @@ impl ParseTo for UserOptions {
             } else {
                 parser.expected(
                     "SUPERUSER | NOSUPERUSER | CREATEDB | NOCREATEDB | LOGIN | NOLOGIN \
-                        | CREATEUSER | NOCREATEUSER | [ENCRYPTED] PASSWORD | NULL",
+                        | CREATEUSER | NOCREATEUSER | [ENCRYPTED] PASSWORD | NULL | OAUTH",
                     token,
                 )?
             }

src/sqlparser/src/keywords.rs

@@ -344,6 +344,7 @@ define_keywords!(
     NULLIF,
     NULLS,
     NUMERIC,
+    OAUTH,
     OBJECT,
     OCCURRENCES_REGEX,
     OCTET_LENGTH,

src/sqlparser/src/parser.rs

@@ -2377,7 +2377,7 @@ impl Parser {
     //     | CREATEDB | NOCREATEDB
     //     | CREATEUSER | NOCREATEUSER
     //     | LOGIN | NOLOGIN
-    //     | [ ENCRYPTED ] PASSWORD 'password' | PASSWORD NULL
+    //     | [ ENCRYPTED ] PASSWORD 'password' | PASSWORD NULL | OAUTH
     fn parse_create_user(&mut self) -> Result<Statement, ParserError> {
         Ok(Statement::CreateUser(CreateUserStatement::parse_to(self)?))
     }

src/utils/pgwire/Cargo.toml

@@ -21,11 +21,14 @@ byteorder = "1.5"
 bytes = "1"
 futures = { version = "0.3", default-features = false, features = ["alloc"] }
 itertools = "0.12"
+jsonwebtoken = "9"
 openssl = "0.10.60"
 panic-message = "0.3"
 parking_lot = "0.12"
+reqwest = { version = "0.11" }
 risingwave_common = { workspace = true }
 risingwave_sqlparser = { workspace = true }
+serde = { version = "1", features = ["derive"] }
 thiserror = "1"
 thiserror-ext = { workspace = true }
 tokio = { version = "0.2", package = "madsim-tokio", features = ["rt", "macros"] }

src/utils/pgwire/src/pg_protocol.rs

@@ -386,8 +386,8 @@ where
 
         match msg {
             FeMessage::Ssl => self.process_ssl_msg().await?,
-            FeMessage::Startup(msg) => self.process_startup_msg(msg)?,
-            FeMessage::Password(msg) => self.process_password_msg(msg)?,
+            FeMessage::Startup(msg) => self.process_startup_msg(msg).await?,
+            FeMessage::Password(msg) => self.process_password_msg(msg).await?,
             FeMessage::Query(query_msg) => self.process_query_msg(query_msg.get_sql()).await?,
             FeMessage::CancelQuery(m) => self.process_cancel_msg(m)?,
             FeMessage::Terminate => self.process_terminate(),
@@ -469,7 +469,7 @@ where
         Ok(())
     }
 
-    fn process_startup_msg(&mut self, msg: FeStartupMessage) -> PsqlResult<()> {
+    async fn process_startup_msg(&mut self, msg: FeStartupMessage) -> PsqlResult<()> {
         let db_name = msg
             .config
             .get("database")
@@ -483,7 +483,8 @@ where
 
         let session = self
             .session_mgr
-            .connect(&db_name, &user_name, self.peer_addr.clone())
+            .connect(db_name, user_name, self.peer_addr.clone())
+            .await
             .map_err(PsqlError::StartupError)?;
 
         let application_name = msg.config.get("application_name");
@@ -508,7 +509,7 @@ where
                     })?;
                 self.ready_for_query()?;
             }
-            UserAuthenticator::ClearText(_) => {
+            UserAuthenticator::ClearText(_) | UserAuthenticator::OAuth(_) => {
                 self.stream
                     .write_no_flush(&BeMessage::AuthenticationCleartextPassword)?;
             }
@@ -523,11 +524,9 @@ where
         Ok(())
     }
 
-    fn process_password_msg(&mut self, msg: FePasswordMessage) -> PsqlResult<()> {
+    async fn process_password_msg(&mut self, msg: FePasswordMessage) -> PsqlResult<()> {
         let authenticator = self.session.as_ref().unwrap().user_authenticator();
-        if !authenticator.authenticate(&msg.password) {
-            return Err(PsqlError::PasswordError);
-        }
+        authenticator.authenticate(&msg.password).await?;
         self.stream.write_no_flush(&BeMessage::AuthenticationOk)?;
         self.stream
             .write_parameter_status_msg_no_flush(&ParameterStatus::default())?;