1
use caps::{CapSet, Capability, CapsHashSet};
2
use nix::unistd::{getgid, getuid, setgid, setgroups, setuid};
3

            
4
#[derive(Debug)]
5
pub enum Error {
6
    CapsRead(caps::errors::CapsError),
7
    CapsUpdate(caps::errors::CapsError),
8
    DropGroups(nix::Error),
9
    SetGid(nix::Error),
10
    SetUid(nix::Error),
11
}
12

            
13
impl std::fmt::Display for Error {
14
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
15
        match self {
16
            Self::CapsRead(e) => write!(f, "Failed to read process capabilities: {e}"),
17
            Self::CapsUpdate(e) => write!(f, "Failed updating process capabilities: {e}"),
18
            Self::DropGroups(e) => write!(f, "Failed to drop supplementary groups: {e}"),
19
            Self::SetGid(e) => write!(f, "Failed to setgid: {e}"),
20
            Self::SetUid(e) => write!(f, "Failed to setuid: {e}"),
21
        }
22
    }
23
}
24

            
25
impl std::error::Error for Error {}
26

            
27
#[derive(Debug, PartialEq)]
28
enum CapabilityState {
29
    // We are either setuid root or the root user
30
    Full,
31
    // File system based capabilities
32
    Partial,
33
    None,
34
}
35

            
36
fn handle_full_capabilities() -> Result<(), Error> {
37
    // First, prepare the capability sets we want to end up with
38
    let mut ipc_lock_caps = CapsHashSet::new();
39
    ipc_lock_caps.insert(Capability::CAP_IPC_LOCK);
40

            
41
    // Clear all capabilities first, but DON'T touch bounding set yet
42
    let empty_caps = CapsHashSet::new();
43
    caps::set(None, CapSet::Effective, &empty_caps).map_err(Error::CapsUpdate)?;
44
    caps::set(None, CapSet::Permitted, &empty_caps).map_err(Error::CapsUpdate)?;
45

            
46
    // Set only CAP_IPC_LOCK in permitted and effective (before identity change)
47
    caps::set(None, CapSet::Permitted, &ipc_lock_caps).map_err(Error::CapsUpdate)?;
48
    caps::set(None, CapSet::Effective, &ipc_lock_caps).map_err(Error::CapsUpdate)?;
49

            
50
    // Drop supplementary groups first
51
    setgroups(&[]).map_err(Error::DropGroups)?;
52

            
53
    // Change to real GID
54
    setgid(getgid()).map_err(Error::SetGid)?;
55

            
56
    // Change to real UID (this should be done last)
57
    setuid(getuid()).map_err(Error::SetUid)?;
58

            
59
    // NOW we can safely clear the bounding set (after identity change)
60
    if let Err(err) = caps::set(None, CapSet::Bounding, &ipc_lock_caps) {
61
        tracing::debug!(
62
            "Could not clear bounding set (may not be supported): {}",
63
            err
64
        );
65
    }
66

            
67
    Ok(())
68
}
69

            
70
fn handle_partial_capabilities() -> Result<(), Error> {
71
    let effective_caps = caps::read(None, CapSet::Effective).map_err(Error::CapsRead)?;
72

            
73
    // Check if we have CAP_IPC_LOCK in effective set
74
    if !effective_caps.contains(&Capability::CAP_IPC_LOCK) {
75
        tracing::warn!("Insufficient process capabilities, insecure memory might get used");
76
    }
77

            
78
    // Check if we have CAP_SETPCAP for bounding set manipulation
79
    let has_setpcap =
80
        caps::has_cap(None, CapSet::Effective, Capability::CAP_SETPCAP).map_err(Error::CapsRead)?;
81

            
82
    // Clear all capabilities first
83
    let empty_caps = CapsHashSet::new();
84
    caps::set(None, CapSet::Effective, &empty_caps).map_err(Error::CapsUpdate)?;
85
    caps::set(None, CapSet::Permitted, &empty_caps).map_err(Error::CapsUpdate)?;
86

            
87
    // Only clear bounding set if we have CAP_SETPCAP
88
    if has_setpcap {
89
        if let Err(err) = caps::set(None, CapSet::Bounding, &empty_caps) {
90
            tracing::warn!("Failed to clear bounding set: {}", err);
91
        }
92
    }
93

            
94
    // Add only CAP_IPC_LOCK to effective and permitted sets
95
    let mut ipc_lock_caps = CapsHashSet::new();
96
    ipc_lock_caps.insert(Capability::CAP_IPC_LOCK);
97

            
98
    caps::set(None, CapSet::Effective, &ipc_lock_caps).map_err(Error::CapsUpdate)?;
99
    caps::set(None, CapSet::Permitted, &ipc_lock_caps).map_err(Error::CapsUpdate)?;
100

            
101
    // Only set bounding set if we have CAP_SETPCAP and cleared it successfully
102
    if has_setpcap {
103
        if let Err(err) = caps::set(None, CapSet::Bounding, &ipc_lock_caps) {
104
            tracing::warn!("Failed to set bounding set: {}", err);
105
        }
106
    }
107

            
108
    Ok(())
109
}
110

            
111
/// Determines the current capability state of the process
112
/// This mirrors the logic from libcap-ng's capng_have_capabilities()
113
fn determine_capability_state() -> Result<CapabilityState, Error> {
114
    let effective_caps = caps::read(None, CapSet::Effective).map_err(Error::CapsRead)?;
115
    let permitted_caps = caps::read(None, CapSet::Permitted).map_err(Error::CapsRead)?;
116
    let bounding_caps = caps::read(None, CapSet::Bounding).map_err(Error::CapsRead)?;
117

            
118
    // Check if we have no capabilities at all
119
    if permitted_caps.is_empty() && effective_caps.is_empty() {
120
        return Ok(CapabilityState::None);
121
    }
122

            
123
    // To match libcap-ng logic more closely, check if we have "most" capabilities
124
    // This is a heuristic - if we have a substantial number of capabilities,
125
    // we're likely in FULL state (setuid root or running as root)
126

            
127
    // Count total unique capabilities across all sets
128
    let mut all_caps = permitted_caps.clone();
129
    all_caps.extend(&effective_caps);
130
    all_caps.extend(&bounding_caps);
131

            
132
    // If we have 10+ capabilities total, likely FULL state
133
    // This matches libcap-ng's heuristic more closely than checking specific caps
134
    if all_caps.len() >= 10 {
135
        Ok(CapabilityState::Full)
136
    } else {
137
        // Otherwise, we have partial/filesystem-based capabilities
138
        Ok(CapabilityState::Partial)
139
    }
140
}
141

            
142
pub fn drop_unnecessary_capabilities() -> Result<(), Error> {
143
    // First, verify we can read capabilities at all (equivalent to CAPNG_FAIL
144
    // check)
145
    if let Err(e) = caps::read(None, CapSet::Effective) {
146
        // Critical error - cannot proceed safely, should abort like C version
147
        tracing::error!("Error getting process capabilities: {:?}, aborting", e);
148
        std::process::exit(1);
149
    }
150

            
151
    match determine_capability_state()? {
152
        CapabilityState::Full => {
153
            // We are either setuid root or the root user
154
            handle_full_capabilities()?;
155
        }
156
        CapabilityState::None => {
157
            tracing::warn!("No process capabilities, insecure memory might get used");
158
            return Ok(());
159
        }
160
        CapabilityState::Partial => {
161
            // File system based capabilities
162
            handle_partial_capabilities()?;
163
        }
164
    }
165

            
166
    Ok(())
167
}