Skip to content

Commit

Permalink
feat: auth by refresh and session tokens.
Browse files Browse the repository at this point in the history
  • Loading branch information
youngsofun committed Aug 8, 2024
1 parent c5531ab commit b506740
Show file tree
Hide file tree
Showing 38 changed files with 1,116 additions and 86 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 5 additions & 0 deletions src/common/exception/src/exception_code.rs
Original file line number Diff line number Diff line change
Expand Up @@ -413,4 +413,9 @@ build_exceptions! {
build_exceptions! {
// A task that already stopped and can not stopped twice.
AlreadyStopped(5002),

SessionTokenExpired(5100),
RefreshTokenExpired(5101),
SessionTokenNotFound(5102),
RefreshTokenNotFound(5103)
}
2 changes: 2 additions & 0 deletions src/meta/app/src/principal/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,8 @@ pub mod tenant_user_ident;
pub mod user_defined_file_format_ident;
pub mod user_setting_ident;
pub mod user_stage_ident;
pub mod user_token;
pub mod user_token_ident;

pub use connection::*;
pub use file_format::*;
Expand Down
31 changes: 31 additions & 0 deletions src/meta/app/src/principal/user_token.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use serde::Deserialize;
use serde::Serialize;

#[derive(
serde::Serialize, serde::Deserialize, Clone, Debug, Eq, PartialEq, num_derive::FromPrimitive,
)]
pub enum TokenType {
Refresh = 1,
Session = 2,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
pub struct QueryTokenInfo {
pub token_type: TokenType,
// used to delete refresh token when close session
pub parent: Option<String>,
}
60 changes: 60 additions & 0 deletions src/meta/app/src/principal/user_token_ident.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::tenant_key::ident::TIdent;

/// Define the meta-service key for a user setting.
pub type TokenIdent = TIdent<Resource>;

pub use kvapi_impl::Resource;

mod kvapi_impl {

use databend_common_meta_kvapi::kvapi;

use crate::principal::user_token::QueryTokenInfo;
use crate::tenant_key::resource::TenantResource;

pub struct Resource;
impl TenantResource for Resource {
const PREFIX: &'static str = "__fd_token";
const TYPE: &'static str = "TokenIdent";
const HAS_TENANT: bool = true;
type ValueType = QueryTokenInfo;
}

impl kvapi::Value for QueryTokenInfo {
fn dependency_keys(&self) -> impl IntoIterator<Item = String> {
[]
}
}
}

#[cfg(test)]
mod tests {
use databend_common_meta_kvapi::kvapi::Key;

use crate::principal::user_token_ident::TokenIdent;
use crate::tenant::Tenant;

#[test]
fn test_setting_ident() {
let tenant = Tenant::new_literal("tenant1");
let ident = TokenIdent::new(tenant.clone(), "test");
assert_eq!("__fd_token/tenant1/test", ident.to_string_key());

let got = TokenIdent::from_str_key(&ident.to_string_key()).unwrap();
assert_eq!(ident, got);
}
}
1 change: 1 addition & 0 deletions src/meta/proto-conv/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,7 @@ mod stage_from_to_protobuf_impl;
mod table_from_to_protobuf_impl;
mod tenant_quota_from_to_protobuf_impl;
mod tident_from_to_protobuf_impl;
mod token_from_to_protobuf_impl;
mod udf_from_to_protobuf_impl;
mod user_from_to_protobuf_impl;
mod util;
Expand Down
57 changes: 57 additions & 0 deletions src/meta/proto-conv/src/token_from_to_protobuf_impl.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! This mod is the key point about compatibility.
//! Everytime update anything in this file, update the `VER` and let the tests pass.
use databend_common_meta_app::principal::user_token as mt;
use databend_common_protos::pb;
use num::FromPrimitive;

use crate::reader_check_msg;
use crate::FromToProto;
use crate::Incompatible;
use crate::MIN_READER_VER;
use crate::VER;

impl FromToProto for mt::QueryTokenInfo {
type PB = pb::TokenInfo;

fn get_pb_ver(p: &Self::PB) -> u64 {
p.ver
}

fn from_pb(p: Self::PB) -> Result<Self, Incompatible>
where Self: Sized {
reader_check_msg(p.ver, p.min_reader_ver)?;

let v = Self {
token_type: FromPrimitive::from_i32(p.token_type).ok_or_else(|| Incompatible {
reason: format!("invalid TokenType: {}", p.token_type),
})?,
parent: p.parent,
};
Ok(v)
}

fn to_pb(&self) -> Result<Self::PB, Incompatible> {
let p = pb::TokenInfo {
ver: VER,
min_reader_ver: MIN_READER_VER,
token_type: self.token_type.clone() as i32,
parent: self.parent.clone(),
};
Ok(p)
}
}
1 change: 1 addition & 0 deletions src/meta/proto-conv/src/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,7 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
(102, "2024-07-11: Add: UserOption add must_change_password, AuthInfo.Password add need_change"),
(103, "2024-07-31: Add: ShareMetaV2"),
(104, "2024-08-02: Add: add share catalog into Catalog meta"),
(105, "2024-08-08: Add: add QueryTokenInfo"),
// Dear developer:
// If you're gonna add a new metadata version, you'll have to add a test for it.
// You could just copy an existing test file(e.g., `../tests/it/v024_table_meta.rs`)
Expand Down
1 change: 1 addition & 0 deletions src/meta/proto-conv/tests/it/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -108,3 +108,4 @@ mod v101_database_meta;
mod v102_user_must_change_password;
mod v103_share_meta_v2;
mod v104_share_catalog;
mod v105_query_token;
44 changes: 44 additions & 0 deletions src/meta/proto-conv/tests/it/v105_query_token.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
// Copyright 2023 Datafuse Labs.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use databend_common_meta_app::principal::user_token::TokenType;
use fastrace::func_name;

use crate::common;

// These bytes are built when a new version in introduced,
// and are kept for backward compatibility test.
//
// *************************************************************
// * These messages should never be updated, *
// * only be added when a new version is added, *
// * or be removed when an old version is no longer supported. *
// *************************************************************
//
// The message bytes are built from the output of `test_pb_from_to()`
#[test]
fn test_v105_query_token_info() -> anyhow::Result<()> {
let query_token_info_v105 = vec![
8, 1, 18, 17, 112, 97, 114, 101, 110, 116, 95, 116, 111, 107, 101, 110, 95, 104, 97, 115,
104, 160, 6, 105, 168, 6, 24,
];

let want = || databend_common_meta_app::principal::user_token::QueryTokenInfo {
token_type: TokenType::Refresh,
parent: Some("parent_token_hash".to_string()),
};

common::test_pb_from_to(func_name!(), want())?;
common::test_load_old(func_name!(), query_token_info_v105.as_slice(), 105, want())
}
33 changes: 33 additions & 0 deletions src/meta/protos/proto/token.proto
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// The identifier of a database by name. Names can be changed.
// There is no guarantee that two get-database request by name will return the
// same instance.

syntax = "proto3";

package databend_proto;

message TokenInfo {
uint64 ver = 100;
uint64 min_reader_ver = 101;

enum TokenType {
Session = 0;
Refresh = 2;
}
TokenType token_type = 1;
optional string parent = 2;
}
3 changes: 3 additions & 0 deletions src/query/management/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ pub mod udf;
mod user;

pub mod errors;
mod token;

pub use cluster::ClusterApi;
pub use cluster::ClusterMgr;
Expand All @@ -46,5 +47,7 @@ pub use setting::SettingApi;
pub use setting::SettingMgr;
pub use stage::StageApi;
pub use stage::StageMgr;
pub use token::TokenApi;
pub use token::TokenMgr;
pub use user::UserApi;
pub use user::UserMgr;
19 changes: 19 additions & 0 deletions src/query/management/src/token/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod token_api;
mod token_mgr;

pub use token_api::TokenApi;
pub use token_mgr::TokenMgr;
31 changes: 31 additions & 0 deletions src/query/management/src/token/token_api.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
// Copyright 2021 Datafuse Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use databend_common_exception::Result;
use databend_common_meta_app::principal::user_token::QueryTokenInfo;

#[async_trait::async_trait]
pub trait TokenApi: Sync + Send {
async fn upsert_token(
&self,
token_hash: &str,
token_info: QueryTokenInfo,
ttl_in_secs: u64,
is_update: bool,
) -> Result<bool>;

async fn get_token(&self, token_hash: &str) -> Result<Option<QueryTokenInfo>>;

async fn drop_token(&self, token_hash: &str) -> Result<()>;
}
Loading

0 comments on commit b506740

Please sign in to comment.