ia-sandbox 0.4.0

use std::error::Error;
use std::ffi::{OsStr, OsString};
use std::fs::{self, OpenOptions};
use std::io::{Read, Write};
use std::path::Path;
use std::result;
use std::str::FromStr;
use std::time::Duration;

use crate::config::{ClearUsage, Limits, SpaceUsage};
use crate::errors::CgroupsError;
use crate::ffi;
use crate::run_info::RunUsage;

type Result<T> = result::Result<T, CgroupsError>;

fn cgroup_write<T1: AsRef<Path>, T2: AsRef<str>>(
    hierarchy_path: &Path,
    file: T1,
    line: T2,
) -> Result<()> {
    let path = hierarchy_path.join(file.as_ref());
    let mut cgroup_file = OpenOptions::new().write(true).open(&path).map_err(|err| {
        CgroupsError::OpenCgroupsFileError {
            hierarchy_path: hierarchy_path.to_path_buf(),
            file: file.as_ref().to_path_buf(),
            error: err.to_string(),
        }
    })?;

    cgroup_file
        .write_all(line.as_ref().as_bytes())
        .map_err(|err| CgroupsError::WriteCgroupsFileError {
            hierarchy_path: hierarchy_path.to_path_buf(),
            file: file.as_ref().to_path_buf(),
            error: err.to_string(),
        })
}

fn cgroup_read<T1: AsRef<Path>, T2: FromStr>(hierarchy_path: &Path, file: T1) -> Result<T2>
where
    <T2 as FromStr>::Err: Error,
{
    let path = hierarchy_path.join(file.as_ref());
    let mut cgroup_file = OpenOptions::new().read(true).open(&path).map_err(|err| {
        CgroupsError::OpenCgroupsFileError {
            hierarchy_path: hierarchy_path.to_path_buf(),
            file: file.as_ref().to_path_buf(),
            error: err.to_string(),
        }
    })?;

    let mut buffer = String::new();
    let _ = cgroup_file.read_to_string(&mut buffer).map_err(|err| {
        CgroupsError::ReadCgroupsFileError {
            hierarchy_path: hierarchy_path.to_path_buf(),
            file: file.as_ref().to_path_buf(),
            error: err.to_string(),
        }
    })?;

    buffer
        .trim()
        .parse::<T2>()
        .map_err(|err| CgroupsError::ParseCgroupsFileError {
            hierarchy_path: hierarchy_path.to_path_buf(),
            file: file.as_ref().to_path_buf(),
            buffer,
            error: err.to_string(),
        })
}

fn cgroup_read_field<T1: AsRef<Path>, T2: FromStr>(
    hierarchy_path: &Path,
    file: T1,
    field: &str,
) -> Result<T2>
where
    <T2 as FromStr>::Err: Error,
{
    let data: String = cgroup_read(hierarchy_path, file.as_ref())?;

    for line in data.lines() {
        let (key, value) = if let Some(split) = line.split_once(' ') {
            split
        } else {
            continue;
        };
        if key == field {
            return value
                .trim()
                .parse::<T2>()
                .map_err(|err| CgroupsError::ParseCgroupsFileError {
                    hierarchy_path: hierarchy_path.to_path_buf(),
                    file: file.as_ref().to_path_buf(),
                    buffer: value.to_string(),
                    error: err.to_string(),
                });
        }
    }

    Err(CgroupsError::MissingFieldError {
        hierarchy_path: hierarchy_path.to_path_buf(),
        file: file.as_ref().to_path_buf(),
        field: field.to_string(),
    })
}

const ISOLATED_CGROUP_NAME: &str = "isolated";
pub(crate) fn clear_cgroup(hierarchy_path: &Path) -> Result<()> {
    let remove = |path: &Path| {
        if path.exists() {
            fs::remove_dir(path).map_err(|err| CgroupsError::InstanceClearError {
                hierarchy_path: path.to_path_buf(),
                error: err.to_string(),
            })
        } else {
            Ok(())
        }
    };

    remove(&hierarchy_path.join(ISOLATED_CGROUP_NAME))?;
    remove(hierarchy_path)
}

