Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(frontend): implement OAuth authentication #13151

Merged
merged 17 commits into from
Mar 5, 2024
Merged
1 change: 1 addition & 0 deletions proto/user.proto
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ message AuthInfo {
PLAINTEXT = 1;
SHA256 = 2;
MD5 = 3;
OAuth = 4;
yezizp2012 marked this conversation as resolved.
Show resolved Hide resolved
}
EncryptionType encryption_type = 1;
bytes encrypted_value = 2;
Expand Down
6 changes: 5 additions & 1 deletion src/frontend/src/handler/alter_user.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use super::RwPgResponse;
use crate::binder::Binder;
use crate::catalog::CatalogError;
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(
Expand Down Expand Up @@ -109,6 +109,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))
Expand Down
3 changes: 2 additions & 1 deletion src/frontend/src/handler/create_user.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use super::RwPgResponse;
use crate::binder::Binder;
use crate::catalog::{CatalogError, DatabaseId};
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(
Expand Down Expand Up @@ -89,6 +89,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(),
}
}

Expand Down
2 changes: 2 additions & 0 deletions src/frontend/src/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -825,6 +825,8 @@ impl SessionManager for SessionManagerImpl {
),
salt,
}
} else if auth_info.encryption_type == EncryptionType::OAuth as i32 {
UserAuthenticator::OAuth
} else {
return Err(Box::new(Error::new(
ErrorKind::Unsupported,
Expand Down
10 changes: 10 additions & 0 deletions src/frontend/src/user/user_authentication.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
}
Expand Down
7 changes: 5 additions & 2 deletions src/sqlparser/src/ast/statement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1053,6 +1053,7 @@ pub enum UserOption {
NoLogin,
EncryptedPassword(AstString),
Password(Option<AstString>),
OAuth,
}

impl fmt::Display for UserOption {
Expand All @@ -1069,6 +1070,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"),
}
}
}
Expand Down Expand Up @@ -1156,10 +1158,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!()
Expand All @@ -1169,7 +1172,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,
)?
}
Expand Down
1 change: 1 addition & 0 deletions src/sqlparser/src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -337,6 +337,7 @@ define_keywords!(
NULLIF,
NULLS,
NUMERIC,
OAUTH,
OBJECT,
OCCURRENCES_REGEX,
OCTET_LENGTH,
Expand Down
2 changes: 1 addition & 1 deletion src/sqlparser/src/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2322,7 +2322,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)?))
}
Expand Down
2 changes: 1 addition & 1 deletion src/utils/pgwire/src/pg_protocol.rs
Original file line number Diff line number Diff line change
Expand Up @@ -387,7 +387,7 @@ where
})?;
self.ready_for_query()?;
}
UserAuthenticator::ClearText(_) => {
UserAuthenticator::ClearText(_) | UserAuthenticator::OAuth => {
self.stream
.write_no_flush(&BeMessage::AuthenticationCleartextPassword)?;
}
Expand Down
6 changes: 6 additions & 0 deletions src/utils/pgwire/src/pg_server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,7 @@ pub enum UserAuthenticator {
encrypted_password: Vec<u8>,
salt: [u8; 4],
},
OAuth,
}

impl UserAuthenticator {
Expand All @@ -126,6 +127,11 @@ impl UserAuthenticator {
UserAuthenticator::Md5WithSalt {
encrypted_password, ..
} => encrypted_password == password,
UserAuthenticator::OAuth => {
// TODO: OAuth authentication happens here.
tracing::info!("OAuth authenticator gets: {}", String::from_utf8_lossy(password));
Rossil2012 marked this conversation as resolved.
Show resolved Hide resolved
true
}
}
}
}
Expand Down
Loading