risingwave_license/
manager.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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/// A feature that's specified in the custom tier.
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
27#[serde(untagged)]
28pub enum MaybeFeature {
29    /// A known feature that exists in the [`Feature`] enum.
30    Feature(Feature),
31    /// An unknown feature. It could be features introduced in future release. We still allow it to
32    /// be here for compatibility purposes.
33    Unknown(String),
34}
35
36/// License tier.
37///
38/// Each enterprise [`Feature`] is available for a specific tier and above.
39#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
40#[serde(rename_all = "snake_case")]
41pub enum Tier {
42    /// Free tier. No feature is available. This is more like a placeholder.
43    Free,
44
45    /// All features available as of 2.5.
46    #[serde(rename = "paid")]
47    AllAsOf2_5,
48
49    /// All features available currently and in the future.
50    All,
51
52    /// Custom tier, with a list of available features.
53    #[serde(untagged)]
54    Custom {
55        name: String,
56        features: Vec<MaybeFeature>,
57    },
58}
59
60impl Tier {
61    /// Get all available features based on the license tier.
62    #[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/// Issuer of the license.
77///
78/// The issuer must be `prod.risingwave.com` in production, and can be `test.risingwave.com` in
79/// development. This will be validated when refreshing the license key.
80#[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/// The content of a license.
93///
94/// We use JSON Web Token (JWT) to represent the license. This struct is the payload.
95///
96/// Prefer calling [`crate::Feature::check_available`] to check the availability of a feature,
97/// other than directly checking the content of the license.
98// TODO(license): Shall we add a version field?
99#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
100#[serde(rename_all = "snake_case")]
101pub struct License {
102    /// Subject of the license.
103    ///
104    /// See <https://tools.ietf.org/html/rfc7519#section-4.1.2>.
105    #[allow(dead_code)]
106    pub sub: String,
107
108    /// Issuer of the license.
109    ///
110    /// See <https://tools.ietf.org/html/rfc7519#section-4.1.1>.
111    #[allow(dead_code)]
112    pub iss: Issuer,
113
114    /// Tier of the license.
115    pub tier: Tier,
116
117    /// Maximum number of compute-node CPU cores allowed to use.
118    pub cpu_core_limit: Option<NonZeroUsize>,
119
120    /// Expiration time in seconds since UNIX epoch.
121    ///
122    /// See <https://tools.ietf.org/html/rfc7519#section-4.1.4>.
123    pub exp: u64,
124}
125
126impl Default for License {
127    /// The default license is a free license that never expires.
128    ///
129    /// Used when `license_key` is unset or invalid.
130    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/// The error type for invalid license key when verifying as JWT.
142#[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
160/// The singleton license manager.
161pub 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    /// Create a new license manager with the default license.
172    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    /// Get the singleton instance of the license manager.
182    pub fn get() -> &'static Self {
183        static INSTANCE: LazyLock<LicenseManager> = LazyLock::new(LicenseManager::new);
184        &INSTANCE
185    }
186
187    /// Refresh the license with the given license key.
188    pub fn refresh(&self, license_key: LicenseKeyRef<'_>) {
189        let license_key = license_key.0;
190        let mut inner = self.inner.write().unwrap();
191
192        // Empty license key means unset. Use the default one here.
193        if license_key.is_empty() {
194            inner.license = Ok(License::default());
195            return;
196        }
197
198        // TODO(license): shall we also validate `nbf`(Not Before)?
199        let mut validation = Validation::new(Algorithm::RS512);
200        // Only accept `prod` issuer in production, so that we can use license keys issued by
201        // the `test` issuer in development without leaking them to production.
202        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    /// Update the cached CPU core count.
220    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    /// Get the current license if it is valid.
226    ///
227    /// Since the license can expire, the returned license should not be cached by the caller.
228    ///
229    /// Prefer calling [`crate::Feature::check_available`] to check the availability of a feature,
230    /// other than directly calling this method and checking the content of the license.
231    pub fn license(&self) -> Result<License, LicenseError> {
232        let inner = self.inner.read().unwrap();
233        let license = inner.license.clone()?;
234
235        // Check the expiration time additionally.
236        if license.exp < jsonwebtoken::get_current_timestamp() {
237            return Err(LicenseError::InvalidKey(
238                jsonwebtoken::errors::ErrorKind::ExpiredSignature.into(),
239            ));
240        }
241
242        // Check the CPU core limit.
243        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// Tests below only work in debug mode.
258#[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        // Default license will be used.
335        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        // "exp": 0
402        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        // "iss": "bad.risingwave.com"
412        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}
risingwave_license/manager.rs

risingwave_license/
manager.rs