pub(crate) fn enter_cgroup(hierarchy_path: &Path, instance_name: &OsStr) -> Result<()> {
    if !hierarchy_path.exists() {
        return Err(CgroupsError::HierarchyMissing(hierarchy_path.to_path_buf()));
    }

    let instance_path = hierarchy_path.join(instance_name);
    if !instance_path.exists() {
        fs::create_dir(&instance_path).map_err(|err| {
            CgroupsError::InstanceHierarchyCreateError {
                hierarchy_path: hierarchy_path.to_path_buf(),
                instance_name: instance_name.to_os_string(),
                error: err.to_string(),
            }
        })?;
    }

    // Nest it lower so that when we unshare the cgroup namespace the
    // process does not have aceess to its limits.
    let isolated_cgroup = instance_path.join(ISOLATED_CGROUP_NAME);

    if !isolated_cgroup.exists() {
        fs::create_dir(&isolated_cgroup).map_err(|err| {
            CgroupsError::InstanceHierarchyCreateError {
                hierarchy_path: isolated_cgroup.to_path_buf(),
                instance_name: OsString::from(ISOLATED_CGROUP_NAME),
                error: err.to_string(),
            }
        })?;
    }

    cgroup_write(
        &isolated_cgroup,
        "cgroup.procs",
        format!("{}\n", ffi::getpid()),
    )
}

const EXTRA_MEMORY_GIVEN: u64 = 4 * 1_024 * 1_024;
pub(crate) fn enter_memory_cgroup(
    instance_path: &Path,
    memory_limit: Option<SpaceUsage>,
) -> Result<()> {
    cgroup_write(instance_path, "memory.swap.max", "0\n")?;
    if let Some(memory_limit) = memory_limit {
        // Assign some extra memory so that we can tell when a killed by signal 9 is actually a
        // memory limit exceeded
        cgroup_write(
            instance_path,
            "memory.max",
            format!("{}\n", memory_limit.as_bytes() + EXTRA_MEMORY_GIVEN),
        )?;
        cgroup_write(
            instance_path,
            "memory.high",
            format!("{}\n", memory_limit.as_bytes()),
        )?;
    }
    Ok(())
}

pub(crate) fn enter_pids_cgroup(instance_path: &Path, pids_limit: Option<usize>) -> Result<()> {
    if let Some(pids_limit) = pids_limit {
        cgroup_write(instance_path, "pids.max", format!("{}\n", pids_limit))
    } else {
        cgroup_write(instance_path, "pids.max", "max\n")
    }
}

pub(crate) fn enter_all_cgroups(
    hierarchy_path: &Path,
    instance_name: &OsStr,
    limits: Limits,
    clear_usage: ClearUsage,
) -> Result<()> {
    let instance_path = &hierarchy_path.join(instance_name);

    if clear_usage == ClearUsage::Yes {
        clear_cgroup(instance_path)?;
        enter_cgroup(hierarchy_path, instance_name)?;
        enter_memory_cgroup(instance_path, limits.memory())?;
        enter_pids_cgroup(instance_path, limits.pids())
    } else {
        enter_cgroup(hierarchy_path, instance_name)
    }
}

pub(crate) fn get_usage(
    hierarchy_path: &Path,
    instance_name: &OsStr,
    wall_time: Duration,
) -> Result<RunUsage> {
    let instance_path = &hierarchy_path.join(instance_name);

    let user_time =
        Duration::from_micros(cgroup_read_field(instance_path, "cpu.stat", "usage_usec")?);

    let memory = SpaceUsage::from_bytes(cgroup_read(instance_path, "memory.peak")?);
    Ok(RunUsage::new(user_time, wall_time, memory))
}