src/main.rs - master - githubmirror/mdevctl

src/main.rs
728 строк24 KB
Jonathon Jongsma
Fix all clippy warnings
25 июл 2025, 20:15
25 июл 2025, 20:15b68becd
//! mdevctl is a utility for managing and persisting devices in the mediated device framework of
//! the Linux kernel.  Mediated devices are sub-devices of a parent device (ex. a vGPU) which can
//! be dynamically created and potentially used by drivers like vfio-mdev for assignment to virtual
//! machines.
//!
//! See `mdevctl help` or the manpage for more information.

use clap::Parser;
use error::Error;
use log::{debug, warn};
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::fmt::Write;
use std::fs;
use std::io::stdout;
use std::path::PathBuf;
use std::vec::Vec;
use uuid::Uuid;

use crate::callouts::*;
use crate::cli::{LsmdevOptions, MdevctlCommands};
use crate::environment::Environment;
use crate::logger::logger;
use crate::mdev::*;

mod callouts;
mod cli;
mod environment;
mod error;
mod logger;
mod mdev;

#[cfg(test)]
mod tests;

/// Format a map of mediated devices into a json string
fn format_json(devices: BTreeMap<String, Vec<MDev>>) -> Result<String, Error> {
    let mut parents = serde_json::map::Map::new();
    for (parentname, children) in devices {
        let mut childrenarray = Vec::new();
        for child in children {
            childrenarray.push(child.to_json(true)?);
        }
        parents.insert(parentname, childrenarray.into());
    }
    // don't serialize an empty object if there are no devices
    let jsonval = match parents.len() {
        0 => serde_json::json!([]),
        _ => serde_json::json!([parents]),
    };
    serde_json::to_string_pretty(&jsonval).map_err(Into::into)
}

/// convert 'define' command arguments into a MDev struct
fn define_command_helper(
    env: &Environment,
    uuid: Option<Uuid>,
    auto: bool,
    parent: Option<String>,
    mdev_type: Option<String>,
    jsonfile: Option<PathBuf>,
    force: bool,
) -> Result<MDev, Error> {
    let uuid_provided = uuid.is_some();
    let uuid = uuid.unwrap_or_else(Uuid::new_v4);
    let mut dev = MDev::new(env, uuid);

    if let Some(jsonfile) = jsonfile {
        let _ = std::fs::File::open(&jsonfile)
            .map_err(|e| Error::IOError(format!("Unable to read file {jsonfile:?}"), e))?;

        if mdev_type.is_some() {
            return Err(Error::InvalidConfiguration(format!(
                "Device type cannot be specified separately from {jsonfile:?}"
            )));
        }

        let parent = parent.ok_or_else(|| {
            Error::InvalidConfiguration(format!(
                "Parent device required to define device via {jsonfile:?}"
            ))
        })?;

        let devs = env.get_defined_devices(Some(&uuid), Some(&parent))?;
        if !devs.is_empty() {
            return Err(Error::DeviceExists(format!(
                "Cowardly refusing to overwrite existing config for {parent}/{uuid}"
            )));
        }

        let filecontents = fs::read_to_string(&jsonfile).map_err(|e| {
            let var_name = format!("Unable to read jsonfile {jsonfile:?}");
            Error::IOError(var_name, e)
        })?;
        let jsonval = serde_json::from_str(&filecontents)?;
        dev.load_from_json(parent, &jsonval)?;
    } else {
        if uuid_provided {
            MDevSysfsData::load_for_mdev(&dev)
                .and_then(|sysfs_data| {
                    if parent.is_none() && mdev_type.is_some() {
                        return Err(Error::InvalidConfiguration(
                            "No parent specified".to_string(),
                        ));
                    }
                    dev.set_sysfs_data(sysfs_data);
                    Ok(())
                })
                .or_else(|e| match e {
                    Error::DeviceNotFound => Ok(()),
                    _ => {
                        if !force {
                            return Err(e);
                        }
                        warn!(
                            "For device {} a sysfs update caused the error: {:?}",
                            dev.uuid, e
                        );
                        Ok(())
                    }
                })?;
        }

        dev.autostart = auto;
        if parent.is_some() {
            dev.parent = parent;
        }
        if mdev_type.is_some() {
            dev.mdev_type = mdev_type;
        }

        if dev.parent.is_none() {
            return Err(Error::InvalidConfiguration(
                "No parent specified".to_string(),
            ));
        }
        if dev.mdev_type.is_none() {
            return Err(Error::InvalidConfiguration("No type specified".to_string()));
        }

        if dev.is_defined() {
            return Err(Error::DeviceExists(format!(
                "Device {} on {} already defined",
                dev.uuid,
                dev.parent()?
            )));
        }
    }

    Ok(dev)
}

