diff --git a/rustfs/src/connect/diagnostics/trace_runtime.rs b/rustfs/src/connect/diagnostics/trace_runtime.rs index dac9b86ac..2428c7e08 100644 --- a/rustfs/src/connect/diagnostics/trace_runtime.rs +++ b/rustfs/src/connect/diagnostics/trace_runtime.rs @@ -306,6 +306,11 @@ pub(crate) async fn request_local_runtime_profile( request: LocalRuntimeProfileRequest, cancel: &CancellationToken, ) -> Result { + let tool = match request.capability.as_str() { + super::profile_cpu::THREAD_PROFILE_CAPABILITY => ProfileTool::Threads, + super::profile_cpu::MEMORY_PROFILE_CAPABILITY => ProfileTool::Memory, + _ => return Err(LocalTraceCaptureError::Protocol), + }; let owner = private_state_owner(state_root)?; let socket_path = state_root.join(SOCKET_FILE); socket_identity(&socket_path, owner)?; @@ -374,7 +379,7 @@ pub(crate) async fn request_local_runtime_profile( } Ok(SignedProfileExport { artifact_uid, - tool: ProfileTool::Threads, + tool, outcome: ProfileOutcome::Succeeded, reason_code: ProfileReasonCode::Complete, archive_bytes, @@ -717,7 +722,12 @@ async fn capture_local_runtime_profile( { return Err(RuntimeErrorCode::LimitExceeded); } - if input.schema_version != 1 || input.capability != super::profile_cpu::THREAD_PROFILE_CAPABILITY { + let tool = match input.capability.as_str() { + super::profile_cpu::THREAD_PROFILE_CAPABILITY => ProfileTool::Threads, + super::profile_cpu::MEMORY_PROFILE_CAPABILITY => ProfileTool::Memory, + _ => return Err(RuntimeErrorCode::InvalidRequest), + }; + if input.schema_version != 1 { return Err(RuntimeErrorCode::InvalidRequest); } let key = load_offline_key(state_root, &input.offline_key_id)?; @@ -749,19 +759,27 @@ async fn capture_local_runtime_profile( 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); + match tool { + ProfileTool::Memory => super::profile_memory::export_memory_profile(&request, &key, cancel) + .await + .map_err(RuntimeErrorCode::from), + ProfileTool::Threads => { + 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) } - encode_signed_profile_export(&request, &result, &key, &cancel) - }) - .await - .map_err(|_| RuntimeErrorCode::CollectionFailed)? - .map_err(RuntimeErrorCode::from) + ProfileTool::Cpu => Err(RuntimeErrorCode::InvalidRequest), + } } async fn run_listener(listener: UnixListener, owner: u32, state_root: PathBuf, shutdown: CancellationToken) { @@ -1338,6 +1356,38 @@ mod tests { assert!(!state.path().join("identity").exists()); } + #[tokio::test] + #[serial] + async fn local_memory_profile_signs_service_capture_with_offline_identity() { + use std::io::Read as _; + let state = tempfile::tempdir().unwrap(); + std::fs::set_permissions(state.path(), std::fs::Permissions::from_mode(0o700)).unwrap(); + let (mut request, key) = runtime_request(state.path()); + request.capability = super::super::profile_cpu::MEMORY_PROFILE_CAPABILITY.to_owned(); + let runtime = spawn_local_trace_capture_runtime(state.path(), &CancellationToken::new()).unwrap(); + let export = super::request_local_runtime_profile(state.path(), request.clone(), &CancellationToken::new()) + .await + .unwrap(); + assert_eq!(export.tool, super::ProfileTool::Memory); + runtime.shutdown().await; + 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!(result["toolId"], "profile.memory"); + assert_eq!(result["outcome"], "SUCCEEDED"); + assert!(result["data"]["allocatedBytes"].is_u64()); + assert!(result["data"]["allocationCount"].is_u64()); + 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() { @@ -1352,6 +1402,13 @@ mod tests { Err(super::RuntimeErrorCode::ConsentRequired) )); let mut invalid = request.clone(); + invalid.capability = super::super::profile_cpu::MEMORY_PROFILE_CAPABILITY.to_owned(); + 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, diff --git a/rustfs/src/startup_entrypoint.rs b/rustfs/src/startup_entrypoint.rs index f2e68a335..0ec730984 100644 --- a/rustfs/src/startup_entrypoint.rs +++ b/rustfs/src/startup_entrypoint.rs @@ -1401,8 +1401,12 @@ async fn execute_connect_profile(options: ConnectProfileOpts) -> Result<()> { }; use rand::{TryRng as _, rngs::SysRng}; - let runtime_scope = - options.tool == ConnectProfileTool::Threads && options.thread_scope == Some(ConnectThreadProfileScope::TokioRuntime); + let runtime_scope = (options.tool == ConnectProfileTool::Threads + && options.thread_scope == Some(ConnectThreadProfileScope::TokioRuntime)) + || (options.tool == ConnectProfileTool::Memory && options.offline_key_id.is_some()); + if options.tool == ConnectProfileTool::Memory && options.thread_scope.is_some() { + return Err(Error::other("--thread-scope is valid only for the threads profile")); + } if !runtime_scope && options.offline_key_id.is_some() { return Err(Error::other("--offline-key-id is valid only for a service runtime profile")); }