py_spy_for_datakit/
python_spy.rs

1use std;
2use std::collections::HashMap;
3#[cfg(all(target_os="linux", unwind))]
4use std::collections::HashSet;
5use std::mem::size_of;
6use std::slice;
7use std::path::Path;
8#[cfg(all(target_os="linux", unwind))]
9use std::iter::FromIterator;
10use regex::Regex;
11#[cfg(windows)]
12use regex::RegexBuilder;
13
14use anyhow::{Error, Result, Context};
15use lazy_static::lazy_static;
16use remoteprocess::{Process, ProcessMemory, Pid, Tid};
17use proc_maps::{get_process_maps, MapRange};
18
19
20use crate::binary_parser::{parse_binary, BinaryInfo};
21use crate::config::{Config, LockingStrategy, LineNo};
22#[cfg(unwind)]
23use crate::native_stack_trace::NativeStack;
24use crate::python_bindings::{pyruntime, v2_7_15, v3_3_7, v3_5_5, v3_6_6, v3_7_0, v3_8_0, v3_9_5, v3_10_0, v3_11_0};
25use crate::python_interpreters::{self, InterpreterState, ThreadState};
26use crate::python_threading::thread_name_lookup;
27use crate::stack_trace::{StackTrace, get_stack_traces, get_stack_trace};
28use crate::version::Version;
29
30/// Lets you retrieve stack traces of a running python program
31pub struct PythonSpy {
32    pub pid: Pid,
33    pub process: Process,
34    pub version: Version,
35    pub interpreter_address: usize,
36    pub threadstate_address: usize,
37    pub python_filename: std::path::PathBuf,
38    pub version_string: String,
39    pub config: Config,
40    #[cfg(unwind)]
41    pub native: Option<NativeStack>,
42    pub short_filenames: HashMap<String, Option<String>>,
43    pub python_thread_ids: HashMap<u64, Tid>,
44    pub python_thread_names: HashMap<u64, String>,
45    #[cfg(target_os="linux")]
46    pub dockerized: bool
47}
48
49fn error_if_gil(config: &Config, version: &Version, msg: &str) -> Result<(), Error> {
50    lazy_static! {
51        static ref WARNED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
52    }
53
54    if config.gil_only {
55        if !WARNED.load(std::sync::atomic::Ordering::Relaxed) {
56            // only print this once
57            eprintln!("Cannot detect GIL holding in version '{}' on the current platform (reason: {})", version, msg);
58            eprintln!("Please open an issue in https://github.com/benfred/py-spy with the Python version and your platform.");
59            WARNED.store(true, std::sync::atomic::Ordering::Relaxed);
60        }
61        Err(format_err!("Cannot detect GIL holding in version '{}' on the current platform (reason: {})", version, msg))
62    } else {
63        warn!("Unable to detect GIL usage: {}", msg);
64        Ok(())
65    }
66}
67
68impl PythonSpy {
69    /// Constructs a new PythonSpy object.
70    pub fn new(pid: Pid, config: &Config) -> Result<PythonSpy, Error> {
71        let process = remoteprocess::Process::new(pid)
72            .context("Failed to open process - check if it is running.")?;
73
74        // get basic process information (memory maps/symbols etc)
75        let python_info = PythonProcessInfo::new(&process)?;
76
77        // lock the process when loading up on freebsd (rather than locking
78        // on every memory read). Needs done after getting python process info
79        // because procmaps also tries to attach w/ ptrace on freebsd
80        #[cfg(target_os="freebsd")]
81        let _lock = process.lock();
82
83        let version = get_python_version(&python_info, &process)?;
84        info!("python version {} detected", version);
85
86        let interpreter_address = get_interpreter_address(&python_info, &process, &version)?;
87        info!("Found interpreter at 0x{:016x}", interpreter_address);
88
89        // lets us figure out which thread has the GIL
90         let threadstate_address = match version {
91             Version{major: 3, minor: 7..=11, ..} => {
92                match python_info.get_symbol("_PyRuntime") {
93                    Some(&addr) => {
94                        if let Some(offset) = pyruntime::get_tstate_current_offset(&version) {
95                            info!("Found _PyRuntime @ 0x{:016x}, getting gilstate.tstate_current from offset 0x{:x}",
96                                addr, offset);
97                            addr as usize + offset
98                        } else {
99                            error_if_gil(config, &version, "unknown pyruntime.gilstate.tstate_current offset")?;
100                            0
101                        }
102                    },
103                    None => {
104                        error_if_gil(config, &version, "failed to find _PyRuntime symbol")?;
105                        0
106                    }
107                }
108             },
109             _ => {
110                 match python_info.get_symbol("_PyThreadState_Current") {
111                    Some(&addr) => {
112                        info!("Found _PyThreadState_Current @ 0x{:016x}", addr);
113                        addr as usize
114                    },
115                    None => {
116                        error_if_gil(config, &version, "failed to find _PyThreadState_Current symbol")?;
117                        0
118                    }
119                }
120             }
121         };
122
123        let version_string = format!("python{}.{}", version.major, version.minor);
124
125        #[cfg(unwind)]
126        let native = if config.native {
127            Some(NativeStack::new(pid, python_info.python_binary, python_info.libpython_binary)?)
128        } else {
129            None
130        };
131
132        Ok(PythonSpy{pid, process, version, interpreter_address, threadstate_address,
133                     python_filename: python_info.python_filename,
134                     version_string,
135                     #[cfg(unwind)]
136                     native,
137                     #[cfg(target_os="linux")]
138                     dockerized: python_info.dockerized,
139                     config: config.clone(),
140                     short_filenames: HashMap::new(),
141                     python_thread_ids: HashMap::new(),
142                     python_thread_names: HashMap::new()})
143    }
144
145    /// Creates a PythonSpy object, retrying up to max_retries times.
146    /// Mainly useful for the case where the process is just started and
147    /// symbols or the python interpreter might not be loaded yet.
148    pub fn retry_new(pid: Pid, config: &Config, max_retries:u64) -> Result<PythonSpy, Error> {
149        let mut retries = 0;
150        loop {
151            let err = match PythonSpy::new(pid, config) {
152                Ok(mut process) => {
153                    // verify that we can load a stack trace before returning success
154                    match process.get_stack_traces() {
155                        Ok(_) => return Ok(process),
156                        Err(err) => err
157                    }
158                },
159                Err(err) => err
160            };
161
162            // If we failed, retry a couple times before returning the last error
163            retries += 1;
164            if retries >= max_retries {
165                return Err(err);
166            }
167            info!("Failed to connect to process, retrying. Error: {}", err);
168            std::thread::sleep(std::time::Duration::from_millis(20));
169        }
170    }
171
172    /// Gets a StackTrace for each thread in the current process
173    pub fn get_stack_traces(&mut self) -> Result<Vec<StackTrace>, Error> {
174        match self.version {
175            // ABI for 2.3/2.4/2.5/2.6/2.7 is compatible for our purpose
176            Version{major: 2, minor: 3..=7, ..} => self._get_stack_traces::<v2_7_15::_is>(),
177            Version{major: 3, minor: 3, ..} => self._get_stack_traces::<v3_3_7::_is>(),
178            // ABI for 3.4 and 3.5 is the same for our purposes
179            Version{major: 3, minor: 4, ..} => self._get_stack_traces::<v3_5_5::_is>(),
180            Version{major: 3, minor: 5, ..} => self._get_stack_traces::<v3_5_5::_is>(),
181            Version{major: 3, minor: 6, ..} => self._get_stack_traces::<v3_6_6::_is>(),
182            Version{major: 3, minor: 7, ..} => self._get_stack_traces::<v3_7_0::_is>(),
183            // v3.8.0a1 to v3.8.0a3 is compatible with 3.7 ABI, but later versions of 3.8.0 aren't
184            Version{major: 3, minor: 8, patch: 0, ..} => {
185                match self.version.release_flags.as_ref() {
186                    "a1" | "a2" | "a3" => self._get_stack_traces::<v3_7_0::_is>(),
187                    _ => self._get_stack_traces::<v3_8_0::_is>()
188                }
189            }
190            Version{major: 3, minor: 8, ..} => self._get_stack_traces::<v3_8_0::_is>(),
191            Version{major: 3, minor: 9, ..} => self._get_stack_traces::<v3_9_5::_is>(),
192            Version{major: 3, minor: 10, ..} => self._get_stack_traces::<v3_10_0::_is>(),
193            Version{major: 3, minor: 11, ..} => self._get_stack_traces::<v3_11_0::_is>(),
194            _ => Err(format_err!("Unsupported version of Python: {}", self.version)),
195        }
196    }
197
198    // implementation of get_stack_traces, where we have a type for the InterpreterState
199    fn _get_stack_traces<I: InterpreterState>(&mut self) -> Result<Vec<StackTrace>, Error> {
200        // Query the OS to get if each thread in the process is running or not
201        let mut thread_activity = HashMap::new();
202        for thread in self.process.threads()?.iter() {
203            let threadid: Tid = thread.id()?;
204            thread_activity.insert(threadid, thread.active()?);
205        }
206
207        // Lock the process if appropriate. Note we have to lock AFTER getting the thread
208        // activity status from the OS (otherwise each thread would report being inactive always).
209        // This has the potential for race conditions (in that the thread activity could change
210        // between getting the status and locking the thread, but seems unavoidable right now
211        let _lock = if self.config.blocking == LockingStrategy::Lock {
212            Some(self.process.lock().context("Failed to suspend process")?)
213        } else {
214            None
215        };
216
217        let gil_thread_id = self._get_gil_threadid::<I>()?;
218
219        // Get the python interpreter, and loop over all the python threads
220        let interp: I = self.process.copy_struct(self.interpreter_address)
221           .context("Failed to copy PyInterpreterState from process")?;
222
223        let mut traces = Vec::new();
224        let mut threads = interp.head();
225        while !threads.is_null() {
226            // Get the stack trace of the python thread
227            let thread = self.process.copy_pointer(threads).context("Failed to copy PyThreadState")?;
228            let mut trace = get_stack_trace(&thread, &self.process, self.config.dump_locals > 0, self.config.lineno)?;
229
230            // Try getting the native thread id
231            let python_thread_id = thread.thread_id();
232
233            // python 3.11+ has the native thread id directly on the PyThreadState object,
234            // so use that if available
235            trace.os_thread_id = thread.native_thread_id();
236
237            // for older versions of python, try using OS specific code to get the native
238            // thread id (doesn' work on freebsd, or on arm/i686 processors on linux)
239            if trace.os_thread_id.is_none() {
240                let mut os_thread_id = self._get_os_thread_id(python_thread_id, &interp)?;
241
242                // linux can see issues where pthread_ids get recycled for new OS threads,
243                // which totally breaks the caching we were doing here. Detect this and retry
244                if let Some(tid) = os_thread_id {
245                    if thread_activity.len() > 0 && !thread_activity.contains_key(&tid) {
246                        info!("clearing away thread id caches, thread {} has exited", tid);
247                        self.python_thread_ids.clear();
248                        self.python_thread_names.clear();
249                        os_thread_id = self._get_os_thread_id(python_thread_id, &interp)?;
250                    }
251                }
252
253                trace.os_thread_id = os_thread_id.map(|id| id as u64);
254            }
255
256            trace.thread_name = self._get_python_thread_name(python_thread_id);
257            trace.owns_gil = trace.thread_id == gil_thread_id;
258
259            // Figure out if the thread is sleeping from the OS if possible
260            trace.active = true;
261            if let Some(id) = trace.os_thread_id {
262                let id = id as Tid;
263                if let Some(active) = thread_activity.get(&id as _) {
264                    trace.active = *active;
265                }
266            }
267
268            // fallback to using a heuristic if we think the thread is still active
269            // Note that on linux the OS thread activity can only be gotten on x86_64
270            // processors and even then seems to be wrong occasionally in thinking 'select'
271            // calls are active (which seems related to the thread locking code,
272            // this problem doesn't seem to happen with the --nonblocking option)
273            // Note: this should be done before the native merging for correct results
274            if trace.active {
275                trace.active = !self._heuristic_is_thread_idle(&trace);
276            }
277
278            // Merge in the native stack frames if necessary
279            #[cfg(unwind)]
280            {
281                if self.config.native {
282                    if let Some(native) = self.native.as_mut() {
283                        let thread_id = trace.os_thread_id.ok_or_else(|| format_err!("failed to get os threadid"))?;
284                        let os_thread = remoteprocess::Thread::new(thread_id as Tid)?;
285                        trace.frames = native.merge_native_thread(&trace.frames, &os_thread)?
286                    }
287                }
288            }
289
290            for frame in &mut trace.frames {
291                frame.short_filename = self.shorten_filename(&frame.filename);
292                if let Some(locals) = frame.locals.as_mut() {
293                    use crate::python_data_access::format_variable;
294                    let max_length = (128 * self.config.dump_locals) as isize;
295                    for local in locals {
296                        let repr = format_variable::<I>(&self.process, &self.version, local.addr, max_length);
297                        local.repr = Some(repr.unwrap_or("?".to_owned()));
298                    }
299                }
300            }
301
302            traces.push(trace);
303
304            // This seems to happen occasionally when scanning BSS addresses for valid interpreters
305            if traces.len() > 4096 {
306                return Err(format_err!("Max thread recursion depth reached"));
307            }
308
309            threads = thread.next();
310        }
311        Ok(traces)
312    }
313
314    // heuristic fallback for determining if a thread is active, used
315    // when we don't have the ability to get the thread information from the OS
316    fn _heuristic_is_thread_idle(&self, trace: &StackTrace) -> bool {
317        let frames = &trace.frames;
318        if frames.is_empty() {
319            // we could have 0 python frames, but still be active running native
320            // code.
321            false
322        } else {
323            let frame = &frames[0];
324            (frame.name == "wait" && frame.filename.ends_with("threading.py")) ||
325            (frame.name == "select" && frame.filename.ends_with("selectors.py")) ||
326            (frame.name == "poll" && (frame.filename.ends_with("asyncore.py") ||
327                                    frame.filename.contains("zmq") ||
328                                    frame.filename.contains("gevent") ||
329                                    frame.filename.contains("tornado")))
330        }
331    }
332
333    #[cfg(windows)]
334    fn _get_os_thread_id<I: InterpreterState>(&mut self, python_thread_id: u64, _interp: &I) -> Result<Option<Tid>, Error> {
335        Ok(Some(python_thread_id as Tid))
336    }
337
338    #[cfg(target_os="macos")]
339    fn _get_os_thread_id<I: InterpreterState>(&mut self, python_thread_id: u64, _interp: &I) -> Result<Option<Tid>, Error> {
340        // If we've already know this threadid, we're good
341        if let Some(thread_id) = self.python_thread_ids.get(&python_thread_id) {
342            return Ok(Some(*thread_id));
343        }
344
345        for thread in self.process.threads()?.iter() {
346            // ok, this is crazy pants. is this 224 constant right?  Is this right for all versions of OSX? how is this determined?
347            // is this correct for all versions of python? Why does this even work?
348            let current_handle = thread.thread_handle()? - 224;
349            self.python_thread_ids.insert(current_handle, thread.id()?);
350        }
351
352        if let Some(thread_id) = self.python_thread_ids.get(&python_thread_id) {
353            return Ok(Some(*thread_id));
354        }
355        Ok(None)
356    }
357
358    #[cfg(all(target_os="linux", not(unwind)))]
359    fn _get_os_thread_id<I: InterpreterState>(&mut self, _python_thread_id: u64, _interp: &I) -> Result<Option<Tid>, Error> {
360        Ok(None)
361    }
362
363    #[cfg(all(target_os="linux", unwind))]
364    fn _get_os_thread_id<I: InterpreterState>(&mut self, python_thread_id: u64, interp: &I) -> Result<Option<Tid>, Error> {
365        // in nonblocking mode, we can't get the threadid reliably (method here requires reading the RBX
366        // register which requires a ptrace attach). fallback to heuristic thread activity here
367        if self.config.blocking == LockingStrategy::NonBlocking {
368            return Ok(None);
369        }
370
371        // likewise this doesn't yet work for profiling processes running inside docker containers from the host os
372        if self.dockerized {
373            return Ok(None);
374        }
375
376        // If we've already know this threadid, we're good
377        if let Some(thread_id) = self.python_thread_ids.get(&python_thread_id) {
378            return Ok(Some(*thread_id));
379        }
380
381        // Get a list of all the python thread ids
382        let mut all_python_threads = HashSet::new();
383        let mut threads = interp.head();
384        while !threads.is_null() {
385            let thread = self.process.copy_pointer(threads).context("Failed to copy PyThreadState")?;
386            let current = thread.thread_id();
387            all_python_threads.insert(current);
388            threads = thread.next();
389        }
390
391        let processed_os_threads: HashSet<Tid> = HashSet::from_iter(self.python_thread_ids.values().map(|x| *x));
392
393        let unwinder = self.process.unwinder()?;
394
395        // Try getting the pthread_id from the native stack registers for threads we haven't looked up yet
396        for thread in self.process.threads()?.iter() {
397            let threadid = thread.id()?;
398            if processed_os_threads.contains(&threadid) {
399                continue;
400            }
401
402            match self._get_pthread_id(&unwinder, &thread, &all_python_threads) {
403                Ok(pthread_id) => {
404                    if pthread_id != 0 {
405                        self.python_thread_ids.insert(pthread_id, threadid);
406                    }
407                },
408                Err(e) => { warn!("Failed to get get_pthread_id for {}: {}", threadid, e); }
409            };
410        }
411
412        // we can't get the python threadid for the main thread from registers,
413        // so instead assign the main threadid (pid) to the missing python thread
414        if !processed_os_threads.contains(&self.pid) {
415            let mut unknown_python_threadids = HashSet::new();
416            for python_thread_id in all_python_threads.iter() {
417                if !self.python_thread_ids.contains_key(python_thread_id) {
418                    unknown_python_threadids.insert(*python_thread_id);
419                }
420            }
421
422            if unknown_python_threadids.len() == 1 {
423                let python_thread_id = *unknown_python_threadids.iter().next().unwrap();
424                self.python_thread_ids.insert(python_thread_id, self.pid);
425            } else {
426                warn!("failed to get python threadid for main thread!");
427            }
428        }
429
430        if let Some(thread_id) = self.python_thread_ids.get(&python_thread_id) {
431            return Ok(Some(*thread_id));
432        }
433        info!("failed looking up python threadid for {}. known python_thread_ids {:?}. all_python_threads {:?}",
434            python_thread_id, self.python_thread_ids, all_python_threads);
435        Ok(None)
436    }
437
438
439    #[cfg(all(target_os="linux", unwind))]
440    pub fn _get_pthread_id(&self, unwinder: &remoteprocess::Unwinder, thread: &remoteprocess::Thread, threadids: &HashSet<u64>) -> Result<u64, Error> {
441        let mut pthread_id = 0;
442
443        let mut cursor = unwinder.cursor(thread)?;
444        while let Some(_) = cursor.next() {
445            // the pthread_id is usually in the top-level frame of the thread, but on some configs
446            // can be 2nd level. Handle this by taking the top-most rbx value that is one of the
447            // pthread_ids we're looking for
448            if let Ok(bx) = cursor.bx() {
449                if bx != 0 && threadids.contains(&bx) {
450                    pthread_id = bx;
451                }
452            }
453        }
454
455        Ok(pthread_id)
456    }
457
458    #[cfg(target_os="freebsd")]
459    fn _get_os_thread_id<I: InterpreterState>(&mut self, _python_thread_id: u64, _interp: &I) -> Result<Option<Tid>, Error> {
460        Ok(None)
461    }
462
463    fn _get_gil_threadid<I: InterpreterState>(&self) -> Result<u64, Error> {
464        // figure out what thread has the GIL by inspecting _PyThreadState_Current
465        if self.threadstate_address > 0 {
466            let addr: usize = self.process.copy_struct(self.threadstate_address)?;
467
468            // if the addr is 0, no thread is currently holding the GIL
469            if addr != 0 {
470                let threadstate: I::ThreadState = self.process.copy_struct(addr)?;
471                return Ok(threadstate.thread_id());
472            }
473        }
474        Ok(0)
475    }
476
477    fn _get_python_thread_name(&mut self, python_thread_id: u64) -> Option<String> {
478        match self.python_thread_names.get(&python_thread_id) {
479            Some(thread_name) => Some(thread_name.clone()),
480            None => {
481                self.python_thread_names = thread_name_lookup(self).unwrap_or_else(|| HashMap::new());
482                self.python_thread_names.get(&python_thread_id).map(|name| name.clone())
483            }
484        }
485    }
486
487    /// We want to display filenames without the boilerplate of the python installation
488    /// directory etc. This function looks only includes paths inside a python
489    /// package or subpackage, and not the path the package is installed at
490    fn shorten_filename(&mut self, filename: &str) -> Option<String> {
491        // if the user requested full filenames, skip shortening
492        if self.config.full_filenames {
493            return Some(filename.to_string());
494        }
495
496        // if we have figured out the short filename already, use it
497        if let Some(short) = self.short_filenames.get(filename) {
498            return short.clone();
499        }
500
501        // on linux the process could be running in docker, access the filename through procfs
502        #[cfg(target_os="linux")]
503        let filename_storage;
504
505        #[cfg(target_os="linux")]
506        let filename = if self.dockerized {
507            filename_storage = format!("/proc/{}/root{}", self.pid, filename);
508            if Path::new(&filename_storage).exists() {
509                &filename_storage
510            } else {
511                filename
512            }
513        } else {
514            filename
515        };
516
517        // only include paths that include an __init__.py
518        let mut path = Path::new(filename);
519        while let Some(parent) = path.parent() {
520            path = parent;
521            if !parent.join("__init__.py").exists() {
522                break;
523            }
524        }
525
526        // remote the parent prefix and convert to an optional string
527        let shortened = Path::new(filename)
528            .strip_prefix(path)
529            .ok()
530            .map(|p| p.to_string_lossy().to_string());
531
532        self.short_filenames.insert(filename.to_owned(), shortened.clone());
533        shortened
534    }
535}
536
537pub fn resolve_python_version(pid: Pid) -> Result<Version, Error> {
538    let process = Process::new(pid)
539        .context("Failed to open process - check if it is running.")?;
540
541    // get basic process information (memory maps/symbols etc)
542    let python_info = PythonProcessInfo::new(&process)?;
543
544    // lock the process when loading up on freebsd (rather than locking
545    // on every memory read). Needs done after getting python process info
546    // because procmaps also tries to attach w/ ptrace on freebsd
547    #[cfg(target_os="freebsd")]
548        let _lock = process.lock();
549
550    return get_python_version(&python_info, &process)
551}
552
553/// Returns the version of python running in the process.
554fn get_python_version(python_info: &PythonProcessInfo, process: &remoteprocess::Process)
555        -> Result<Version, Error> {
556    // If possible, grab the sys.version string from the processes memory (mac osx).
557    if let Some(&addr) = python_info.get_symbol("Py_GetVersion.version") {
558        info!("Getting version from symbol address");
559        if let Ok(bytes) = process.copy(addr as usize, 128) {
560            if let Ok(version) = Version::scan_bytes(&bytes) {
561                return Ok(version);
562            }
563        }
564    }
565
566    // otherwise get version info from scanning BSS section for sys.version string
567    if let Some(ref pb) = python_info.python_binary {
568        info!("Getting version from python binary BSS");
569        let bss = process.copy(pb.bss_addr as usize,
570                               pb.bss_size as usize)?;
571        match Version::scan_bytes(&bss) {
572            Ok(version) => return Ok(version),
573            Err(err) => info!("Failed to get version from BSS section: {}", err)
574        }
575    }
576
577    // try again if there is a libpython.so
578    if let Some(ref libpython) = python_info.libpython_binary {
579        info!("Getting version from libpython BSS");
580        let bss = process.copy(libpython.bss_addr as usize,
581                               libpython.bss_size as usize)?;
582        match Version::scan_bytes(&bss) {
583            Ok(version) => return Ok(version),
584            Err(err) => info!("Failed to get version from libpython BSS section: {}", err)
585        }
586    }
587
588    // the python_filename might have the version encoded in it (/usr/bin/python3.5 etc).
589    // try reading that in (will miss patch level on python, but that shouldn't matter)
590    info!("Trying to get version from path: {}", python_info.python_filename.display());
591    let path = Path::new(&python_info.python_filename);
592    if let Some(python) = path.file_name() {
593        if let Some(python) = python.to_str() {
594            if python.starts_with("python") {
595                let tokens: Vec<&str> = python[6..].split('.').collect();
596                if tokens.len() >= 2 {
597                    if let (Ok(major), Ok(minor)) = (tokens[0].parse::<u64>(), tokens[1].parse::<u64>()) {
598                        return Ok(Version{major, minor, patch:0, release_flags: "".to_owned()})
599                    }
600                }
601            }
602        }
603    }
604    Err(format_err!("Failed to find python version from target process"))
605}
606
607fn get_interpreter_address(python_info: &PythonProcessInfo,
608                           process: &remoteprocess::Process,
609                           version: &Version) -> Result<usize, Error> {
610    // get the address of the main PyInterpreterState object from loaded symbols if we can
611    // (this tends to be faster than scanning through the bss section)
612    match version {
613        Version{major: 3, minor: 7..=11, ..} => {
614            if let Some(&addr) = python_info.get_symbol("_PyRuntime") {
615                let addr = process.copy_struct(addr as usize + pyruntime::get_interp_head_offset(&version))?;
616
617                // Make sure the interpreter addr is valid before returning
618                match check_interpreter_addresses(&[addr], &python_info.maps, process, version) {
619                    Ok(addr) => return Ok(addr),
620                    Err(_) => { warn!("Interpreter address from _PyRuntime symbol is invalid {:016x}", addr); }
621                };
622            }
623        },
624        _ => {
625            if let Some(&addr) = python_info.get_symbol("interp_head") {
626                let addr = process.copy_struct(addr as usize)?;
627                match check_interpreter_addresses(&[addr], &python_info.maps, process, version) {
628                    Ok(addr) => return Ok(addr),
629                    Err(_) => { warn!("Interpreter address from interp_head symbol is invalid {:016x}", addr); }
630                };
631            }
632        }
633    };
634    info!("Failed to get interp_head from symbols, scanning BSS section from main binary");
635
636    // try scanning the BSS section of the binary for things that might be the interpreterstate
637    let err =
638        if let Some(ref pb) = python_info.python_binary {
639            match get_interpreter_address_from_binary(pb, &python_info.maps, process, version) {
640                Ok(addr) => return Ok(addr),
641                err => Some(err)
642            }
643        } else {
644            None
645        };
646    // Before giving up, try again if there is a libpython.so
647    if let Some(ref lpb) = python_info.libpython_binary {
648        info!("Failed to get interpreter from binary BSS, scanning libpython BSS");
649        match get_interpreter_address_from_binary(lpb, &python_info.maps, process, version) {
650            Ok(addr) => return Ok(addr),
651            lib_err => err.unwrap_or(lib_err)
652        }
653    } else {
654        err.expect("Both python and libpython are invalid.")
655    }
656}
657
658fn get_interpreter_address_from_binary(binary: &BinaryInfo,
659                                       maps: &[MapRange],
660                                       process: &remoteprocess::Process,
661                                       version: &Version) -> Result<usize, Error> {
662    // We're going to scan the BSS/data section for things, and try to narrowly scan things that
663    // look like pointers to PyinterpreterState
664    let bss = process.copy(binary.bss_addr as usize, binary.bss_size as usize)?;
665
666    #[allow(clippy::cast_ptr_alignment)]
667    let addrs = unsafe { slice::from_raw_parts(bss.as_ptr() as *const usize, bss.len() / size_of::<usize>()) };
668    check_interpreter_addresses(addrs, maps, process, version)
669}
670
671// Checks whether a block of memory (from BSS/.data etc) contains pointers that are pointing
672// to a valid PyInterpreterState
673fn check_interpreter_addresses(addrs: &[usize],
674                               maps: &[MapRange],
675                               process: &remoteprocess::Process,
676                               version: &Version) -> Result<usize, Error> {
677    // On windows, we can't just check if a pointer is valid by looking to see if it points
678    // to something in the virtual memory map. Brute-force it instead
679    #[cfg(windows)]
680    fn maps_contain_addr(_: usize, _: &[MapRange]) -> bool { true }
681
682    #[cfg(not(windows))]
683    use proc_maps::maps_contain_addr;
684
685    // This function does all the work, but needs a type of the interpreter
686    fn check<I>(addrs: &[usize],
687                maps: &[MapRange],
688                process: &remoteprocess::Process) -> Result<usize, Error>
689            where I: python_interpreters::InterpreterState {
690        for &addr in addrs {
691            if maps_contain_addr(addr, maps) {
692                // this address points to valid memory. try loading it up as a PyInterpreterState
693                // to further check
694                let interp: I = match process.copy_struct(addr) {
695                    Ok(interp) => interp,
696                    Err(_) => continue
697                };
698
699                // get the pythreadstate pointer from the interpreter object, and if it is also
700                // a valid pointer then load it up.
701                let threads = interp.head();
702                if maps_contain_addr(threads as usize, maps) {
703                    // If the threadstate points back to the interpreter like we expect, then
704                    // this is almost certainly the address of the intrepreter
705                    let thread = match process.copy_pointer(threads) {
706                        Ok(thread) => thread,
707                        Err(_) => continue
708                    };
709
710                    // as a final sanity check, try getting the stack_traces, and only return if this works
711                    if thread.interp() as usize == addr && get_stack_traces(&interp, process, LineNo::NoLine).is_ok() {
712                        return Ok(addr);
713                    }
714                }
715            }
716        }
717        Err(format_err!("Failed to find a python interpreter in the .data section"))
718    }
719
720    // different versions have different layouts, check as appropriate
721    match version {
722        Version{major: 2, minor: 3..=7, ..} => check::<v2_7_15::_is>(addrs, maps, process),
723        Version{major: 3, minor: 3, ..} => check::<v3_3_7::_is>(addrs, maps, process),
724        Version{major: 3, minor: 4..=5, ..} => check::<v3_5_5::_is>(addrs, maps, process),
725        Version{major: 3, minor: 6, ..} => check::<v3_6_6::_is>(addrs, maps, process),
726        Version{major: 3, minor: 7, ..} => check::<v3_7_0::_is>(addrs, maps, process),
727        Version{major: 3, minor: 8, patch: 0, ..} => {
728            match version.release_flags.as_ref() {
729                "a1" | "a2" | "a3" => check::<v3_7_0::_is>(addrs, maps, process),
730                _ => check::<v3_8_0::_is>(addrs, maps, process)
731            }
732        },
733        Version{major: 3, minor: 8, ..} => check::<v3_8_0::_is>(addrs, maps, process),
734        Version{major: 3, minor: 9, ..} => check::<v3_9_5::_is>(addrs, maps, process),
735        Version{major: 3, minor: 10, ..} => check::<v3_10_0::_is>(addrs, maps, process),
736        Version{major: 3, minor: 11, ..} => check::<v3_11_0::_is>(addrs, maps, process),
737        _ => Err(format_err!("Unsupported version of Python: {}", version))
738    }
739}
740
741/// Holds information about the python process: memory map layout, parsed binary info
742/// for python /libpython etc.
743pub struct PythonProcessInfo {
744    python_binary: Option<BinaryInfo>,
745    // if python was compiled with './configure --enabled-shared', code/symbols will
746    // be in a libpython.so file instead of the executable. support that.
747    libpython_binary: Option<BinaryInfo>,
748    maps: Vec<MapRange>,
749    python_filename: std::path::PathBuf,
750    #[cfg(target_os="linux")]
751    dockerized: bool,
752}
753
754impl PythonProcessInfo {
755    fn new(process: &remoteprocess::Process) -> Result<PythonProcessInfo, Error> {
756        let filename = process.exe()
757            .context("Failed to get process executable name. Check that the process is running.")?;
758
759        #[cfg(windows)]
760        let filename = filename.to_lowercase();
761
762        #[cfg(windows)]
763        let is_python_bin = |pathname: &str| pathname.to_lowercase() == filename;
764
765        #[cfg(not(windows))]
766        let is_python_bin = |pathname: &str| pathname == filename;
767
768        // get virtual memory layout
769        let maps = get_process_maps(process.pid)?;
770        info!("Got virtual memory maps from pid {}:", process.pid);
771        for map in &maps {
772            debug!("map: {:016x}-{:016x} {}{}{} {}", map.start(), map.start() + map.size(),
773                if map.is_read() {'r'} else {'-'}, if map.is_write() {'w'} else {'-'}, if map.is_exec() {'x'} else {'-'},
774                map.filename().unwrap_or(&std::path::PathBuf::from("")).display());
775        }
776
777        // parse the main python binary
778        let (python_binary, python_filename) = {
779            // Get the memory address for the executable by matching against virtual memory maps
780            let map = maps.iter()
781                .find(|m| {
782                    if let Some(pathname) = m.filename() {
783                        if let Some(pathname) = pathname.to_str() {
784                            return is_python_bin(pathname) && m.is_exec();
785                        }
786                    }
787                    false
788                });
789
790            let map = match map {
791                Some(map) => map,
792                None => {
793                    warn!("Failed to find '{}' in virtual memory maps, falling back to first map region", filename);
794                    // If we failed to find the executable in the virtual memory maps, just take the first file we find
795                    // sometimes on windows get_process_exe returns stale info =( https://github.com/benfred/py-spy/issues/40
796                    // and on all operating systems I've tried, the exe is the first region in the maps
797                    &maps.first().ok_or_else(|| format_err!("Failed to get virtual memory maps from process"))?
798                }
799            };
800
801            let filename = std::path::PathBuf::from(filename);
802
803            // TODO: consistent types? u64 -> usize? for map.start etc
804            #[allow(unused_mut)]
805            let python_binary = parse_binary(process.pid, &filename, map.start() as u64, map.size() as u64, true)
806                .and_then(|mut pb| {
807                    // windows symbols are stored in separate files (.pdb), load
808                    #[cfg(windows)]
809                    {
810                        get_windows_python_symbols(process.pid, &filename, map.start() as u64)
811                            .map(|symbols| { pb.symbols.extend(symbols); pb })
812                            .map_err(|err| err.into())
813                    }
814
815                    // For OSX, need to adjust main binary symbols by subtracting _mh_execute_header
816                    // (which we've added to by map.start already, so undo that here)
817                    #[cfg(target_os = "macos")]
818                    {
819                        let offset = pb.symbols["_mh_execute_header"] - map.start() as u64;
820                        for address in pb.symbols.values_mut() {
821                            *address -= offset;
822                        }
823
824                        if pb.bss_addr != 0 {
825                            pb.bss_addr -= offset;
826                        }
827                    }
828
829                    #[cfg(not(windows))]
830                    Ok(pb)
831                });
832
833            (python_binary, filename.clone())
834        };
835
836        // likewise handle libpython for python versions compiled with --enabled-shared
837        let libpython_binary = {
838            let libmap = maps.iter()
839                .find(|m| {
840                    if let Some(pathname) = m.filename() {
841                        if let Some(pathname) = pathname.to_str() {
842                            return is_python_lib(pathname) && m.is_exec();
843                        }
844                    }
845                    false
846                });
847
848            let mut libpython_binary: Option<BinaryInfo> = None;
849            if let Some(libpython) = libmap {
850                if let Some(filename) = &libpython.filename() {
851                    info!("Found libpython binary @ {}", filename.display());
852                    #[allow(unused_mut)]
853                    let mut parsed = parse_binary(process.pid, filename, libpython.start() as u64, libpython.size() as u64, false)?;
854                    #[cfg(windows)]
855                    parsed.symbols.extend(get_windows_python_symbols(process.pid, filename, libpython.start() as u64)?);
856                    libpython_binary = Some(parsed);
857                }
858            }
859
860            // On OSX, it's possible that the Python library is a dylib loaded up from the system
861            // framework (like /System/Library/Frameworks/Python.framework/Versions/2.7/Python)
862            // In this case read in the dyld_info information and figure out the filename from there
863            #[cfg(target_os = "macos")]
864            {
865                if libpython_binary.is_none() {
866                    use proc_maps::mac_maps::get_dyld_info;
867                    let dyld_infos = get_dyld_info(process.pid)?;
868
869                    for dyld in &dyld_infos {
870                        let segname = unsafe { std::ffi::CStr::from_ptr(dyld.segment.segname.as_ptr()) };
871                        debug!("dyld: {:016x}-{:016x} {:10} {}",
872                            dyld.segment.vmaddr, dyld.segment.vmaddr + dyld.segment.vmsize,
873                            segname.to_string_lossy(), dyld.filename.display());
874                    }
875
876                    let python_dyld_data = dyld_infos.iter()
877                        .find(|m| {
878                            if let Some(filename) = m.filename.to_str() {
879                                return is_python_framework(filename) &&
880                                      m.segment.segname[0..7] == [95, 95, 68, 65, 84, 65, 0];
881                            }
882                            false
883                        });
884
885
886                    if let Some(libpython) = python_dyld_data {
887                        info!("Found libpython binary from dyld @ {}", libpython.filename.display());
888
889                        let mut binary = parse_binary(process.pid, &libpython.filename, libpython.segment.vmaddr, libpython.segment.vmsize, false)?;
890
891                        // TODO: bss addr offsets returned from parsing binary are wrong
892                        // (assumes data section isn't split from text section like done here).
893                        // BSS occurs somewhere in the data section, just scan that
894                        // (could later tighten this up to look at segment sections too)
895                        binary.bss_addr = libpython.segment.vmaddr;
896                        binary.bss_size = libpython.segment.vmsize;
897                        libpython_binary = Some(binary);
898                    }
899                }
900            }
901
902            libpython_binary
903        };
904
905        // If we have a libpython binary - we can tolerate failures on parsing the main python binary.
906        let python_binary = match libpython_binary {
907            None => Some(python_binary.context("Failed to parse python binary")?),
908            _ => python_binary.ok(),
909        };
910
911        #[cfg(target_os="linux")]
912        let dockerized = is_dockerized(process.pid).unwrap_or(false);
913
914        Ok(PythonProcessInfo{python_binary, libpython_binary, maps, python_filename,
915                             #[cfg(target_os="linux")]
916                             dockerized
917        })
918    }
919
920    pub fn get_symbol(&self, symbol: &str) -> Option<&u64> {
921        if let Some(ref pb) = self.python_binary {
922            if let Some(addr) = pb.symbols.get(symbol) {
923                info!("got symbol {} (0x{:016x}) from python binary", symbol, addr);
924                return Some(addr);
925            }
926        }
927
928        if let Some(ref binary) = self.libpython_binary {
929            if let Some(addr) = binary.symbols.get(symbol) {
930                info!("got symbol {} (0x{:016x}) from libpython binary", symbol, addr);
931                return Some(addr);
932            }
933        }
934        None
935    }
936}
937
938#[cfg(target_os="linux")]
939fn is_dockerized(pid: Pid) -> Result<bool, Error> {
940    let self_mnt = std::fs::read_link("/proc/self/ns/mnt")?;
941    let target_mnt = std::fs::read_link(&format!("/proc/{}/ns/mnt", pid))?;
942    Ok(self_mnt != target_mnt)
943}
944
945// We can't use goblin to parse external symbol files (like in a separate .pdb file) on windows,
946// So use the win32 api to load up the couple of symbols we need on windows. Note:
947// we still can get export's from the PE file
948#[cfg(windows)]
949pub fn get_windows_python_symbols(pid: Pid, filename: &Path, offset: u64) -> std::io::Result<HashMap<String, u64>> {
950    use proc_maps::win_maps::SymbolLoader;
951
952    let handler = SymbolLoader::new(pid)?;
953    let _module = handler.load_module(filename)?; // need to keep this module in scope
954
955    let mut ret = HashMap::new();
956
957    // currently we only need a subset of symbols, and enumerating the symbols is
958    // expensive (via SymEnumSymbolsW), so rather than load up all symbols like we
959    // do for goblin, just load the the couple we need directly.
960    for symbol in ["_PyThreadState_Current", "interp_head", "_PyRuntime"].iter() {
961        if let Ok((base, addr)) = handler.address_from_name(symbol) {
962            // If we have a module base (ie from PDB), need to adjust by the offset
963            // otherwise seems like we can take address directly
964            let addr = if base == 0 { addr } else { offset + addr - base };
965            ret.insert(String::from(*symbol), addr);
966        }
967    }
968
969    Ok(ret)
970}
971
972#[cfg(any(target_os="linux", target_os="freebsd"))]
973pub fn is_python_lib(pathname: &str) -> bool {
974    lazy_static! {
975        static ref RE: Regex = Regex::new(r"/libpython\d.\d\d?(m|d|u)?.so").unwrap();
976    }
977    RE.is_match(pathname)
978}
979
980#[cfg(target_os="macos")]
981pub fn is_python_lib(pathname: &str) -> bool {
982    lazy_static! {
983        static ref RE: Regex = Regex::new(r"/libpython\d.\d\d?(m|d|u)?.(dylib|so)$").unwrap();
984    }
985    RE.is_match(pathname) || is_python_framework(pathname)
986}
987
988#[cfg(windows)]
989pub fn is_python_lib(pathname: &str) -> bool {
990    lazy_static! {
991        static ref RE: Regex = RegexBuilder::new(r"\\python\d\d\d?(m|d|u)?.dll$").case_insensitive(true).build().unwrap();
992    }
993    RE.is_match(pathname)
994}
995
996#[cfg(target_os="macos")]
997pub fn is_python_framework(pathname: &str) -> bool {
998    pathname.ends_with("/Python")  &&
999    !pathname.contains("Python.app")
1000}
1001
1002#[cfg(test)]
1003mod tests {
1004    use super::*;
1005
1006    #[cfg(target_os="macos")]
1007    #[test]
1008    fn test_is_python_lib() {
1009        assert!(is_python_lib("~/Anaconda2/lib/libpython2.7.dylib"));
1010
1011        // python lib configured with --with-pydebug (flag: d)
1012        assert!(is_python_lib("/lib/libpython3.4d.dylib"));
1013
1014        // configured --with-pymalloc (flag: m)
1015        assert!(is_python_lib("/usr/local/lib/libpython3.8m.dylib"));
1016
1017        // python2 configured with --with-wide-unicode (flag: u)
1018        assert!(is_python_lib("./libpython2.7u.dylib"));
1019
1020        assert!(!is_python_lib("/libboost_python.dylib"));
1021        assert!(!is_python_lib("/lib/heapq.cpython-36m-darwin.dylib"));
1022    }
1023
1024    #[cfg(any(target_os="linux", target_os="freebsd"))]
1025    #[test]
1026    fn test_is_python_lib() {
1027        // libpython bundled by pyinstaller https://github.com/benfred/py-spy/issues/42
1028        assert!(is_python_lib("/tmp/_MEIOqzg01/libpython2.7.so.1.0"));
1029
1030        // test debug/malloc/unicode flags
1031        assert!(is_python_lib("./libpython2.7.so"));
1032        assert!(is_python_lib("/usr/lib/libpython3.4d.so"));
1033        assert!(is_python_lib("/usr/local/lib/libpython3.8m.so"));
1034        assert!(is_python_lib("/usr/lib/libpython2.7u.so"));
1035
1036        // don't blindly match libraries with python in the name (boost_python etc)
1037        assert!(!is_python_lib("/usr/lib/libboost_python.so"));
1038        assert!(!is_python_lib("/usr/lib/x86_64-linux-gnu/libboost_python-py27.so.1.58.0"));
1039        assert!(!is_python_lib("/usr/lib/libboost_python-py35.so"));
1040
1041    }
1042
1043    #[cfg(windows)]
1044    #[test]
1045    fn test_is_python_lib() {
1046        assert!(is_python_lib("C:\\Users\\test\\AppData\\Local\\Programs\\Python\\Python37\\python37.dll"));
1047        // .NET host via https://github.com/pythonnet/pythonnet
1048        assert!(is_python_lib("C:\\Users\\test\\AppData\\Local\\Programs\\Python\\Python37\\python37.DLL"));
1049    }
1050
1051
1052    #[cfg(target_os="macos")]
1053    #[test]
1054    fn test_python_frameworks() {
1055        // homebrew v2
1056        assert!(!is_python_framework("/usr/local/Cellar/python@2/2.7.15_1/Frameworks/Python.framework/Versions/2.7/Resources/Python.app/Contents/MacOS/Python"));
1057        assert!(is_python_framework("/usr/local/Cellar/python@2/2.7.15_1/Frameworks/Python.framework/Versions/2.7/Python"));
1058
1059        // System python from osx 10.13.6 (high sierra)
1060        assert!(!is_python_framework("/System/Library/Frameworks/Python.framework/Versions/2.7/Resources/Python.app/Contents/MacOS/Python"));
1061        assert!(is_python_framework("/System/Library/Frameworks/Python.framework/Versions/2.7/Python"));
1062
1063        // pyenv 3.6.6 with OSX framework enabled (https://github.com/benfred/py-spy/issues/15)
1064        // env PYTHON_CONFIGURE_OPTS="--enable-framework" pyenv install 3.6.6
1065        assert!(is_python_framework("/Users/ben/.pyenv/versions/3.6.6/Python.framework/Versions/3.6/Python"));
1066        assert!(!is_python_framework("/Users/ben/.pyenv/versions/3.6.6/Python.framework/Versions/3.6/Resources/Python.app/Contents/MacOS/Python"));
1067
1068        // single file pyinstaller
1069        assert!(is_python_framework("/private/var/folders/3x/qy479lpd1fb2q88lc9g4d3kr0000gn/T/_MEI2Akvi8/Python"));
1070    }
1071}