/// Implementation of the `mdevctl define` command
fn define_command(
    env: &Environment,
    uuid: Option<Uuid>,
    auto: bool,
    parent: Option<String>,
    mdev_type: Option<String>,
    jsonfile: Option<PathBuf>,
    force: bool,
) -> Result<(), Error> {
    debug!("Defining mdev {uuid:?}");

    let mut dev = define_command_helper(env, uuid, auto, parent, mdev_type, jsonfile, force)?;

    /*
        Call Callout::get_attributes() when defining an active device without a config file.
        This function allows callout script to acquire device-specific attributes from sysfs,
        and populate the attrs field correspondingly before the device is defined in the system.
        The device config file will contain the same attributes that were used to start this device。
    */
    let mut c = callout(&mut dev)?;
    c.invoke(Action::Define, force, |c| {
        if c.dev.active {
            let attrs = c.get_attributes()?;
            c.dev.add_attributes(&attrs)?;
        }
        c.dev.define()?;
        Ok(())
    })
    .map(|_| {
        if uuid.is_none() {
            println!("{}", dev.uuid.hyphenated());
        }
    })
}

/// Implementation of the `mdevctl undefine` command
fn undefine_command(
    env: &Environment,
    uuid: Uuid,
    parent: Option<String>,
    force: bool,
) -> Result<(), Error> {
    debug!("Undefining mdev {uuid:?}");
    let mut failed = false;
    let devs = env.get_defined_devices(Some(&uuid), parent.as_ref())?;
    if devs.is_empty() {
        return Err(Error::DeviceNotFound);
    }
    for (_, mut children) in devs {
        for child in children.iter_mut() {
            let mut c = callout(child)?;
            if let Err(e) = c.invoke(Action::Undefine, force, |c| {
                c.dev.undefine()?;
                Ok(())
            }) {
                failed = true;
                warn!(
                    "Undefine of {} on parent {} failed with error: {}",
                    c.dev.uuid,
                    c.dev.parent()?,
                    e
                )
            }
        }
    }
    if failed {
        return Err(Error::System("Undefine failed".to_string()));
    }
    Ok(())
}

