lightswitch 0.7.0

use std::collections::HashMap;
use std::path::PathBuf;
use std::time::Instant;
use tracing::debug;

use lightswitch_object::BuildId;
use lightswitch_object::ElfLoad;
use lightswitch_object::ExecutableId;
use lightswitch_object::Runtime;

use crate::util::FileId;

pub type Pid = i32;

/// What type of mapping we are dealing with.
#[derive(Debug, Clone, PartialEq)]
pub enum ExecutableMappingType {
    /// An object file that got loaded from disk.
    FileBacked,
    /// Not file backed, typically produced by a JIT runtime.
    Anonymous,
    /// Special mapping to optimise certain system calls.
    Vdso,
    Kernel,
}

#[derive(Debug, PartialEq, Clone)]
pub enum ProcessStatus {
    Running { evicted: bool },
    Exited,
}

impl ProcessStatus {
    /// Checks if the process is still running as far as we know and that we
    /// haven't evicted any of the BPF structures for it.
    pub fn is_running_and_loaded(&self) -> bool {
        matches!(self, ProcessStatus::Running { evicted: false })
    }

    /// Sets a process as evicted. This means that the userspace structures are
    /// left as-is but the BPF structures have been unloaded.
    pub fn evict(&mut self) {
        *self = ProcessStatus::Running { evicted: true };
    }

    /// Sets a process as running with the full BPF structures loaded.
    pub fn running(&mut self) {
        *self = ProcessStatus::Running { evicted: false };
    }
}

#[derive(Debug, Clone)]
pub struct ProcessInfo {
    pub comm: String,
    pub status: ProcessStatus,
    pub mappings: ExecutableMappings,
    /// If there's at least one 32 bit executable linked in this process.
    pub compatibility_mode: bool,
    pub last_used: Instant,
    pub file_id: FileId,
}

/// Stores information for a executable mapping with all
/// the information we need to do everything symbolization
/// related.
#[derive(Debug, Clone)]
pub struct ExecutableMapping {
    pub executable_id: ExecutableId,
    pub kind: ExecutableMappingType,
    pub start_addr: u64,
    pub end_addr: u64,
    // kaslr info etc etc
    pub offset: u64,
    pub load_address: u64,
    pub soft_delete: bool,
}

#[derive(Debug, Clone)]
pub struct ExecutableMappings(pub Vec<ExecutableMapping>);

impl ExecutableMappings {
    /// Find the executable mapping a given virtual address falls into.
    pub fn for_address(&self, virtual_address: &u64) -> Option<&ExecutableMapping> {
        self.0
            .iter()
            .find(|&mapping| (mapping.start_addr..mapping.end_addr).contains(virtual_address))
    }
}

impl ExecutableMapping {
    /// Soft delete a mapping. We don't want to delete it straight away as we
    /// might need it for a bit longer for normalization and / or local
    /// symbolization.
    pub fn mark_as_deleted(
        &mut self,
        object_files: &mut HashMap<ExecutableId, ObjectFileInfo>,
    ) -> bool {
        // The executable mapping can be removed at a later time, and function
        // might be called multiple times. To avoid this, we keep track
        // of whether this mapping has been soft deleted.
        if self.soft_delete {
            return false;
        }
        self.soft_delete = true;

        if let Some(object_file) = object_files.get_mut(&self.executable_id) {
            // Object files are also soft deleted, so do not try to decrease the
            // reference count if it's already zero.
            if object_file.references == 0 {
                return false;
            }

            object_file.references -= 1;

            if object_file.references == 0 {
                debug!(
                    "object file with path {} can be deleted",
                    object_file.path.display()
                );
                return true;
            }

            debug_assert!(
                object_file.references >= 0,
                "Reference count for {} is negative: {}",
                object_file.path.display(),
                object_file.references,
            );
        }
        false
    }
}

pub struct ObjectFileInfo {
    pub build_id: Option<BuildId>,
    pub path: PathBuf,
    pub elf_load_segments: Vec<ElfLoad>,
    pub is_dyn: bool,
    pub references: i64,
    pub native_unwind_info_size: Option<u64>,
    pub is_vdso: bool,
    pub runtime: Runtime,
    pub is_64: bool,
}

impl Clone for ObjectFileInfo {
    fn clone(&self) -> Self {
        ObjectFileInfo {
            build_id: self.build_id.clone(),
            path: self.path.clone(),
            elf_load_segments: self.elf_load_segments.clone(),
            is_dyn: self.is_dyn,
            references: self.references,
            native_unwind_info_size: self.native_unwind_info_size,
            is_vdso: self.is_vdso,
            runtime: self.runtime.clone(),
            is_64: self.is_64,
        }
    }
}

impl ObjectFileInfo {
    /// For a virtual address return the offset within the object file. This is
    /// necessary for off-host symbolization. In order to do this we must check
    /// every `PT_LOAD` segment.
    pub fn normalized_address(
        &self,
        virtual_address: u64,
        mapping: &ExecutableMapping,
    ) -> Option<u64> {
        if mapping.kind == ExecutableMappingType::Kernel {
            return Some(virtual_address - mapping.offset);
        }

        let offset = virtual_address - mapping.start_addr + mapping.offset;

        for segment in &self.elf_load_segments {
            let address_range = segment.p_offset..(segment.p_offset + segment.p_filesz);
            if address_range.contains(&offset) {
                return Some(offset - segment.p_offset + segment.p_vaddr);
            }
        }

        None
    }
}

/// Returns the path for the procfs file referencing the mapped file which can
/// be read from even after the executable gets deleted and can't be accessed
/// via its path anymore.
pub(crate) fn opened_exe_path(pid: Pid, start_address: u64, end_address: u64) -> PathBuf {
    PathBuf::from(format!(
        "/proc/{pid}/map_files/{:x}-{:x}",
        start_address, end_address
    ))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_address_normalization() {
        let mut object_file_info = ObjectFileInfo {
            build_id: None,
            path: "/".into(),
            elf_load_segments: vec![],
            is_dyn: false,
            references: 0,
            native_unwind_info_size: None,
            is_vdso: false,
            runtime: Runtime::CLike,
            is_64: true,
        };

        let mapping = ExecutableMapping {
            executable_id: ExecutableId(0x0),
            kind: ExecutableMappingType::FileBacked,
            start_addr: 0x100,
            end_addr: 0x100 + 100,
            offset: 0x0,
            load_address: 0x0,
            soft_delete: false,
        };

        // no elf segments
        assert!(
            object_file_info
                .normalized_address(0x110, &mapping)
                .is_none()
        );

        // matches an elf segment
        object_file_info.elf_load_segments = vec![ElfLoad {
            p_offset: 0x1,
            p_vaddr: 0x0,
            p_filesz: 0x20,
        }];
        assert_eq!(
            object_file_info.normalized_address(0x110, &mapping),
            Some(0xF)
        );
        // does not match any elf segments
        object_file_info.elf_load_segments = vec![ElfLoad {
            p_offset: 0x0,
            p_vaddr: 0x0,
            p_filesz: 0x5,
        }];
        assert!(
            object_file_info
                .normalized_address(0x110, &mapping)
                .is_none()
        );
    }
}