Skip to main content

otel_bootstrap/
profiling.rs

1#![cfg(feature = "profiling")]
2
3use std::error::Error;
4use std::sync::OnceLock;
5
6#[cfg(feature = "profiling-bridge-pyroscope-rs")]
7use opentelemetry::trace::TraceContextExt;
8
9/// Validate that a pyroscope endpoint targets only loopback (per ADR platform/0203 AC1).
10/// Allowed: 127.0.0.1, ::1, localhost, unix socket paths.
11/// Rejects routable addresses to prevent unauthenticated plaintext profile data leaving the pod.
12fn validate_pyroscope_endpoint(endpoint: &str) -> Result<(), Box<dyn Error>> {
13    use url::Url;
14
15    // Unix socket paths are allowed
16    if endpoint.starts_with("unix://") {
17        return Ok(());
18    }
19
20    // HTTP/HTTPS endpoints must target loopback
21    if endpoint.starts_with("http://") || endpoint.starts_with("https://") {
22        let url = Url::parse(endpoint)?;
23
24        // Reject endpoints with userinfo (user:pass@host) to prevent redirect attacks
25        if !url.username().is_empty() || url.password().is_some() {
26            return Err(format!(
27                "pyroscope endpoint must not contain userinfo; got: {endpoint} (ADR platform/0203 AC1)"
28            ).into());
29        }
30
31        let host = url.host_str().unwrap_or("");
32
33        match host {
34            "127.0.0.1" | "::1" | "[::1]" | "localhost" => Ok(()),
35            _ => Err(format!(
36                "pyroscope endpoint must target loopback (127.0.0.1, ::1, localhost, or unix socket); \
37                 got: {endpoint} (ADR platform/0203 AC1)"
38            ).into()),
39        }
40    } else {
41        Err(
42            format!("pyroscope endpoint must be http://, https://, or unix://; got: {endpoint}")
43                .into(),
44        )
45    }
46}
47
48/// Profiling bridge handle. Owns the active profiling agent and ensures
49/// graceful shutdown on drop.
50pub struct ProfilingHandle {
51    #[cfg(feature = "profiling-bridge-pyroscope-rs")]
52    agent: Option<pyroscope::PyroscopeAgent<pyroscope::pyroscope::PyroscopeAgentRunning>>,
53}
54
55#[cfg(feature = "profiling-bridge-pyroscope-rs")]
56impl Drop for ProfilingHandle {
57    fn drop(&mut self) {
58        if let Some(agent) = self.agent.take() {
59            let _ = agent.stop();
60        }
61    }
62}
63
64#[cfg(feature = "profiling-bridge-pyroscope-rs")]
65type BoxedTagFn = Box<dyn Fn(String, String) -> pyroscope::Result<()> + Send + Sync>;
66
67/// Module-level storage for profiling tag functions (add_tag, remove_tag)
68/// obtained from the running pyroscope agent.
69#[cfg(feature = "profiling-bridge-pyroscope-rs")]
70static PROFILING_TAG_FNS: OnceLock<(BoxedTagFn, BoxedTagFn)> = OnceLock::new();
71
72/// Guards against starting more than one profiling agent per process.
73/// The `pprof` backend keeps a single process-wide profiler guard, so a
74/// second concurrent agent would fail to start; subsequent calls are
75/// treated as no-ops rather than errors.
76#[cfg(feature = "profiling-bridge-pyroscope-rs")]
77static PROFILING_STARTED: OnceLock<()> = OnceLock::new();
78
79/// Start the pyroscope profiling bridge.
80///
81/// The bridge pushes profiles over plain HTTP/loopback to a local SPIFFE-terminating
82/// sidecar (or an already-mTLS'd endpoint reachable without client-side TLS material).
83/// pyroscope-rs hardcodes its own HTTP client internally with no hook
84/// for custom TLS/identity, so in-process mTLS is not possible; the sidecar carries
85/// the workload identity upstream.
86///
87/// **Temporary exception** (Tracks #40): This bridge is a sunset-bound interim implementation
88/// pending a native Rust OTLP profiles exporter. See ADR platform/0202 and issue #40.
89#[cfg(feature = "profiling-bridge-pyroscope-rs")]
90pub(crate) fn start_pyroscope_bridge(
91    service_name: &str,
92    pyroscope_endpoint: &str,
93) -> Result<Option<ProfilingHandle>, Box<dyn Error>> {
94    use pyroscope::backend::{BackendConfig, PprofConfig, pprof_backend};
95
96    // Validate endpoint targets loopback only (ADR platform/0203 AC1)
97    validate_pyroscope_endpoint(pyroscope_endpoint)?;
98
99    // Only one profiling agent may run per process (the `pprof` backend holds
100    // a single process-wide profiler guard); ignore subsequent start attempts.
101    if PROFILING_STARTED.set(()).is_err() {
102        return Ok(None);
103    }
104
105    let agent = pyroscope::pyroscope::PyroscopeAgentBuilder::new(
106        pyroscope_endpoint,
107        service_name,
108        100,
109        "pyroscope-rs",
110        env!("CARGO_PKG_VERSION"),
111        pprof_backend(PprofConfig { sample_rate: 100 }, BackendConfig::default()),
112    )
113    .build()?
114    .start()?;
115
116    let (add_tag, remove_tag) = agent.tag_wrapper();
117    PROFILING_TAG_FNS
118        .set((Box::new(add_tag), Box::new(remove_tag)))
119        .ok();
120
121    Ok(Some(ProfilingHandle { agent: Some(agent) }))
122}
123
124/// No-op bridge for when profiling is enabled but the pyroscope feature is not.
125#[cfg(all(feature = "profiling", not(feature = "profiling-bridge-pyroscope-rs")))]
126pub(crate) fn start_pyroscope_bridge(
127    _service_name: &str,
128    _pyroscope_endpoint: &str,
129) -> Result<Option<ProfilingHandle>, Box<dyn Error>> {
130    Ok(None)
131}
132
133/// Tracing layer that tags active span enter/exit with trace_id and span_id
134/// in the running pyroscope agent.
135#[cfg(feature = "profiling-bridge-pyroscope-rs")]
136pub struct ProfilingTagLayer;
137
138#[cfg(feature = "profiling-bridge-pyroscope-rs")]
139impl<S> tracing_subscriber::Layer<S> for ProfilingTagLayer
140where
141    S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
142{
143    fn on_enter(&self, _id: &tracing::span::Id, _ctx: tracing_subscriber::layer::Context<'_, S>) {
144        if let Some((add_tag, _)) = PROFILING_TAG_FNS.get() {
145            let cx = opentelemetry::Context::current();
146            let span_ref = cx.span();
147            let span_context = span_ref.span_context();
148            if span_context.is_valid() {
149                let trace_id = span_context.trace_id();
150                let span_id = span_context.span_id();
151                let _ = add_tag("trace_id".to_string(), format!("{trace_id:x}"));
152                let _ = add_tag("span_id".to_string(), format!("{span_id:x}"));
153            }
154        }
155    }
156
157    fn on_exit(&self, _id: &tracing::span::Id, _ctx: tracing_subscriber::layer::Context<'_, S>) {
158        if let Some((_, remove_tag)) = PROFILING_TAG_FNS.get() {
159            let cx = opentelemetry::Context::current();
160            let span_ref = cx.span();
161            let span_context = span_ref.span_context();
162            if span_context.is_valid() {
163                let trace_id = span_context.trace_id();
164                let span_id = span_context.span_id();
165                let _ = remove_tag("trace_id".to_string(), format!("{trace_id:x}"));
166                let _ = remove_tag("span_id".to_string(), format!("{span_id:x}"));
167            }
168        }
169    }
170}
171
172#[cfg(all(test, feature = "profiling-bridge-pyroscope-rs"))]
173mod tests {
174    use super::*;
175
176    #[test]
177    fn start_bridge_with_nonexistent_server() {
178        let result = start_pyroscope_bridge("test-svc", "http://localhost:4040");
179        assert!(
180            result.is_ok(),
181            "pyroscope agent start() is lazy and does not eagerly connect"
182        );
183        if let Ok(Some(_handle)) = result {
184            // Bridge is active
185        }
186    }
187
188    #[test]
189    fn start_bridge_multiple_times_ignores_second() {
190        let result1 = start_pyroscope_bridge("test-svc-1", "http://localhost:4040");
191        assert!(result1.is_ok());
192        let result2 = start_pyroscope_bridge("test-svc-2", "http://localhost:4041");
193        assert!(result2.is_ok());
194        // Second call is a no-op: the `pprof` backend only supports one
195        // process-wide profiler guard, so the bridge returns `Ok(None)`.
196        assert!(result2.unwrap().is_none());
197    }
198
199    #[test]
200    fn validate_endpoint_accepts_loopback_ipv4() {
201        assert!(validate_pyroscope_endpoint("http://127.0.0.1:4040").is_ok());
202    }
203
204    #[test]
205    fn validate_endpoint_accepts_loopback_ipv6() {
206        // IPv6 literals in a URL authority must be bracketed (RFC 3986 §3.2.2).
207        assert!(validate_pyroscope_endpoint("http://[::1]:4040").is_ok());
208    }
209
210    #[test]
211    fn validate_endpoint_accepts_localhost() {
212        assert!(validate_pyroscope_endpoint("http://localhost:4040").is_ok());
213    }
214
215    #[test]
216    fn validate_endpoint_accepts_https_loopback() {
217        assert!(validate_pyroscope_endpoint("https://127.0.0.1:4040").is_ok());
218    }
219
220    #[test]
221    fn validate_endpoint_rejects_routable_ipv4() {
222        assert!(validate_pyroscope_endpoint("http://10.0.0.1:4040").is_err());
223    }
224
225    #[test]
226    fn validate_endpoint_rejects_userinfo_bypass() {
227        // Userinfo bypass: attacker tries to use loopback as userinfo but target evil.com
228        assert!(validate_pyroscope_endpoint("http://127.0.0.1:4040@evil.com/").is_err());
229    }
230
231    #[test]
232    fn validate_endpoint_rejects_userinfo_with_password() {
233        assert!(validate_pyroscope_endpoint("http://user:pass@localhost:4040").is_err());
234    }
235
236    #[test]
237    fn validate_endpoint_rejects_unix_socket_check() {
238        assert!(validate_pyroscope_endpoint("unix:///var/run/profiling.sock").is_ok());
239    }
240}
241
242#[cfg(all(
243    test,
244    feature = "profiling",
245    not(feature = "profiling-bridge-pyroscope-rs")
246))]
247mod tests_no_bridge {
248    use super::*;
249
250    #[test]
251    fn start_bridge_returns_none() {
252        let result = start_pyroscope_bridge("test-svc", "http://localhost:4040");
253        assert!(result.is_ok());
254        if let Ok(handle) = result {
255            assert!(handle.is_none());
256        }
257    }
258}