From d2c4383ce074fc374fbc1939c69a865685b8e194 Mon Sep 17 00:00:00 2001 From: Chris Date: Tue, 15 Sep 2026 01:53:05 +0800 Subject: [PATCH] Execute authenticated top.rpc service jobs (#7881) feat(connect): execute top.rpc service jobs --- rustfs/src/connect/diagnostics/job.rs | 155 +++++++++++++++++++++++++- rustfs/src/server/http.rs | 147 +++++++++++++++++++++++- 2 files changed, 293 insertions(+), 9 deletions(-) diff --git a/rustfs/src/connect/diagnostics/job.rs b/rustfs/src/connect/diagnostics/job.rs index eb52bdfc2..1fceab07d 100644 --- a/rustfs/src/connect/diagnostics/job.rs +++ b/rustfs/src/connect/diagnostics/job.rs @@ -35,10 +35,10 @@ use super::{ MAX_TOP_EXPORT_VALIDITY, NETWORK_CAPABILITY, NETWORK_SCHEMA_VERSION, NetworkOutcome, NetworkPerformanceError, NetworkPerformanceRequest, NetworkProvenance, NetworkReasonCode, PROFILE_SCHEMA_VERSION, ProfileCaptureRequest, ProfileOutcome, ProfileProvenance, THREAD_PROFILE_CAPABILITY, TOP_API_CAPABILITY, TOP_CLASSIFICATION, TOP_LOCKS_CAPABILITY, - TOP_SCHEMA_VERSION, ThreadProfileScope, TopApiOperation, TopCaptureLimits, TopCaptureRequest, TopCaptureScope, TopOutcome, - capture_cpu_profile, capture_thread_profile, capture_top_api, capture_top_locks, encode_signed_profile_export, measure_drive, - measure_network, runtime_network_peer_aliases, sign_drive_export, sign_network_export, sign_top_export, - sign_top_export_with_nonce, + TOP_RPC_CAPABILITY, TOP_SCHEMA_VERSION, ThreadProfileScope, TopApiOperation, TopCaptureLimits, TopCaptureRequest, + TopCaptureScope, TopOutcome, capture_cpu_profile, capture_thread_profile, capture_top_api, capture_top_locks, + capture_top_rpc, encode_signed_profile_export, measure_drive, measure_network, runtime_network_peer_aliases, + sign_drive_export, sign_network_export, sign_top_export, sign_top_export_with_nonce, }; use crate::connect::DeviceIdentity; @@ -49,6 +49,7 @@ const PERFORMANCE_DRIVE_JOB_TYPE: &str = "performance.drive"; const PERFORMANCE_NETWORK_JOB_TYPE: &str = "performance.network"; const TOP_API_JOB_TYPE: &str = "top.api"; const TOP_LOCKS_JOB_TYPE: &str = "top.locks"; +const TOP_RPC_JOB_TYPE: &str = "top.rpc"; pub const DIAGNOSTIC_JOB_SIGNATURE_DOMAIN: &[u8] = b"rustfs-connect-agent-job-v1\0"; const MAX_JOB_LIFETIME_SECONDS: i64 = 1_800; const MAX_FUTURE_SKEW_SECONDS: i64 = 300; @@ -67,6 +68,8 @@ const MAX_TOP_API_CPU_MILLIS: u64 = 5_000; const MIN_TOP_API_MEMORY_BYTES: u64 = 1_048_576; const MAX_TOP_LOCKS_CPU_MILLIS: u64 = 5_000; const MIN_TOP_LOCKS_MEMORY_BYTES: u64 = 1_048_576; +const MAX_TOP_RPC_CPU_MILLIS: u64 = 5_000; +const MIN_TOP_RPC_MEMORY_BYTES: u64 = 1_048_576; #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum DiagnosticJobKind { @@ -76,6 +79,7 @@ enum DiagnosticJobKind { PerformanceNetwork, TopApi, TopLocks, + TopRpc, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -454,6 +458,11 @@ impl DiagnosticJobEnvelope { { return Err(DiagnosticJobError::LimitExceeded); } + if kind == DiagnosticJobKind::TopRpc + && (self.limits.max_cpu_millis > MAX_TOP_RPC_CPU_MILLIS || self.limits.max_memory_bytes < MIN_TOP_RPC_MEMORY_BYTES) + { + return Err(DiagnosticJobError::LimitExceeded); + } Ok(()) } @@ -481,6 +490,11 @@ impl DiagnosticJobEnvelope { { Ok(DiagnosticJobKind::TopLocks) } + (TOP_RPC_JOB_TYPE, [capability], version) + if capability == TOP_RPC_CAPABILITY && version == u16::from(TOP_SCHEMA_VERSION) => + { + Ok(DiagnosticJobKind::TopRpc) + } _ => Err(DiagnosticJobError::Unsupported), } } @@ -514,6 +528,7 @@ pub async fn execute_diagnostic_job( } DiagnosticJobKind::TopApi => execute_top_api_job(envelope, nonce, identity, provenance, cancel).await, DiagnosticJobKind::TopLocks => execute_top_locks_job(envelope, identity, provenance, cancel).await, + DiagnosticJobKind::TopRpc => execute_top_rpc_job(envelope, nonce, identity, provenance, cancel).await, } } @@ -961,6 +976,70 @@ async fn execute_top_locks_job( }) } +async fn execute_top_rpc_job( + envelope: DiagnosticJobEnvelope, + nonce: [u8; 32], + identity: &DeviceIdentity, + provenance: ProfileProvenance, + cancel: &CancellationToken, +) -> Result { + let expire = parse_time(&envelope.expire_time)?; + let consent_expire = parse_time(&envelope.parameters.consent_expires_at)?; + let request = TopCaptureRequest { + scope: TopCaptureScope { + organization_name: envelope.organization_name, + cluster_name: envelope.cluster_name, + device_name: envelope.device_name, + run_uid: envelope.job_id.clone(), + artifact_uid: envelope.parameters.artifact_uid, + policy_revision: envelope.parameters.consent_policy_revision, + run_expires_at_unix: expire.timestamp(), + executable_sha256: provenance.executable_sha256().to_owned(), + build_features: provenance.build_features().to_vec(), + consent: LocalTopConsent { + uid: envelope.parameters.consent_uid, + tool_id: TOP_RPC_JOB_TYPE.to_owned(), + classification: TOP_CLASSIFICATION.to_owned(), + active: true, + expires_at_unix: consent_expire.timestamp(), + }, + }, + limits: TopCaptureLimits { + max_duration_millis: envelope.parameters.duration_millis, + max_working_memory_bytes: envelope.limits.max_memory_bytes, + max_cpu_millis: envelope.limits.max_cpu_millis, + ..TopCaptureLimits::default() + }, + window: Duration::from_millis(envelope.parameters.duration_millis), + export_validity: MAX_TOP_EXPORT_VALIDITY, + }; + let result = capture_top_rpc(&request, cancel).await.map_err(top_capture_failure)?; + let outcome = result.outcome.as_str().to_owned(); + let reason = result.reason_code.as_str().to_owned(); + if !matches!(result.outcome, TopOutcome::Succeeded | TopOutcome::Partial) { + return Ok(DiagnosticJobExecution { + job_id: envelope.job_id, + outcome, + reason, + artifact_uid: None, + artifact_sha256: None, + artifact_bytes: None, + }); + } + let export = sign_top_export_with_nonce(&request, &result, identity, cancel, nonce).map_err(top_export_failure)?; + if export.archive_bytes.len() > usize::try_from(envelope.limits.max_output_bytes).unwrap_or(usize::MAX) { + return Err(DiagnosticJobError::LimitExceeded); + } + Ok(DiagnosticJobExecution { + job_id: envelope.job_id, + outcome, + reason, + artifact_uid: Some(export.artifact_uid), + artifact_sha256: Some(export.archive_sha256), + artifact_bytes: Some(export.archive_bytes), + }) +} + fn capture_failure(error: super::ProfileError) -> DiagnosticJobError { match error { super::ProfileError::Cancelled => DiagnosticJobError::Cancelled, @@ -1402,6 +1481,74 @@ mod tests { ); } + #[test] + fn accepts_only_the_bounded_top_rpc_capability_pair() { + let mut top = envelope(); + top.job_type = TOP_RPC_JOB_TYPE.to_owned(); + top.required_capabilities = vec![TOP_RPC_CAPABILITY.to_owned()]; + top.limits.max_cpu_millis = MAX_TOP_RPC_CPU_MILLIS; + let (top, signer) = signed_envelope(top); + signer + .verify(&top, &target(&top), "2030-01-01T00:00:10Z".parse().expect("time")) + .expect("valid top.rpc job"); + + let mut mismatched = top.clone(); + mismatched.required_capabilities = vec![TOP_LOCKS_CAPABILITY.to_owned()]; + assert_eq!( + signer.verify(&mismatched, &target(&mismatched), "2030-01-01T00:00:10Z".parse().expect("time")), + Err(DiagnosticJobError::Unsupported) + ); + + let mut unbounded = top; + unbounded.limits.max_cpu_millis += 1; + assert_eq!( + signer.verify(&unbounded, &target(&unbounded), "2030-01-01T00:00:10Z".parse().expect("time")), + Err(DiagnosticJobError::LimitExceeded) + ); + } + + #[tokio::test] + #[serial_test::serial] + async fn top_rpc_job_captures_service_process_events_and_exports_a_signed_artifact() { + use rustfs_common::trace_bus::{ + TelemetryTraceEvent, TelemetryTraceOperation, TelemetryTraceStatus, telemetry_trace_emit, + }; + + let mut top = envelope(); + top.job_type = TOP_RPC_JOB_TYPE.to_owned(); + top.required_capabilities = vec![TOP_RPC_CAPABILITY.to_owned()]; + top.limits.max_cpu_millis = MAX_TOP_RPC_CPU_MILLIS; + top.parameters.duration_millis = 50; + + let emit = async { + tokio::time::sleep(Duration::from_millis(10)).await; + assert!(telemetry_trace_emit(|| { + TelemetryTraceEvent::new( + TelemetryTraceOperation::InternalRpc, + Duration::from_micros(37), + TelemetryTraceStatus::Ok, + ) + })); + }; + let identity = DeviceIdentity::generate(); + let cancellation = CancellationToken::new(); + let execute = execute_diagnostic_job( + VerifiedDiagnosticJob { + envelope: top, + nonce: [7_u8; 32], + }, + &identity, + ProfileProvenance::new("a".repeat(40), "b".repeat(64), "1.0.0", vec![]), + &cancellation, + ); + let (execution, ()) = tokio::join!(execute, emit); + let execution = execution.expect("top.rpc job should execute"); + + assert_eq!(execution.outcome, "SUCCEEDED"); + assert_eq!(execution.reason, "COMPLETE"); + assert!(execution.artifact_bytes.is_some_and(|bytes| !bytes.is_empty())); + } + #[test] fn rejects_tampering_cross_device_replay_and_expiry() { let (envelope, signer) = signed(); diff --git a/rustfs/src/server/http.rs b/rustfs/src/server/http.rs index 8e1addadd..64d1cf399 100644 --- a/rustfs/src/server/http.rs +++ b/rustfs/src/server/http.rs @@ -57,7 +57,10 @@ use hyper_util::{ use metrics::{counter, gauge, histogram}; use opentelemetry::global; use opentelemetry::trace::TraceContextExt; -use rustfs_common::GlobalReadiness; +use rustfs_common::{ + GlobalReadiness, + trace_bus::{TelemetryTraceEvent, TelemetryTraceOperation, TelemetryTraceStatus, telemetry_trace_emit}, +}; use rustfs_io_metrics::internode_metrics::{ INTERNODE_OPERATION_GRPC_OTHER, INTERNODE_OPERATION_GRPC_READ_ALL, INTERNODE_OPERATION_GRPC_READ_MULTIPLE, INTERNODE_OPERATION_GRPC_WRITE_ALL, INTERNODE_TRANSPORT_BACKEND_GRPC, global_internode_metrics, @@ -84,7 +87,7 @@ use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::task::{Context, Poll}; -use std::time::Duration; +use std::time::{Duration, Instant}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tonic::service::Routes; @@ -258,6 +261,93 @@ struct RpcRequestPathService { inner: S, } +struct RpcCompletionBody { + inner: B, + started_at: Instant, + header_status: Option, + complete: bool, +} + +impl RpcCompletionBody { + fn new(inner: B, started_at: Instant, header_status: Option) -> Self { + Self { + inner, + started_at, + header_status, + complete: false, + } + } + + fn complete(&mut self, status: TelemetryTraceStatus) { + if self.complete { + return; + } + self.complete = true; + telemetry_trace_emit(|| { + TelemetryTraceEvent::new(TelemetryTraceOperation::InternalRpc, self.started_at.elapsed(), status) + }); + } +} + +impl http_body::Body for RpcCompletionBody +where + B: http_body::Body + Unpin, +{ + type Data = Bytes; + type Error = B::Error; + + fn is_end_stream(&self) -> bool { + self.complete || self.inner.is_end_stream() + } + + fn poll_frame( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + match Pin::new(&mut self.inner).poll_frame(cx) { + Poll::Ready(Some(Ok(frame))) => { + if let Some(trailers) = frame.trailers_ref() { + let status = grpc_telemetry_status(trailers).unwrap_or(TelemetryTraceStatus::Error); + self.complete(status); + } + Poll::Ready(Some(Ok(frame))) + } + Poll::Ready(Some(Err(error))) => { + self.complete(TelemetryTraceStatus::Error); + Poll::Ready(Some(Err(error))) + } + Poll::Ready(None) => { + let status = self.header_status.unwrap_or(TelemetryTraceStatus::Error); + self.complete(status); + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } + + fn size_hint(&self) -> http_body::SizeHint { + self.inner.size_hint() + } +} + +impl Drop for RpcCompletionBody { + fn drop(&mut self) { + if !self.complete { + self.complete(self.header_status.unwrap_or(TelemetryTraceStatus::Error)); + } + } +} + +fn grpc_telemetry_status(headers: &HeaderMap) -> Option { + headers.get("grpc-status").map(|status| { + if status == "0" { + TelemetryTraceStatus::Ok + } else { + TelemetryTraceStatus::Error + } + }) +} + impl RpcRequestPathService { fn new(inner: S) -> Self { Self { inner } @@ -270,9 +360,9 @@ where S::Error: Send + 'static, S::Future: Send + 'static, B: Send + 'static, - ResBody: Send + 'static, + ResBody: http_body::Body + Unpin + Send + 'static, { - type Response = Response; + type Response = Response>; type Error = S::Error; type Future = Pin> + Send>>; @@ -281,6 +371,7 @@ where } fn call(&mut self, mut req: HttpRequest) -> Self::Future { + let started_at = Instant::now(); let target = RpcRequestTarget { uri: req.uri().clone(), method: req.method().clone(), @@ -300,7 +391,10 @@ where if let Some(headers) = response_headers { response.headers_mut().extend(headers); } - Ok(response) + let header_status = grpc_telemetry_status(response.headers()); + let (parts, body) = response.into_parts(); + let tracked = RpcCompletionBody::new(body, started_at, header_status); + Ok(Response::from_parts(parts, tracked)) }) } } @@ -3434,6 +3528,49 @@ mod tests { assert_eq!(captured.method, Method::POST); } + #[tokio::test] + #[serial_test::serial] + async fn rpc_completion_waits_for_the_grpc_stream_trailer() { + use http_body_util::StreamBody; + use rustfs_common::trace_bus::subscribe_telemetry_trace_events; + use tokio_stream::wrappers::ReceiverStream; + + let mut subscription = subscribe_telemetry_trace_events(); + let (tx, rx) = mpsc::channel::, Infallible>>(2); + let mut body = RpcCompletionBody::new(StreamBody::new(ReceiverStream::new(rx)), Instant::now(), None); + + tx.send(Ok(Frame::data(Bytes::from_static(b"rpc-data")))) + .await + .expect("response data frame should send"); + let frame = body + .frame() + .await + .expect("response data frame") + .expect("response body should remain valid"); + assert!(frame.is_data()); + assert!(matches!(subscription.try_recv(), Err(tokio::sync::broadcast::error::TryRecvError::Empty))); + + let mut trailers = HeaderMap::new(); + trailers.insert("grpc-status", HeaderValue::from_static("0")); + tx.send(Ok(Frame::trailers(trailers))) + .await + .expect("response trailer should send"); + let frame = body + .frame() + .await + .expect("response trailer frame") + .expect("response body should remain valid"); + assert!(frame.is_trailers()); + + let event = tokio::time::timeout(Duration::from_secs(1), subscription.recv()) + .await + .expect("RPC completion event should arrive") + .expect("telemetry source should remain open"); + assert_eq!(event.operation, TelemetryTraceOperation::InternalRpc); + assert_eq!(event.status, TelemetryTraceStatus::Ok); + assert!(event.duration > Duration::ZERO); + } + #[tokio::test] #[serial_test::serial] async fn rpc_auth_binds_post_method_authority_and_exact_path() {