From 855251cdd854f1bf258e5471d061d575b03b6f27 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Thu, 27 Jun 2024 12:02:43 +0800 Subject: [PATCH] feat: introduce license manager and feature tiers (#17396) Signed-off-by: Bugen Zhao Co-authored-by: xxchan --- .typos.toml | 1 + Cargo.lock | 24 +- Cargo.toml | 2 + e2e_test/batch/catalog/pg_settings.slt.part | 1 + e2e_test/error_ui/simple/license.slt | 59 ++++ proto/expr.proto | 1 + proto/meta.proto | 1 + src/common/Cargo.toml | 1 + src/common/src/config.rs | 11 +- src/common/src/lib.rs | 2 +- src/common/src/system_param/common.rs | 11 +- src/common/src/system_param/local_manager.rs | 1 + src/common/src/system_param/mod.rs | 16 +- src/common/src/system_param/reader.rs | 4 + src/config/docs.md | 1 + src/config/example.toml | 1 + src/expr/impl/src/scalar/mod.rs | 1 + src/expr/impl/src/scalar/test_license.rs | 27 ++ src/frontend/src/binder/expr/function.rs | 1 + src/frontend/src/expr/pure.rs | 1 + .../src/optimizer/plan_expr_visitor/strong.rs | 1 + src/license/Cargo.toml | 28 ++ src/license/src/feature.rs | 127 ++++++++ src/license/src/key.pub | 9 + src/license/src/lib.rs | 21 ++ src/license/src/manager.rs | 293 ++++++++++++++++++ 26 files changed, 637 insertions(+), 9 deletions(-) create mode 100644 e2e_test/error_ui/simple/license.slt create mode 100644 src/expr/impl/src/scalar/test_license.rs create mode 100644 src/license/Cargo.toml create mode 100644 src/license/src/feature.rs create mode 100644 src/license/src/key.pub create mode 100644 src/license/src/lib.rs create mode 100644 src/license/src/manager.rs diff --git a/.typos.toml b/.typos.toml index 567904f5c319b..a7b5570bb766d 100644 --- a/.typos.toml +++ b/.typos.toml @@ -25,6 +25,7 @@ extend-exclude = [ "src/sqlparser/tests/testdata/", "src/frontend/planner_test/tests/testdata", "src/tests/sqlsmith/tests/freeze", + "src/license/src/manager.rs", "Cargo.lock", "**/Cargo.toml", "**/go.mod", diff --git a/Cargo.lock b/Cargo.lock index 0e8e629e4d55f..80cfce6eb09dc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3425,15 +3425,16 @@ dependencies = [ [[package]] name = "curve25519-dalek" -version = "4.1.3" +version = "4.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be" +checksum = "0a677b8922c94e01bdbb12126b0bc852f00447528dee1782229af9c720c3f348" dependencies = [ "cfg-if", "cpufeatures", "curve25519-dalek-derive", "digest", "fiat-crypto", + "platforms", "rustc_version 0.4.0", "subtle", "zeroize", @@ -9251,6 +9252,12 @@ version = "0.3.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964" +[[package]] +name = "platforms" +version = "3.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4503fa043bf02cee09a9582e9554b4c6403b2ef55e4612e96561d294419429f8" + [[package]] name = "plotters" version = "0.3.5" @@ -10853,6 +10860,7 @@ dependencies = [ "risingwave_common_metrics", "risingwave_common_proc_macro", "risingwave_error", + "risingwave_license", "risingwave_pb", "rust_decimal", "rusty-fork", @@ -11667,6 +11675,18 @@ dependencies = [ "tracing", ] +[[package]] +name = "risingwave_license" +version = "1.11.0-alpha" +dependencies = [ + "expect-test", + "jsonwebtoken", + "serde", + "thiserror", + "thiserror-ext", + "tracing", +] + [[package]] name = "risingwave_mem_table_spill_test" version = "1.11.0-alpha" diff --git a/Cargo.toml b/Cargo.toml index da795196dfbc9..3fe6abc8787ba 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,6 +25,7 @@ members = [ "src/frontend/planner_test", "src/java_binding", "src/jni_core", + "src/license", "src/meta", "src/meta/dashboard", "src/meta/model_v2", @@ -217,6 +218,7 @@ risingwave_frontend = { path = "./src/frontend" } risingwave_hummock_sdk = { path = "./src/storage/hummock_sdk" } risingwave_hummock_test = { path = "./src/storage/hummock_test" } risingwave_hummock_trace = { path = "./src/storage/hummock_trace" } +risingwave_license = { path = "./src/license" } risingwave_mem_table_spill_test = { path = "./src/stream/spill_test" } risingwave_meta = { path = "./src/meta" } risingwave_meta_dashboard = { path = "./src/meta/dashboard" } diff --git a/e2e_test/batch/catalog/pg_settings.slt.part b/e2e_test/batch/catalog/pg_settings.slt.part index a3fd0edc65227..f405cc71c2c0d 100644 --- a/e2e_test/batch/catalog/pg_settings.slt.part +++ b/e2e_test/batch/catalog/pg_settings.slt.part @@ -13,6 +13,7 @@ postmaster backup_storage_url postmaster barrier_interval_ms postmaster checkpoint_frequency postmaster enable_tracing +postmaster license_key postmaster max_concurrent_creating_streaming_jobs postmaster pause_on_next_bootstrap user application_name diff --git a/e2e_test/error_ui/simple/license.slt b/e2e_test/error_ui/simple/license.slt new file mode 100644 index 0000000000000..22bd4c60cf689 --- /dev/null +++ b/e2e_test/error_ui/simple/license.slt @@ -0,0 +1,59 @@ +# Set the license key to a free tier key. +statement ok +ALTER SYSTEM SET license_key TO 'eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjo5OTk5OTk5OTk5fQ.ALC3Kc9LI6u0S-jeMB1YTxg1k8Azxwvc750ihuSZgjA_e1OJC9moxMvpLrHdLZDzCXHjBYi0XJ_1lowmuO_0iPEuPqN5AFpDV1ywmzJvGmMCMtw3A2wuN7hhem9OsWbwe6lzdwrefZLipyo4GZtIkg5ZdwGuHzm33zsM-X5gl_Ns4P6axHKiorNSR6nTAyA6B32YVET_FAM2YJQrXqpwA61wn1XLfarZqpdIQyJ5cgyiC33BFBlUL3lcRXLMLeYe6TjYGeV4K63qARCjM9yeOlsRbbW5ViWeGtR2Yf18pN8ysPXdbaXm_P_IVhl3jCTDJt9ctPh6pUCbkt36FZqO9A'; + +query error +SELECT rw_test_paid_tier(); +---- +db error: ERROR: Failed to run the query + +Caused by these errors (recent errors listed first): + 1: Expr error + 2: error while evaluating expression `test_paid_tier()` + 3: feature TestPaid is only available for tier Paid and above, while the current tier is Free + +Hint: You may want to set a license key with `ALTER SYSTEM SET license_key = '...';` command. + + +# Set the license key to an invalid key. +statement ok +ALTER SYSTEM SET license_key TO 'invalid'; + +query error +SELECT rw_test_paid_tier(); +---- +db error: ERROR: Failed to run the query + +Caused by these errors (recent errors listed first): + 1: Expr error + 2: error while evaluating expression `test_paid_tier()` + 3: feature TestPaid is not available due to license error + 4: invalid license key + 5: InvalidToken + + +# Set the license key to empty. This demonstrates the default behavior in production, i.e., free tier. +statement ok +ALTER SYSTEM SET license_key TO ''; + +query error +SELECT rw_test_paid_tier(); +---- +db error: ERROR: Failed to run the query + +Caused by these errors (recent errors listed first): + 1: Expr error + 2: error while evaluating expression `test_paid_tier()` + 3: feature TestPaid is only available for tier Paid and above, while the current tier is Free + +Hint: You may want to set a license key with `ALTER SYSTEM SET license_key = '...';` command. + + +# Set the license key to default. In debug mode, this will set the license key to a paid tier key. +statement ok +ALTER SYSTEM SET license_key TO DEFAULT; + +query T +SELECT rw_test_paid_tier(); +---- +t diff --git a/proto/expr.proto b/proto/expr.proto index ada2159bb80ae..3bb54fae6f0bb 100644 --- a/proto/expr.proto +++ b/proto/expr.proto @@ -286,6 +286,7 @@ message ExprNode { // ------------------------ // Internal functions VNODE = 1101; + TEST_PAID_TIER = 1102; // Non-deterministic functions PROCTIME = 2023; PG_SLEEP = 2024; diff --git a/proto/meta.proto b/proto/meta.proto index 4a67dd0455e5c..e4068a0b8cd58 100644 --- a/proto/meta.proto +++ b/proto/meta.proto @@ -603,6 +603,7 @@ message SystemParams { optional string wasm_storage_url = 14 [deprecated = true]; optional bool enable_tracing = 15; optional bool use_new_object_prefix_strategy = 16; + optional string license_key = 17; } message GetSystemParamsRequest {} diff --git a/src/common/Cargo.toml b/src/common/Cargo.toml index 9fd9415cb7ccd..ae6f67faf3aac 100644 --- a/src/common/Cargo.toml +++ b/src/common/Cargo.toml @@ -87,6 +87,7 @@ risingwave_common_estimate_size = { workspace = true } risingwave_common_metrics = { path = "./metrics" } risingwave_common_proc_macro = { workspace = true } risingwave_error = { workspace = true } +risingwave_license = { workspace = true } risingwave_pb = { workspace = true } rust_decimal = { version = "1", features = ["db-postgres", "maths"] } rw_iter_util = { workspace = true } diff --git a/src/common/src/config.rs b/src/common/src/config.rs index a554e220ec632..2579395b59ee2 100644 --- a/src/common/src/config.rs +++ b/src/common/src/config.rs @@ -2278,6 +2278,13 @@ pub struct CompactionConfig { mod tests { use super::*; + fn default_config_for_docs() -> RwConfig { + let mut config = RwConfig::default(); + // Set `license_key` to empty to avoid showing the test-only license key in the docs. + config.system.license_key = Some("".to_owned()); + config + } + /// This test ensures that `config/example.toml` is up-to-date with the default values specified /// in this file. Developer should run `./risedev generate-example-config` to update it if this /// test fails. @@ -2287,7 +2294,7 @@ mod tests { # Check detailed comments in src/common/src/config.rs"; let actual = expect_test::expect_file!["../../config/example.toml"]; - let default = toml::to_string(&RwConfig::default()).expect("failed to serialize"); + let default = toml::to_string(&default_config_for_docs()).expect("failed to serialize"); let expected = format!("{HEADER}\n\n{default}"); actual.assert_eq(&expected); @@ -2328,7 +2335,7 @@ mod tests { .collect(); let toml_doc: BTreeMap = - toml::from_str(&toml::to_string(&RwConfig::default()).unwrap()).unwrap(); + toml::from_str(&toml::to_string(&default_config_for_docs()).unwrap()).unwrap(); toml_doc.into_iter().for_each(|(name, value)| { set_default_values("".to_string(), name, value, &mut configs); }); diff --git a/src/common/src/lib.rs b/src/common/src/lib.rs index 1cbb2d837aa78..05ca777ea7d39 100644 --- a/src/common/src/lib.rs +++ b/src/common/src/lib.rs @@ -75,12 +75,12 @@ pub mod field_generator; pub mod hash; pub mod log; pub mod memory; -pub use risingwave_common_metrics as metrics; pub use risingwave_common_metrics::{ monitor, register_guarded_gauge_vec_with_registry, register_guarded_histogram_vec_with_registry, register_guarded_int_counter_vec_with_registry, register_guarded_int_gauge_vec_with_registry, }; +pub use {risingwave_common_metrics as metrics, risingwave_license as license}; pub mod lru; pub mod opts; pub mod range; diff --git a/src/common/src/system_param/common.rs b/src/common/src/system_param/common.rs index d8ff741399533..7aec7d14f012b 100644 --- a/src/common/src/system_param/common.rs +++ b/src/common/src/system_param/common.rs @@ -12,13 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. +use risingwave_license::LicenseManager; + use super::diff::SystemParamsDiff; use super::reader::SystemParamsReader; use crate::util::tracing::layer::toggle_otel_layer; /// Node-independent handler for system parameter changes. -/// -/// Currently, it is only used to enable or disable the distributed tracing layer. #[derive(Debug)] pub struct CommonHandler; @@ -32,8 +32,13 @@ impl CommonHandler { /// Handle the change of system parameters. pub fn handle_change(&self, diff: &SystemParamsDiff) { + // Toggle the distributed tracing layer. if let Some(enabled) = diff.enable_tracing { - toggle_otel_layer(enabled) + toggle_otel_layer(enabled); + } + // Refresh the license key. + if let Some(key) = diff.license_key.as_ref() { + LicenseManager::get().refresh(key); } } } diff --git a/src/common/src/system_param/local_manager.rs b/src/common/src/system_param/local_manager.rs index 5040d30f811d0..8a81959219704 100644 --- a/src/common/src/system_param/local_manager.rs +++ b/src/common/src/system_param/local_manager.rs @@ -48,6 +48,7 @@ impl LocalSystemParamsManager { let this = Self::new_inner(initial_params.clone()); // Spawn a task to run the common handler. + // TODO(bugen): this may be spawned multiple times under standalone deployment, though idempotent. tokio::spawn({ let mut rx = this.tx.subscribe(); async move { diff --git a/src/common/src/system_param/mod.rs b/src/common/src/system_param/mod.rs index c8382d35fee81..aa6f207ead1a9 100644 --- a/src/common/src/system_param/mod.rs +++ b/src/common/src/system_param/mod.rs @@ -30,6 +30,7 @@ use std::ops::RangeBounds; use std::str::FromStr; use paste::paste; +use risingwave_license::TEST_PAID_LICENSE_KEY; use risingwave_pb::meta::PbSystemParams; use self::diff::SystemParamsDiff; @@ -60,6 +61,15 @@ impl_param_value!(u64); impl_param_value!(f64); impl_param_value!(String => &'a str); +/// Set the default value of `license_key` to [`TEST_PAID_LICENSE_KEY`] in debug mode. +fn default_license_key() -> String { + if cfg!(debug_assertions) { + TEST_PAID_LICENSE_KEY.to_owned() + } else { + "".to_owned() + } +} + /// Define all system parameters here. /// /// To match all these information, write the match arm as follows: @@ -87,7 +97,8 @@ macro_rules! for_all_params { { 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.", }, { enable_tracing, bool, Some(false), true, "Whether to enable distributed tracing.", }, - { use_new_object_prefix_strategy, bool, None, false, "Whether to split object prefix.", }, + { use_new_object_prefix_strategy, bool, None, false, "Whether to split object prefix.", }, + { license_key, String, Some(default_license_key()), true, "The license key to activate enterprise features.", }, } }; } @@ -148,6 +159,8 @@ macro_rules! def_default { pub mod default { use std::sync::LazyLock; + use super::*; + for_all_params!(def_default_opt); for_all_params!(def_default); } @@ -444,6 +457,7 @@ mod tests { (PAUSE_ON_NEXT_BOOTSTRAP_KEY, "false"), (ENABLE_TRACING_KEY, "true"), (USE_NEW_OBJECT_PREFIX_STRATEGY_KEY, "false"), + (LICENSE_KEY_KEY, "foo"), ("a_deprecated_param", "foo"), ]; diff --git a/src/common/src/system_param/reader.rs b/src/common/src/system_param/reader.rs index 9a2c6e49534af..2ef13ff3f5509 100644 --- a/src/common/src/system_param/reader.rs +++ b/src/common/src/system_param/reader.rs @@ -168,4 +168,8 @@ where .enable_tracing .unwrap_or_else(default::enable_tracing) } + + fn license_key(&self) -> &str { + self.inner().license_key.as_deref().unwrap_or_default() + } } diff --git a/src/config/docs.md b/src/config/docs.md index a52ce9202a3b3..f93700f5c5e65 100644 --- a/src/config/docs.md +++ b/src/config/docs.md @@ -163,6 +163,7 @@ This page is automatically generated by `./risedev generate-example-config` | checkpoint_frequency | There will be a checkpoint for every n barriers. | 1 | | data_directory | Remote directory for storing data and metadata objects. | | | enable_tracing | Whether to enable distributed tracing. | false | +| license_key | The license key to activate enterprise features. | "" | | max_concurrent_creating_streaming_jobs | Max number of concurrent creating streaming jobs. | 1 | | parallel_compact_size_mb | The size of parallel task for one compact/flush job. | 512 | | pause_on_next_bootstrap | Whether to pause all data sources on next bootstrap. | false | diff --git a/src/config/example.toml b/src/config/example.toml index a708fed3b84bf..afc41b5647da3 100644 --- a/src/config/example.toml +++ b/src/config/example.toml @@ -238,3 +238,4 @@ bloom_false_positive = 0.001 max_concurrent_creating_streaming_jobs = 1 pause_on_next_bootstrap = false enable_tracing = false +license_key = "" diff --git a/src/expr/impl/src/scalar/mod.rs b/src/expr/impl/src/scalar/mod.rs index c4e7990de133b..fbf9b512ea86d 100644 --- a/src/expr/impl/src/scalar/mod.rs +++ b/src/expr/impl/src/scalar/mod.rs @@ -84,6 +84,7 @@ pub use to_jsonb::*; mod encrypt; mod external; mod inet; +mod test_license; mod to_timestamp; mod translate; mod trigonometric; diff --git a/src/expr/impl/src/scalar/test_license.rs b/src/expr/impl/src/scalar/test_license.rs new file mode 100644 index 0000000000000..de1afd5ed616b --- /dev/null +++ b/src/expr/impl/src/scalar/test_license.rs @@ -0,0 +1,27 @@ +// Copyright 2024 RisingWave 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 risingwave_common::license::Feature; +use risingwave_expr::{function, ExprError, Result}; + +/// A function that checks if the `TestPaid` feature is available. +/// +/// It's mainly for testing purposes only. +#[function("test_paid_tier() -> boolean")] +pub fn test_paid_tier() -> Result { + Feature::TestPaid + .check_available() + .map_err(|e| ExprError::Internal(anyhow::Error::from(e)))?; + Ok(true) +} diff --git a/src/frontend/src/binder/expr/function.rs b/src/frontend/src/binder/expr/function.rs index b9a3b27825abf..95466021863e4 100644 --- a/src/frontend/src/binder/expr/function.rs +++ b/src/frontend/src/binder/expr/function.rs @@ -1426,6 +1426,7 @@ impl Binder { ("pg_is_in_recovery", raw_literal(ExprImpl::literal_bool(false))), // internal ("rw_vnode", raw_call(ExprType::Vnode)), + ("rw_test_paid_tier", raw_call(ExprType::TestPaidTier)), // for testing purposes // TODO: choose which pg version we should return. ("version", raw_literal(ExprImpl::literal_varchar(current_cluster_version()))), // non-deterministic diff --git a/src/frontend/src/expr/pure.rs b/src/frontend/src/expr/pure.rs index b404fb3408df0..fe87eb6c2280c 100644 --- a/src/frontend/src/expr/pure.rs +++ b/src/frontend/src/expr/pure.rs @@ -258,6 +258,7 @@ impl ExprVisitor for ImpureAnalyzer { } // expression output is not deterministic Type::Vnode + | Type::TestPaidTier | Type::Proctime | Type::PgSleep | Type::PgSleepFor diff --git a/src/frontend/src/optimizer/plan_expr_visitor/strong.rs b/src/frontend/src/optimizer/plan_expr_visitor/strong.rs index 55d8e3a18adf4..0cf1cd0a07a35 100644 --- a/src/frontend/src/optimizer/plan_expr_visitor/strong.rs +++ b/src/frontend/src/optimizer/plan_expr_visitor/strong.rs @@ -291,6 +291,7 @@ impl Strong { | ExprType::JsonbToRecord | ExprType::JsonbSet | ExprType::Vnode + | ExprType::TestPaidTier | ExprType::Proctime | ExprType::PgSleep | ExprType::PgSleepFor diff --git a/src/license/Cargo.toml b/src/license/Cargo.toml new file mode 100644 index 0000000000000..47e00228626b8 --- /dev/null +++ b/src/license/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "risingwave_license" +description = "License validation and feature gating for RisingWave" +version = { workspace = true } +edition = { workspace = true } +homepage = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +repository = { workspace = true } + +[package.metadata.cargo-machete] +ignored = ["workspace-hack"] + +[package.metadata.cargo-udeps.ignore] +normal = ["workspace-hack"] + +[dependencies] +jsonwebtoken = "9" +serde = { version = "1", features = ["derive"] } +thiserror = "1" +thiserror-ext = { workspace = true } +tracing = "0.1" + +[dev-dependencies] +expect-test = "1" + +[lints] +workspace = true diff --git a/src/license/src/feature.rs b/src/license/src/feature.rs new file mode 100644 index 0000000000000..302538cc3ecc3 --- /dev/null +++ b/src/license/src/feature.rs @@ -0,0 +1,127 @@ +// Copyright 2024 RisingWave 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 thiserror::Error; + +use super::{License, LicenseKeyError, LicenseManager, Tier}; + +/// Define all features that are available based on the tier of the license. +/// +/// # Define a new feature +/// +/// To add a new feature, add a new entry below following the same pattern as the existing ones. +/// +/// Check the definition of [`Tier`] for all available tiers. Note that normally there's no need to +/// add a feature with the minimum tier of `Free`, as you can directly write the code without +/// gating it with a feature check. +/// +/// # Check the availability of a feature +/// +/// To check the availability of a feature during runtime, call the method +/// [`check_available`](Feature::check_available) on the feature. If the feature is not available, +/// an error of type [`FeatureNotAvailable`] will be returned and you should handle it properly, +/// generally by returning an error to the user. +/// +/// # Feature availability in tests +/// +/// In tests with `debug_assertions` enabled, a license key of the paid (maximum) tier is set by +/// default. As a result, all features are available in tests. To test the behavior when a feature +/// is not available, you can manually set a license key with a lower tier. Check the e2e test cases +/// under `error_ui` for examples. +macro_rules! for_all_features { + ($macro:ident) => { + $macro! { + // name min tier doc + { TestPaid, Paid, "A dummy feature that's only available on paid tier for testing purposes." }, + } + }; +} + +macro_rules! def_feature { + ($({ $name:ident, $min_tier:ident, $doc:literal },)*) => { + /// A set of features that are available based on the tier of the license. + /// + /// To define a new feature, add a new entry in the macro [`for_all_features`]. + #[derive(Clone, Copy, Debug)] + pub enum Feature { + $( + #[doc = concat!($doc, "\n\nAvailable for tier `", stringify!($min_tier), "` and above.")] + $name, + )* + } + + impl Feature { + /// Minimum tier required to use this feature. + fn min_tier(self) -> Tier { + match self { + $( + Self::$name => Tier::$min_tier, + )* + } + } + } + }; +} + +for_all_features!(def_feature); + +/// The error type for feature not available due to license. +#[derive(Debug, Error)] +pub enum FeatureNotAvailable { + #[error( + "feature {:?} is only available for tier {:?} and above, while the current tier is {:?}\n\n\ + Hint: You may want to set a license key with `ALTER SYSTEM SET license_key = '...';` command.", + feature, feature.min_tier(), current_tier, + )] + InsufficientTier { + feature: Feature, + current_tier: Tier, + }, + + #[error("feature {feature:?} is not available due to license error")] + LicenseError { + feature: Feature, + source: LicenseKeyError, + }, +} + +impl Feature { + /// Check whether the feature is available based on the current license. + pub fn check_available(self) -> Result<(), FeatureNotAvailable> { + match LicenseManager::get().license() { + Ok(license) => { + if license.tier >= self.min_tier() { + Ok(()) + } else { + Err(FeatureNotAvailable::InsufficientTier { + feature: self, + current_tier: license.tier, + }) + } + } + Err(error) => { + // If there's a license key error, we still try against the default license first + // to see if the feature is available for free. + if License::default().tier >= self.min_tier() { + Ok(()) + } else { + Err(FeatureNotAvailable::LicenseError { + feature: self, + source: error, + }) + } + } + } + } +} diff --git a/src/license/src/key.pub b/src/license/src/key.pub new file mode 100644 index 0000000000000..8c57095eb3e39 --- /dev/null +++ b/src/license/src/key.pub @@ -0,0 +1,9 @@ +-----BEGIN PUBLIC KEY----- +MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAv9mzpTfLkQc8S2R6rgRG +nX/uBtXYcieUWV1ppLOW0iOfOu3Ao8FzYY2RIaLrJ0Wi3I7qktBV+lMnn2ncqD73 +1Rfo/RD4/M8WGm93+MLVTNEw5Wn1FNf9M8lC1Zo3iPlssVbVYlIJZDkZZ5phF3qo +PHvs3zm0XtglF3JfAzo0ecmfh+q3Jjjc2j5Ceu+Ngcm3wJV4aVSvK6C802e6ONbq +OwwMCie960hdpVHArBpeI6FbPFMSR4f9tCB6gSIP4sQmXFUPmvw2Khc9cjFc/QpP +WzmCpyvUeK0XhV4LYPA2+vQ4Ui5I8de/igR2JUANUNOJ7vP3kLc57g8DgxTtE75/ +LQIDAQAB +-----END PUBLIC KEY----- diff --git a/src/license/src/lib.rs b/src/license/src/lib.rs new file mode 100644 index 0000000000000..0e641be9789b1 --- /dev/null +++ b/src/license/src/lib.rs @@ -0,0 +1,21 @@ +// Copyright 2024 RisingWave 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. + +#![feature(lazy_cell)] + +mod feature; +mod manager; + +pub use feature::*; +pub use manager::*; diff --git a/src/license/src/manager.rs b/src/license/src/manager.rs new file mode 100644 index 0000000000000..f815a145ba01d --- /dev/null +++ b/src/license/src/manager.rs @@ -0,0 +1,293 @@ +// Copyright 2024 RisingWave 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 std::sync::{LazyLock, RwLock}; + +use jsonwebtoken::{Algorithm, DecodingKey, Validation}; +use serde::Deserialize; +use thiserror::Error; +use thiserror_ext::AsReport; + +/// License tier. +/// +/// Each enterprise [`Feature`](super::Feature) is available for a specific tier and above. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Tier { + /// Free tier. + /// + /// This is more like a placeholder. If a feature is available for the free tier, there's no + /// need to add it to the [`Feature`](super::Feature) enum at all. + Free, + + /// Paid tier. + // TODO(license): Add more tiers if needed. + Paid, +} + +/// Issuer of the license. +/// +/// The issuer must be `prod.risingwave.com` in production, and can be `test.risingwave.com` in +/// development. This will be validated when refreshing the license key. +#[derive(Debug, Clone, Deserialize)] +pub enum Issuer { + #[serde(rename = "prod.risingwave.com")] + Prod, + + #[serde(rename = "test.risingwave.com")] + Test, + + #[serde(untagged)] + Unknown(String), +} + +/// The content of a license. +/// +/// We use JSON Web Token (JWT) to represent the license. This struct is the payload. +// TODO(license): Shall we add a version field? +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(super) struct License { + /// Subject of the license. + /// + /// See . + #[allow(dead_code)] + pub sub: String, + + /// Issuer of the license. + /// + /// See . + #[allow(dead_code)] + pub iss: Issuer, + + /// Tier of the license. + pub tier: Tier, + + /// Expiration time in seconds since UNIX epoch. + /// + /// See . + pub exp: u64, +} + +impl Default for License { + /// The default license is a free license that never expires. + /// + /// Used when `license_key` is unset or invalid. + fn default() -> Self { + Self { + sub: "default".to_owned(), + tier: Tier::Free, + iss: Issuer::Prod, + exp: u64::MAX, + } + } +} + +/// The error type for invalid license key when verifying as JWT. +#[derive(Debug, Clone, Error)] +#[error("invalid license key")] +pub struct LicenseKeyError(#[source] jsonwebtoken::errors::Error); + +struct Inner { + license: Result, +} + +/// The singleton license manager. +pub struct LicenseManager { + inner: RwLock, +} + +static PUBLIC_KEY: LazyLock = LazyLock::new(|| { + DecodingKey::from_rsa_pem(include_bytes!("key.pub")) + .expect("invalid public key for license validation") +}); + +impl LicenseManager { + /// Create a new license manager with the default license. + fn new() -> Self { + Self { + inner: RwLock::new(Inner { + license: Ok(License::default()), + }), + } + } + + /// Get the singleton instance of the license manager. + pub fn get() -> &'static Self { + static INSTANCE: LazyLock = LazyLock::new(LicenseManager::new); + &INSTANCE + } + + /// Refresh the license with the given license key. + pub fn refresh(&self, license_key: &str) { + let mut inner = self.inner.write().unwrap(); + + // Empty license key means unset. Use the default one here. + if license_key.is_empty() { + inner.license = Ok(License::default()); + return; + } + + // TODO(license): shall we also validate `nbf`(Not Before)? + let mut validation = Validation::new(Algorithm::RS512); + // Only accept `prod` issuer in production, so that we can use license keys issued by + // the `test` issuer in development without leaking them to production. + validation.set_issuer(&[ + "prod.risingwave.com", + #[cfg(debug_assertions)] + "test.risingwave.com", + ]); + + inner.license = match jsonwebtoken::decode(license_key, &PUBLIC_KEY, &validation) { + Ok(data) => Ok(data.claims), + Err(error) => Err(LicenseKeyError(error)), + }; + + match &inner.license { + Ok(license) => tracing::info!(?license, "license refreshed"), + Err(error) => tracing::warn!(error = %error.as_report(), "invalid license key"), + } + } + + /// Get the current license if it is valid. + /// + /// Since the license can expire, the returned license should not be cached by the caller. + pub(super) fn license(&self) -> Result { + let license = self.inner.read().unwrap().license.clone()?; + + // Check the expiration time additionally. + if license.exp < jsonwebtoken::get_current_timestamp() { + return Err(LicenseKeyError( + jsonwebtoken::errors::ErrorKind::ExpiredSignature.into(), + )); + } + + Ok(license) + } +} + +/// A license key with the paid tier that only works in tests. +pub const TEST_PAID_LICENSE_KEY: &str = + "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\ + eyJzdWIiOiJydy10ZXN0IiwidGllciI6InBhaWQiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjo5OTk5OTk5OTk5fQ.\ + c6Gmb6xh3dBDYX_4cOnHUbwRXJbUCM7W3mrJA77nLC5FkoOLpGstzvQ7qfnPVBu412MFtKRDvh-Lk8JwG7pVa0WLw16DeHTtVHxZukMTZ1Q_ciZ1xKeUx_pwUldkVzv6c9j99gNqPSyTjzOXTdKlidBRLer2zP0v3Lf-ZxnMG0tEcIbTinTb3BNCtAQ8bwBSRP-X48cVTWafjaZxv_zGiJT28uV3bR6jwrorjVB4VGvqhsJi6Fd074XOmUlnOleoAtyzKvjmGC5_FvnL0ztIe_I0z_pyCMfWpyJ_J4C7rCP1aVWUImyoowLmVDA-IKjclzOW5Fvi0wjXsc6OckOc_A"; + +// Tests below only work in debug mode. +#[cfg(debug_assertions)] +#[cfg(test)] +mod tests { + use expect_test::expect; + + use super::*; + + fn do_test(key: &str, expect: expect_test::Expect) { + let manager = LicenseManager::new(); + manager.refresh(key); + + match manager.license() { + Ok(license) => expect.assert_debug_eq(&license), + Err(error) => expect.assert_eq(&error.to_report_string()), + } + } + + #[test] + fn test_paid_license_key() { + do_test( + TEST_PAID_LICENSE_KEY, + expect![[r#" + License { + sub: "rw-test", + iss: Test, + tier: Paid, + exp: 9999999999, + } + "#]], + ); + } + + #[test] + fn test_free_license_key() { + const KEY: &str = + "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\ + eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjo5OTk5OTk5OTk5fQ.\ + ALC3Kc9LI6u0S-jeMB1YTxg1k8Azxwvc750ihuSZgjA_e1OJC9moxMvpLrHdLZDzCXHjBYi0XJ_1lowmuO_0iPEuPqN5AFpDV1ywmzJvGmMCMtw3A2wuN7hhem9OsWbwe6lzdwrefZLipyo4GZtIkg5ZdwGuHzm33zsM-X5gl_Ns4P6axHKiorNSR6nTAyA6B32YVET_FAM2YJQrXqpwA61wn1XLfarZqpdIQyJ5cgyiC33BFBlUL3lcRXLMLeYe6TjYGeV4K63qARCjM9yeOlsRbbW5ViWeGtR2Yf18pN8ysPXdbaXm_P_IVhl3jCTDJt9ctPh6pUCbkt36FZqO9A"; + + do_test( + KEY, + expect![[r#" + License { + sub: "rw-test", + iss: Test, + tier: Free, + exp: 9999999999, + } + "#]], + ); + } + + #[test] + fn test_empty_license_key() { + // Default license will be used. + do_test( + "", + expect![[r#" + License { + sub: "default", + iss: Prod, + tier: Free, + exp: 18446744073709551615, + } + "#]], + ); + } + + #[test] + fn test_invalid_license_key() { + const KEY: &str = "invalid"; + + do_test(KEY, expect!["invalid license key: InvalidToken"]); + } + + #[test] + fn test_expired_license_key() { + // "exp": 0 + const KEY: &str = + "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\ + eyJzdWIiOiJydy10ZXN0IiwidGllciI6InBhaWQiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjowfQ.\ + TyYmoT5Gw9-FN7DWDbeg3myW8g_3Xlc90i4M9bGuPf2WLv9zRMJy2r9J7sl1BO7t6F1uGgyrvNxsVRVZ2XF_WAs6uNlluYBnd4Cqvsj6Xny1XJCCo8II3RIea-ZlRjp6tc1saaoe-_eTtqDH8NIIWe73vVtBeBTBU4zAiN2vCtU_Si2XuoTLBKJMIjtn0HjLNhb6-DX2P3SCzp75tMyWzr49qcsBgratyKdu_v2kqBM1qw_dTaRg2ZeNNO6scSOBwu4YHHJTL4nUaZO2yEodI_OKUztIPLYuO2A33Fb5OE57S7LTgSzmxZLf7e23Vrck7Os14AfBQr7p9ncUeyIXhA"; + + do_test(KEY, expect!["invalid license key: ExpiredSignature"]); + } + + #[test] + fn test_invalid_issuer() { + // "iss": "bad.risingwave.com" + const KEY: &str = + "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\ + eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJiYWQucmlzaW5nd2F2ZS5jb20iLCJleHAiOjk5OTk5OTk5OTl9.\ + SUbDJTri902FbGgIoe5L3LG4edTXoR42BQCIu_KLyW41OK47bMnD2aK7JggyJmWyGtN7b_596hxM9HjU58oQtHePUo_zHi5li5IcRaMi8gqHae7CJGqOGAUo9vYOWCP5OjEuDfozJhpgcHBLzDRnSwYnWhLKtsrzb3UcpOXEqRVK7EDShBNx6kNqfYs2LlFI7ASsgFRLhoRuOTR5LeVDjj6NZfkZGsdMe1VyrODWoGT9kcAF--hBpUd1ZJ5mZ67A0_948VPFBYDbDPcTRnw1-5MvdibO-jKX49rJ0rlPXcAbqKPE_yYUaqUaORUzb3PaPgCT_quO9PWPuAFIgAb_fg"; + + do_test(KEY, expect!["invalid license key: InvalidIssuer"]); + } + + #[test] + fn test_invalid_signature() { + const KEY: &str = + "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\ + eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjo5OTk5OTk5OTk5fQ.\ + InvalidSignatureoe5L3LG4edTXoR42BQCIu_KLyW41OK47bMnD2aK7JggyJmWyGtN7b_596hxM9HjU58oQtHePUo_zHi5li5IcRaMi8gqHae7CJGqOGAUo9vYOWCP5OjEuDfozJhpgcHBLzDRnSwYnWhLKtsrzb3UcpOXEqRVK7EDShBNx6kNqfYs2LlFI7ASsgFRLhoRuOTR5LeVDjj6NZfkZGsdMe1VyrODWoGT9kcAF--hBpUd1ZJ5mZ67A0_948VPFBYDbDPcTRnw1-5MvdibO-jKX49rJ0rlPXcAbqKPE_yYUaqUaORUzb3PaPgCT_quO9PWPuAFIgAb_fg"; + + do_test(KEY, expect!["invalid license key: InvalidSignature"]); + } +}