/// Implementation of the `mdevctl modify` command
#[allow(clippy::too_many_arguments)]
fn modify_command(
    env: &Environment,
    uuid: Uuid,
    parent: Option<String>,
    mdev_type: Option<String>,
    addattr: Option<String>,
    delattr: bool,
    index: Option<u32>,
    value: Option<String>,
    auto: bool,
    manual: bool,
    live: bool,
    defined: bool,
    jsonfile: Option<PathBuf>,
    force: bool,
) -> Result<(), Error> {
    debug!("Modifying mdev {uuid:?}");
    if live {
        if mdev_type.is_some() {
            return Err(Error::Unsupported(
                "'type' cannot be changed on active mdev".to_string(),
            ));
        }
        if auto {
            return Err(Error::Unsupported(
                "'auto' cannot be changed on active mdev".to_string(),
            ));
        }
        if manual {
            return Err(Error::Unsupported(
                "'manual' cannot be changed on active mdev".to_string(),
            ));
        }
        let mut act_dev = env.get_active_device(uuid, parent.as_ref())?;
        if let Some(f) = jsonfile {
            let act_parent = act_dev.parent.clone().ok_or_else(|| {
                Error::InvalidConfiguration(
                    "Parent device required to modify device via json file".to_string(),
                )
            })?;
            let json_dev = MDev::new_from_jsonfile(env, uuid, act_parent, f)?;
            if json_dev.mdev_type != act_dev.mdev_type {
                return Err(Error::Unsupported(
                    "'type' cannot be changed on active mdev".to_string(),
                ));
            }
            if json_dev.parent != act_dev.parent {
                return Err(Error::Unsupported(
                    "'parent' cannot be changed on active mdev".to_string(),
                ));
            }
            act_dev = json_dev;
        } else {
            return Err(Error::InvalidConfiguration(
                "'live' option must be used with 'jsonfile' option".to_string(),
            ));
        }

        if defined {
            // live and stored modify - defined dev config exists and types match
            let def_dev = env.get_defined_device(uuid, act_dev.parent.as_ref())?;
            if def_dev.mdev_type != act_dev.mdev_type {
                return Err(Error::InvalidConfiguration(
                    "'type' of active and defined mdev does not match".to_string(),
                ));
            }

            let mut c = callout(&mut act_dev)?;
            debug!("mdev device used for live update '{:?}'", c.dev);
            return c.invoke_modify_live().and_then(|_| {
                c.invoke(Action::Modify, force, |c| {
                    c.dev.write_config()?;
                    Ok(())
                })
            });
        }
        // live modify only
        callout(&mut act_dev)?.invoke_modify_live()
    } else {
        let mut dev: MDev;
        // stored configuration modify
        if let Some(f) = jsonfile {
            let parent = parent.ok_or_else(|| {
                Error::InvalidConfiguration(
                    "Parent device required to modify device via json file".to_string(),
                )
            })?;
            dev = MDev::new_from_jsonfile(env, uuid, parent, f)?;
        } else {
            dev = env.get_defined_device(uuid, parent.as_ref())?;
            if mdev_type.is_some() {
                dev.mdev_type = mdev_type;
            }
            if auto && manual {
                return Err(Error::InvalidConfiguration(
                    "'auto' and 'manual' are mutually exclusive".to_string(),
                ));
            }
            if auto {
                dev.autostart = true;
            } else if manual {
                dev.autostart = false;
            }
        }

        let index = index.map(|n| n as usize);
        match addattr {
            Some(attr) => match value {
                None => {
                    return Err(Error::InvalidConfiguration(
                        "No attribute value provided".to_string(),
                    ))
                }
                Some(v) => dev.add_attribute(attr, v, index)?,
            },
            None => {
                if delattr {
                    dev.delete_attribute(index)?;
                }
            }
        }
        callout(&mut dev)?.invoke(Action::Modify, force, |c| {
            c.dev.write_config()?;
            Ok(())
        })
    }
}

/// convert 'start' command arguments into a MDev struct
fn start_command_helper(
    env: &Environment,
    uuid: Option<Uuid>,
    parent: Option<String>,
    mdev_type: Option<String>,
    jsonfile: Option<PathBuf>,
    force: bool,
) -> Result<MDev, Error> {
    debug!("Starting device '{uuid:?}'");
    let mut dev: Option<MDev> = None;
    match jsonfile {
        Some(fname) => {
            let contents = fs::read_to_string(&fname)
                .map_err(|e| Error::IOError(format!("Unable to read jsonfile {fname:?}"), e))?;
            let val = serde_json::from_str(&contents)?;

            if mdev_type.is_some() {
                return Err(Error::InvalidConfiguration(
                    "Device type cannot be specified separately from json file".to_string(),
                ));
            }

            let parent = parent.ok_or_else(|| {
                Error::InvalidConfiguration(
                    "Parent device required to start device via json file".to_string(),
                )
            })?;

            let mut d = MDev::new(env, uuid.unwrap_or_else(Uuid::new_v4));
            d.load_from_json(parent, &val)?;
            dev = Some(d);
        }
        _ => {
            // if the user specified a uuid, check to see if they're referring to a defined device
            if let Some(uuid) = uuid {
                let devs = env.get_defined_devices(Some(&uuid), parent.as_ref())?;
                let n = devs.values().flatten().count();
                match n.cmp(&1) {
                    Ordering::Greater => {
                        return Err(Error::DeviceState(
                            "Multiple definitions found. Specify a parent.".to_string(),
                            uuid,
                            None,
                        ));
                    }
                    Ordering::Equal => {
                        // FIXME: use into_values() to consume the iterator and avoid cloning below
                        // when we can require rust 1.54.0
                        let d = devs.values().flatten().next();
                        if let Some(d) = d {
                            // See https://github.com/mdevctl/mdevctl/issues/38
                            // If a user specifies the uuid (and optional parent) of a defined device
                            if mdev_type.is_some() && mdev_type != d.mdev_type {
                                return Err(Error::DeviceExists(format!(
                                    "Device {} exists on parent {} with type {}",
                                    d.uuid,
                                    d.parent()?,
                                    d.mdev_type()?
                                )));
                            } else {
                                dev = Some(d.clone());
                            }
                        }
                    }
                    _ => (),
                }
            }

            if dev.is_none() {
                let mut d = MDev::new(env, uuid.unwrap_or_else(Uuid::new_v4));
                d.parent = parent;
                d.mdev_type = mdev_type;
                dev = Some(d);
            }

            if let Some(ref d) = dev {
                if d.mdev_type.is_some() && d.parent.is_none() {
                    return Err(Error::InvalidConfiguration(
                        "can't provide type without parent".to_string(),
                    ));
                }
                if d.mdev_type.is_none() || d.parent.is_none() {
                    return Err(Error::InvalidConfiguration(
                        "Device is insufficiently specified".to_string(),
                    ));
                }
            }
        }
    }
    let mut dev = dev.ok_or_else(|| Error::System("Unknown error".to_string()))?;

    callout(&mut dev)?.invoke(Action::Start, force, |c| {
        c.dev.start()?;
        Ok(())
    })?;
    Ok(dev)
}

