Upload diagnostic job results through support bundles (#7831)

feat(connect): upload diagnostic job results
This commit is contained in:
Chris
2026-09-14 12:58:00 +08:00
committed by GitHub
parent f59e6be338
commit c842eed84c
3 changed files with 206 additions and 28 deletions
+26
View File
@@ -79,6 +79,14 @@ pub struct DiagnosticJobSignature {
value: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct DiagnosticJobAuthorization {
actor_type: String,
actor_name: String,
request_id: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct DiagnosticJobEnvelope {
@@ -93,6 +101,7 @@ pub struct DiagnosticJobEnvelope {
pub expire_time: String,
pub nonce: String,
pub required_capabilities: Vec<String>,
pub authorization: DiagnosticJobAuthorization,
pub limits: DiagnosticJobLimits,
pub parameters: DiagnosticJobParameters,
pub signature: DiagnosticJobSignature,
@@ -112,6 +121,7 @@ struct UnsignedDiagnosticJob<'a> {
expire_time: &'a str,
nonce: &'a str,
required_capabilities: &'a [String],
authorization: &'a DiagnosticJobAuthorization,
limits: &'a DiagnosticJobLimits,
parameters: &'a DiagnosticJobParameters,
}
@@ -285,6 +295,7 @@ impl DiagnosticJobEnvelope {
expire_time: &self.expire_time,
nonce: &self.nonce,
required_capabilities: &self.required_capabilities,
authorization: &self.authorization,
limits: &self.limits,
parameters: &self.parameters,
}
@@ -315,6 +326,10 @@ impl DiagnosticJobEnvelope {
if !uuid7(&self.job_id)
|| !uuid7(&self.parameters.artifact_uid)
|| !uuid7(&self.parameters.consent_uid)
|| self.authorization.actor_type != "BROWSER_USER"
|| !self.authorization.actor_name.strip_prefix("users/").is_some_and(uuid7)
|| !Uuid::parse_str(&self.authorization.request_id)
.is_ok_and(|value| value.get_version_num() == 4 && value.to_string() == self.authorization.request_id)
|| self.parameters.consent_policy_revision == 0
{
return Err(DiagnosticJobError::Invalid);
@@ -452,6 +467,11 @@ mod tests {
expire_time: "2030-01-01T00:00:30Z".to_owned(),
nonce: URL_SAFE_NO_PAD.encode_to_string([7_u8; 32]),
required_capabilities: vec![CPU_PROFILE_CAPABILITY.to_owned()],
authorization: DiagnosticJobAuthorization {
actor_type: "BROWSER_USER".to_owned(),
actor_name: "users/018cc251-f400-7abc-8def-0123456789ab".to_owned(),
request_id: "123e4567-e89b-42d3-a456-426614174001".to_owned(),
},
limits: DiagnosticJobLimits {
timeout_seconds: 30,
max_output_bytes: 524_288,
@@ -521,6 +541,12 @@ mod tests {
signer.verify(&envelope, &target(&envelope), "2030-01-01T00:00:30Z".parse().expect("time")),
Err(DiagnosticJobError::Expired)
);
let (mut actor_tampered, signer) = signed();
actor_tampered.authorization.actor_name.push('0');
assert_eq!(
signer.verify(&actor_tampered, &target(&actor_tampered), "2030-01-01T00:00:10Z".parse().expect("time")),
Err(DiagnosticJobError::Invalid)
);
}
#[test]
+178 -28
View File
@@ -22,7 +22,6 @@ use std::time::Duration;
#[cfg(unix)]
use std::os::unix::fs::{MetadataExt as _, OpenOptionsExt as _, PermissionsExt as _};
use base64_simd::URL_SAFE_NO_PAD;
use chrono::Utc;
use serde::{Deserialize, Serialize};
use sha2::{Digest as _, Sha256};
@@ -34,6 +33,7 @@ use super::{
execute_diagnostic_job,
};
use crate::connect::config::HeartbeatConfig;
use crate::connect::report_upload::ReportUploadClient;
use crate::connect::telemetry::{TelemetryDelivery, TelemetryTransport};
const PROTOCOL_VERSION: &str = "v1";
@@ -56,6 +56,7 @@ pub(crate) struct DiagnosticJobRuntime {
enum JobState {
Active,
Completed(DiagnosticJobExecution),
Uploaded(UploadedDiagnosticJobResult),
Delivered,
}
@@ -64,6 +65,16 @@ struct JobStateStore {
directory: PathBuf,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
struct UploadedDiagnosticJobResult {
job_id: String,
outcome: String,
reason: String,
artifact_name: Option<String>,
artifact_sha256: Option<String>,
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct DiagnosticJobResultRequest<'a> {
@@ -72,10 +83,8 @@ struct DiagnosticJobResultRequest<'a> {
job_id: &'a str,
outcome: &'a str,
reason: &'a str,
artifact_uid: Option<&'a str>,
artifact_name: Option<&'a str>,
artifact_sha256: Option<&'a str>,
artifact_encoding: Option<&'static str>,
artifact: Option<String>,
}
impl DiagnosticJobRuntime {
@@ -106,6 +115,16 @@ impl DiagnosticJobRuntime {
match states.get(&job_id) {
Some(JobState::Active) => return,
Some(JobState::Completed(result)) => Some(result.clone()),
Some(JobState::Uploaded(result)) => {
let result = result.clone();
drop(states);
let runtime = self.clone();
let cancel = shutdown.child_token();
tokio::spawn(async move {
runtime.deliver_uploaded(result, &cancel).await;
});
return;
}
Some(JobState::Delivered) => return,
None => match self.store.load(&job_id) {
Ok(Some(JobState::Completed(result))) => {
@@ -119,6 +138,19 @@ impl DiagnosticJobRuntime {
states.insert(job_id, JobState::Delivered);
return;
}
Ok(Some(JobState::Uploaded(result))) => {
if !valid_uploaded_result(&result, &job_id) {
return;
}
states.insert(job_id.clone(), JobState::Uploaded(result.clone()));
drop(states);
let runtime = self.clone();
let cancel = shutdown.child_token();
tokio::spawn(async move {
runtime.deliver_uploaded(result, &cancel).await;
});
return;
}
Ok(Some(JobState::Active)) => {
let result = failed_execution(&job_id, "INTERRUPTED");
if self.store.save(&job_id, &JobState::Completed(result.clone())).is_err() {
@@ -159,14 +191,21 @@ impl DiagnosticJobRuntime {
result
}
};
if deliver_result(&runtime.config, &result, &cancel).await
&& runtime.store.save(&job_id, &JobState::Delivered).is_ok()
&& let Ok(mut states) = runtime.states.lock()
{
states.insert(job_id, JobState::Delivered);
if let Some(uploaded) = prepare_result(&runtime, result, &cancel).await {
runtime.deliver_uploaded(uploaded, &cancel).await;
}
});
}
async fn deliver_uploaded(&self, result: UploadedDiagnosticJobResult, cancel: &CancellationToken) {
let job_id = result.job_id.clone();
if deliver_result(&self.config, &result, cancel).await
&& self.store.save(&job_id, &JobState::Delivered).is_ok()
&& let Ok(mut states) = self.states.lock()
{
states.insert(job_id, JobState::Delivered);
}
}
}
impl JobStateStore {
@@ -372,6 +411,13 @@ fn enabled_build_features() -> Vec<String> {
}
fn failed_execution(job_id: &str, reason: &'static str) -> DiagnosticJobExecution {
let reason = if reason == "CANCELLED" {
"CANCELLED"
} else if reason == "LIMIT_EXCEEDED" {
"LIMIT_EXCEEDED"
} else {
"COLLECTION_FAILED"
};
DiagnosticJobExecution {
job_id: job_id.to_owned(),
outcome: if reason == "CANCELLED" { "CANCELLED" } else { "FAILED" }.to_owned(),
@@ -382,7 +428,64 @@ fn failed_execution(job_id: &str, reason: &'static str) -> DiagnosticJobExecutio
}
}
async fn deliver_result(config: &HeartbeatConfig, result: &DiagnosticJobExecution, cancel: &CancellationToken) -> bool {
async fn prepare_result(
runtime: &DiagnosticJobRuntime,
result: DiagnosticJobExecution,
cancel: &CancellationToken,
) -> Option<UploadedDiagnosticJobResult> {
let uploaded = if let (Some(bytes), true) =
(result.artifact_bytes.as_ref(), matches!(result.outcome.as_str(), "SUCCEEDED" | "PARTIAL"))
{
if runtime.store.ensure_directory().is_err() {
return None;
}
let path = runtime
.store
.directory
.join(format!(".{}.{}.artifact", result.job_id, Uuid::new_v4()));
let mut options = OpenOptions::new();
options.write(true).create_new(true);
#[cfg(unix)]
options.mode(0o600);
let mut file = options.open(&path).ok()?;
if file.write_all(bytes).and_then(|()| file.sync_all()).is_err() {
let _ = fs::remove_file(&path);
return None;
}
drop(file);
let receipt = match ReportUploadClient::new(runtime.config.clone(), Duration::from_secs(15 * 60)) {
Ok(client) => client.upload(&path, cancel).await.ok(),
Err(_) => None,
};
let _ = fs::remove_file(path);
let receipt = receipt?;
UploadedDiagnosticJobResult {
job_id: result.job_id,
outcome: result.outcome,
reason: result.reason,
artifact_name: Some(receipt.name),
artifact_sha256: Some(receipt.declared_sha256),
}
} else {
UploadedDiagnosticJobResult {
job_id: result.job_id,
outcome: result.outcome,
reason: result.reason,
artifact_name: None,
artifact_sha256: None,
}
};
let state = JobState::Uploaded(uploaded.clone());
if runtime.store.save(&uploaded.job_id, &state).is_err() {
return None;
}
if let Ok(mut states) = runtime.states.lock() {
states.insert(uploaded.job_id.clone(), state);
}
Some(uploaded)
}
async fn deliver_result(config: &HeartbeatConfig, result: &UploadedDiagnosticJobResult, cancel: &CancellationToken) -> bool {
let Ok(transport) = TelemetryTransport::new(config.clone()) else {
return false;
};
@@ -393,13 +496,8 @@ async fn deliver_result(config: &HeartbeatConfig, result: &DiagnosticJobExecutio
job_id: &result.job_id,
outcome: &result.outcome,
reason: &result.reason,
artifact_uid: result.artifact_uid.as_deref(),
artifact_name: result.artifact_name.as_deref(),
artifact_sha256: result.artifact_sha256.as_deref(),
artifact_encoding: result.artifact_bytes.as_ref().map(|_| "base64url"),
artifact: result
.artifact_bytes
.as_ref()
.map(|bytes| URL_SAFE_NO_PAD.encode_to_string(bytes)),
};
for attempt in 0..DELIVERY_ATTEMPTS {
if cancel.is_cancelled() {
@@ -421,14 +519,38 @@ async fn deliver_result(config: &HeartbeatConfig, result: &DiagnosticJobExecutio
false
}
fn valid_execution(result: &DiagnosticJobExecution, job_id: &str) -> bool {
let valid_outcome = matches!(result.outcome.as_str(), "SUCCEEDED" | "PARTIAL" | "FAILED" | "UNSUPPORTED" | "CANCELLED");
let valid_reason = !result.reason.is_empty()
&& result.reason.len() <= 64
&& result
.reason
fn valid_uploaded_result(result: &UploadedDiagnosticJobResult, job_id: &str) -> bool {
result.job_id == job_id
&& valid_terminal_fields(&result.outcome, &result.reason)
&& match (&result.artifact_name, &result.artifact_sha256) {
(Some(name), Some(digest)) => {
matches!(result.outcome.as_str(), "SUCCEEDED" | "PARTIAL")
&& name.len() <= 512
&& name.starts_with("organizations/")
&& lower_hex(digest, 64)
}
(None, None) => matches!(result.outcome.as_str(), "FAILED" | "UNSUPPORTED" | "CANCELLED"),
_ => false,
}
}
fn valid_terminal_fields(outcome: &str, reason: &str) -> bool {
matches!(outcome, "SUCCEEDED" | "PARTIAL" | "FAILED" | "UNSUPPORTED" | "CANCELLED")
&& !reason.is_empty()
&& reason.len() <= 64
&& reason
.bytes()
.all(|byte| byte.is_ascii_uppercase() || byte.is_ascii_digit() || byte == b'_');
.all(|byte| byte.is_ascii_uppercase() || byte.is_ascii_digit() || byte == b'_')
}
fn lower_hex(value: &str, length: usize) -> bool {
value.len() == length
&& value
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
}
fn valid_execution(result: &DiagnosticJobExecution, job_id: &str) -> bool {
let valid_artifact = match (&result.artifact_uid, &result.artifact_sha256, &result.artifact_bytes) {
(Some(uid), Some(digest), Some(bytes)) => {
Uuid::parse_str(uid).is_ok_and(|value| value.get_version_num() == 7 && value.to_string() == *uid)
@@ -439,7 +561,7 @@ fn valid_execution(result: &DiagnosticJobExecution, job_id: &str) -> bool {
(None, None, None) => matches!(result.outcome.as_str(), "FAILED" | "CANCELLED") && result.reason != "COMPLETE",
_ => false,
};
result.job_id == job_id && valid_outcome && valid_reason && valid_artifact
result.job_id == job_id && valid_terminal_fields(&result.outcome, &result.reason) && valid_artifact
}
#[cfg(test)]
@@ -465,14 +587,21 @@ mod tests {
job_id: "018cc251-f400-7abc-8def-0123456789ab",
outcome: "FAILED",
reason: "CANCELLED",
artifact_uid: None,
artifact_name: None,
artifact_sha256: None,
artifact_encoding: None,
artifact: None,
};
let value = serde_json::to_value(request).expect("result request");
assert_eq!(value["jobId"], "018cc251-f400-7abc-8def-0123456789ab");
for forbidden in ["command", "script", "path", "sql", "arguments"] {
for forbidden in [
"command",
"script",
"path",
"sql",
"arguments",
"artifact",
"artifactUid",
"artifactEncoding",
] {
assert!(value.get(forbidden).is_none());
}
}
@@ -498,6 +627,27 @@ mod tests {
assert_eq!(loaded, expected);
assert!(valid_execution(&loaded, &job_id));
let uploaded = UploadedDiagnosticJobResult {
job_id: job_id.clone(),
outcome: "SUCCEEDED".to_owned(),
reason: "COMPLETE".to_owned(),
artifact_name: Some(format!(
"organizations/{}/clusters/{}/supportBundles/{}",
Uuid::now_v7(),
Uuid::now_v7(),
Uuid::now_v7()
)),
artifact_sha256: Some("a".repeat(64)),
};
store
.save(&job_id, &JobState::Uploaded(uploaded.clone()))
.expect("uploaded state");
let loaded = match store.load(&job_id).expect("load state") {
Some(JobState::Uploaded(result)) => result,
_ => panic!("uploaded state expected"),
};
assert!(valid_uploaded_result(&loaded, &job_id));
store.save(&job_id, &JobState::Delivered).expect("delivered state");
assert!(matches!(store.load(&job_id), Ok(Some(JobState::Delivered))));
}
+2
View File
@@ -108,6 +108,7 @@ impl PendingHeartbeat {
DiagnosticCollectionPolicy::policy_sync_capability(),
ENVIRONMENT_CAPABILITY,
"jobs",
super::diagnostics::CPU_PROFILE_CAPABILITY,
])
&& self.sequence <= MAX_SEQUENCE
&& self.coarse_node_summary.is_valid()
@@ -262,6 +263,7 @@ impl HeartbeatStateStore {
];
if self.job_capable {
capabilities.push("jobs".to_owned());
capabilities.push(super::diagnostics::CPU_PROFILE_CAPABILITY.to_owned());
}
let pending = PendingHeartbeat {
protocol_version: PROTOCOL_VERSION.to_owned(),