Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 9 additions & 24 deletions dstack/kms/src/main_service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -200,16 +200,14 @@ pub(crate) enum KeyRelease {
AdminHandover,
AppKey,
AppCert,
TempCa,
}

impl KeyRelease {
const ALL: [Self; 5] = [
const ALL: [Self; 4] = [
Self::RpcHandover,
Self::AdminHandover,
Self::AppKey,
Self::AppCert,
Self::TempCa,
];

/// The listener and operation labels.
Expand All @@ -219,7 +217,6 @@ impl KeyRelease {
Self::AdminHandover => ("admin", "kms_export_keys"),
Self::AppKey => ("rpc", "app_get_key"),
Self::AppCert => ("rpc", "app_get_cert"),
Self::TempCa => ("rpc", "temp_ca"),
}
}
}
Expand Down Expand Up @@ -506,28 +503,25 @@ impl RpcHandler {
}

/// Count a key release and log it, regardless of the metrics setting.
/// `boot_info` is the authorized caller, if the operation identifies one.
fn audit_key_release(
&self,
release: KeyRelease,
boot_info: Option<&BootInfo>,
boot_info: &BootInfo,
requested: &str,
returned: &str,
cached: bool,
) {
self.state.metrics().record_key_release(release);
let hex_or_unknown = |bytes: &[u8]| match bytes {
let instance_id = match &boot_info.instance_id[..] {
[] => "UNKNOWN".to_string(),
bytes => hex::encode(bytes),
id => hex::encode(id),
};
let app_id = hex_or_unknown(boot_info.map_or(&[], |info| &info.app_id[..]));
let instance_id = hex_or_unknown(boot_info.map_or(&[], |info| &info.instance_id[..]));
let (listener, operation) = release.labels();
info!(
target: "dstack_kms::audit",
listener,
operation,
app_id = %app_id,
app_id = %hex::encode(&boot_info.app_id),
instance_id = %instance_id,
remote = ?self.remote_endpoint,
requested_epoch = requested,
Expand Down Expand Up @@ -555,7 +549,7 @@ impl RpcHandler {
let info = self.ensure_kms_allowed(&request.vm_config).await?;
ensure_key_release_allowed(&info, (&self.state.config).into())?;
let returned = format!("root 0..={}", self.state.keyring.latest());
self.audit_key_release(release, Some(&info), "all", &returned, cached);
self.audit_key_release(release, &info, "all", &returned, cached);
Ok(KmsKeyResponse {
temp_ca_key: self.state.temp_ca.key.serialize_pem(),
keys: self
Expand Down Expand Up @@ -816,7 +810,7 @@ impl KmsRpc for RpcHandler {
let epoch = |epoch: Option<u32>| epoch.map_or("current".into(), |e| e.to_string());
self.audit_key_release(
KeyRelease::AppKey,
Some(&boot_info),
&boot_info,
&format!(
"root {} app {}",
epoch(request.root_epoch),
Expand Down Expand Up @@ -960,7 +954,6 @@ impl KmsRpc for RpcHandler {
}

async fn get_temp_ca_cert(self) -> Result<GetTempCaCertResponse> {
let cached = self.self_allowed_cached(Operation::KmsServe);
let epoch = self.ensure_self_allowed(Operation::KmsServe).await?;
if self.state.config.enforce_self_authorization {
let boot_info = self
Expand All @@ -971,18 +964,10 @@ impl KmsRpc for RpcHandler {
.context("missing authorized KMS boot info")?;
ensure_key_release_allowed(&boot_info, (&self.state.config).into())?;
}
let ca_cert = self.state.epoch_root(epoch)?.ca.pem_cert.clone();
self.audit_key_release(
KeyRelease::TempCa,
None,
"current",
&format!("root {epoch}"),
cached,
);
Ok(GetTempCaCertResponse {
temp_ca_cert: self.state.temp_ca.pem_cert.clone(),
temp_ca_key: self.state.temp_ca.key.serialize_pem(),
ca_cert,
ca_cert: self.state.epoch_root(epoch)?.ca.pem_cert.clone(),
})
}

Expand Down Expand Up @@ -1038,7 +1023,7 @@ impl KmsRpc for RpcHandler {
.context("Failed to sign certificate")?;
self.audit_key_release(
KeyRelease::AppCert,
Some(&app_info.boot_info),
&app_info.boot_info,
"current",
&format!("root {} app {}", app_info.root_epoch, app_info.app_epoch),
cached,
Expand Down
Loading