/// Implementation of the `mdevctl start` command
fn start_command(
    env: &Environment,
    uuid: Option<Uuid>,
    parent: Option<String>,
    mdev_type: Option<String>,
    jsonfile: Option<PathBuf>,
    force: bool,
) -> Result<(), Error> {
    let dev = start_command_helper(env, uuid, parent, mdev_type, jsonfile, force)?;

    if uuid.is_none() {
        println!("{}", dev.uuid.hyphenated());
    }
    Ok(())
}

/// Implementation of the `mdevctl stop` command
fn stop_command(env: &Environment, uuid: Uuid, force: bool) -> Result<(), Error> {
    debug!("Stopping '{uuid}'");
    let mut dev = MDev::new(env, uuid);
    match MDevSysfsData::load_for_mdev(&dev) {
        Ok(sysfs_data) => dev.set_sysfs_data(sysfs_data),
        Err(Error::DeviceNotFound) => {
            return Err(Error::DeviceState(
                "device is not active".to_string(),
                uuid,
                None,
            ))
        }
        Err(e) => {
            if !force {
                return Err(e);
            }
            warn!(
                "For device {} a sysfs update caused the error: {:?}",
                dev.uuid, e
            );
        }
    };

    callout(&mut dev)?.invoke(Action::Stop, force, |c| {
        c.dev.stop()?;
        Ok(())
    })
}

/// Implementation of the `mdevctl list` command
fn list_command(
    env: &Environment,
    defined: bool,
    dumpjson: bool,
    verbose: bool,
    uuid: Option<Uuid>,
    parent: Option<String>,
    output: &mut dyn std::io::Write,
) -> Result<(), Error> {
    let mut devices: BTreeMap<String, Vec<MDev>>;
    if defined {
        devices = env.get_defined_devices(uuid.as_ref(), parent.as_ref())?;
    } else {
        devices = env.get_active_devices(uuid.as_ref(), parent.as_ref())?;
    }

    // ensure that devices are sorted in a stable order
    for v in devices.values_mut() {
        v.sort_by_key(|e| e.uuid);
    }

    match dumpjson {
        true => {
            // if specified to a single device, output such that it can be piped into a config
            // file, else print entire heirarchy
            if uuid.is_none() || devices.values().flatten().count() > 1 {
                output.write(format_json(devices)?.as_bytes())
            } else {
                let jsonval = match devices.values().next() {
                    Some(children) => children
                        .first()
                        .ok_or_else(|| Error::System("Failed to get device".to_string()))?
                        .to_json(false)?,
                    None => serde_json::json!([]),
                };
                output.write(serde_json::to_string_pretty(&jsonval)?.as_bytes())
            }
        }
        false => {
            let ft = match defined {
                true => FormatType::Defined,
                false => FormatType::Active,
            };
            output.write(
                devices
                    .values()
                    // convert child vector into an iterator over the vector's elements
                    .flat_map(|v| v.iter())
                    // convert MDev elements to a text representation, filtering out errors
                    .flat_map(|d| d.to_text(ft, verbose))
                    .collect::<String>()
                    .as_bytes(),
            )
        }
    }
    .map(|_| ())
    .map_err(|e| Error::IOError("Failed to write data".to_string(), e))
}

