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::NonZeroU64;
16use std::sync::{LazyLock, RwLock};
17
18use jsonwebtoken::{Algorithm, DecodingKey, Validation};
19use risingwave_pb::common::ClusterResource;
20use serde::{Deserialize, Serialize};
21use thiserror::Error;
22use thiserror_ext::AsReport;
23
24use crate::{Feature, LicenseKeyRef};
25
26/// A feature that's specified in the custom tier.
27#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
28#[serde(untagged)]
29pub enum MaybeFeature {
30    /// A known feature that exists in the [`Feature`] enum.
31    Feature(Feature),
32    /// An unknown feature. It could be features introduced in future release. We still allow it to
33    /// be here for compatibility purposes.
34    Unknown(String),
35}
36
37/// License tier.
38///
39/// Each enterprise [`Feature`] is available for a specific tier and above.
40#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
41#[serde(rename_all = "snake_case")]
42pub enum Tier {
43    /// Free tier. No feature is available. This is more like a placeholder.
44    Free,
45
46    /// All features available as of 2.5.
47    #[serde(rename = "paid")]
48    AllAsOf2_5,
49
50    /// All features available currently and in the future.
51    All,
52
53    /// Custom tier, with a list of available features.
54    #[serde(untagged)]
55    Custom {
56        name: String,
57        features: Vec<MaybeFeature>,
58    },
59}
60
61impl Tier {
62    /// Get all available features based on the license tier.
63    #[auto_enums::auto_enum(Iterator)]
64    pub fn available_features(&self) -> impl Iterator<Item = Feature> {
65        match self {
66            Tier::Free => std::iter::empty(),
67            Tier::AllAsOf2_5 => Feature::all_as_of_2_5().iter().copied(),
68            Tier::All => Feature::all().iter().copied(),
69            Tier::Custom { features, .. } => features.iter().filter_map(|feature| match feature {
70                MaybeFeature::Feature(feature) => Some(*feature),
71                MaybeFeature::Unknown(_) => None,
72            }),
73        }
74    }
75
76    pub fn name(&self) -> &str {
77        match self {
78            Tier::Free => "free",
79            Tier::AllAsOf2_5 => "paid",
80            Tier::All => "all",
81            Tier::Custom { name, .. } => name,
82        }
83    }
84}
85
86/// Issuer of the license.
87///
88/// The issuer must be `prod.risingwave.com` in production, and can be `test.risingwave.com` in
89/// development. This will be validated when refreshing the license key.
90#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
91pub enum Issuer {
92    #[serde(rename = "prod.risingwave.com")]
93    Prod,
94
95    #[serde(rename = "test.risingwave.com")]
96    Test,
97
98    #[serde(untagged)]
99    Unknown(String),
100}
101
102/// The content of a license.
103///
104/// We use JSON Web Token (JWT) to represent the license. This struct is the payload.
105///
106/// Prefer calling [`crate::Feature::check_available`] to check the availability of a feature,
107/// other than directly checking the content of the license.
108// TODO(license): Shall we add a version field?
109#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
110#[serde(rename_all = "snake_case")]
111pub struct License {
112    /// Subject of the license.
113    ///
114    /// See <https://tools.ietf.org/html/rfc7519#section-4.1.2>.
115    #[allow(dead_code)]
116    pub sub: String,
117
118    /// Issuer of the license.
119    ///
120    /// See <https://tools.ietf.org/html/rfc7519#section-4.1.1>.
121    #[allow(dead_code)]
122    pub iss: Issuer,
123
124    /// Tier of the license.
125    pub tier: Tier,
126
127    /// Maximum number of RWU allowed to use.
128    ///
129    /// 1 RWU corresponds to 4 GiB memory and 1 CPU core.
130    /// See <https://docs.risingwave.com/cloud/pricing#risingwave-unit-rwu> for more details.
131    #[serde(alias = "cpu_core_limit")]
132    pub rwu_limit: Option<NonZeroU64>,
133
134    /// Expiration time in seconds since UNIX epoch.
135    ///
136    /// See <https://tools.ietf.org/html/rfc7519#section-4.1.4>.
137    pub exp: u64,
138}
139
140impl License {
141    /// Return the CPU core limit based on the RWU limit in the license.
142    pub fn cpu_core_limit(&self) -> Option<u64> {
143        self.rwu_limit.map(|limit| limit.get())
144    }
145
146    /// Return the memory limit (in bytes) based on the RWU limit in the license.
147    pub fn memory_limit(&self) -> Option<u64> {
148        // 4GB per RWU
149        const MEMORY_PER_RWU: u64 = 4 * 1024 * 1024 * 1024;
150
151        self.rwu_limit.map(
152            |limit| (limit.get() + 1) * MEMORY_PER_RWU - 1, // allow some margin
153        )
154    }
155
156    /// Check whether the given cluster resource exceeds the limit in the license.
157    pub fn check_cluster_resource(&self, resource: ClusterResource) -> Result<(), LicenseError> {
158        let ClusterResource {
159            total_memory_bytes,
160            total_cpu_cores,
161        } = resource;
162
163        if let Some(limit) = self.cpu_core_limit()
164            && total_cpu_cores > limit
165        {
166            return Err(LicenseError::CpuLimitExceeded {
167                limit,
168                actual: total_cpu_cores,
169            });
170        }
171
172        #[cfg(not(madsim))] // skip checking memory limit in simulation tests
173        if let Some(limit) = self.memory_limit()
174            && total_memory_bytes > limit
175        {
176            return Err(LicenseError::MemoryLimitExceeded {
177                limit,
178                actual: total_memory_bytes,
179            });
180        }
181
182        Ok(())
183    }
184}
185
186impl Default for License {
187    /// The default license is a free license that never expires.
188    ///
189    /// Used when `license_key` is unset or invalid.
190    fn default() -> Self {
191        Self {
192            sub: "default".to_owned(),
193            tier: Tier::Free,
194            iss: Issuer::Prod,
195            rwu_limit: None,
196            exp: u64::MAX,
197        }
198    }
199}
200
201/// The error type for invalid license key when verifying as JWT.
202#[derive(Debug, Clone, Error)]
203pub enum LicenseError {
204    #[error("invalid license key")]
205    InvalidKey(#[source] jsonwebtoken::errors::Error),
206
207    #[error(
208        "a valid license key is set, but it is currently not effective because the CPU core in the cluster \
209        ({actual}) exceeds the maximum allowed by the license key ({limit}); \
210        consider removing some nodes or acquiring a new license key with a higher limit"
211    )]
212    CpuLimitExceeded { limit: u64, actual: u64 },
213
214    #[error(
215        "a valid license key is set, but it is currently not effective because the memory in the cluster \
216        ({actual}) exceeds the maximum allowed by the license key ({limit}); \
217        consider removing some nodes or acquiring a new license key with a higher limit",
218        actual = humansize::format_size(*actual, humansize::BINARY),
219        limit = humansize::format_size(*limit, humansize::BINARY),
220    )]
221    MemoryLimitExceeded { limit: u64, actual: u64 },
222}
223
224struct Inner {
225    license: Result<License, LicenseError>,
226    cached_cluster_resource: ClusterResource,
227}
228
229/// The singleton license manager.
230pub struct LicenseManager {
231    inner: RwLock<Inner>,
232}
233
234static PUBLIC_KEY: LazyLock<DecodingKey> = LazyLock::new(|| {
235    DecodingKey::from_rsa_pem(include_bytes!("key.pub"))
236        .expect("invalid public key for license validation")
237});
238
239impl LicenseManager {
240    /// Create a new license manager with the default license.
241    pub(crate) fn new() -> Self {
242        Self {
243            inner: RwLock::new(Inner {
244                license: Ok(License::default()),
245                cached_cluster_resource: ClusterResource {
246                    total_cpu_cores: 0,
247                    total_memory_bytes: 0,
248                },
249            }),
250        }
251    }
252
253    /// Get the singleton instance of the license manager.
254    pub fn get() -> &'static Self {
255        static INSTANCE: LazyLock<LicenseManager> = LazyLock::new(LicenseManager::new);
256        &INSTANCE
257    }
258
259    /// Refresh the license with the given license key.
260    pub fn refresh(&self, license_key: LicenseKeyRef<'_>) {
261        let license_key = license_key.0;
262        let mut inner = self.inner.write().unwrap();
263
264        // Empty license key means unset. Use the default one here.
265        if license_key.is_empty() {
266            inner.license = Ok(License::default());
267            return;
268        }
269
270        // TODO(license): shall we also validate `nbf`(Not Before)?
271        let mut validation = Validation::new(Algorithm::RS512);
272        // Only accept `prod` issuer in production, so that we can use license keys issued by
273        // the `test` issuer in development without leaking them to production.
274        validation.set_issuer(&[
275            "prod.risingwave.com",
276            #[cfg(debug_assertions)]
277            "test.risingwave.com",
278        ]);
279
280        inner.license = match jsonwebtoken::decode(license_key, &PUBLIC_KEY, &validation) {
281            Ok(data) => Ok(data.claims),
282            Err(error) => Err(LicenseError::InvalidKey(error)),
283        };
284
285        match &inner.license {
286            Ok(license) => tracing::info!(?license, "license refreshed"),
287            Err(error) => tracing::warn!(error = %error.as_report(), "invalid license key"),
288        }
289    }
290
291    /// Update the cached cluster resource.
292    pub fn update_cluster_resource(&self, resource: ClusterResource) {
293        let mut inner = self.inner.write().unwrap();
294        inner.cached_cluster_resource = resource;
295    }
296
297    /// Get the current license if it is valid.
298    ///
299    /// Since the license can expire, the returned license should not be cached by the caller.
300    ///
301    /// Prefer calling [`crate::Feature::check_available`] to check the availability of a feature,
302    /// other than directly calling this method and checking the content of the license.
303    pub fn license(&self) -> Result<License, LicenseError> {
304        let inner = self.inner.read().unwrap();
305        let license = inner.license.clone()?;
306
307        // Check the expiration time additionally.
308        if license.exp < jsonwebtoken::get_current_timestamp() {
309            return Err(LicenseError::InvalidKey(
310                jsonwebtoken::errors::ErrorKind::ExpiredSignature.into(),
311            ));
312        }
313
314        // Check the resource limit.
315        license.check_cluster_resource(inner.cached_cluster_resource)?;
316
317        Ok(license)
318    }
319}
320
321// Tests below only work in debug mode.
322#[cfg(debug_assertions)]
323#[cfg(test)]
324mod tests {
325    use expect_test::expect;
326
327    use super::*;
328    use crate::{LicenseKey, PROD_ALL_4_CORE_LICENSE_KEY_CONTENT, TEST_ALL_LICENSE_KEY_CONTENT};
329
330    fn do_test(key: &str, expect: expect_test::Expect) -> LicenseManager {
331        let manager = LicenseManager::new();
332        manager.refresh(LicenseKey(key));
333
334        match manager.license() {
335            Ok(license) => expect.assert_debug_eq(&license),
336            Err(error) => expect.assert_eq(&error.to_report_string()),
337        }
338
339        manager
340    }
341
342    #[test]
343    fn test_all_license_key() {
344        do_test(
345            TEST_ALL_LICENSE_KEY_CONTENT,
346            expect![[r#"
347                License {
348                    sub: "rw-test-all",
349                    iss: Test,
350                    tier: All,
351                    rwu_limit: None,
352                    exp: 10000627200,
353                }
354            "#]],
355        );
356    }
357
358    #[test]
359    fn test_prod_all_4_core_license_key() {
360        do_test(
361            PROD_ALL_4_CORE_LICENSE_KEY_CONTENT,
362            expect![[r#"
363                License {
364                    sub: "rw-default-all-4-core",
365                    iss: Prod,
366                    tier: All,
367                    rwu_limit: Some(
368                        4,
369                    ),
370                    exp: 10000627200,
371                }
372            "#]],
373        );
374    }
375
376    #[test]
377    fn test_free_license_key() {
378        const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
379          eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjo5OTk5OTk5OTk5fQ.\
380          ALC3Kc9LI6u0S-jeMB1YTxg1k8Azxwvc750ihuSZgjA_e1OJC9moxMvpLrHdLZDzCXHjBYi0XJ_1lowmuO_0iPEuPqN5AFpDV1ywmzJvGmMCMtw3A2wuN7hhem9OsWbwe6lzdwrefZLipyo4GZtIkg5ZdwGuHzm33zsM-X5gl_Ns4P6axHKiorNSR6nTAyA6B32YVET_FAM2YJQrXqpwA61wn1XLfarZqpdIQyJ5cgyiC33BFBlUL3lcRXLMLeYe6TjYGeV4K63qARCjM9yeOlsRbbW5ViWeGtR2Yf18pN8ysPXdbaXm_P_IVhl3jCTDJt9ctPh6pUCbkt36FZqO9A";
381
382        do_test(
383            KEY,
384            expect![[r#"
385                License {
386                    sub: "rw-test",
387                    iss: Test,
388                    tier: Free,
389                    rwu_limit: None,
390                    exp: 9999999999,
391                }
392            "#]],
393        );
394    }
395
396    #[test]
397    fn test_empty_license_key() {
398        // Default license will be used.
399        do_test(
400            "",
401            expect![[r#"
402                License {
403                    sub: "default",
404                    iss: Prod,
405                    tier: Free,
406                    rwu_limit: None,
407                    exp: 18446744073709551615,
408                }
409            "#]],
410        );
411    }
412
413    #[test]
414    fn test_custom_tier_license_key() {
415        const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
416          eyJzdWIiOiJydy11bml0LXRlc3QiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjoxMDAwMDYyNzIwMCwiaWF0IjoxNzUxODY5OTI3LCJ0aWVyIjp7Im5hbWUiOiJzZWNyZXQtb25seSIsImZlYXR1cmVzIjpbIlNlY3JldE1hbmFnZW1lbnQiXX19.\
417          iZKNyAHxz1l24Qh4F9wauuQEtDPLAeOmW2m_ttlew2ENmP_XcGWCx_O8r50NXRDaKv66z-ibteWVL5XOUqaVgJw9EnCyfVuFNoKtFQu18Vt0l52Yw2zNh3iHNQFKwHuCUki5FlOHw-K57a5f414-gzfwAwQkO1bAYoyAFhtoX6QQ2jdbxctFg0NxTQqpjnP-h0k2myZ_IhxA8fKrQIAqBj5Y8tGljuxjTLJpoK7X0ESXNh7bhA8njEf2Hm4QCymI1uo8OYRaR1Siw87r0aZykw9wW15Q8VK58VxIQLS7b7gQOmDToBjJt9yIF3MT6YMfqMX_l3Dtn9bOS_htVd1bjQ";
418
419        let manager = do_test(
420            KEY,
421            expect![[r#"
422                License {
423                    sub: "rw-unit-test",
424                    iss: Test,
425                    tier: Custom {
426                        name: "secret-only",
427                        features: [
428                            Feature(
429                                SecretManagement,
430                            ),
431                        ],
432                    },
433                    rwu_limit: None,
434                    exp: 10000627200,
435                }
436            "#]],
437        );
438
439        let tier = manager.license().unwrap().tier;
440        assert_eq!(
441            tier.available_features().collect::<Vec<_>>(),
442            vec![Feature::SecretManagement]
443        );
444        assert!(
445            Feature::SecretManagement
446                .check_available_with(&manager)
447                .is_ok()
448        );
449        assert!(
450            Feature::BigQuerySink
451                .check_available_with(&manager)
452                .is_err()
453        );
454    }
455
456    #[test]
457    fn test_invalid_license_key() {
458        const KEY: &str = "invalid";
459
460        do_test(KEY, expect!["invalid license key: InvalidToken"]);
461    }
462
463    #[test]
464    fn test_expired_license_key() {
465        // "exp": 0
466        const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
467          eyJzdWIiOiJydy10ZXN0IiwidGllciI6InBhaWQiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjowfQ.\
468          TyYmoT5Gw9-FN7DWDbeg3myW8g_3Xlc90i4M9bGuPf2WLv9zRMJy2r9J7sl1BO7t6F1uGgyrvNxsVRVZ2XF_WAs6uNlluYBnd4Cqvsj6Xny1XJCCo8II3RIea-ZlRjp6tc1saaoe-_eTtqDH8NIIWe73vVtBeBTBU4zAiN2vCtU_Si2XuoTLBKJMIjtn0HjLNhb6-DX2P3SCzp75tMyWzr49qcsBgratyKdu_v2kqBM1qw_dTaRg2ZeNNO6scSOBwu4YHHJTL4nUaZO2yEodI_OKUztIPLYuO2A33Fb5OE57S7LTgSzmxZLf7e23Vrck7Os14AfBQr7p9ncUeyIXhA";
469
470        do_test(KEY, expect!["invalid license key: ExpiredSignature"]);
471    }
472
473    #[test]
474    fn test_invalid_issuer() {
475        // "iss": "bad.risingwave.com"
476        const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
477          eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJiYWQucmlzaW5nd2F2ZS5jb20iLCJleHAiOjk5OTk5OTk5OTl9.\
478          SUbDJTri902FbGgIoe5L3LG4edTXoR42BQCIu_KLyW41OK47bMnD2aK7JggyJmWyGtN7b_596hxM9HjU58oQtHePUo_zHi5li5IcRaMi8gqHae7CJGqOGAUo9vYOWCP5OjEuDfozJhpgcHBLzDRnSwYnWhLKtsrzb3UcpOXEqRVK7EDShBNx6kNqfYs2LlFI7ASsgFRLhoRuOTR5LeVDjj6NZfkZGsdMe1VyrODWoGT9kcAF--hBpUd1ZJ5mZ67A0_948VPFBYDbDPcTRnw1-5MvdibO-jKX49rJ0rlPXcAbqKPE_yYUaqUaORUzb3PaPgCT_quO9PWPuAFIgAb_fg";
479
480        do_test(KEY, expect!["invalid license key: InvalidIssuer"]);
481    }
482
483    #[test]
484    fn test_invalid_signature() {
485        const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
486         eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjo5OTk5OTk5OTk5fQ.\
487         InvalidSignatureoe5L3LG4edTXoR42BQCIu_KLyW41OK47bMnD2aK7JggyJmWyGtN7b_596hxM9HjU58oQtHePUo_zHi5li5IcRaMi8gqHae7CJGqOGAUo9vYOWCP5OjEuDfozJhpgcHBLzDRnSwYnWhLKtsrzb3UcpOXEqRVK7EDShBNx6kNqfYs2LlFI7ASsgFRLhoRuOTR5LeVDjj6NZfkZGsdMe1VyrODWoGT9kcAF--hBpUd1ZJ5mZ67A0_948VPFBYDbDPcTRnw1-5MvdibO-jKX49rJ0rlPXcAbqKPE_yYUaqUaORUzb3PaPgCT_quO9PWPuAFIgAb_fg";
488
489        do_test(KEY, expect!["invalid license key: InvalidSignature"]);
490    }
491}
risingwave_license/manager.rs

risingwave_license/
manager.rs
