feat(connect): capture service runtime profiles over local IPC (#8168)

feat(connect): export service runtime profiles over local IPC
This commit is contained in:
Chris
2026-09-28 00:49:45 +08:00
committed by GitHub
parent e33542b0c2
commit acebc9649d
8 changed files with 749 additions and 91 deletions
+4
View File
@@ -1043,6 +1043,10 @@ pub struct ConnectProfileOpts {
#[arg(long = "thread-scope", value_enum)]
pub thread_scope: Option<ConnectThreadProfileScope>,
/// SHA-256 of the enrolled offline public key; required for service runtime capture
#[arg(long = "offline-key-id")]
pub offline_key_id: Option<String>,
/// Negotiated producer schema version
#[arg(long = "schema-version", default_value_t = 1)]
pub schema_version: u16,
@@ -385,7 +385,7 @@ fn target_and_identity(config: &HeartbeatConfig) -> Result<(DiagnosticJobTarget,
))
}
async fn executable_provenance() -> Result<ProfileProvenance, &'static str> {
pub(super) async fn executable_provenance() -> Result<ProfileProvenance, &'static str> {
let digest = tokio::task::spawn_blocking(hash_current_executable)
.await
.map_err(|_| "PROVENANCE_FAILED")?
+3 -1
View File
@@ -131,6 +131,7 @@ pub use perf_site_replication::{
SiteReplicationTargetResult, measure_site_replication, read_protected_site_replication_credential,
save_signed_site_replication_export, sign_site_replication_export, validate_site_replication_limits,
};
pub(crate) use profile_cpu::LocalRuntimeProfileRequest;
pub use profile_cpu::{
CPU_PROFILE_CAPABILITY, LocalProfileConsent, MAX_PROFILE_DURATION, MEMORY_PROFILE_CAPABILITY, PROFILE_SCHEMA_VERSION,
ProfileCaptureRequest, ProfileData, ProfileError, ProfileOutcome, ProfileProvenance, ProfileReasonCode, ProfileResult,
@@ -171,5 +172,6 @@ pub use trace_record::{
};
pub use trace_replay::{LocallyReviewedTraceArtifact, ReplayedTrace, TraceReplayError, replay_trace, replay_trace_result};
pub(crate) use trace_runtime::{
LocalTraceCaptureError, LocalTraceCaptureRuntime, request_local_trace_capture, spawn_local_trace_capture_runtime,
LocalTraceCaptureError, LocalTraceCaptureRuntime, request_local_runtime_profile, request_local_trace_capture,
spawn_local_trace_capture_runtime,
};
@@ -244,6 +244,27 @@ pub struct LocalProfileConsent {
pub confirmed: bool,
}
/// Owner-authorized request for the running service, never a CLI process profile.
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub(crate) struct LocalRuntimeProfileRequest {
pub offline_key_id: String,
pub organization_name: String,
pub cluster_name: String,
pub device_name: String,
pub run_uid: String,
pub artifact_uid: String,
pub schema_version: u16,
pub capability: String,
pub consent_uid: String,
pub policy_revision: u64,
pub consent_expires_at_unix: i64,
pub acknowledge_l3: bool,
pub expires_at_unix: i64,
pub duration_millis: u64,
pub sample_period_micros: u64,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ProfileCaptureRequest {
pub organization_name: String,
+603 -24
View File
@@ -12,14 +12,28 @@
// See the License for the specific language governing permissions and
// limitations under the License.
//! Owner-only local transport between the telemetry CLI and the running server.
//! Owner-only local transport between diagnostic CLI commands and the running server.
//!
//! Version 1 accepts only TRACE_RECORD and RUNTIME_PROFILE. Runtime requests select
//! an existing offline key by SPKI digest; this is not proof of Connect enrollment.
//! The receiver checks enrollment, target ownership and consent at import. The
//! server owns provenance, nonce generation, capture and signing; the CLI receives
//! only a bounded archive and saves it without uploading. No paths cross this IPC.
//! The write half stays open during collection. EOF cancels the server collector;
//! cancellation is acknowledged only after joining it and releasing its lease.
use std::io;
use std::os::unix::fs::{FileTypeExt as _, MetadataExt as _, PermissionsExt as _};
use std::path::Path;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use super::profile_cpu::{
LocalProfileConsent, LocalRuntimeProfileRequest, MAX_ARCHIVE_BYTES, MAX_PROFILE_DURATION, ProfileCaptureRequest,
ProfileError, ProfileOutcome, ProfileReasonCode, ProfileTool, SignedProfileExport, encode_signed_profile_export,
};
use base64_simd::URL_SAFE_NO_PAD;
use serde::{Deserialize, Serialize};
use sha2::{Digest as _, Sha256};
use thiserror::Error;
use tokio::io::{AsyncBufReadExt as _, AsyncReadExt as _, AsyncWriteExt as _, BufReader};
use tokio::net::{UnixListener, UnixStream};
@@ -30,7 +44,8 @@ use super::{LocalTelemetryConsent, TelemetryProducerError, TraceRecordCapture, T
const SOCKET_FILE: &str = "telemetry-record.sock";
const PROTOCOL_VERSION: u16 = 1;
const MAX_REQUEST_BYTES: u64 = 1_024;
const MAX_REQUEST_BYTES: u64 = 8_192;
const MAX_RUNTIME_RESPONSE_BYTES: u64 = (MAX_ARCHIVE_BYTES as u64).div_ceil(3) * 4 + 1_024;
const MAX_RESPONSE_BYTES: u64 = 300_000;
const MAX_CONNECTIONS: usize = 8;
const REQUEST_TIMEOUT: Duration = Duration::from_secs(1);
@@ -49,6 +64,10 @@ pub(crate) enum LocalTraceCaptureError {
Io(#[source] io::Error),
#[error(transparent)]
Producer(#[from] TelemetryProducerError),
#[error("runtime profile request rejected: {0}")]
RuntimeProfile(String),
#[error("runtime profile cancellation was not acknowledged")]
CancellationUnconfirmed,
}
pub(crate) struct LocalTraceCaptureRuntime {
@@ -72,12 +91,23 @@ impl Drop for LocalTraceCaptureRuntime {
}
#[derive(Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
struct CaptureRequest {
protocol_version: u16,
consent_expires_at_unix: i64,
duration_millis: u64,
max_spans: usize,
#[serde(
deny_unknown_fields,
tag = "operation",
rename_all = "SCREAMING_SNAKE_CASE",
rename_all_fields = "camelCase"
)]
enum CaptureRequest {
TraceRecord {
protocol_version: u16,
consent_expires_at_unix: i64,
duration_millis: u64,
max_spans: usize,
},
RuntimeProfile {
protocol_version: u16,
request: LocalRuntimeProfileRequest,
},
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
@@ -124,10 +154,27 @@ impl From<ProducerErrorCode> for TelemetryProducerError {
}
#[derive(Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields, tag = "status", rename_all = "SCREAMING_SNAKE_CASE")]
#[serde(
deny_unknown_fields,
tag = "status",
rename_all = "SCREAMING_SNAKE_CASE",
rename_all_fields = "camelCase"
)]
enum CaptureResponse {
Ok { capture: TraceRecordCapture },
Error { code: ProducerErrorCode },
Ok {
capture: TraceRecordCapture,
},
Error {
code: ProducerErrorCode,
},
RuntimeOk {
archive_base64: String,
archive_sha256: String,
artifact_uid: String,
},
RuntimeError {
code: RuntimeErrorCode,
},
}
pub(crate) fn spawn_local_trace_capture_runtime(
@@ -140,8 +187,9 @@ pub(crate) fn spawn_local_trace_capture_runtime(
let socket_identity = socket_identity(&socket_path, owner)?;
let shutdown = parent_shutdown.child_token();
let task_shutdown = shutdown.clone();
let state_root = state_root.to_path_buf();
let task = tokio::spawn(async move {
run_listener(listener, owner, task_shutdown).await;
run_listener(listener, owner, state_root, task_shutdown).await;
remove_own_socket(&socket_path, socket_identity);
});
Ok(LocalTraceCaptureRuntime {
@@ -164,8 +212,11 @@ pub(crate) async fn request_local_trace_capture(
io::ErrorKind::NotFound | io::ErrorKind::ConnectionRefused => LocalTraceCaptureError::RuntimeUnavailable,
_ => LocalTraceCaptureError::Io(error),
})?;
if !stream.peer_cred().is_ok_and(|credentials| credentials.uid() == owner) {
return Err(LocalTraceCaptureError::StateSecurity);
}
let (reader, mut writer) = stream.into_split();
let request = CaptureRequest {
let request = CaptureRequest::TraceRecord {
protocol_version: PROTOCOL_VERSION,
consent_expires_at_unix,
duration_millis: u64::try_from(limits.duration.as_millis()).map_err(|_| TelemetryProducerError::InvalidDuration)?,
@@ -173,7 +224,7 @@ pub(crate) async fn request_local_trace_capture(
};
let mut encoded = serde_json::to_vec(&request).map_err(|_| LocalTraceCaptureError::Protocol)?;
encoded.push(b'\n');
if encoded.len() as u64 > MAX_REQUEST_BYTES {
if encoded.len() > 1_024 {
return Err(LocalTraceCaptureError::Protocol);
}
writer.write_all(&encoded).await.map_err(LocalTraceCaptureError::Io)?;
@@ -198,10 +249,287 @@ pub(crate) async fn request_local_trace_capture(
match response {
CaptureResponse::Ok { capture } => Ok(capture),
CaptureResponse::Error { code } => Err(TelemetryProducerError::from(code).into()),
_ => Err(LocalTraceCaptureError::Protocol),
}
}
async fn run_listener(listener: UnixListener, owner: u32, shutdown: CancellationToken) {
#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
enum RuntimeErrorCode {
InvalidRequest,
IdentityUnavailable,
ConsentRequired,
ConsentExpired,
Expired,
LimitExceeded,
Busy,
Cancelled,
TimedOut,
CollectionFailed,
}
impl From<ProfileError> for RuntimeErrorCode {
fn from(error: ProfileError) -> Self {
match error {
ProfileError::ConsentRequired => Self::ConsentRequired,
ProfileError::ConsentExpired => Self::ConsentExpired,
ProfileError::Expired => Self::Expired,
ProfileError::LimitExceeded => Self::LimitExceeded,
ProfileError::Busy => Self::Busy,
ProfileError::Cancelled => Self::Cancelled,
ProfileError::TimedOut => Self::TimedOut,
ProfileError::InvalidRequest | ProfileError::UnsupportedCapability | ProfileError::UnsupportedVersion => {
Self::InvalidRequest
}
_ => Self::CollectionFailed,
}
}
}
pub(crate) async fn request_local_runtime_profile(
state_root: &Path,
request: LocalRuntimeProfileRequest,
cancel: &CancellationToken,
) -> Result<SignedProfileExport, LocalTraceCaptureError> {
let owner = private_state_owner(state_root)?;
let socket_path = state_root.join(SOCKET_FILE);
socket_identity(&socket_path, owner)?;
let stream = UnixStream::connect(&socket_path).await.map_err(LocalTraceCaptureError::Io)?;
if !stream.peer_cred().is_ok_and(|credentials| credentials.uid() == owner) {
return Err(LocalTraceCaptureError::StateSecurity);
}
let artifact_uid = request.artifact_uid.clone();
let message = CaptureRequest::RuntimeProfile {
protocol_version: PROTOCOL_VERSION,
request,
};
let mut bytes = serde_json::to_vec(&message).map_err(|_| LocalTraceCaptureError::Protocol)?;
bytes.push(b'\n');
if bytes.len() as u64 > MAX_REQUEST_BYTES {
return Err(LocalTraceCaptureError::Protocol);
}
let (reader, mut writer) = stream.into_split();
tokio::time::timeout(REQUEST_TIMEOUT, writer.write_all(&bytes))
.await
.map_err(|_| LocalTraceCaptureError::Protocol)?
.map_err(LocalTraceCaptureError::Io)?;
// Keep the write half open: EOF is the server's cancellation signal.
let response = async {
let mut bytes = Vec::new();
reader
.take(MAX_RUNTIME_RESPONSE_BYTES + 1)
.read_to_end(&mut bytes)
.await
.map_err(LocalTraceCaptureError::Io)?;
if bytes.is_empty() || bytes.len() as u64 > MAX_RUNTIME_RESPONSE_BYTES {
return Err(LocalTraceCaptureError::Protocol);
}
serde_json::from_slice::<CaptureResponse>(&bytes).map_err(|_| LocalTraceCaptureError::Protocol)
};
let response = tokio::time::timeout(MAX_PROFILE_DURATION + Duration::from_secs(5), response);
tokio::pin!(response);
let response = tokio::select! {
biased;
_ = cancel.cancelled() => {
writer.shutdown().await.map_err(|_| LocalTraceCaptureError::CancellationUnconfirmed)?;
// The service closes the response only after the blocking collector has joined.
let acknowledged = matches!(tokio::time::timeout(Duration::from_secs(2), &mut response).await,
Ok(Ok(Ok(CaptureResponse::RuntimeError { .. } | CaptureResponse::RuntimeOk { .. }))));
if !acknowledged { return Err(LocalTraceCaptureError::CancellationUnconfirmed); }
return Err(LocalTraceCaptureError::RuntimeProfile("CANCELLED".to_owned()));
}
response = &mut response => response.map_err(|_| LocalTraceCaptureError::Protocol)??,
};
match response {
CaptureResponse::RuntimeOk {
archive_base64,
archive_sha256,
artifact_uid: returned_uid,
} => {
let archive_bytes = URL_SAFE_NO_PAD
.decode_to_vec(&archive_base64)
.map_err(|_| LocalTraceCaptureError::Protocol)?;
if archive_bytes.is_empty()
|| archive_bytes.len() > MAX_ARCHIVE_BYTES
|| returned_uid != artifact_uid
|| URL_SAFE_NO_PAD.encode_to_string(&archive_bytes) != archive_base64
|| hex_simd::encode_to_string(Sha256::digest(&archive_bytes), hex_simd::AsciiCase::Lower) != archive_sha256
{
return Err(LocalTraceCaptureError::Protocol);
}
Ok(SignedProfileExport {
artifact_uid,
tool: ProfileTool::Threads,
outcome: ProfileOutcome::Succeeded,
reason_code: ProfileReasonCode::Complete,
archive_bytes,
archive_sha256,
})
}
CaptureResponse::RuntimeError { code } => Err(LocalTraceCaptureError::RuntimeProfile(format!("{code:?}"))),
_ => Err(LocalTraceCaptureError::Protocol),
}
}
async fn handle_runtime_profile(
mut reader: BufReader<tokio::net::unix::OwnedReadHalf>,
mut writer: tokio::net::unix::OwnedWriteHalf,
state_root: &Path,
protocol_version: u16,
request: LocalRuntimeProfileRequest,
shutdown: CancellationToken,
) {
let cancel = shutdown.child_token();
let capture = capture_local_runtime_profile(state_root, protocol_version, request, &cancel);
tokio::pin!(capture);
let mut unexpected = [0_u8; 1];
let result = tokio::select! {
biased;
_ = shutdown.cancelled() => { cancel.cancel(); let _ = capture.await; return; }
_ = reader.read(&mut unexpected) => { cancel.cancel(); let _ = capture.await; Err(RuntimeErrorCode::Cancelled) }
result = &mut capture => result,
};
let response = match result {
Ok(export) => CaptureResponse::RuntimeOk {
archive_base64: URL_SAFE_NO_PAD.encode_to_string(&export.archive_bytes),
archive_sha256: export.archive_sha256,
artifact_uid: export.artifact_uid,
},
Err(code) => CaptureResponse::RuntimeError { code },
};
if let Ok(bytes) = serde_json::to_vec(&response)
&& bytes.len() as u64 <= MAX_RUNTIME_RESPONSE_BYTES
{
let _ = tokio::time::timeout(REQUEST_TIMEOUT, async {
writer.write_all(&bytes).await?;
writer.shutdown().await
})
.await;
}
}
async fn capture_local_runtime_profile(
state_root: &Path,
protocol_version: u16,
input: LocalRuntimeProfileRequest,
cancel: &CancellationToken,
) -> Result<SignedProfileExport, RuntimeErrorCode> {
use rand::{TryRng as _, rngs::SysRng};
if protocol_version != PROTOCOL_VERSION
|| input.offline_key_id.len() != 64
|| !input
.offline_key_id
.bytes()
.all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
{
return Err(RuntimeErrorCode::InvalidRequest);
}
if cancel.is_cancelled() {
return Err(RuntimeErrorCode::Cancelled);
}
if !input.acknowledge_l3 || input.policy_revision == 0 {
return Err(RuntimeErrorCode::ConsentRequired);
}
let now = unix_now().map_err(|_| RuntimeErrorCode::CollectionFailed)?;
if input.consent_expires_at_unix <= now || input.expires_at_unix > input.consent_expires_at_unix {
return Err(RuntimeErrorCode::ConsentExpired);
}
if input.expires_at_unix <= now || input.expires_at_unix.saturating_sub(now) > super::profile_cpu::MAX_VALIDITY_SECONDS {
return Err(RuntimeErrorCode::Expired);
}
if input.duration_millis == 0
|| input.duration_millis > 30_000
|| input.sample_period_micros == 0
|| input.sample_period_micros > input.duration_millis * 1_000
{
return Err(RuntimeErrorCode::LimitExceeded);
}
if input.schema_version != 1 || input.capability != super::profile_cpu::THREAD_PROFILE_CAPABILITY {
return Err(RuntimeErrorCode::InvalidRequest);
}
let owner = private_state_owner(state_root).map_err(|_| RuntimeErrorCode::IdentityUnavailable)?;
let key_root = state_root.join("offline");
let key_directory = std::fs::symlink_metadata(&key_root).map_err(|_| RuntimeErrorCode::IdentityUnavailable)?;
// Existing offline enrollment creates this subdirectory with the process umask.
// The state root is 0700; require the child to remain owned and non-writable by others.
if !key_directory.is_dir()
|| key_directory.file_type().is_symlink()
|| key_directory.uid() != owner
|| key_directory.permissions().mode() & 0o022 != 0
{
return Err(RuntimeErrorCode::IdentityUnavailable);
}
// IdentityStore::load follows symlinks; this IPC boundary must reject them before reading.
use std::os::unix::fs::OpenOptionsExt as _;
let path = crate::connect::OfflineKeyStore::new(state_root).key_path();
let file = std::fs::OpenOptions::new()
.read(true)
.custom_flags(libc::O_NOFOLLOW | libc::O_CLOEXEC | libc::O_NONBLOCK)
.open(path)
.map_err(|_| RuntimeErrorCode::IdentityUnavailable)?;
let metadata = file.metadata().map_err(|_| RuntimeErrorCode::IdentityUnavailable)?;
if !metadata.is_file()
|| metadata.uid() != owner
|| metadata.permissions().mode() & 0o7777 != 0o600
|| metadata.len() == 0
|| metadata.len() > 4_096
{
return Err(RuntimeErrorCode::IdentityUnavailable);
}
let mut der = zeroize::Zeroizing::new(Vec::new());
std::io::Read::read_to_end(&mut std::io::Read::take(file, 4_097), &mut der)
.map_err(|_| RuntimeErrorCode::IdentityUnavailable)?;
if der.len() > 4_096 {
return Err(RuntimeErrorCode::IdentityUnavailable);
}
let key = crate::connect::DeviceIdentity::from_pkcs8_der(&der).map_err(|_| RuntimeErrorCode::IdentityUnavailable)?;
if hex_simd::encode_to_string(Sha256::digest(key.public_key_der()), hex_simd::AsciiCase::Lower) != input.offline_key_id {
return Err(RuntimeErrorCode::IdentityUnavailable);
}
let provenance = super::job_delivery::executable_provenance()
.await
.map_err(|_| RuntimeErrorCode::CollectionFailed)?;
let mut nonce = [0_u8; 32];
SysRng
.try_fill_bytes(&mut nonce)
.map_err(|_| RuntimeErrorCode::CollectionFailed)?;
let request = ProfileCaptureRequest {
organization_name: input.organization_name,
cluster_name: input.cluster_name,
device_name: input.device_name,
run_uid: input.run_uid,
artifact_uid: input.artifact_uid,
schema_version: input.schema_version,
capability: input.capability,
consent: LocalProfileConsent {
consent_uid: input.consent_uid,
policy_revision: input.policy_revision,
expires_at_unix: input.consent_expires_at_unix,
confirmed: input.acknowledge_l3,
},
produced_at_unix: unix_now().map_err(|_| RuntimeErrorCode::CollectionFailed)?,
expires_at_unix: input.expires_at_unix,
nonce,
duration: Duration::from_millis(input.duration_millis),
sample_period: Duration::from_micros(input.sample_period_micros),
provenance,
};
let metrics = tokio::runtime::Handle::current().metrics();
let cancel = cancel.clone();
// Always await this task, including after disconnect, so the lease is released before acknowledgement.
tokio::task::spawn_blocking(move || {
let result = super::profile_threads::capture_runtime_profile(&request, &metrics, MAX_PROFILE_DURATION, &cancel)?;
if result.outcome() != ProfileOutcome::Succeeded {
return Err(ProfileError::CollectionFailed);
}
encode_signed_profile_export(&request, &result, &key, &cancel)
})
.await
.map_err(|_| RuntimeErrorCode::CollectionFailed)?
.map_err(RuntimeErrorCode::from)
}
async fn run_listener(listener: UnixListener, owner: u32, state_root: PathBuf, shutdown: CancellationToken) {
let mut connections = JoinSet::new();
loop {
tokio::select! {
@@ -213,7 +541,7 @@ async fn run_listener(listener: UnixListener, owner: u32, shutdown: Cancellation
if connections.len() < MAX_CONNECTIONS
&& stream.peer_cred().is_ok_and(|credentials| credentials.uid() == owner)
{
connections.spawn(handle_connection(stream, shutdown.clone()));
connections.spawn(handle_connection(stream, state_root.clone(), shutdown.clone()));
}
}
Err(_) => break,
@@ -223,7 +551,7 @@ async fn run_listener(listener: UnixListener, owner: u32, shutdown: Cancellation
while connections.join_next().await.is_some() {}
}
async fn handle_connection(stream: UnixStream, shutdown: CancellationToken) {
async fn handle_connection(stream: UnixStream, state_root: PathBuf, shutdown: CancellationToken) {
let (reader, mut writer) = stream.into_split();
let mut reader = BufReader::new(reader);
let request = tokio::select! {
@@ -234,17 +562,31 @@ async fn handle_connection(stream: UnixStream, shutdown: CancellationToken) {
Ok(Err(_)) | Err(_) => return,
},
};
let (protocol_version, consent_expires_at_unix, duration_millis, max_spans) = match request {
CaptureRequest::RuntimeProfile {
protocol_version,
request,
} => {
handle_runtime_profile(reader, writer, &state_root, protocol_version, request, shutdown).await;
return;
}
CaptureRequest::TraceRecord {
protocol_version,
consent_expires_at_unix,
duration_millis,
max_spans,
} => (protocol_version, consent_expires_at_unix, duration_millis, max_spans),
};
let limits = TraceRecordLimits {
duration: Duration::from_millis(request.duration_millis),
max_spans: request.max_spans,
duration: Duration::from_millis(duration_millis),
max_spans,
};
let capture = async {
if request.protocol_version != PROTOCOL_VERSION {
if protocol_version != PROTOCOL_VERSION {
return Err(TelemetryProducerError::SourceUnavailable);
}
limits.validate()?;
let remaining = request
.consent_expires_at_unix
let remaining = consent_expires_at_unix
.checked_sub(unix_now().map_err(|_| TelemetryProducerError::ConsentExpired)?)
.and_then(|seconds| u64::try_from(seconds).ok())
.filter(|seconds| *seconds > 0)
@@ -294,7 +636,11 @@ async fn read_request(reader: &mut BufReader<tokio::net::unix::OwnedReadHalf>) -
if bytes.is_empty() || bytes.len() as u64 > MAX_REQUEST_BYTES || bytes.pop() != Some(b'\n') {
return Err(LocalTraceCaptureError::Protocol);
}
serde_json::from_slice(&bytes).map_err(|_| LocalTraceCaptureError::Protocol)
let request = serde_json::from_slice(&bytes).map_err(|_| LocalTraceCaptureError::Protocol)?;
if matches!(request, CaptureRequest::TraceRecord { .. }) && bytes.len() + 1 > 1_024 {
return Err(LocalTraceCaptureError::Protocol);
}
Ok(request)
}
fn private_state_owner(path: &Path) -> Result<u32, LocalTraceCaptureError> {
@@ -523,4 +869,237 @@ mod tests {
.expect("server capture stops after disconnect");
runtime.shutdown().await;
}
use sha2::Digest as _;
fn runtime_request(state: &std::path::Path) -> (super::LocalRuntimeProfileRequest, crate::connect::DeviceIdentity) {
let key = crate::connect::OfflineKeyStore::new(state).load_or_create().unwrap();
let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() as i64;
let organization = "organizations/019e3ae0-0000-7000-8000-000000000021";
let cluster = format!("{organization}/clusters/019e3ae0-0000-7000-8000-000000000022");
(
super::LocalRuntimeProfileRequest {
offline_key_id: hex_simd::encode_to_string(
sha2::Sha256::digest(key.public_key_der()),
hex_simd::AsciiCase::Lower,
),
organization_name: organization.to_owned(),
cluster_name: cluster.clone(),
device_name: format!("{cluster}/clusterDevices/019e3ae0-0000-7000-8000-000000000023"),
run_uid: "019e3ae0-0000-7000-8000-000000000024".to_owned(),
artifact_uid: "019e3ae0-0000-7000-8000-000000000025".to_owned(),
schema_version: 1,
capability: "profile.threads@1".to_owned(),
consent_uid: "019e3ae0-0000-7000-8000-000000000026".to_owned(),
policy_revision: 1,
consent_expires_at_unix: now + 120,
acknowledge_l3: true,
expires_at_unix: now + 60,
duration_millis: 20,
sample_period_micros: 1000,
},
key,
)
}
#[tokio::test]
#[serial]
async fn local_runtime_profile_uses_server_workers_and_offline_signature() {
use std::io::Read as _;
let state = tempfile::tempdir().unwrap();
std::fs::set_permissions(state.path(), std::fs::Permissions::from_mode(0o700)).unwrap();
let (request, key) = runtime_request(state.path());
let state_root = state.path().to_path_buf();
let stop = CancellationToken::new();
let server_stop = stop.clone();
let (ready, wait) = tokio::sync::oneshot::channel();
let server = std::thread::spawn(move || {
tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.enable_all()
.build()
.unwrap()
.block_on(async {
let runtime = spawn_local_trace_capture_runtime(&state_root, &server_stop).unwrap();
ready.send(()).unwrap();
server_stop.cancelled().await;
runtime.shutdown().await;
});
});
wait.await.unwrap();
let export = super::request_local_runtime_profile(state.path(), request.clone(), &CancellationToken::new())
.await
.unwrap();
stop.cancel();
server.join().unwrap();
let mut zip = zip::ZipArchive::new(std::io::Cursor::new(&export.archive_bytes)).unwrap();
let result: serde_json::Value = serde_json::from_reader(zip.by_name("result.json").unwrap()).unwrap();
assert_eq!(tokio::runtime::Handle::current().metrics().num_workers(), 1);
assert_eq!(result["data"]["workerCount"], 2);
assert_eq!(result["data"]["samples"].as_array().unwrap().len(), 2);
assert_eq!(result["provenance"]["sourceCommit"], crate::version::build::COMMIT_HASH);
let mut binary = std::fs::File::open(std::env::current_exe().unwrap()).unwrap();
let mut hasher = sha2::Sha256::new();
let mut buffer = [0_u8; 64 * 1024];
loop {
let read = binary.read(&mut buffer).unwrap();
if read == 0 {
break;
}
hasher.update(&buffer[..read]);
}
let binary_hash = hasher.finalize();
assert_eq!(
result["provenance"]["executableSha256"],
hex_simd::encode_to_string(binary_hash, hex_simd::AsciiCase::Lower)
);
let mut envelope = Vec::new();
zip.by_name("envelope.json").unwrap().read_to_end(&mut envelope).unwrap();
let metadata: serde_json::Value = serde_json::from_slice(&envelope).unwrap();
assert_eq!(metadata["classification"], "L3");
assert_eq!(metadata["deviceKeyId"], request.offline_key_id);
let signature: serde_json::Value = serde_json::from_reader(zip.by_name("envelope.sig").unwrap()).unwrap();
let mut signed = b"rustfs-diagnostic-envelope-v1\0".to_vec();
signed.extend_from_slice(&envelope);
assert!(key.verifies_pending_registration_state(&signed, signature["value"].as_str().unwrap()));
assert!(!state.path().join("identity").exists());
}
#[tokio::test]
#[serial]
async fn local_runtime_profile_rejects_invalid_consent_identity_and_protocol() {
let state = tempfile::tempdir().unwrap();
std::fs::set_permissions(state.path(), std::fs::Permissions::from_mode(0o700)).unwrap();
let (request, _) = runtime_request(state.path());
let cancel = CancellationToken::new();
let mut invalid = request.clone();
invalid.acknowledge_l3 = false;
assert!(matches!(
super::capture_local_runtime_profile(state.path(), 1, invalid, &cancel).await,
Err(super::RuntimeErrorCode::ConsentRequired)
));
let mut invalid = request.clone();
invalid.consent_expires_at_unix = 1;
assert!(matches!(
super::capture_local_runtime_profile(state.path(), 1, invalid, &cancel).await,
Err(super::RuntimeErrorCode::ConsentExpired)
));
let mut invalid = request.clone();
invalid.duration_millis = 30_001;
assert!(matches!(
super::capture_local_runtime_profile(state.path(), 1, invalid, &cancel).await,
Err(super::RuntimeErrorCode::LimitExceeded)
));
let mut invalid = request.clone();
invalid.offline_key_id = "0".repeat(64);
assert!(matches!(
super::capture_local_runtime_profile(state.path(), 1, invalid, &cancel).await,
Err(super::RuntimeErrorCode::IdentityUnavailable)
));
let mut value = serde_json::to_value(&request).unwrap();
value["provenance"] = serde_json::json!({});
assert!(serde_json::from_value::<super::LocalRuntimeProfileRequest>(value).is_err());
assert!(serde_json::from_str::<super::CaptureRequest>(r#"{"operation":"TRACE_RECORD","protocolVersion":1,"protocolVersion":1,"consentExpiresAtUnix":1,"durationMillis":1,"maxSpans":1}"#).is_err());
let key_dir = state.path().join("offline");
std::fs::set_permissions(&key_dir, std::fs::Permissions::from_mode(0o777)).unwrap();
assert!(matches!(
super::capture_local_runtime_profile(state.path(), 1, request.clone(), &cancel).await,
Err(super::RuntimeErrorCode::IdentityUnavailable)
));
std::fs::set_permissions(&key_dir, std::fs::Permissions::from_mode(0o755)).unwrap();
let path = crate::connect::OfflineKeyStore::new(state.path()).key_path();
let real = state.path().join("original-key");
std::fs::rename(&path, &real).unwrap();
std::os::unix::fs::symlink(&real, &path).unwrap();
assert!(matches!(
super::capture_local_runtime_profile(state.path(), 1, request.clone(), &cancel).await,
Err(super::RuntimeErrorCode::IdentityUnavailable)
));
std::fs::remove_file(&path).unwrap();
assert!(matches!(
super::capture_local_runtime_profile(state.path(), 1, request, &cancel).await,
Err(super::RuntimeErrorCode::IdentityUnavailable)
));
assert!(!path.exists());
}
#[tokio::test]
#[serial]
async fn local_runtime_profile_cancellation_waits_for_lease_release() {
let state = tempfile::tempdir().unwrap();
std::fs::set_permissions(state.path(), std::fs::Permissions::from_mode(0o700)).unwrap();
let (mut request, _) = runtime_request(state.path());
request.duration_millis = 5_000;
let runtime = spawn_local_trace_capture_runtime(state.path(), &CancellationToken::new()).unwrap();
let cancel = CancellationToken::new();
let task_cancel = cancel.clone();
let state_root = state.path().to_path_buf();
let task = tokio::spawn(async move { super::request_local_runtime_profile(&state_root, request, &task_cancel).await });
tokio::time::timeout(Duration::from_secs(10), async {
loop {
if super::super::profile_cpu::CollectorLease::acquire().is_err() {
break;
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.unwrap();
cancel.cancel();
assert!(matches!(task.await.unwrap(), Err(LocalTraceCaptureError::RuntimeProfile(code)) if code == "CANCELLED"));
let lease = super::super::profile_cpu::CollectorLease::acquire().expect("client cancellation joins collector");
drop(lease);
runtime.shutdown().await;
}
#[tokio::test]
#[serial]
async fn local_runtime_profile_does_not_claim_unacknowledged_cancellation() {
let state = tempfile::tempdir().unwrap();
std::fs::set_permissions(state.path(), std::fs::Permissions::from_mode(0o700)).unwrap();
let (request, _) = runtime_request(state.path());
let owner = super::private_state_owner(state.path()).unwrap();
let listener = super::bind_listener(&state.path().join(super::SOCKET_FILE), owner).unwrap();
let (ready, received) = tokio::sync::oneshot::channel();
let stop = CancellationToken::new();
let peer_stop = stop.clone();
let peer = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let (reader, _writer) = stream.into_split();
let mut reader = tokio::io::BufReader::new(reader);
super::read_request(&mut reader).await.unwrap();
ready.send(()).unwrap();
peer_stop.cancelled().await;
});
let cancel = CancellationToken::new();
let task_cancel = cancel.clone();
let state_root = state.path().to_path_buf();
let task = tokio::spawn(async move { super::request_local_runtime_profile(&state_root, request, &task_cancel).await });
received.await.unwrap();
cancel.cancel();
assert!(matches!(task.await.unwrap(), Err(LocalTraceCaptureError::CancellationUnconfirmed)));
stop.cancel();
peer.await.unwrap();
}
#[tokio::test]
async fn local_capture_wire_keeps_separate_request_limits() {
use tokio::io::AsyncWriteExt as _;
for (size, accepted) in [(1_024, true), (1_025, false), (8_193, false)] {
let (mut sender, receiver) = tokio::net::UnixStream::pair().unwrap();
let mut bytes = serde_json::to_vec(&super::CaptureRequest::TraceRecord {
protocol_version: 1,
consent_expires_at_unix: 1,
duration_millis: 1,
max_spans: 1,
})
.unwrap();
bytes.resize(size - 1, b' ');
bytes.push(b'\n');
let writer = tokio::spawn(async move {
sender.write_all(&bytes).await.unwrap();
});
let (reader, _) = receiver.into_split();
let result = super::read_request(&mut tokio::io::BufReader::new(reader)).await;
assert_eq!(result.is_ok(), accepted, "wire size {size}");
writer.await.unwrap();
}
}
}
@@ -55,3 +55,11 @@ pub(crate) async fn request_local_trace_capture(
) -> Result<TraceRecordCapture, LocalTraceCaptureError> {
Err(LocalTraceCaptureError::RuntimeUnavailable)
}
pub(crate) async fn request_local_runtime_profile(
_state_root: &Path,
_request: super::profile_cpu::LocalRuntimeProfileRequest,
_cancel: &CancellationToken,
) -> Result<super::profile_cpu::SignedProfileExport, LocalTraceCaptureError> {
Err(LocalTraceCaptureError::RuntimeUnavailable)
}
+2 -1
View File
@@ -129,7 +129,8 @@ pub use diagnostics::{
evaluate_network_window, save_signed_top_export, sign_top_export,
};
pub(crate) use diagnostics::{
LocalTraceCaptureError, LocalTraceCaptureRuntime, request_local_trace_capture, spawn_local_trace_capture_runtime,
LocalTraceCaptureError, LocalTraceCaptureRuntime, request_local_runtime_profile, request_local_trace_capture,
spawn_local_trace_capture_runtime,
};
pub use environment::{
ENVIRONMENT_CAPABILITY, ENVIRONMENT_SCHEMA_VERSION, EnvironmentCollectionRequest, EnvironmentError, EnvironmentExportRequest,
+107 -64
View File
@@ -1343,81 +1343,124 @@ async fn execute_connect_profile(options: ConnectProfileOpts) -> Result<()> {
};
use rand::{TryRng as _, rngs::SysRng};
let key = IdentityStore::new(options.state_dir.join("identity"))
.load()
.map_err(Error::other)?
.ok_or_else(|| Error::other("connect profile requires an enrolled device identity"))?;
let executable_sha256 = hash_current_executable()?;
let produced_at_unix = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(Error::other)
.and_then(|duration| i64::try_from(duration.as_secs()).map_err(Error::other))?;
let mut nonce = [0_u8; 32];
SysRng.try_fill_bytes(&mut nonce).map_err(Error::other)?;
let request = ProfileCaptureRequest {
organization_name: options.organization,
cluster_name: options.cluster,
device_name: options.device,
run_uid: options.run_uid,
artifact_uid: options.artifact_uid,
schema_version: options.schema_version,
capability: options.capability,
consent: LocalProfileConsent {
let runtime_scope =
options.tool == ConnectProfileTool::Threads && options.thread_scope == Some(ConnectThreadProfileScope::TokioRuntime);
if !runtime_scope && options.offline_key_id.is_some() {
return Err(Error::other("--offline-key-id is valid only for a service runtime profile"));
}
let cancel = tokio_util::sync::CancellationToken::new();
let export = if runtime_scope {
let offline_key_id = options
.offline_key_id
.ok_or_else(|| Error::other("--offline-key-id must explicitly select an existing offline identity"))?;
let input = crate::connect::diagnostics::LocalRuntimeProfileRequest {
offline_key_id,
organization_name: options.organization,
cluster_name: options.cluster,
device_name: options.device,
run_uid: options.run_uid,
artifact_uid: options.artifact_uid,
schema_version: options.schema_version,
capability: options.capability,
consent_uid: options.consent_uid,
policy_revision: options.policy_revision,
expires_at_unix: options.consent_expires_at_unix,
confirmed: options.acknowledge_l3,
},
produced_at_unix,
expires_at_unix: options.expires_at_unix,
nonce,
duration: Duration::from_millis(options.duration_millis),
sample_period: Duration::from_micros(options.sample_period_micros),
provenance: ProfileProvenance::new(
crate::version::build::COMMIT_HASH,
executable_sha256,
env!("CARGO_PKG_VERSION"),
enabled_build_features(),
),
};
let cancel = tokio_util::sync::CancellationToken::new();
let export = {
let capture = async {
match options.tool {
ConnectProfileTool::Cpu => {
if options.thread_scope.is_some() {
return Err(Error::other("--thread-scope is valid only for the threads profile"));
}
export_cpu_profile(&request, &key, &cancel).await.map_err(Error::other)
}
ConnectProfileTool::Memory => {
if options.thread_scope.is_some() {
return Err(Error::other("--thread-scope is valid only for the threads profile"));
}
export_memory_profile(&request, &key, &cancel).await.map_err(Error::other)
}
ConnectProfileTool::Threads => {
let scope = match options.thread_scope {
Some(ConnectThreadProfileScope::TokioRuntime) => ThreadProfileScope::TokioRuntime,
Some(ConnectThreadProfileScope::NativeThreads) => ThreadProfileScope::NativeThreads,
None => return Err(Error::other("--thread-scope is required for the threads profile")),
};
export_thread_profile(&request, scope, &key, &cancel)
.await
.map_err(Error::other)
}
}
consent_expires_at_unix: options.consent_expires_at_unix,
acknowledge_l3: options.acknowledge_l3,
expires_at_unix: options.expires_at_unix,
duration_millis: options.duration_millis,
sample_period_micros: options.sample_period_micros,
};
let capture = crate::connect::request_local_runtime_profile(&options.state_dir, input, &cancel);
tokio::pin!(capture);
tokio::select! {
biased;
signal = tokio::signal::ctrl_c() => {
signal.map_err(Error::other)?;
cancel.cancel();
return Err(Error::other("profile collection cancelled"));
return match capture.await {
Err(error) => Err(Error::other(error)),
Ok(_) => Err(Error::other("profile collection cancelled")),
};
}
result = capture.as_mut() => result?,
result = &mut capture => result.map_err(Error::other)?,
}
} else {
let key = IdentityStore::new(options.state_dir.join("identity"))
.load()
.map_err(Error::other)?
.ok_or_else(|| Error::other("connect profile requires an enrolled device identity"))?;
let executable_sha256 = hash_current_executable()?;
let produced_at_unix = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(Error::other)
.and_then(|duration| i64::try_from(duration.as_secs()).map_err(Error::other))?;
let mut nonce = [0_u8; 32];
SysRng.try_fill_bytes(&mut nonce).map_err(Error::other)?;
let request = ProfileCaptureRequest {
organization_name: options.organization,
cluster_name: options.cluster,
device_name: options.device,
run_uid: options.run_uid,
artifact_uid: options.artifact_uid,
schema_version: options.schema_version,
capability: options.capability,
consent: LocalProfileConsent {
consent_uid: options.consent_uid,
policy_revision: options.policy_revision,
expires_at_unix: options.consent_expires_at_unix,
confirmed: options.acknowledge_l3,
},
produced_at_unix,
expires_at_unix: options.expires_at_unix,
nonce,
duration: Duration::from_millis(options.duration_millis),
sample_period: Duration::from_micros(options.sample_period_micros),
provenance: ProfileProvenance::new(
crate::version::build::COMMIT_HASH,
executable_sha256,
env!("CARGO_PKG_VERSION"),
enabled_build_features(),
),
};
let export = {
let capture = async {
match options.tool {
ConnectProfileTool::Cpu => {
if options.thread_scope.is_some() {
return Err(Error::other("--thread-scope is valid only for the threads profile"));
}
export_cpu_profile(&request, &key, &cancel).await.map_err(Error::other)
}
ConnectProfileTool::Memory => {
if options.thread_scope.is_some() {
return Err(Error::other("--thread-scope is valid only for the threads profile"));
}
export_memory_profile(&request, &key, &cancel).await.map_err(Error::other)
}
ConnectProfileTool::Threads => {
let scope = match options.thread_scope {
Some(ConnectThreadProfileScope::TokioRuntime) => ThreadProfileScope::TokioRuntime,
Some(ConnectThreadProfileScope::NativeThreads) => ThreadProfileScope::NativeThreads,
None => return Err(Error::other("--thread-scope is required for the threads profile")),
};
export_thread_profile(&request, scope, &key, &cancel)
.await
.map_err(Error::other)
}
}
};
tokio::pin!(capture);
tokio::select! {
biased;
signal = tokio::signal::ctrl_c() => {
signal.map_err(Error::other)?;
cancel.cancel();
return Err(Error::other("profile collection cancelled"));
}
result = capture.as_mut() => result?,
}
};
export
};
let tool = export.tool;
let outcome = export.outcome;