/// convert 'types' command arguments into a text output
fn types_command(
    env: &Environment,
    parent: Option<String>,
    dumpjson: bool,
    output: &mut dyn std::io::Write,
) -> Result<(), Error> {
    let types = env.get_supported_types(parent)?;
    debug!("{types:?}");
    if dumpjson {
        let mut parents = serde_json::map::Map::new();
        for (parent, children) in types {
            let mut childarray = Vec::new();
            for child in children {
                childarray.push(child.to_json());
            }
            parents.insert(parent, childarray.into());
        }

        let jsonval = match parents.len() {
            0 => serde_json::json!([]),
            _ => serde_json::json!([parents]),
        };
        output.write(serde_json::to_string_pretty(&jsonval)?.as_bytes())
    } else {
        let mut text: String = Default::default();
        for (parent, children) in types {
            let _ = writeln!(text, "{parent}");
            for child in children {
                let _ = writeln!(text, "  {}", child.typename);
                let _ = writeln!(
                    text,
                    "    Available instances: {}",
                    child.available_instances
                );
                let _ = writeln!(text, "    Device API: {}", child.device_api);
                if !child.name.is_empty() {
                    let _ = writeln!(text, "    Name: {}", child.name);
                }
                if !child.description.is_empty() {
                    let _ = writeln!(text, "    Description: {}", child.description);
                }
            }
        }
        output.write(text.as_bytes())
    }
    .map(|_| ())
    .map_err(|e| Error::IOError("Unable to write output".to_string(), e))
}

/// Implementation of the `start-parent-mdevs` command
fn start_parent_mdevs_command(env: &Environment, parent: String) -> Result<(), Error> {
    let mut devs = env.get_defined_devices(None, Some(&parent))?;
    if devs.is_empty() {
        // nothing to do
        return Ok(());
    }

    if devs.len() != 1 {
        return Err(Error::InvalidConfiguration(
            "More than one parent found".to_string(),
        ));
    };

    for (_, children) in devs.iter_mut() {
        for child in children {
            if child.autostart {
                debug!("Autostarting {:?}", child.uuid);
                if let Err(e) = callout(child)?.invoke(Action::Start, false, |c| {
                    c.dev.start()?;
                    Ok(())
                }) {
                    warn!("{e}");
                }
            }
        }
    }
    Ok(())
}

/// parse command line arguments and dispatch to command-specific functions
fn main() -> Result<(), Error> {
    logger().init();
    debug!("Starting up");

    let env = Environment::new("/".to_string());
    debug!("{env:?}");

    // make sure the environment is sane
    env.self_check()?;

    // check if we're running as the symlink executable 'lsmdev'. If so, just execute the 'list'
    // command directly
    let exe = std::env::args_os()
        .next()
        .ok_or_else(|| Error::System("Failed to get the executable name".to_string()))?;
    match exe.to_str() {
        Some(val) if val.ends_with("lsmdev") => {
            debug!("running as 'lsmdev'");
            let opts = LsmdevOptions::parse();
            list_command(
                &env,
                opts.defined,
                opts.dumpjson,
                opts.verbose,
                opts.uuid,
                opts.parent,
                &mut stdout(),
            )
        }
        _ => match MdevctlCommands::parse() {
            MdevctlCommands::Define {
                uuid,
                auto,
                parent,
                mdev_type,
                jsonfile,
                force,
            } => define_command(&env, uuid, auto, parent, mdev_type, jsonfile, force),
            MdevctlCommands::Undefine {
                uuid,
                parent,
                force,
            } => undefine_command(&env, uuid, parent, force),
            MdevctlCommands::Modify {
                uuid,
                parent,
                mdev_type,
                addattr,
                delattr,
                index,
                value,
                auto,
                manual,
                live,
                defined,
                jsonfile,
                force,
            } => modify_command(
                &env, uuid, parent, mdev_type, addattr, delattr, index, value, auto, manual, live,
                defined, jsonfile, force,
            ),
            MdevctlCommands::Start {
                uuid,
                parent,
                mdev_type,
                jsonfile,
                force,
            } => start_command(&env, uuid, parent, mdev_type, jsonfile, force),
            MdevctlCommands::Stop { uuid, force } => stop_command(&env, uuid, force),
            MdevctlCommands::List(list) => list_command(
                &env,
                list.defined,
                list.dumpjson,
                list.verbose,
                list.uuid,
                list.parent,
                &mut stdout(),
            ),
            MdevctlCommands::Types { parent, dumpjson } => {
                types_command(&env, parent, dumpjson, &mut stdout())
            }
            MdevctlCommands::StartParentMdevs { parent } => {
                start_parent_mdevs_command(&env, parent)
            }
        },
    }
}