1use std::num::NonZeroUsize;
16use std::sync::{LazyLock, RwLock};
17
18use jsonwebtoken::{Algorithm, DecodingKey, Validation};
19use serde::{Deserialize, Serialize};
20use thiserror::Error;
21use thiserror_ext::AsReport;
22
23use crate::{Feature, LicenseKeyRef};
24
25#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
27#[serde(untagged)]
28pub enum MaybeFeature {
29 Feature(Feature),
31 Unknown(String),
34}
35
36#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
40#[serde(rename_all = "snake_case")]
41pub enum Tier {
42 Free,
44
45 #[serde(rename = "paid")]
47 AllAsOf2_5,
48
49 All,
51
52 #[serde(untagged)]
54 Custom {
55 name: String,
56 features: Vec<MaybeFeature>,
57 },
58}
59
60impl Tier {
61 #[auto_enums::auto_enum(Iterator)]
63 pub fn available_features(&self) -> impl Iterator<Item = Feature> {
64 match self {
65 Tier::Free => std::iter::empty(),
66 Tier::AllAsOf2_5 => Feature::all_as_of_2_5().iter().copied(),
67 Tier::All => Feature::all().iter().copied(),
68 Tier::Custom { features, .. } => features.iter().filter_map(|feature| match feature {
69 MaybeFeature::Feature(feature) => Some(*feature),
70 MaybeFeature::Unknown(_) => None,
71 }),
72 }
73 }
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
81pub enum Issuer {
82 #[serde(rename = "prod.risingwave.com")]
83 Prod,
84
85 #[serde(rename = "test.risingwave.com")]
86 Test,
87
88 #[serde(untagged)]
89 Unknown(String),
90}
91
92#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
100#[serde(rename_all = "snake_case")]
101pub struct License {
102 #[allow(dead_code)]
106 pub sub: String,
107
108 #[allow(dead_code)]
112 pub iss: Issuer,
113
114 pub tier: Tier,
116
117 pub cpu_core_limit: Option<NonZeroUsize>,
119
120 pub exp: u64,
124}
125
126impl Default for License {
127 fn default() -> Self {
131 Self {
132 sub: "default".to_owned(),
133 tier: Tier::Free,
134 iss: Issuer::Prod,
135 cpu_core_limit: None,
136 exp: u64::MAX,
137 }
138 }
139}
140
141#[derive(Debug, Clone, Error)]
143pub enum LicenseError {
144 #[error("invalid license key")]
145 InvalidKey(#[source] jsonwebtoken::errors::Error),
146
147 #[error(
148 "a valid license key is set, but it is currently not effective because the CPU core in the cluster \
149 ({actual}) exceeds the maximum allowed by the license key ({limit}); \
150 consider removing some nodes or acquiring a new license key with a higher limit"
151 )]
152 CpuCoreLimitExceeded { limit: NonZeroUsize, actual: usize },
153}
154
155struct Inner {
156 license: Result<License, LicenseError>,
157 cached_cpu_core_count: usize,
158}
159
160pub struct LicenseManager {
162 inner: RwLock<Inner>,
163}
164
165static PUBLIC_KEY: LazyLock<DecodingKey> = LazyLock::new(|| {
166 DecodingKey::from_rsa_pem(include_bytes!("key.pub"))
167 .expect("invalid public key for license validation")
168});
169
170impl LicenseManager {
171 pub(crate) fn new() -> Self {
173 Self {
174 inner: RwLock::new(Inner {
175 license: Ok(License::default()),
176 cached_cpu_core_count: 0,
177 }),
178 }
179 }
180
181 pub fn get() -> &'static Self {
183 static INSTANCE: LazyLock<LicenseManager> = LazyLock::new(LicenseManager::new);
184 &INSTANCE
185 }
186
187 pub fn refresh(&self, license_key: LicenseKeyRef<'_>) {
189 let license_key = license_key.0;
190 let mut inner = self.inner.write().unwrap();
191
192 if license_key.is_empty() {
194 inner.license = Ok(License::default());
195 return;
196 }
197
198 let mut validation = Validation::new(Algorithm::RS512);
200 validation.set_issuer(&[
203 "prod.risingwave.com",
204 #[cfg(debug_assertions)]
205 "test.risingwave.com",
206 ]);
207
208 inner.license = match jsonwebtoken::decode(license_key, &PUBLIC_KEY, &validation) {
209 Ok(data) => Ok(data.claims),
210 Err(error) => Err(LicenseError::InvalidKey(error)),
211 };
212
213 match &inner.license {
214 Ok(license) => tracing::info!(?license, "license refreshed"),
215 Err(error) => tracing::warn!(error = %error.as_report(), "invalid license key"),
216 }
217 }
218
219 pub fn update_cpu_core_count(&self, cpu_core_count: usize) {
221 let mut inner = self.inner.write().unwrap();
222 inner.cached_cpu_core_count = cpu_core_count;
223 }
224
225 pub fn license(&self) -> Result<License, LicenseError> {
232 let inner = self.inner.read().unwrap();
233 let license = inner.license.clone()?;
234
235 if license.exp < jsonwebtoken::get_current_timestamp() {
237 return Err(LicenseError::InvalidKey(
238 jsonwebtoken::errors::ErrorKind::ExpiredSignature.into(),
239 ));
240 }
241
242 let actual_cpu_core = inner.cached_cpu_core_count;
244 if let Some(limit) = license.cpu_core_limit
245 && actual_cpu_core > limit.get()
246 {
247 return Err(LicenseError::CpuCoreLimitExceeded {
248 limit,
249 actual: actual_cpu_core,
250 });
251 }
252
253 Ok(license)
254 }
255}
256
257#[cfg(debug_assertions)]
259#[cfg(test)]
260mod tests {
261 use expect_test::expect;
262
263 use super::*;
264 use crate::{LicenseKey, PROD_ALL_4_CORE_LICENSE_KEY_CONTENT, TEST_ALL_LICENSE_KEY_CONTENT};
265
266 fn do_test(key: &str, expect: expect_test::Expect) -> LicenseManager {
267 let manager = LicenseManager::new();
268 manager.refresh(LicenseKey(key));
269
270 match manager.license() {
271 Ok(license) => expect.assert_debug_eq(&license),
272 Err(error) => expect.assert_eq(&error.to_report_string()),
273 }
274
275 manager
276 }
277
278 #[test]
279 fn test_all_license_key() {
280 do_test(
281 TEST_ALL_LICENSE_KEY_CONTENT,
282 expect![[r#"
283 License {
284 sub: "rw-test-all",
285 iss: Test,
286 tier: All,
287 cpu_core_limit: None,
288 exp: 10000627200,
289 }
290 "#]],
291 );
292 }
293
294 #[test]
295 fn test_prod_all_4_core_license_key() {
296 do_test(
297 PROD_ALL_4_CORE_LICENSE_KEY_CONTENT,
298 expect![[r#"
299 License {
300 sub: "rw-default-all-4-core",
301 iss: Prod,
302 tier: All,
303 cpu_core_limit: Some(
304 4,
305 ),
306 exp: 10000627200,
307 }
308 "#]],
309 );
310 }
311
312 #[test]
313 fn test_free_license_key() {
314 const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
315 eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjo5OTk5OTk5OTk5fQ.\
316 ALC3Kc9LI6u0S-jeMB1YTxg1k8Azxwvc750ihuSZgjA_e1OJC9moxMvpLrHdLZDzCXHjBYi0XJ_1lowmuO_0iPEuPqN5AFpDV1ywmzJvGmMCMtw3A2wuN7hhem9OsWbwe6lzdwrefZLipyo4GZtIkg5ZdwGuHzm33zsM-X5gl_Ns4P6axHKiorNSR6nTAyA6B32YVET_FAM2YJQrXqpwA61wn1XLfarZqpdIQyJ5cgyiC33BFBlUL3lcRXLMLeYe6TjYGeV4K63qARCjM9yeOlsRbbW5ViWeGtR2Yf18pN8ysPXdbaXm_P_IVhl3jCTDJt9ctPh6pUCbkt36FZqO9A";
317
318 do_test(
319 KEY,
320 expect![[r#"
321 License {
322 sub: "rw-test",
323 iss: Test,
324 tier: Free,
325 cpu_core_limit: None,
326 exp: 9999999999,
327 }
328 "#]],
329 );
330 }
331
332 #[test]
333 fn test_empty_license_key() {
334 do_test(
336 "",
337 expect![[r#"
338 License {
339 sub: "default",
340 iss: Prod,
341 tier: Free,
342 cpu_core_limit: None,
343 exp: 18446744073709551615,
344 }
345 "#]],
346 );
347 }
348
349 #[test]
350 fn test_custom_tier_license_key() {
351 const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
352 eyJzdWIiOiJydy11bml0LXRlc3QiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjoxMDAwMDYyNzIwMCwiaWF0IjoxNzUxODY5OTI3LCJ0aWVyIjp7Im5hbWUiOiJzZWNyZXQtb25seSIsImZlYXR1cmVzIjpbIlNlY3JldE1hbmFnZW1lbnQiXX19.\
353 iZKNyAHxz1l24Qh4F9wauuQEtDPLAeOmW2m_ttlew2ENmP_XcGWCx_O8r50NXRDaKv66z-ibteWVL5XOUqaVgJw9EnCyfVuFNoKtFQu18Vt0l52Yw2zNh3iHNQFKwHuCUki5FlOHw-K57a5f414-gzfwAwQkO1bAYoyAFhtoX6QQ2jdbxctFg0NxTQqpjnP-h0k2myZ_IhxA8fKrQIAqBj5Y8tGljuxjTLJpoK7X0ESXNh7bhA8njEf2Hm4QCymI1uo8OYRaR1Siw87r0aZykw9wW15Q8VK58VxIQLS7b7gQOmDToBjJt9yIF3MT6YMfqMX_l3Dtn9bOS_htVd1bjQ";
354
355 let manager = do_test(
356 KEY,
357 expect![[r#"
358 License {
359 sub: "rw-unit-test",
360 iss: Test,
361 tier: Custom {
362 name: "secret-only",
363 features: [
364 Feature(
365 SecretManagement,
366 ),
367 ],
368 },
369 cpu_core_limit: None,
370 exp: 10000627200,
371 }
372 "#]],
373 );
374
375 let tier = manager.license().unwrap().tier;
376 assert_eq!(
377 tier.available_features().collect::<Vec<_>>(),
378 vec![Feature::SecretManagement]
379 );
380 assert!(
381 Feature::SecretManagement
382 .check_available_with(&manager)
383 .is_ok()
384 );
385 assert!(
386 Feature::BigQuerySink
387 .check_available_with(&manager)
388 .is_err()
389 );
390 }
391
392 #[test]
393 fn test_invalid_license_key() {
394 const KEY: &str = "invalid";
395
396 do_test(KEY, expect!["invalid license key: InvalidToken"]);
397 }
398
399 #[test]
400 fn test_expired_license_key() {
401 const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
403 eyJzdWIiOiJydy10ZXN0IiwidGllciI6InBhaWQiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjowfQ.\
404 TyYmoT5Gw9-FN7DWDbeg3myW8g_3Xlc90i4M9bGuPf2WLv9zRMJy2r9J7sl1BO7t6F1uGgyrvNxsVRVZ2XF_WAs6uNlluYBnd4Cqvsj6Xny1XJCCo8II3RIea-ZlRjp6tc1saaoe-_eTtqDH8NIIWe73vVtBeBTBU4zAiN2vCtU_Si2XuoTLBKJMIjtn0HjLNhb6-DX2P3SCzp75tMyWzr49qcsBgratyKdu_v2kqBM1qw_dTaRg2ZeNNO6scSOBwu4YHHJTL4nUaZO2yEodI_OKUztIPLYuO2A33Fb5OE57S7LTgSzmxZLf7e23Vrck7Os14AfBQr7p9ncUeyIXhA";
405
406 do_test(KEY, expect!["invalid license key: ExpiredSignature"]);
407 }
408
409 #[test]
410 fn test_invalid_issuer() {
411 const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
413 eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJiYWQucmlzaW5nd2F2ZS5jb20iLCJleHAiOjk5OTk5OTk5OTl9.\
414 SUbDJTri902FbGgIoe5L3LG4edTXoR42BQCIu_KLyW41OK47bMnD2aK7JggyJmWyGtN7b_596hxM9HjU58oQtHePUo_zHi5li5IcRaMi8gqHae7CJGqOGAUo9vYOWCP5OjEuDfozJhpgcHBLzDRnSwYnWhLKtsrzb3UcpOXEqRVK7EDShBNx6kNqfYs2LlFI7ASsgFRLhoRuOTR5LeVDjj6NZfkZGsdMe1VyrODWoGT9kcAF--hBpUd1ZJ5mZ67A0_948VPFBYDbDPcTRnw1-5MvdibO-jKX49rJ0rlPXcAbqKPE_yYUaqUaORUzb3PaPgCT_quO9PWPuAFIgAb_fg";
415
416 do_test(KEY, expect!["invalid license key: InvalidIssuer"]);
417 }
418
419 #[test]
420 fn test_invalid_signature() {
421 const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
422 eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjo5OTk5OTk5OTk5fQ.\
423 InvalidSignatureoe5L3LG4edTXoR42BQCIu_KLyW41OK47bMnD2aK7JggyJmWyGtN7b_596hxM9HjU58oQtHePUo_zHi5li5IcRaMi8gqHae7CJGqOGAUo9vYOWCP5OjEuDfozJhpgcHBLzDRnSwYnWhLKtsrzb3UcpOXEqRVK7EDShBNx6kNqfYs2LlFI7ASsgFRLhoRuOTR5LeVDjj6NZfkZGsdMe1VyrODWoGT9kcAF--hBpUd1ZJ5mZ67A0_948VPFBYDbDPcTRnw1-5MvdibO-jKX49rJ0rlPXcAbqKPE_yYUaqUaORUzb3PaPgCT_quO9PWPuAFIgAb_fg";
424
425 do_test(KEY, expect!["invalid license key: InvalidSignature"]);
426 }
427}