Skip to main content

mz_license_keys/
lib.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use anyhow::{Context, anyhow, bail};
11use jsonwebtoken::{Algorithm, DecodingKey, TokenData, Validation};
12use serde::{Deserialize, Serialize};
13
14#[cfg(feature = "signing")]
15mod signing;
16#[cfg(feature = "signing")]
17pub use signing::{get_pubkey_pem, make_intermediate_jwt, make_license_key};
18
19const ISSUER: &str = "Materialize, Inc.";
20// list of public keys which are allowed to validate license keys. this is a
21// list to allow for key rotation if necessary.
22const PUBLIC_KEYS: &[&str] = &[
23    include_str!("license_keys/production.pub"),
24    // Signs the intermediate JWTs whose keys sign BYOC license keys.
25    include_str!("license_keys/intermediate_root.pub"),
26];
27// keys which we have issued but need to be revoked before their expiration
28// (due to being accidentally exposed or similar). holds both license key and
29// intermediate JWT ids.
30const REVOKED_KEYS: &[&str] = &["eddaf004-dc1e-48cf-9cc1-41d1543d940a"];
31const INTERMEDIATE_TYP: &str = "mz-intermediate+jwt";
32// License key header field carrying the intermediate JWT that certifies the
33// key the license key was signed with.
34const INTERMEDIATE_HEADER: &str = "mz_intermediate";
35
36#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
37#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
38pub enum ExpirationBehavior {
39    Warn,
40    DisableClusterCreation,
41    Disable,
42}
43
44#[derive(Debug, Clone)]
45pub struct ValidatedLicenseKey {
46    pub id: String,
47    pub organization: String,
48    pub environment_id: String,
49    pub expiration: u64,
50    pub not_before: u64,
51
52    pub max_credit_consumption_rate: f64,
53    pub allow_credit_consumption_override: bool,
54    pub expiration_behavior: ExpirationBehavior,
55    pub expired: bool,
56    /// Optional feature flags / third-party integrations enabled for this key
57    /// (e.g. `"ory"` to permit pulling images through the OCI registry proxy).
58    /// Empty for keys that predate this field.
59    pub entitlements: Vec<String>,
60}
61
62impl ValidatedLicenseKey {
63    pub fn for_tests() -> Self {
64        Self {
65            id: "".to_string(),
66            organization: "".to_string(),
67            environment_id: "".to_string(),
68            expiration: 0,
69            not_before: 0,
70
71            max_credit_consumption_rate: 999999.0,
72            allow_credit_consumption_override: true,
73            expiration_behavior: ExpirationBehavior::Warn,
74            expired: false,
75            entitlements: Vec::new(),
76        }
77    }
78
79    pub fn disabled() -> Self {
80        Self {
81            id: "".to_string(),
82            organization: "".to_string(),
83            environment_id: "".to_string(),
84            expiration: 0,
85            not_before: 0,
86
87            max_credit_consumption_rate: 999999.0,
88            allow_credit_consumption_override: true,
89            expiration_behavior: ExpirationBehavior::Warn,
90            expired: false,
91            entitlements: Vec::new(),
92        }
93    }
94
95    /// Returns true if `entitlement` is present in this key.
96    pub fn has_entitlement(&self, entitlement: &str) -> bool {
97        self.entitlements.iter().any(|e| e == entitlement)
98    }
99
100    pub fn max_credit_consumption_rate(&self) -> Option<f64> {
101        if self.expired
102            && matches!(
103                self.expiration_behavior,
104                ExpirationBehavior::DisableClusterCreation | ExpirationBehavior::Disable
105            )
106        {
107            Some(0.0)
108        } else if self.allow_credit_consumption_override {
109            None
110        } else {
111            Some(self.max_credit_consumption_rate)
112        }
113    }
114}
115
116impl Default for ValidatedLicenseKey {
117    fn default() -> Self {
118        Self {
119            id: "".to_string(),
120            organization: "".to_string(),
121            environment_id: "".to_string(),
122            expiration: 0,
123            not_before: 0,
124
125            max_credit_consumption_rate: 24.0,
126            allow_credit_consumption_override: false,
127            expiration_behavior: ExpirationBehavior::Disable,
128            expired: false,
129            entitlements: Vec::new(),
130        }
131    }
132}
133
134/// Validates `license_key` against the root public keys, falling back to the
135/// intermediate JWT embedded in its header.
136///
137/// An intermediate JWT is signed by a root key and carries another public key
138/// (see `make_intermediate_jwt`). A license key signed by that public key is
139/// accepted while the intermediate JWT is not revoked. The result expires at
140/// the earlier of the license key's and the intermediate JWT's expiration.
141pub fn validate(license_key: &str) -> anyhow::Result<ValidatedLicenseKey> {
142    validate_with_roots(license_key, PUBLIC_KEYS)
143}
144
145fn validate_with_roots(
146    license_key: &str,
147    root_pubkeys: &[&str],
148) -> anyhow::Result<ValidatedLicenseKey> {
149    let mut errs = Vec::new();
150    for (i, pubkey) in root_pubkeys.iter().enumerate() {
151        match validate_with_pubkey(license_key, pubkey) {
152            Ok(key) => return Ok(key),
153            Err(e) => errs.push(format!("root key {i}: {e:#}")),
154        }
155    }
156
157    // The header is read before any signature is checked. That is safe
158    // because the intermediate JWT is verified against the root keys before
159    // its public key is used.
160    let intermediate = jsonwebtoken::decode_header(license_key)
161        .map_err(anyhow::Error::from)
162        .and_then(|header| Ok(header.extras.get::<String>(INTERMEDIATE_HEADER)?));
163    match intermediate {
164        Ok(Some(intermediate_jwt)) => {
165            let res = validate_intermediate(&intermediate_jwt, root_pubkeys)
166                .context("invalid intermediate JWT")
167                .and_then(|intermediate| {
168                    let mut key = validate_with_pubkey(license_key, &intermediate.public_key)?;
169                    // NOTE: An expired intermediate JWT does not fail
170                    // validation, so the license key's `expiration_behavior`
171                    // decides what happens. This means its expiration does not
172                    // bound a leaked intermediate private key, which can mint
173                    // `Warn` keys. Only `REVOKED_KEYS` stops that.
174                    key.expired |= intermediate.expired;
175                    key.expiration = key.expiration.min(intermediate.expiration);
176                    Ok(key)
177                });
178            match res {
179                Ok(key) => return Ok(key),
180                Err(e) => errs.push(format!("intermediate JWT: {e:#}")),
181            }
182        }
183        Ok(None) => {}
184        Err(e) => errs.push(format!("intermediate JWT: {e:#}")),
185    }
186
187    if errs.is_empty() {
188        bail!("no public key found");
189    }
190    bail!("license key validation failed:\n  {}", errs.join("\n  "))
191}
192
193struct ValidatedIntermediate {
194    /// PEM of the public key that signs license keys.
195    public_key: String,
196    expiration: u64,
197    expired: bool,
198}
199
200/// Returns the claims of `intermediate_jwt` if a root key signed it.
201fn validate_intermediate(
202    intermediate_jwt: &str,
203    root_pubkeys: &[&str],
204) -> anyhow::Result<ValidatedIntermediate> {
205    let mut validation = Validation::new(Algorithm::PS256);
206    validation.set_required_spec_claims(&["exp", "nbf", "iss", "sub"]);
207    validation.set_issuer(&[ISSUER]);
208    validation.validate_exp = true;
209    validation.validate_nbf = true;
210
211    let mut err = None;
212    for root in root_pubkeys {
213        let key = DecodingKey::from_rsa_pem(root.as_bytes())?;
214        let res = match jsonwebtoken::decode::<IntermediatePayload>(
215            intermediate_jwt,
216            &key,
217            &validation,
218        ) {
219            Err(e) if matches!(e.kind(), jsonwebtoken::errors::ErrorKind::ExpiredSignature) => {
220                let mut validation = validation.clone();
221                validation.validate_exp = false;
222                jsonwebtoken::decode(intermediate_jwt, &key, &validation).map(|d| (d, true))
223            }
224            res => res.map(|d| (d, false)),
225        };
226        match res {
227            Ok((decoded, expired)) => {
228                // The typ check keeps a root-signed license key from being
229                // accepted as an intermediate JWT, and vice versa.
230                if decoded.header.typ.as_deref() != Some(INTERMEDIATE_TYP) {
231                    bail!("invalid intermediate JWT header type");
232                }
233                if decoded.claims.version != 1 {
234                    bail!("invalid intermediate JWT version");
235                }
236                if decoded.claims.iat < decoded.claims.nbf {
237                    bail!("intermediate JWT issued before it is valid");
238                }
239                if decoded.claims.iat > decoded.claims.exp {
240                    bail!("intermediate JWT issued after it expired");
241                }
242                if REVOKED_KEYS.contains(&decoded.claims.jti.as_str()) {
243                    bail!("revoked intermediate JWT");
244                }
245                return Ok(ValidatedIntermediate {
246                    public_key: decoded.claims.public_key,
247                    expiration: decoded.claims.exp,
248                    expired,
249                });
250            }
251            Err(e) => err = Some(e.into()),
252        }
253    }
254    Err(err.unwrap_or_else(|| anyhow!("no public key found")))
255}
256
257fn validate_with_pubkey(
258    license_key: &str,
259    pubkey_pem: &str,
260) -> anyhow::Result<ValidatedLicenseKey> {
261    // don't just read the version out of the payload before verifying it,
262    // trusting unsigned data to determine how to verify the signature is a
263    // bad idea. instead, just try validating it as each version
264    // independently, and if the signature is valid, only then check to
265    // ensure that the version matches what we validated.
266
267    // try current version first, so we can prefer that for error messages
268    let res = validate_with_pubkey_v1(license_key, pubkey_pem);
269    let err = match res {
270        Ok(key) => return Ok(key),
271        Err(e) => e,
272    };
273
274    let previous_versions: Vec<Box<dyn Fn() -> anyhow::Result<ValidatedLicenseKey>>> = vec![
275        // add to this if/when we add new versions
276        // for example,
277        // Box::new(|| validate_with_pubkey_v1(license_key, pubkey_pem, environment_id)),
278    ];
279    for validator in previous_versions {
280        if let Ok(key) = validator() {
281            return Ok(key);
282        }
283    }
284
285    Err(err)
286}
287
288#[derive(Debug, Clone, Serialize, Deserialize)]
289struct Payload {
290    sub: String,
291    exp: u64,
292    nbf: u64,
293    iss: String,
294    aud: String,
295    iat: u64,
296    jti: String,
297
298    version: u64,
299    max_credit_consumption_rate: f64,
300    #[serde(default, skip_serializing_if = "is_default")]
301    allow_credit_consumption_override: bool,
302    expiration_behavior: ExpirationBehavior,
303    // Defaulted + skipped-when-empty so keys issued before entitlements
304    // existed continue to validate and we don't bloat keys that don't need
305    // any entitlements.
306    #[serde(default, skip_serializing_if = "Vec::is_empty")]
307    entitlements: Vec<String>,
308}
309
310#[derive(Debug, Clone, Serialize, Deserialize)]
311struct IntermediatePayload {
312    /// Human readable name of the intermediate key, e.g. the stack it serves.
313    sub: String,
314    exp: u64,
315    nbf: u64,
316    iss: String,
317    iat: u64,
318    jti: String,
319
320    version: u64,
321    public_key: String,
322}
323
324fn validate_with_pubkey_v1(
325    license_key: &str,
326    pubkey_pem: &str,
327) -> anyhow::Result<ValidatedLicenseKey> {
328    let mut validation = Validation::new(Algorithm::PS256);
329    validation.set_required_spec_claims(&["exp", "nbf", "aud", "iss", "sub"]);
330    validation.set_issuer(&[ISSUER]);
331    validation.validate_exp = true;
332    validation.validate_nbf = true;
333    validation.validate_aud = false;
334
335    let key = DecodingKey::from_rsa_pem(pubkey_pem.as_bytes())?;
336
337    let (jwt, expired): (TokenData<Payload>, _) =
338        jsonwebtoken::decode(license_key, &key, &validation).map_or_else(
339            |e| {
340                if matches!(e.kind(), jsonwebtoken::errors::ErrorKind::ExpiredSignature) {
341                    validation.validate_exp = false;
342                    Ok((jsonwebtoken::decode(license_key, &key, &validation)?, true))
343                } else {
344                    Err::<_, anyhow::Error>(e.into())
345                }
346            },
347            |jwt| Ok((jwt, false)),
348        )?;
349
350    if jwt.header.typ.as_deref() != Some("JWT") {
351        bail!("invalid jwt header type");
352    }
353
354    if jwt.claims.version != 1 {
355        bail!("invalid license key version");
356    }
357
358    if !(jwt.claims.nbf..=jwt.claims.exp).contains(&jwt.claims.iat) {
359        bail!("invalid issuance time");
360    }
361
362    if REVOKED_KEYS.contains(&jwt.claims.jti.as_str()) {
363        bail!("revoked license key");
364    }
365
366    Ok(ValidatedLicenseKey {
367        id: jwt.claims.jti,
368        organization: jwt.claims.sub,
369        environment_id: jwt.claims.aud,
370        expiration: jwt.claims.exp,
371        not_before: jwt.claims.nbf,
372
373        max_credit_consumption_rate: jwt.claims.max_credit_consumption_rate,
374        allow_credit_consumption_override: jwt.claims.allow_credit_consumption_override,
375        expiration_behavior: jwt.claims.expiration_behavior,
376        expired,
377        entitlements: jwt.claims.entitlements,
378    })
379}
380
381fn is_default<T: PartialEq + Eq + Default>(val: &T) -> bool {
382    *val == T::default()
383}
384
385#[cfg(test)]
386mod tests {
387    use aws_lc_rs::encoding::AsDer;
388    use aws_lc_rs::rsa::{KeyPair, KeySize};
389    use aws_lc_rs::signature::KeyPair as _;
390    use jsonwebtoken::{EncodingKey, Header};
391
392    use super::*;
393
394    /// Returns `(private_pem, public_pem)`.
395    fn keypair() -> (String, String) {
396        let key = KeyPair::generate(KeySize::Rsa2048).unwrap();
397        let private: aws_lc_rs::encoding::Pkcs8V1Der = key.as_der().unwrap();
398        let public: aws_lc_rs::encoding::PublicKeyX509Der = key.public_key().as_der().unwrap();
399        (
400            pem::Pem::new("PRIVATE KEY", private.as_ref()).to_string(),
401            pem::Pem::new("PUBLIC KEY", public.as_ref()).to_string(),
402        )
403    }
404
405    fn sign(
406        private_pem: &str,
407        typ: &str,
408        intermediate_jwt: Option<&str>,
409        claims: &impl Serialize,
410    ) -> String {
411        let mut header = Header::new(Algorithm::PS256);
412        header.typ = Some(typ.to_string());
413        if let Some(intermediate_jwt) = intermediate_jwt {
414            header.extras.insert(INTERMEDIATE_HEADER, intermediate_jwt);
415        }
416        let key = EncodingKey::from_rsa_pem(private_pem.as_bytes()).unwrap();
417        jsonwebtoken::encode(&header, claims, &key).unwrap()
418    }
419
420    fn now() -> u64 {
421        std::time::SystemTime::now()
422            .duration_since(std::time::UNIX_EPOCH)
423            .unwrap()
424            .as_secs()
425    }
426
427    fn license_payload() -> Payload {
428        let now = now();
429        Payload {
430            exp: now + 3600,
431            nbf: now,
432            iat: now,
433            ..sample_payload(Vec::new())
434        }
435    }
436
437    fn intermediate_payload(public_key: String, exp: u64) -> IntermediatePayload {
438        let now = now();
439        IntermediatePayload {
440            sub: "test-stack".to_string(),
441            exp,
442            nbf: now - 7200,
443            iss: ISSUER.to_string(),
444            iat: now - 7200,
445            jti: "intermediate-1".to_string(),
446            version: 1,
447            public_key,
448        }
449    }
450
451    #[mz_ore::test]
452    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function
453    fn intermediate_validation() {
454        let (root_priv, root_pub) = keypair();
455        let (int_priv, int_pub) = keypair();
456        let (other_priv, _) = keypair();
457        let roots = &[root_pub.as_str()];
458        let int_key = |intermediate_jwt: &str| {
459            sign(&int_priv, "JWT", Some(intermediate_jwt), &license_payload())
460        };
461
462        let int_jwt = sign(
463            &root_priv,
464            INTERMEDIATE_TYP,
465            None,
466            &intermediate_payload(int_pub.clone(), now() + 3600),
467        );
468        let root_key = sign(&root_priv, "JWT", None, &license_payload());
469
470        let root_validated = validate_with_roots(&root_key, roots).unwrap();
471        assert_eq!(root_validated.id, "jti-1");
472        validate_with_roots(&int_key(&int_jwt), roots).unwrap();
473        assert!(
474            validate_with_roots(&sign(&int_priv, "JWT", None, &license_payload()), roots).is_err(),
475            "intermediate-signed key needs its intermediate JWT"
476        );
477
478        let untrusted_jwt = sign(
479            &other_priv,
480            INTERMEDIATE_TYP,
481            None,
482            &intermediate_payload(int_pub.clone(), now() + 3600),
483        );
484        assert!(
485            validate_with_roots(&int_key(&untrusted_jwt), roots).is_err(),
486            "intermediate JWT must be signed by a root key"
487        );
488
489        let mut revoked = license_payload();
490        revoked.jti = REVOKED_KEYS[0].to_string();
491        let revoked_key = sign(&root_priv, "JWT", Some(&int_jwt), &revoked);
492        let err = validate_with_roots(&revoked_key, roots)
493            .unwrap_err()
494            .to_string();
495        assert!(
496            err.contains("root key 0: revoked license key") && err.contains("intermediate JWT: "),
497            "error must report every attempt: {err}"
498        );
499
500        let expired_exp = now() - 3600;
501        let expired_jwt = sign(
502            &root_priv,
503            INTERMEDIATE_TYP,
504            None,
505            &intermediate_payload(int_pub.clone(), expired_exp),
506        );
507        let validated = validate_with_roots(&int_key(&expired_jwt), roots).unwrap();
508        assert!(
509            validated.expired,
510            "expired intermediate JWT must expire the key"
511        );
512        assert_eq!(validated.expiration, expired_exp);
513
514        let wrong_typ_jwt = sign(
515            &root_priv,
516            "JWT",
517            None,
518            &intermediate_payload(int_pub, now() + 3600),
519        );
520        assert!(
521            validate_with_roots(&int_key(&wrong_typ_jwt), roots).is_err(),
522            "intermediate JWT must have the intermediate typ"
523        );
524
525        assert!(
526            validate_with_roots(&int_key(&root_key), roots).is_err(),
527            "license key must not be usable as an intermediate JWT"
528        );
529        let self_embedded = sign(
530            &int_priv,
531            INTERMEDIATE_TYP,
532            Some(&int_jwt),
533            &intermediate_payload(String::new(), now() + 3600),
534        );
535        assert!(
536            validate_with_roots(&self_embedded, roots).is_err(),
537            "intermediate JWT must not be usable as a license key"
538        );
539    }
540
541    #[mz_ore::test]
542    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function
543    fn public_keys_parse() {
544        for (i, pem) in PUBLIC_KEYS.iter().enumerate() {
545            DecodingKey::from_rsa_pem(pem.as_bytes())
546                .unwrap_or_else(|e| panic!("PUBLIC_KEYS[{i}] is not an RSA public key: {e}"));
547        }
548    }
549
550    fn sample_payload(entitlements: Vec<String>) -> Payload {
551        Payload {
552            sub: "org-1".to_string(),
553            exp: 200,
554            nbf: 100,
555            iss: ISSUER.to_string(),
556            aud: "env-1".to_string(),
557            iat: 100,
558            jti: "jti-1".to_string(),
559            version: 1,
560            max_credit_consumption_rate: 10.0,
561            allow_credit_consumption_override: false,
562            expiration_behavior: ExpirationBehavior::Warn,
563            entitlements,
564        }
565    }
566
567    #[mz_ore::test]
568    fn entitlements_roundtrip_through_payload() {
569        let payload = sample_payload(vec!["ory".to_string(), "foo".to_string()]);
570        let json = serde_json::to_string(&payload).unwrap();
571        assert!(
572            json.contains(r#""entitlements":["ory","foo"]"#),
573            "payload should serialize entitlements: {json}"
574        );
575
576        let decoded: Payload = serde_json::from_str(&json).unwrap();
577        assert_eq!(decoded.entitlements, vec!["ory", "foo"]);
578    }
579
580    #[mz_ore::test]
581    fn empty_entitlements_omitted_from_payload() {
582        let payload = sample_payload(Vec::new());
583        let json = serde_json::to_string(&payload).unwrap();
584        // Skipping the field on empty keeps issued JWTs the same shape they
585        // were before entitlements existed, so old + new validators agree.
586        assert!(
587            !json.contains("entitlements"),
588            "empty entitlements should be skipped: {json}"
589        );
590    }
591
592    #[mz_ore::test]
593    fn legacy_payload_without_entitlements_decodes() {
594        // Pre-DEP-130 keys have no `entitlements` field. They must still
595        // deserialize, with entitlements defaulting to empty.
596        let legacy = serde_json::json!({
597            "sub": "org-1",
598            "exp": 200,
599            "nbf": 100,
600            "iss": ISSUER,
601            "aud": "env-1",
602            "iat": 100,
603            "jti": "jti-1",
604            "version": 1,
605            "max_credit_consumption_rate": 10.0,
606            "expiration_behavior": "Warn",
607        });
608        let decoded: Payload = serde_json::from_value(legacy).unwrap();
609        assert!(decoded.entitlements.is_empty());
610    }
611}