risingwave_license/
manager.rs

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

risingwave_license/
manager.rs
