1use 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.";
20const PUBLIC_KEYS: &[&str] = &[
23 include_str!("license_keys/production.pub"),
24 include_str!("license_keys/intermediate_root.pub"),
26];
27const REVOKED_KEYS: &[&str] = &["eddaf004-dc1e-48cf-9cc1-41d1543d940a"];
31const INTERMEDIATE_TYP: &str = "mz-intermediate+jwt";
32const 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 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 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
134pub 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 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 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 public_key: String,
196 expiration: u64,
197 expired: bool,
198}
199
200fn 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 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 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 ];
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 #[serde(default, skip_serializing_if = "Vec::is_empty")]
307 entitlements: Vec<String>,
308}
309
310#[derive(Debug, Clone, Serialize, Deserialize)]
311struct IntermediatePayload {
312 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 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)] 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)] 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 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 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}