s390-tools/rust/pvimg/examples/create-sehdr/main.rs

// SPDX-License-Identifier: MIT
//
// Copyright IBM Corp. 2024

#![allow(missing_docs)]
use std::fmt::Display;
use std::fs::{File, OpenOptions};
use std::io::{BufReader, Read, Write};
use std::path::PathBuf;
use std::str::FromStr;

use anyhow::{anyhow, Context, Error};
use clap::{Parser, ValueEnum, ValueHint};
use log::info;
use pv::misc::{decode_hex, open_file, read_file, read_hkd, try_parse_u64};
use pv::request::{HostKey, SymKeyType};
use pv::Result;
use pvimg::misc::PSW;
use pvimg::secured_comp::{ComponentTrait, Layout, SecuredComponentBuilder};
use pvimg::uvdata::{BuilderTrait, SeHdrBuilder, SeHdrControlFlags, SeHdrVersion, SeTarget};
use utils::{AtomicFile, AtomicFileOperation, HexSlice, PvLogger, VerbosityOptions};

/// Converts the hexstring into a byte vector.
///
/// # Errors
///
/// Raises an error if a non-hex character was found or the length was not a
/// multiple of two.
pub fn decode_hex_str<S: AsRef<str>>(s: S) -> Result<Vec<u8>> {
    let hex_str = s.as_ref();
    let hex_value = if hex_str.starts_with("0x") {
        hex_str.split_at(2).1
    } else {
        hex_str
    };

    Ok(decode_hex(hex_value)?)
}

fn decode_u64_hex_str(s: &str) -> Result<u64> {
    Ok(try_parse_u64(s, "The")?)
}

impl FromStr for ComponentArg {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let parts: Vec<_> = s.split(',').collect();
        if parts.len() != 3 {
            return Err(anyhow!("Invalid component format."));
        }
        let path = parts[0].into();
        let addr = try_parse_u64(parts[1], "Invalid address")?;
        let mut tweak =
            decode_hex_str(parts[2]).with_context(|| format!("Invalid tweak {}", parts[2]))?;
        if tweak.len() > SymKeyType::AES_256_XTS_TWEAK_LEN {
            return Err(anyhow!(
                "Invalid tweak because the length of {} is greater than the expected {}.",
                tweak.len(),
                SymKeyType::AES_256_XTS_TWEAK_LEN
            ));
        }
        tweak.resize(SymKeyType::AES_256_XTS_TWEAK_LEN, 0x0);
        Ok(Self { path, addr, tweak })
    }
}

#[derive(Debug, Clone)]
struct ComponentArg {
    path: PathBuf,
    addr: u64,
    tweak: Vec<u8>,
}

impl Display for ComponentArg {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "component\n\
             Path ......: {:}\n\
            Address ...: {:#0x}\n\
            Tweak .....: {:#}",
            self.path.display(),
            self.addr,
            HexSlice::from(&self.tweak),
        )
    }
}

/// SE-header version selection
#[derive(Debug, Clone, Copy, ValueEnum)]
enum SeHdrVersionArg {
    /// SE-header version 1
    #[value(name = "1")]
    V1,
    /// SE-header version 2
    #[value(name = "2")]
    V2,
}

impl From<SeHdrVersionArg> for SeHdrVersion {
    fn from(arg: SeHdrVersionArg) -> Self {
        match arg {
            SeHdrVersionArg::V1 => SeHdrVersion::V1,
            SeHdrVersionArg::V2 => SeHdrVersion::V2,
        }
    }
}

impl SeHdrVersionArg {
    /// Detect the SE header version from the keys.
    /// Returns V2 if any key is hybrid, otherwise V1.
    pub fn detect<K: AsRef<[HostKey]>>(keys: K) -> Self {
        if keys.as_ref().iter().any(|k| !k.is_hybrid()) {
            Self::V1
        } else {
            Self::V2
        }
    }
}

/// Create a Secure Execution header.
#[derive(Parser, Debug)]
pub struct Args {
    /// Use FILE as the component, ADDR as the component address, and TWEAK as
    /// the component tweak.
    ///
    /// ADDR and TWEAK must be a hex-string. TWEAK is right padded with zero
    /// bytes if the given tweak is not large enough. Can be specified multiple
    /// times and must be used at least once.
    #[arg(short, long = "component", required = true, value_name = "FILE,ADDR,TWEAK", value_hint = ValueHint::FilePath)]
    components: Vec<ComponentArg>,

    /// Use FILE as a host key document.
    ///
    /// Can be specified multiple times and must be used at least once.
    #[arg(short = 'k', long = "host-key", required = true)]
    pub host_key_documents: Vec<PathBuf>,

