Skip to content

Commit a9ab58e

Browse files
JoshLindJoshLind
committed
[Prover] Add simple metrics for JWT attributes.
1 parent bea3613 commit a9ab58e

File tree

3 files changed

+72
-3
lines changed

3 files changed

+72
-3
lines changed

‎keyless-common/src/input_processing/jwt.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,10 @@ impl JwtParts {
129129
String::from(&self.payload)
130130
}
131131

132+
pub fn header_undecoded(&self) -> String {
133+
String::from(&self.header)
134+
}
135+
132136
pub fn header_undecoded_with_dot(&self) -> String {
133137
String::from(&self.header) + "."
134138
}
@@ -150,6 +154,10 @@ impl JwtParts {
150154
pub fn signature(&self) -> Result<RsaSignature> {
151155
RsaSignature::from_b64(&self.signature)
152156
}
157+
158+
pub fn signature_undecoded(&self) -> String {
159+
String::from(&self.signature)
160+
}
153161
}
154162

155163
/// Struct representing the unsigned parts of a JWT with padding

‎prover-service/src/metrics.rs‎

Lines changed: 58 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
use crate::external_resources::prover_config::ProverServiceConfig;
44
use crate::request_handler::handler::is_known_path;
5+
use crate::request_handler::types::VerifiedInput;
56
use aptos_logger::{error, info, warn};
67
use aptos_metrics_core::{
78
exponential_buckets, register_histogram_vec, register_int_counter_vec, Encoder, HistogramVec,
@@ -17,8 +18,6 @@ use std::net::SocketAddr;
1718
use std::sync::Arc;
1819
use std::time::{Duration, Instant};
1920

20-
// TODO: sanity check and expand these metrics!
21-
2221
// Constants for the metrics endpoint and response type
2322
const METRICS_ENDPOINT: &str = "/metrics";
2423
const PLAIN_CONTENT_TYPE: &str = "text/plain";
@@ -39,6 +38,16 @@ pub const PROVER_RESPONSE_GENERATION_LABEL: &str = "prover_response_generation";
3938
pub const VALIDATE_PROVE_REQUEST_LABEL: &str = "validate_prove_request";
4039
pub const WITNESS_GENERATION_LABEL: &str = "witness_generation";
4140

41+
// Useful metric labels for JWT attribute sizes
42+
const JWT_HEADER_SIZE: &str = "jwt_header_size";
43+
const JWT_PAYLOAD_SIZE: &str = "jwt_payload_size";
44+
const JWT_SIGNATURE_SIZE: &str = "jwt_signature_size";
45+
const JWT_ISS_SIZE: &str = "jwt_iss_size";
46+
const JWT_NONCE_SIZE: &str = "jwt_nonce_size";
47+
const JWT_SUB_SIZE: &str = "jwt_sub_size";
48+
const JWT_EMAIL_SIZE: &str = "jwt_email_size";
49+
const JWT_AUD_SIZE: &str = "jwt_aud_size";
50+
4251
// Invalid request path label
4352
const INVALID_PATH: &str = "invalid-path";
4453

@@ -61,6 +70,14 @@ static LATENCY_BUCKETS: Lazy<Vec<f64>> = Lazy::new(|| {
6170
.unwrap()
6271
});
6372

73+
// Buckets for tracking sizes (1 byte to 256 KB)
74+
static SIZE_BUCKETS: Lazy<Vec<f64>> = Lazy::new(|| {
75+
exponential_buckets(
76+
/*start=*/ 1.0, /*factor=*/ 2.0, /*count=*/ 19,
77+
)
78+
.unwrap()
79+
});
80+
6481
// Counter for the number of prover metrics in various states
6582
pub static NUM_TOTAL_METRICS: Lazy<IntCounterVec> = Lazy::new(|| {
6683
register_int_counter_vec!(
@@ -93,6 +110,17 @@ static REQUEST_HANDLING_SECONDS: Lazy<HistogramVec> = Lazy::new(|| {
93110
.unwrap()
94111
});
95112

113+
// Histogram for tracking the attribute sizes of JWT requests
114+
static REQUEST_JWT_ATTRIBUTE_SIZES: Lazy<HistogramVec> = Lazy::new(|| {
115+
register_histogram_vec!(
116+
"keyless_prover_service_request_jwt_attribute_sizes",
117+
"Sizes of request JWT attributes",
118+
&["attribute"],
119+
SIZE_BUCKETS.clone()
120+
)
121+
.unwrap()
122+
});
123+
96124
/// Handles incoming HTTP requests for the metrics server
97125
async fn handle_metrics_request(
98126
request: hyper::Request<Body>,
@@ -227,3 +255,31 @@ pub fn update_request_handling_metrics(
227255
])
228256
.observe(elapsed.as_secs_f64());
229257
}
258+
259+
/// Updates the JWT attribute size metrics for the given attribute and size
260+
fn update_jwt_attribute_size_metrics(attribute: &str, size_bytes: usize) {
261+
REQUEST_JWT_ATTRIBUTE_SIZES
262+
.with_label_values(&[attribute])
263+
.observe(size_bytes as f64);
264+
}
265+
266+
/// Updates the JWT attribute metrics based on the verified input
267+
pub fn update_jwt_attribute_metrics(verified_input: &VerifiedInput) {
268+
// Update the JWT parts metrics
269+
let jwt_parts = &verified_input.jwt_parts;
270+
update_jwt_attribute_size_metrics(JWT_HEADER_SIZE, jwt_parts.header_undecoded().len());
271+
update_jwt_attribute_size_metrics(JWT_PAYLOAD_SIZE, jwt_parts.payload_undecoded().len());
272+
update_jwt_attribute_size_metrics(JWT_SIGNATURE_SIZE, jwt_parts.signature_undecoded().len());
273+
274+
// Update the JWT field metrics
275+
let jwt_payload = &verified_input.jwt.payload;
276+
update_jwt_attribute_size_metrics(JWT_ISS_SIZE, jwt_payload.iss.len());
277+
update_jwt_attribute_size_metrics(JWT_NONCE_SIZE, jwt_payload.nonce.len());
278+
if let Some(sub) = &jwt_payload.sub {
279+
update_jwt_attribute_size_metrics(JWT_SUB_SIZE, sub.len());
280+
}
281+
if let Some(email) = &jwt_payload.email {
282+
update_jwt_attribute_size_metrics(JWT_EMAIL_SIZE, email.len());
283+
}
284+
update_jwt_attribute_size_metrics(JWT_AUD_SIZE, jwt_payload.aud.len());
285+
}

‎prover-service/src/request_handler/prover_handler.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,12 @@ pub async fn hande_prove_request(
4747
// Validate the input request
4848
let verified_input =
4949
match validate_prove_request_input(&prover_service_state, &prove_request_input).await {
50-
Ok(verified_input) => verified_input,
50+
Ok(verified_input) => {
51+
// Update the JWT attribute metrics
52+
metrics::update_jwt_attribute_metrics(&verified_input);
53+
54+
verified_input
55+
}
5156
Err(error) => {
5257
let error_string =
5358
format!("Failed to validate prove request input! Error: {}", error);

0 commit comments

Comments
 (0)