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::LicenseKeyRef;
24
25/// License tier.
26///
27/// Each enterprise [`Feature`](super::Feature) is available for a specific tier and above.
28#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
29#[serde(rename_all = "snake_case")]
30pub enum Tier {
31    /// Free tier.
32    ///
33    /// This is more like a placeholder. If a feature is available for the free tier, there's no
34    /// need to add it to the [`Feature`](super::Feature) enum at all.
35    Free,
36
37    /// Paid tier.
38    // TODO(license): Add more tiers if needed.
39    Paid,
40}
41
42/// Issuer of the license.
43///
44/// The issuer must be `prod.risingwave.com` in production, and can be `test.risingwave.com` in
45/// development. This will be validated when refreshing the license key.
46#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
47pub enum Issuer {
48    #[serde(rename = "prod.risingwave.com")]
49    Prod,
50
51    #[serde(rename = "test.risingwave.com")]
52    Test,
53
54    #[serde(untagged)]
55    Unknown(String),
56}
57
58/// The content of a license.
59///
60/// We use JSON Web Token (JWT) to represent the license. This struct is the payload.
61///
62/// Prefer calling [`crate::Feature::check_available`] to check the availability of a feature,
63/// other than directly checking the content of the license.
64// TODO(license): Shall we add a version field?
65#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
66#[serde(rename_all = "snake_case")]
67pub struct License {
68    /// Subject of the license.
69    ///
70    /// See <https://tools.ietf.org/html/rfc7519#section-4.1.2>.
71    #[allow(dead_code)]
72    pub sub: String,
73
74    /// Issuer of the license.
75    ///
76    /// See <https://tools.ietf.org/html/rfc7519#section-4.1.1>.
77    #[allow(dead_code)]
78    pub iss: Issuer,
79
80    /// Tier of the license.
81    pub tier: Tier,
82
83    /// Maximum number of compute-node CPU cores allowed to use. Typically used for the paid tier.
84    pub cpu_core_limit: Option<NonZeroUsize>,
85
86    /// Expiration time in seconds since UNIX epoch.
87    ///
88    /// See <https://tools.ietf.org/html/rfc7519#section-4.1.4>.
89    pub exp: u64,
90}
91
92impl Default for License {
93    /// The default license is a free license that never expires.
94    ///
95    /// Used when `license_key` is unset or invalid.
96    fn default() -> Self {
97        Self {
98            sub: "default".to_owned(),
99            tier: Tier::Free,
100            iss: Issuer::Prod,
101            cpu_core_limit: None,
102            exp: u64::MAX,
103        }
104    }
105}
106
107/// The error type for invalid license key when verifying as JWT.
108#[derive(Debug, Clone, Error)]
109pub enum LicenseError {
110    #[error("invalid license key")]
111    InvalidKey(#[source] jsonwebtoken::errors::Error),
112
113    #[error(
114        "the license key is currently not effective because the CPU core in the cluster \
115        ({actual}) exceeds the maximum allowed by the license key ({limit}); \
116        consider removing some nodes or acquiring a new license key with a higher limit"
117    )]
118    CpuCoreLimitExceeded { limit: NonZeroUsize, actual: usize },
119}
120
121struct Inner {
122    license: Result<License, LicenseError>,
123    cached_cpu_core_count: usize,
124}
125
126/// The singleton license manager.
127pub struct LicenseManager {
128    inner: RwLock<Inner>,
129}
130
131static PUBLIC_KEY: LazyLock<DecodingKey> = LazyLock::new(|| {
132    DecodingKey::from_rsa_pem(include_bytes!("key.pub"))
133        .expect("invalid public key for license validation")
134});
135
136impl LicenseManager {
137    /// Create a new license manager with the default license.
138    pub(crate) fn new() -> Self {
139        Self {
140            inner: RwLock::new(Inner {
141                license: Ok(License::default()),
142                cached_cpu_core_count: 0,
143            }),
144        }
145    }
146
147    /// Get the singleton instance of the license manager.
148    pub fn get() -> &'static Self {
149        static INSTANCE: LazyLock<LicenseManager> = LazyLock::new(LicenseManager::new);
150        &INSTANCE
151    }
152
153    /// Refresh the license with the given license key.
154    pub fn refresh(&self, license_key: LicenseKeyRef<'_>) {
155        let license_key = license_key.0;
156        let mut inner = self.inner.write().unwrap();
157
158        // Empty license key means unset. Use the default one here.
159        if license_key.is_empty() {
160            inner.license = Ok(License::default());
161            return;
162        }
163
164        // TODO(license): shall we also validate `nbf`(Not Before)?
165        let mut validation = Validation::new(Algorithm::RS512);
166        // Only accept `prod` issuer in production, so that we can use license keys issued by
167        // the `test` issuer in development without leaking them to production.
168        validation.set_issuer(&[
169            "prod.risingwave.com",
170            #[cfg(debug_assertions)]
171            "test.risingwave.com",
172        ]);
173
174        inner.license = match jsonwebtoken::decode(license_key, &PUBLIC_KEY, &validation) {
175            Ok(data) => Ok(data.claims),
176            Err(error) => Err(LicenseError::InvalidKey(error)),
177        };
178
179        match &inner.license {
180            Ok(license) => tracing::info!(?license, "license refreshed"),
181            Err(error) => tracing::warn!(error = %error.as_report(), "invalid license key"),
182        }
183    }
184
185    /// Update the cached CPU core count.
186    pub fn update_cpu_core_count(&self, cpu_core_count: usize) {
187        let mut inner = self.inner.write().unwrap();
188        inner.cached_cpu_core_count = cpu_core_count;
189    }
190
191    /// Get the current license if it is valid.
192    ///
193    /// Since the license can expire, the returned license should not be cached by the caller.
194    ///
195    /// Prefer calling [`crate::Feature::check_available`] to check the availability of a feature,
196    /// other than directly calling this method and checking the content of the license.
197    pub fn license(&self) -> Result<License, LicenseError> {
198        let inner = self.inner.read().unwrap();
199        let license = inner.license.clone()?;
200
201        // Check the expiration time additionally.
202        if license.exp < jsonwebtoken::get_current_timestamp() {
203            return Err(LicenseError::InvalidKey(
204                jsonwebtoken::errors::ErrorKind::ExpiredSignature.into(),
205            ));
206        }
207
208        // Check the CPU core limit.
209        let actual_cpu_core = inner.cached_cpu_core_count;
210        if let Some(limit) = license.cpu_core_limit
211            && actual_cpu_core > limit.get()
212        {
213            return Err(LicenseError::CpuCoreLimitExceeded {
214                limit,
215                actual: actual_cpu_core,
216            });
217        }
218
219        Ok(license)
220    }
221}
222
223// Tests below only work in debug mode.
224#[cfg(debug_assertions)]
225#[cfg(test)]
226mod tests {
227    use expect_test::expect;
228
229    use super::*;
230    use crate::{LicenseKey, TEST_PAID_LICENSE_KEY_CONTENT};
231
232    fn do_test(key: &str, expect: expect_test::Expect) {
233        let manager = LicenseManager::new();
234        manager.refresh(LicenseKey(key));
235
236        match manager.license() {
237            Ok(license) => expect.assert_debug_eq(&license),
238            Err(error) => expect.assert_eq(&error.to_report_string()),
239        }
240    }
241
242    #[test]
243    fn test_paid_license_key() {
244        do_test(
245            TEST_PAID_LICENSE_KEY_CONTENT,
246            expect![[r#"
247                License {
248                    sub: "rw-test",
249                    iss: Test,
250                    tier: Paid,
251                    cpu_core_limit: None,
252                    exp: 9999999999,
253                }
254            "#]],
255        );
256    }
257
258    #[test]
259    fn test_free_license_key() {
260        const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
261          eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjo5OTk5OTk5OTk5fQ.\
262          ALC3Kc9LI6u0S-jeMB1YTxg1k8Azxwvc750ihuSZgjA_e1OJC9moxMvpLrHdLZDzCXHjBYi0XJ_1lowmuO_0iPEuPqN5AFpDV1ywmzJvGmMCMtw3A2wuN7hhem9OsWbwe6lzdwrefZLipyo4GZtIkg5ZdwGuHzm33zsM-X5gl_Ns4P6axHKiorNSR6nTAyA6B32YVET_FAM2YJQrXqpwA61wn1XLfarZqpdIQyJ5cgyiC33BFBlUL3lcRXLMLeYe6TjYGeV4K63qARCjM9yeOlsRbbW5ViWeGtR2Yf18pN8ysPXdbaXm_P_IVhl3jCTDJt9ctPh6pUCbkt36FZqO9A";
263
264        do_test(
265            KEY,
266            expect![[r#"
267                License {
268                    sub: "rw-test",
269                    iss: Test,
270                    tier: Free,
271                    cpu_core_limit: None,
272                    exp: 9999999999,
273                }
274            "#]],
275        );
276    }
277
278    #[test]
279    fn test_empty_license_key() {
280        // Default license will be used.
281        do_test(
282            "",
283            expect![[r#"
284                License {
285                    sub: "default",
286                    iss: Prod,
287                    tier: Free,
288                    cpu_core_limit: None,
289                    exp: 18446744073709551615,
290                }
291            "#]],
292        );
293    }
294
295    #[test]
296    fn test_invalid_license_key() {
297        const KEY: &str = "invalid";
298
299        do_test(KEY, expect!["invalid license key: InvalidToken"]);
300    }
301
302    #[test]
303    fn test_expired_license_key() {
304        // "exp": 0
305        const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
306          eyJzdWIiOiJydy10ZXN0IiwidGllciI6InBhaWQiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjowfQ.\
307          TyYmoT5Gw9-FN7DWDbeg3myW8g_3Xlc90i4M9bGuPf2WLv9zRMJy2r9J7sl1BO7t6F1uGgyrvNxsVRVZ2XF_WAs6uNlluYBnd4Cqvsj6Xny1XJCCo8II3RIea-ZlRjp6tc1saaoe-_eTtqDH8NIIWe73vVtBeBTBU4zAiN2vCtU_Si2XuoTLBKJMIjtn0HjLNhb6-DX2P3SCzp75tMyWzr49qcsBgratyKdu_v2kqBM1qw_dTaRg2ZeNNO6scSOBwu4YHHJTL4nUaZO2yEodI_OKUztIPLYuO2A33Fb5OE57S7LTgSzmxZLf7e23Vrck7Os14AfBQr7p9ncUeyIXhA";
308
309        do_test(KEY, expect!["invalid license key: ExpiredSignature"]);
310    }
311
312    #[test]
313    fn test_invalid_issuer() {
314        // "iss": "bad.risingwave.com"
315        const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
316          eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJiYWQucmlzaW5nd2F2ZS5jb20iLCJleHAiOjk5OTk5OTk5OTl9.\
317          SUbDJTri902FbGgIoe5L3LG4edTXoR42BQCIu_KLyW41OK47bMnD2aK7JggyJmWyGtN7b_596hxM9HjU58oQtHePUo_zHi5li5IcRaMi8gqHae7CJGqOGAUo9vYOWCP5OjEuDfozJhpgcHBLzDRnSwYnWhLKtsrzb3UcpOXEqRVK7EDShBNx6kNqfYs2LlFI7ASsgFRLhoRuOTR5LeVDjj6NZfkZGsdMe1VyrODWoGT9kcAF--hBpUd1ZJ5mZ67A0_948VPFBYDbDPcTRnw1-5MvdibO-jKX49rJ0rlPXcAbqKPE_yYUaqUaORUzb3PaPgCT_quO9PWPuAFIgAb_fg";
318
319        do_test(KEY, expect!["invalid license key: InvalidIssuer"]);
320    }
321
322    #[test]
323    fn test_invalid_signature() {
324        const KEY: &str = "eyJhbGciOiJSUzUxMiIsInR5cCI6IkpXVCJ9.\
325         eyJzdWIiOiJydy10ZXN0IiwidGllciI6ImZyZWUiLCJpc3MiOiJ0ZXN0LnJpc2luZ3dhdmUuY29tIiwiZXhwIjo5OTk5OTk5OTk5fQ.\
326         InvalidSignatureoe5L3LG4edTXoR42BQCIu_KLyW41OK47bMnD2aK7JggyJmWyGtN7b_596hxM9HjU58oQtHePUo_zHi5li5IcRaMi8gqHae7CJGqOGAUo9vYOWCP5OjEuDfozJhpgcHBLzDRnSwYnWhLKtsrzb3UcpOXEqRVK7EDShBNx6kNqfYs2LlFI7ASsgFRLhoRuOTR5LeVDjj6NZfkZGsdMe1VyrODWoGT9kcAF--hBpUd1ZJ5mZ67A0_948VPFBYDbDPcTRnw1-5MvdibO-jKX49rJ0rlPXcAbqKPE_yYUaqUaORUzb3PaPgCT_quO9PWPuAFIgAb_fg";
327
328        do_test(KEY, expect!["invalid license key: InvalidSignature"]);
329    }
330}
risingwave_license/manager.rs

risingwave_license/
manager.rs