    /// Plain control flags. Must be a hex value.
    #[arg(long, default_value = "0x10000000")]
    pub pcf: String,

    /// Secret control flags. Must be a hex value.
    #[arg(long, default_value = "0x0")]
    pub scf: String,

    /// PSW address. Must be a hex value.
    #[arg(long, default_value = "0x10000", value_parser=decode_u64_hex_str)]
    pub psw_addr: u64,

    /// PSW mask. Must be a hex value.
    #[arg(long, default_value = "0x0000000180000000", value_parser=decode_u64_hex_str)]
    pub psw_mask: u64,

    /// Customer communication key (CCK) file path.
    #[arg(long)]
    pub cck: Option<PathBuf>,

    /// Secure Execution header output location.
    #[arg(short, long)]
    pub output: PathBuf,

    /// SE-header version to build
    #[arg(long, value_enum)]
    version: Option<SeHdrVersionArg>,

    #[clap(flatten)]
    pub verbosity: VerbosityOptions,
}

#[derive(Debug)]
pub struct Comp {
    pub reader: BufReader<File>,
}

impl Read for Comp {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.reader.read(buf)
    }
}

enum CompType {
    Dummy = 1,
}

impl ComponentTrait<CompType> for Comp {
    fn secure_mode(&self) -> bool {
        true
    }

    fn kind(&self) -> CompType {
        CompType::Dummy
    }
}

static LOGGER: PvLogger = PvLogger;

fn main() -> anyhow::Result<()> {
    let mut args = Args::parse();
    let log_level = args.verbosity.to_level_filter();
    LOGGER
        .start(log_level)
        .with_context(|| "Failed to set-up logger")?;

    info!("# Preparing components");
    let mut layout = Layout::new(0x0, SecuredComponentBuilder::COMPONENT_ALIGNMENT_V1)?;

    // Don't store the prepared components anywhere as we're only interested in
    // the hashes.
    let mut writer = std::io::empty();
    let mut secure_comp_builer = SecuredComponentBuilder::new_v1(false)?;

    // Sort components by address in ascending order
    args.components.sort_by_key(|a| a.addr);
    for component_arg in args.components {
        info!("## Preparing {}", component_arg);
        let mut comp = Comp {
            reader: BufReader::new(open_file(&component_arg.path)?),
        };

        let comp_addr = component_arg.addr;
        let _ = secure_comp_builer
            .prepare_and_insert_as_secure_component(
                &mut writer,
                &mut layout,
                &mut comp,
                comp_addr,
                component_arg.tweak,
            )
            .with_context(|| {
                format!(
                    "Failed to prepare component '{}'",
                    component_arg.path.display()
                )
            })?;
    }

    info!("\n# Creating Secure Execution Header");
    let addr = args.psw_addr;
    let mask = args.psw_mask;
    let mut target_pub_keys = vec![];
    for hkd_path in args.host_key_documents {
        info!(
            "Use the file '{}' as a host key document",
            hkd_path.display()
        );
        let cert = read_hkd(&hkd_path)?;
        target_pub_keys.push(cert);
    }
    let version: SeHdrVersion = args
        .version
        .unwrap_or(SeHdrVersionArg::detect(&target_pub_keys))
        .into();
    let target = SeTarget::from_se_hdr_version(version);
    let pcf = SeHdrControlFlags::from_u64(try_parse_u64(&args.pcf, "pcf")?, target, true);
    let scf = SeHdrControlFlags::from_u64(try_parse_u64(&args.scf, "scf")?, target, false);
    info!("SE-header version ...: {}", version);
    let mut builder = SeHdrBuilder::new(version, PSW { addr, mask }, secure_comp_builer.finish()?)?;
    builder.add_hostkeys(&target_pub_keys)?;

    info!(
        "PSW addr ............: {addr:#018x}\n\
         PSW mask ............: {mask:#018x}\n\
         PCF .................: {pcf}\n\
         SCF .................: {scf}"
    );
    builder.with_pcf(&pcf)?;
    builder.with_scf(&scf)?;
    if let Some(cck) = args.cck {
        info!("CCK ................: {}", cck.display());
        builder
            .with_cck(read_file(&cck, "CCK")?.into())
            .with_context(|| format!("Invalid CCK in '{}'", &cck.display()))?;
    }

    let mut output = AtomicFile::new(args.output, &mut OpenOptions::new())?;
    output.write_all(&builder.build()?.as_bytes()?)?;
    Ok(output.finish(AtomicFileOperation::Replace)?)
}