kms: add prometheus metrics endpoint

Author: kvinwang

kms/src/main.rs

@@ -10,8 +10,8 @@ use ra_rpc::rocket_helper::QuoteVerifier;
 use rocket::{
     fairing::AdHoc,
     figment::{providers::Serialized, Figment},
-    response::content::RawHtml,
-    Shutdown,
+    response::content::{RawHtml, RawText},
+    Shutdown, State,
 };
 use tracing::{info, warn};
 
@@ -77,6 +77,34 @@ async fn run_onboard_service(kms_config: KmsConfig, figment: Figment) -> Result<
     Ok(())
 }
 
+#[rocket::get("/metrics")]
+fn metrics(state: &State<KmsState>) -> RawText<String> {
+    RawText(state.metrics().render_prometheus())
+}
+
+// Count only RPCs whose primary job is to verify caller/app attestation.
+// Recording in a response fairing also catches failures that happen before
+// RpcHandler is constructed, such as malformed RA-TLS attestation.
+fn is_attestation_rpc_path(path: &str) -> bool {
+    let Some(method) = path.strip_prefix("/prpc/") else {
+        return false;
+    };
+    let method = method.trim_start_matches("KMS.");
+    matches!(method, "GetAppKey" | "GetKmsKey" | "SignCert")
+}
+
+fn record_attestation_metrics(req: &rocket::Request<'_>, res: &rocket::Response<'_>) {
+    if !is_attestation_rpc_path(req.uri().path().as_str()) {
+        return;
+    }
+    let Some(state) = req.rocket().state::<KmsState>() else {
+        return;
+    };
+    state
+        .metrics()
+        .record_attestation_request(res.status().code >= 400);
+}
+
 #[rocket::main]
 async fn main() -> Result<()> {
     {
@@ -119,10 +147,15 @@ async fn main() -> Result<()> {
                 res.set_raw_header("X-App-Version", app_version());
             })
         }))
+        .attach(AdHoc::on_response(
+            "Record KMS attestation metrics",
+            |req, res| Box::pin(async move { record_attestation_metrics(req, res) }),
+        ))
         .mount(
             "/prpc",
             ra_rpc::prpc_routes!(KmsState, RpcHandler, trim: "KMS."),
         )
+        .mount("/", rocket::routes![metrics])
         .manage(state);
 
     let verifier = QuoteVerifier::new(pccs_url);

kms/src/main_service.rs

@@ -4,7 +4,10 @@
 
 use std::{
     path::{Path, PathBuf},
-    sync::Arc,
+    sync::{
+        atomic::{AtomicU64, Ordering},
+        Arc,
+    },
 };
 
 use anyhow::{bail, Context, Result};
@@ -57,6 +60,38 @@ pub struct KmsStateInner {
     temp_ca_key: String,
     verifier: CvmVerifier,
     self_boot_info: OnceCell<BootInfo>,
+    metrics: KmsMetrics,
+}
+
+#[derive(Default)]
+pub(crate) struct KmsMetrics {
+    attestation_requests_total: AtomicU64,
+    attestation_failures_total: AtomicU64,
+}
+
+impl KmsMetrics {
+    pub(crate) fn record_attestation_request(&self, failed: bool) {
+        self.attestation_requests_total
+            .fetch_add(1, Ordering::Relaxed);
+        if failed {
+            self.attestation_failures_total
+                .fetch_add(1, Ordering::Relaxed);
+        }
+    }
+
+    pub(crate) fn render_prometheus(&self) -> String {
+        let attestation_requests_total = self.attestation_requests_total.load(Ordering::Relaxed);
+        let attestation_failures_total = self.attestation_failures_total.load(Ordering::Relaxed);
+
+        format!(
+            "# HELP dstack_kms_attestation_requests_total Total number of KMS attestation requests.\n\
+             # TYPE dstack_kms_attestation_requests_total counter\n\
+             dstack_kms_attestation_requests_total {attestation_requests_total}\n\
+             # HELP dstack_kms_attestation_failures_total Total number of failed KMS attestation requests.\n\
+             # TYPE dstack_kms_attestation_failures_total counter\n\
+             dstack_kms_attestation_failures_total {attestation_failures_total}\n"
+        )
+    }
 }
 
 impl KmsState {
@@ -77,7 +112,9 @@ impl KmsState {
             config.pccs_url.clone(),
         );
         if !config.enforce_self_authorization {
-            warn!("self-authorization is disabled; trusted RPCs will not be gated by KMS self-attestation - do not use in production TEE deployments");
+            warn!(
+                "self-authorization is disabled; trusted RPCs will not be gated by KMS self-attestation - do not use in production TEE deployments"
+            );
         }
         Ok(Self {
             inner: Arc::new(KmsStateInner {
@@ -88,9 +125,14 @@ impl KmsState {
                 temp_ca_key,
                 verifier,
                 self_boot_info: OnceCell::new(),
+                metrics: KmsMetrics::default(),
             }),
         })
     }
+
+    pub(crate) fn metrics(&self) -> &KmsMetrics {
+        &self.inner.metrics
+    }
 }
 
 pub struct RpcHandler {