mz_license_keys/
signing.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 std::time::{Duration, SystemTime, UNIX_EPOCH};
11
12use anyhow::anyhow;
13use aws_lc_rs::digest;
14use aws_sdk_kms::{
15    primitives::Blob,
16    types::{MessageType, SigningAlgorithmSpec},
17};
18use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
19use jsonwebtoken::{Algorithm, Header};
20use pem::Pem;
21use serde::Serialize;
22use uuid::Uuid;
23
24use crate::{
25    ExpirationBehavior, INTERMEDIATE_HEADER, INTERMEDIATE_TYP, ISSUER, IntermediatePayload, Payload,
26};
27
28const VERSION: u64 = 1;
29
30pub async fn get_pubkey_pem(client: &aws_sdk_kms::Client, key_id: &str) -> anyhow::Result<String> {
31    let pubkey = get_pubkey(client, key_id).await?;
32    let pem = Pem::new("PUBLIC KEY", pubkey);
33    Ok(pem.to_string())
34}
35
36/// Signs a license key with `key_id`.
37///
38/// When `key_id` is an intermediate key, `intermediate_jwt` must be the
39/// intermediate JWT for it, which is embedded in the license key's header.
40#[allow(clippy::too_many_arguments)]
41pub async fn make_license_key(
42    client: &aws_sdk_kms::Client,
43    key_id: &str,
44    intermediate_jwt: Option<String>,
45    validity: Duration,
46    organization_id: String,
47    environment_id: String,
48    max_credit_consumption_rate: f64,
49    allow_credit_consumption_override: bool,
50    expiration_behavior: ExpirationBehavior,
51    entitlements: Vec<String>,
52) -> anyhow::Result<String> {
53    let now = SystemTime::now();
54    let expiration = now + validity;
55    let payload = Payload {
56        sub: organization_id,
57        exp: format_time(&expiration),
58        nbf: format_time(&now),
59        iss: ISSUER.to_string(),
60        aud: environment_id,
61        iat: format_time(&now),
62        jti: Uuid::new_v4().to_string(),
63        version: VERSION,
64        max_credit_consumption_rate,
65        allow_credit_consumption_override,
66        expiration_behavior,
67        entitlements,
68    };
69    let mut header = Header::new(Algorithm::PS256);
70    header.typ = Some("JWT".to_string());
71    if let Some(intermediate_jwt) = intermediate_jwt {
72        header.extras.insert(INTERMEDIATE_HEADER, intermediate_jwt);
73    }
74    sign_jwt(client, key_id, &header, &payload).await
75}
76
77/// Signs an intermediate JWT with the root key `root_key_id` that allows license
78/// keys signed by the private half of `public_key_pem` to validate.
79pub async fn make_intermediate_jwt(
80    client: &aws_sdk_kms::Client,
81    root_key_id: &str,
82    validity: Duration,
83    name: String,
84    public_key_pem: String,
85) -> anyhow::Result<String> {
86    let now = SystemTime::now();
87    let expiration = now + validity;
88    let payload = IntermediatePayload {
89        sub: name,
90        exp: format_time(&expiration),
91        nbf: format_time(&now),
92        iss: ISSUER.to_string(),
93        iat: format_time(&now),
94        jti: Uuid::new_v4().to_string(),
95        version: VERSION,
96        public_key: public_key_pem,
97    };
98    let mut header = Header::new(Algorithm::PS256);
99    header.typ = Some(INTERMEDIATE_TYP.to_string());
100    sign_jwt(client, root_key_id, &header, &payload).await
101}
102
103async fn sign_jwt(
104    client: &aws_sdk_kms::Client,
105    key_id: &str,
106    header: &Header,
107    payload: &impl Serialize,
108) -> anyhow::Result<String> {
109    let headers = URL_SAFE_NO_PAD.encode(serde_json::to_string(header).unwrap().as_bytes());
110    let payload = URL_SAFE_NO_PAD.encode(serde_json::to_string(payload).unwrap().as_bytes());
111
112    let signing_string = format!("{}.{}", headers, payload);
113    let signature = URL_SAFE_NO_PAD.encode(sign(client, key_id, signing_string.as_bytes()).await?);
114
115    Ok(format!("{}.{}", signing_string, signature))
116}
117
118async fn get_pubkey(client: &aws_sdk_kms::Client, key_id: &str) -> anyhow::Result<Vec<u8>> {
119    if let Some(pubkey) = client
120        .get_public_key()
121        .key_id(key_id)
122        .send()
123        .await?
124        .public_key
125    {
126        Ok(pubkey.into_inner())
127    } else {
128        Err(anyhow!("failed to get pubkey"))
129    }
130}
131
132async fn sign(
133    client: &aws_sdk_kms::Client,
134    key_id: &str,
135    message: &[u8],
136) -> anyhow::Result<Vec<u8>> {
137    let hash = digest::digest(&digest::SHA256, message);
138    let digest = hash.as_ref().to_vec();
139
140    if let Some(sig) = client
141        .sign()
142        .key_id(key_id)
143        .signing_algorithm(SigningAlgorithmSpec::RsassaPssSha256)
144        .message_type(MessageType::Digest)
145        .message(Blob::new(digest))
146        .send()
147        .await?
148        .signature
149    {
150        Ok(sig.into_inner())
151    } else {
152        Err(anyhow!("failed to get signature"))
153    }
154}
155
156fn format_time(t: &SystemTime) -> u64 {
157    t.duration_since(UNIX_EPOCH).unwrap().as_secs()
158}
mz_license_keys/signing.rs

mz_license_keys/
signing.rs
