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 {}

src/common/src/system_param/mod.rs

@@ -73,21 +73,23 @@ impl_param_value!(String => &'a str);
 macro_rules! for_all_params {
     ($macro:ident) => {
         $macro! {
-            // name                                     type    default value                               mut?    doc
-            { barrier_interval_ms,                      u32,    Some(1000_u32),                             true,   "The interval of periodic barrier.", },
-            { checkpoint_frequency,                     u64,    Some(1_u64),                                true,   "There will be a checkpoint for every n barriers.", },
-            { sstable_size_mb,                          u32,    Some(256_u32),                              false,  "Target size of the Sstable.", },
-            { parallel_compact_size_mb,                 u32,    Some(512_u32),                              false,  "", },
-            { block_size_kb,                            u32,    Some(64_u32),                               false,  "Size of each block in bytes in SST.", },
-            { bloom_false_positive,                     f64,    Some(0.001_f64),                            false,  "False positive probability of bloom filter.", },
-            { state_store,                              String, None,                                       false,  "", },
-            { data_directory,                           String, None,                                       false,  "Remote directory for storing data and metadata objects.", },
-            { backup_storage_url,                       String, None,                                       true,   "Remote storage url for storing snapshots.", },
-            { backup_storage_directory,                 String, None,                                       true,   "Remote directory for storing snapshots.", },
-            { max_concurrent_creating_streaming_jobs,   u32,    Some(1_u32),                                true,   "Max number of concurrent creating streaming jobs.", },
-            { 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.", },
+            // name                                     type    default value                                               mut?    doc
+            { barrier_interval_ms,                      u32,    Some(1000_u32),                                             true,   "The interval of periodic barrier.", },
+            { checkpoint_frequency,                     u64,    Some(1_u64),                                                true,   "There will be a checkpoint for every n barriers.", },
+            { sstable_size_mb,                          u32,    Some(256_u32),                                              false,  "Target size of the Sstable.", },
+            { parallel_compact_size_mb,                 u32,    Some(512_u32),                                              false,  "", },
+            { block_size_kb,                            u32,    Some(64_u32),                                               false,  "Size of each block in bytes in SST.", },
+            { bloom_false_positive,                     f64,    Some(0.001_f64),                                            false,  "False positive probability of bloom filter.", },
+            { state_store,                              String, None,                                                       false,  "", },
+            { data_directory,                           String, None,                                                       false,  "Remote directory for storing data and metadata objects.", },
+            { backup_storage_url,                       String, None,                                                       true,   "Remote storage url for storing snapshots.", },
+            { backup_storage_directory,                 String, None,                                                       true,   "Remote directory for storing snapshots.", },
+            { max_concurrent_creating_streaming_jobs,   u32,    Some(1_u32),                                                true,   "Max number of concurrent creating streaming jobs.", },
+            { 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.", },
+            // TODO: modify default value
+            { oauth_jwks_url,                           String, Some("https://auth-static.confluent.io/jwks".to_string()),  true,   "Url to get JSON Web Key Set(JWKS) for oauth authentication.", },
         }
     };
 }
@@ -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,11 @@ where
             .as_ref()
             .unwrap_or(&default::WASM_STORAGE_URL)
     }
+
+    fn oauth_jwks_url(&self) -> &str {
+        self.inner()
+            .oauth_jwks_url
+            .as_ref()
+            .unwrap_or(&default::OAUTH_JWKS_URL)
+    }
 }

src/config/docs.md

@@ -151,3 +151,4 @@ This page is automatically generated by `./risedev generate-example-config`
 | sstable_size_mb | Target size of the Sstable. | 256 |
 | state_store |  |  |
 | wasm_storage_url |  | "fs://.risingwave/data" |
+| oauth_jwks_url | Url to get JSON Web Key Set(JWKS) for oauth authentication. | https://auth-static.confluent.io/jwks |

src/config/example.toml

@@ -197,3 +197,4 @@ max_concurrent_creating_streaming_jobs = 1
 pause_on_next_bootstrap = false
 wasm_storage_url = "fs://.risingwave/data"
 enable_tracing = false
+oauth_jwks_url = "https://auth-static.confluent.io/jwks"

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::SystemParamsReader;
 use risingwave_common::telemetry::manager::TelemetryManager;
 use risingwave_common::telemetry::telemetry_env_enabled;
 use risingwave_common::types::DataType;
@@ -1089,6 +1090,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/user/user_authentication.rs

@@ -24,7 +24,7 @@ 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.
+/// Build `AuthInfo` for `OAuth`.
 #[inline(always)]
 pub fn build_oauth_info() -> Option<AuthInfo> {
     Some(AuthInfo {

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

@@ -387,7 +387,7 @@ 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::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(),
@@ -523,11 +523,11 @@ 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, Arc::clone(self.session.as_ref().unwrap()))
+            .await?;
         self.stream.write_no_flush(&BeMessage::AuthenticationOk)?;
         self.stream
             .write_parameter_status_msg_no_flush(&ParameterStatus::default())?;

src/utils/pgwire/src/pg_server.rs

@@ -12,20 +12,25 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::future::Future;
 use std::io;
 use std::result::Result;
+use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Instant;
 
 use bytes::Bytes;
+use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Validation};
 use parking_lot::Mutex;
+use risingwave_common::system_param::reader::{SystemParamsRead, SystemParamsReader};
 use risingwave_common::types::DataType;
 use risingwave_sqlparser::ast::Statement;
+use serde::Deserialize;
 use thiserror_ext::AsReport;
 use tokio::io::{AsyncRead, AsyncWrite};
 
-use crate::error::PsqlResult;
+use crate::error::{PsqlError, PsqlResult};
 use crate::net::{AddressRef, Listener};
 use crate::pg_field_descriptor::PgFieldDescriptor;
 use crate::pg_message::TransactionStatus;
@@ -107,6 +112,10 @@ pub trait Session: Send + Sync {
 
     fn user_authenticator(&self) -> &UserAuthenticator;
 
+    fn get_system_params(
+        &self,
+    ) -> impl Future<Output = Result<SystemParamsReader, BoxedError>> + Send;
+
     fn id(&self) -> SessionId;
 
     fn set_config(&self, key: &str, value: String) -> Result<(), BoxedError>;
@@ -158,20 +167,69 @@ pub enum UserAuthenticator {
     OAuth,
 }
 
+#[derive(Debug, Deserialize)]
+struct Jwks {
+    keys: Vec<Jwk>,
+}
+
+#[allow(dead_code)]
+#[derive(Debug, Deserialize)]
+struct Jwk {
+    kid: String,
+    alg: String,
+    n: String,
+    e: String,
+}
+
+async fn fetch_jwks(url: &str) -> Result<Jwks, reqwest::Error> {
+    let resp: Jwks = reqwest::get(url).await?.json().await?;
+    Ok(resp)
+}
+
+async fn validate_jwt(jwt: &str, jwks_url: &str) -> Result<bool, BoxedError> {
+    let header = decode_header(jwt)?;
+    let jwks = fetch_jwks(jwks_url).await?;
+
+    let kid = header.kid.ok_or("kid not found in jwt header")?;
+    let jwk = jwks
+        .keys
+        .into_iter()
+        .find(|k| k.kid == kid)
+        .ok_or("kid not found in jwks")?;
+
+    let decoding_key = DecodingKey::from_rsa_components(&jwk.n, &jwk.e)?;
+    let validation = Validation::new(Algorithm::from_str(&jwk.alg)?);
+
+    Ok(decode::<HashMap<String, String>>(jwt, &decoding_key, &validation).is_ok())
+}
+
 impl UserAuthenticator {
-    pub fn authenticate(&self, password: &[u8]) -> bool {
-        match self {
+    pub async fn authenticate(
+        &self,
+        password: &[u8],
+        session: Arc<impl Session>,
+    ) -> PsqlResult<()> {
+        let success = match self {
             UserAuthenticator::None => true,
             UserAuthenticator::ClearText(text) => password == text,
             UserAuthenticator::Md5WithSalt {
                 encrypted_password, ..
             } => encrypted_password == password,
             UserAuthenticator::OAuth => {
-                // TODO: OAuth authentication happens here.
-                tracing::info!("OAuth authenticator gets: {}", String::from_utf8_lossy(password));
-                true
+                let system_params_reader = session
+                    .get_system_params()
+                    .await
+                    .map_err(PsqlError::StartupError)?;
+                let oauth_jwks_url = system_params_reader.oauth_jwks_url();
+                validate_jwt(&String::from_utf8_lossy(password), oauth_jwks_url)
+                    .await
+                    .map_err(PsqlError::StartupError)?
             }
+        };
+        if !success {
+            return Err(PsqlError::PasswordError);
         }
+        Ok(())
     }
 }
 
@@ -239,6 +297,7 @@ mod tests {
     use bytes::Bytes;
     use futures::stream::BoxStream;
     use futures::StreamExt;
+    use risingwave_common::system_param::reader::SystemParamsReader;
     use risingwave_common::types::DataType;
     use risingwave_sqlparser::ast::Statement;
     use tokio_postgres::NoTls;
@@ -361,6 +420,10 @@ mod tests {
             &UserAuthenticator::None
         }
 
+        async fn get_system_params(&self) -> Result<SystemParamsReader, BoxedError> {
+            Ok(SystemParamsReader::new(Default::default()))
+        }
+
         fn id(&self) -> SessionId {
             (0, 0)
         }