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_pprofrs::{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::PyroscopeAgent::builder(pyroscope_endpoint, service_name)
106        .backend(pprof_backend(PprofConfig::new().sample_rate(100)))
107        .build()?
108        .start()?;
109
110    let (add_tag, remove_tag) = agent.tag_wrapper();
111    PROFILING_TAG_FNS
112        .set((Box::new(add_tag), Box::new(remove_tag)))
113        .ok();
114
115    Ok(Some(ProfilingHandle { agent: Some(agent) }))
116}
117
118/// No-op bridge for when profiling is enabled but the pyroscope feature is not.
119#[cfg(all(feature = "profiling", not(feature = "profiling-bridge-pyroscope-rs")))]
120pub(crate) fn start_pyroscope_bridge(
121    _service_name: &str,
122    _pyroscope_endpoint: &str,
123) -> Result<Option<ProfilingHandle>, Box<dyn Error>> {
124    Ok(None)
125}
126
127/// Tracing layer that tags active span enter/exit with trace_id and span_id
128/// in the running pyroscope agent.
129#[cfg(feature = "profiling-bridge-pyroscope-rs")]
130pub struct ProfilingTagLayer;
131
132#[cfg(feature = "profiling-bridge-pyroscope-rs")]
133impl<S> tracing_subscriber::Layer<S> for ProfilingTagLayer
134where
135    S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
136{
137    fn on_enter(&self, _id: &tracing::span::Id, _ctx: tracing_subscriber::layer::Context<'_, S>) {
138        if let Some((add_tag, _)) = PROFILING_TAG_FNS.get() {
139            let cx = opentelemetry::Context::current();
140            let span_ref = cx.span();
141            let span_context = span_ref.span_context();
142            if span_context.is_valid() {
143                let trace_id = span_context.trace_id();
144                let span_id = span_context.span_id();
145                let _ = add_tag("trace_id".to_string(), format!("{trace_id:x}"));
146                let _ = add_tag("span_id".to_string(), format!("{span_id:x}"));
147            }
148        }
149    }
150
151    fn on_exit(&self, _id: &tracing::span::Id, _ctx: tracing_subscriber::layer::Context<'_, S>) {
152        if let Some((_, remove_tag)) = PROFILING_TAG_FNS.get() {
153            let cx = opentelemetry::Context::current();
154            let span_ref = cx.span();
155            let span_context = span_ref.span_context();
156            if span_context.is_valid() {
157                let trace_id = span_context.trace_id();
158                let span_id = span_context.span_id();
159                let _ = remove_tag("trace_id".to_string(), format!("{trace_id:x}"));
160                let _ = remove_tag("span_id".to_string(), format!("{span_id:x}"));
161            }
162        }
163    }
164}
165
166#[cfg(all(test, feature = "profiling-bridge-pyroscope-rs"))]
167mod tests {
168    use super::*;
169
170    #[test]
171    fn start_bridge_with_nonexistent_server() {
172        let result = start_pyroscope_bridge("test-svc", "http://localhost:4040");
173        assert!(
174            result.is_ok(),
175            "pyroscope agent start() is lazy and does not eagerly connect"
176        );
177        if let Ok(Some(_handle)) = result {
178            // Bridge is active
179        }
180    }
181
182    #[test]
183    fn start_bridge_multiple_times_ignores_second() {
184        let result1 = start_pyroscope_bridge("test-svc-1", "http://localhost:4040");
185        assert!(result1.is_ok());
186        let result2 = start_pyroscope_bridge("test-svc-2", "http://localhost:4041");
187        assert!(result2.is_ok());
188        // Second call is a no-op: the `pprof` backend only supports one
189        // process-wide profiler guard, so the bridge returns `Ok(None)`.
190        assert!(result2.unwrap().is_none());
191    }
192
193    #[test]
194    fn validate_endpoint_accepts_loopback_ipv4() {
195        assert!(validate_pyroscope_endpoint("http://127.0.0.1:4040").is_ok());
196    }
197
198    #[test]
199    fn validate_endpoint_accepts_loopback_ipv6() {
200        // IPv6 literals in a URL authority must be bracketed (RFC 3986 §3.2.2).
201        assert!(validate_pyroscope_endpoint("http://[::1]:4040").is_ok());
202    }
203
204    #[test]
205    fn validate_endpoint_accepts_localhost() {
206        assert!(validate_pyroscope_endpoint("http://localhost:4040").is_ok());
207    }
208
209    #[test]
210    fn validate_endpoint_accepts_https_loopback() {
211        assert!(validate_pyroscope_endpoint("https://127.0.0.1:4040").is_ok());
212    }
213
214    #[test]
215    fn validate_endpoint_rejects_routable_ipv4() {
216        assert!(validate_pyroscope_endpoint("http://10.0.0.1:4040").is_err());
217    }
218
219    #[test]
220    fn validate_endpoint_rejects_userinfo_bypass() {
221        // Userinfo bypass: attacker tries to use loopback as userinfo but target evil.com
222        assert!(validate_pyroscope_endpoint("http://127.0.0.1:4040@evil.com/").is_err());
223    }
224
225    #[test]
226    fn validate_endpoint_rejects_userinfo_with_password() {
227        assert!(validate_pyroscope_endpoint("http://user:pass@localhost:4040").is_err());
228    }
229
230    #[test]
231    fn validate_endpoint_rejects_unix_socket_check() {
232        assert!(validate_pyroscope_endpoint("unix:///var/run/profiling.sock").is_ok());
233    }
234}
235
236#[cfg(all(
237    test,
238    feature = "profiling",
239    not(feature = "profiling-bridge-pyroscope-rs")
240))]
241mod tests_no_bridge {
242    use super::*;
243
244    #[test]
245    fn start_bridge_returns_none() {
246        let result = start_pyroscope_bridge("test-svc", "http://localhost:4040");
247        assert!(result.is_ok());
248        if let Ok(handle) = result {
249            assert!(handle.is_none());
250        }
251    }
252}