tasign 0.1.1

//! `plain.bin`：ELF 签名原数据。
//!
//! 在 `std` 与 `kernel-verify` 下均可用（`no_std` + `alloc` + `goblin`）。
//! 默认 API 使用 [`crate::limits::VerifyLimits::recommended`]；内核可改用 `*_with_limits` 收紧策略。

#[cfg(not(feature = "std"))]
extern crate alloc;
#[cfg(not(feature = "std"))]
use alloc::string::{String, ToString};
#[cfg(not(feature = "std"))]
use alloc::vec;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;

use goblin::elf::header::{ET_DYN, ET_EXEC, ET_REL};
use goblin::elf::program_header::PT_LOAD;
use goblin::elf::section_header::{SHT_NOBITS, SHT_NULL};
use goblin::elf::Elf;
use thiserror::Error;

use crate::limits::{LimitsError, VerifyLimits};

pub const TA_SIGNATURE_SECTION: &str = ".ta_signature";
pub const SHSTRTAB_SECTION: &str = ".shstrtab";
pub const SYMTAB_SECTION: &str = ".symtab";
pub const STRTAB_SECTION: &str = ".strtab";

#[derive(Debug, Error)]
pub enum PlainError {
    /// goblin 在 `no_std` 下未实现 `core::error::Error`，不能对 `#[from] goblin::error::Error` 使用 thiserror，故存字符串。
    #[error("ELF 解析失败: {0}")]
    Goblin(String),
    #[error("ELF 中已存在节 `{0}`，请先删除 `.ta_signature` 后再签名")]
    TaSignatureSectionExists(String),
    #[error("ELF 中不存在节 `{0}`")]
    TaSignatureSectionMissing(String),
    #[error("节 `{0}` 无文件内容（SHT_NOBITS）")]
    TaSignatureSectionNobits(String),
    #[error("节 `{0}` 为空")]
    TaSignatureSectionEmpty(String),
    #[error("无效的节名字符串表")]
    InvalidShstrtab,
    #[cfg(feature = "std")]
    #[error("ELF 重写失败: {0}")]
    ElfRewrite(String),
    #[error("节 `{name}` 的文件范围越界: offset={offset} size={size} file_len={file_len}")]
    SectionOutOfBounds {
        name: String,
        offset: u64,
        size: u64,
        file_len: usize,
    },
    #[error("{0}")]
    ResourceLimit(#[from] LimitsError),
}

impl From<goblin::error::Error> for PlainError {
    fn from(e: goblin::error::Error) -> Self {
        PlainError::Goblin(e.to_string())
    }
}

/// 若存在名为 `.ta_signature` 的节（`SHT_NULL` 除外），返回 `true`。
pub fn has_ta_signature_section(elf: &Elf<'_>) -> Result<bool, PlainError> {
    let shstrtab = &elf.shdr_strtab;
    for sh in &elf.section_headers {
        if sh.sh_type == SHT_NULL {
            continue;
        }
        let name = shstrtab.get_at(sh.sh_name).ok_or(PlainError::InvalidShstrtab)?;
        if name == TA_SIGNATURE_SECTION {
            return Ok(true);
        }
    }
    Ok(false)
}

/// 构造签名原数据 `plain.bin`（与 README 一致）。**要求** ELF 中**不存在** `.ta_signature`。
pub fn build_plain_bin(bytes: &[u8]) -> Result<Vec<u8>, PlainError> {
    build_plain_bin_with_limits(bytes, &VerifyLimits::default())
}

/// 与 [`build_plain_bin`] 相同，但使用自定义 [`VerifyLimits`]。
pub fn build_plain_bin_with_limits(bytes: &[u8], limits: &VerifyLimits) -> Result<Vec<u8>, PlainError> {
    limits.check_elf_input_len(bytes.len())?;
    let elf = Elf::parse(bytes)?;
    limits.check_section_header_count(elf.section_headers.len())?;
    limits.check_program_header_count(elf.program_headers.len())?;
    if has_ta_signature_section(&elf)? {
        return Err(PlainError::TaSignatureSectionExists(
            TA_SIGNATURE_SECTION.into(),
        ));
    }
    build_plain_from_elf(&elf, bytes, false, limits)
}

/// 对已嵌入 `.ta_signature` 的 ELF 计算签名原数据 `plain.bin`（会忽略 `.ta_signature` 本身）。
pub fn plain_bytes_from_signed_elf(bytes: &[u8]) -> Result<Vec<u8>, PlainError> {
    plain_bytes_from_signed_elf_with_limits(bytes, &VerifyLimits::default())
}

/// 与 [`plain_bytes_from_signed_elf`] 相同，但使用自定义 [`VerifyLimits`]。
pub fn plain_bytes_from_signed_elf_with_limits(
    bytes: &[u8],
    limits: &VerifyLimits,
) -> Result<Vec<u8>, PlainError> {
    limits.check_elf_input_len(bytes.len())?;
    let elf = Elf::parse(bytes)?;
    limits.check_section_header_count(elf.section_headers.len())?;
    limits.check_program_header_count(elf.program_headers.len())?;
    build_plain_from_elf(&elf, bytes, true, limits)
}

/// 返回 `.ta_signature` 节在文件中的原始字节（通常为 PKCS#7 DER）。
pub fn ta_signature_section_bytes<'a>(bytes: &'a [u8]) -> Result<&'a [u8], PlainError> {
    ta_signature_section_bytes_with_limits(bytes, &VerifyLimits::default())
}

/// 与 [`ta_signature_section_bytes`] 相同，但校验节长度不超过 `limits.max_ta_signature_section_bytes`。
pub fn ta_signature_section_bytes_with_limits<'a>(
    bytes: &'a [u8],
    limits: &VerifyLimits,
) -> Result<&'a [u8], PlainError> {
    limits.check_elf_input_len(bytes.len())?;
    let elf = Elf::parse(bytes)?;
    limits.check_section_header_count(elf.section_headers.len())?;
    limits.check_program_header_count(elf.program_headers.len())?;
    let shstrtab = &elf.shdr_strtab;
    for sh in &elf.section_headers {
        if sh.sh_type == SHT_NULL {
            continue;
        }
        let name = shstrtab.get_at(sh.sh_name).ok_or(PlainError::InvalidShstrtab)?;
        if name != TA_SIGNATURE_SECTION {
            continue;
        }
        if sh.sh_type == SHT_NOBITS {
            return Err(PlainError::TaSignatureSectionNobits(
                TA_SIGNATURE_SECTION.into(),
            ));
        }
        let off = sh.sh_offset as usize;
        let sz = sh.sh_size as usize;
        if sz == 0 {
            return Err(PlainError::TaSignatureSectionEmpty(
                TA_SIGNATURE_SECTION.into(),
            ));
        }
        limits.check_ta_signature_payload_len(sz)?;
        let end = off.checked_add(sz).ok_or(PlainError::InvalidShstrtab)?;
        if end > bytes.len() {
            return Err(PlainError::SectionOutOfBounds {
                name: TA_SIGNATURE_SECTION.into(),
                offset: sh.sh_offset,
                size: sh.sh_size,
                file_len: bytes.len(),
            });
        }
        return Ok(&bytes[off..end]);
    }
    Err(PlainError::TaSignatureSectionMissing(
        TA_SIGNATURE_SECTION.into(),
    ))
}

fn extend_plain(out: &mut Vec<u8>, data: &[u8], limits: &VerifyLimits) -> Result<(), PlainError> {
    let new_len = out
        .len()
        .checked_add(data.len())
        .ok_or(LimitsError {
            context: "plain 输出长度",
            max: limits.max_plain_output_bytes,
            got: usize::MAX,
        })?;
    if new_len > limits.max_plain_output_bytes {
        return Err(LimitsError {
            context: "plain 输出长度",
            max: limits.max_plain_output_bytes,
            got: new_len,
        }
        .into());
    }
    out.extend_from_slice(data);
    Ok(())
}

fn build_plain_from_elf(
    elf: &Elf<'_>,
    bytes: &[u8],
    allow_ta_signature: bool,
    limits: &VerifyLimits,
) -> Result<Vec<u8>, PlainError> {
    match elf.header.e_type {
        ET_EXEC | ET_DYN => build_plain_from_load_segments(elf, bytes, limits),
        ET_REL => build_plain_from_sections(elf, bytes, allow_ta_signature, limits),
        _ => build_plain_from_sections(elf, bytes, allow_ta_signature, limits),
    }
}

fn build_plain_from_load_segments(
    elf: &Elf<'_>,
    bytes: &[u8],
    limits: &VerifyLimits,
) -> Result<Vec<u8>, PlainError> {
    let mut out = Vec::new();

    // ET_EXEC/ET_DYN：plain = PHDR + merge(sort(PT_LOAD[p_offset, p_offset+p_filesz))。
    let phnum = elf.header.e_phnum as usize;
    let phent = elf.header.e_phentsize as usize;
    let phoff = elf.header.e_phoff as usize;
    if phnum > 0 {
        let ph_len = phnum.saturating_mul(phent);
        let ph_end = phoff.checked_add(ph_len).ok_or(PlainError::InvalidShstrtab)?;
        if ph_end > bytes.len() {
            return Err(PlainError::SectionOutOfBounds {
                name: "<program headers>".into(),
                offset: phoff as u64,
                size: ph_len as u64,
                file_len: bytes.len(),
            });
        }
        extend_plain(&mut out, &bytes[phoff..ph_end], limits)?;
    }
    let ehsize = elf.header.e_ehsize as usize;
    let ph_exclude = if phnum > 0 {
        let ph_len = phnum.saturating_mul(phent);
        Some((phoff, phoff + ph_len))
    } else {
        None
    };

    let mut ranges: Vec<(usize, usize)> = Vec::new();
    for ph in &elf.program_headers {
        if ph.p_type != PT_LOAD || ph.p_filesz == 0 {
            continue;
        }
        let start = ph.p_offset as usize;
        let len = ph.p_filesz as usize;
        let end = start.checked_add(len).ok_or(PlainError::SectionOutOfBounds {
            name: "<PT_LOAD>".into(),
            offset: ph.p_offset,
            size: ph.p_filesz,
            file_len: bytes.len(),
        })?;
        if end > bytes.len() {
            return Err(PlainError::SectionOutOfBounds {
                name: "<PT_LOAD>".into(),
                offset: ph.p_offset,
                size: ph.p_filesz,
                file_len: bytes.len(),
            });
        }
        ranges.push((start, end));
    }
    ranges.sort_by_key(|(s, _)| *s);

    let mut merged: Vec<(usize, usize)> = Vec::new();
    for (s, e) in ranges {
        if let Some((_, last_e)) = merged.last_mut() {
            if s <= *last_e {
                if e > *last_e {
                    *last_e = e;
                }
                continue;
            }
        }
        merged.push((s, e));
    }
    for (s, e) in merged {
        // 规则细化：
        // - ELF header 不纳入 plain；
        // - PHDR 已单独拼接，PT_LOAD 片段中不重复纳入。
        let mut parts = vec![(s, e)];
        parts = subtract_range(parts, (0, ehsize));
        if let Some(ph) = ph_exclude {
            parts = subtract_range(parts, ph);
        }
        for (ps, pe) in parts {
            extend_plain(&mut out, &bytes[ps..pe], limits)?;
        }
    }

    Ok(out)
}

fn subtract_range(ranges: Vec<(usize, usize)>, cut: (usize, usize)) -> Vec<(usize, usize)> {
    let (cut_s, cut_e) = cut;
    if cut_s >= cut_e {
        return ranges;
    }
    let mut out = Vec::new();
    for (s, e) in ranges {
        if e <= cut_s || s >= cut_e {
            out.push((s, e));
            continue;
        }
        if s < cut_s {
            out.push((s, cut_s));
        }
        if e > cut_e {
            out.push((cut_e, e));
        }
    }
    out
}

fn build_plain_from_sections(
    elf: &Elf<'_>,
    bytes: &[u8],
    allow_ta_signature: bool,
    limits: &VerifyLimits,
) -> Result<Vec<u8>, PlainError> {
    let mut out = Vec::new();

    // ET_REL：保留 PHDR（通常为空）+ 现有 section 过滤规则。
    let phnum = elf.header.e_phnum as usize;
    let phent = elf.header.e_phentsize as usize;
    let phoff = elf.header.e_phoff as usize;
    if phnum > 0 {
        let ph_len = phnum.saturating_mul(phent);
        let ph_end = phoff.checked_add(ph_len).ok_or(PlainError::InvalidShstrtab)?;
        if ph_end > bytes.len() {
            return Err(PlainError::SectionOutOfBounds {
                name: "<program headers>".into(),
                offset: phoff as u64,
                size: ph_len as u64,
                file_len: bytes.len(),
            });
        }
        extend_plain(&mut out, &bytes[phoff..ph_end], limits)?;
    }

    let shstrtab = &elf.shdr_strtab;

    for sh in &elf.section_headers {
        let name = shstrtab.get_at(sh.sh_name).ok_or(PlainError::InvalidShstrtab)?;

        if name == TA_SIGNATURE_SECTION {
            if !allow_ta_signature {
                return Err(PlainError::TaSignatureSectionExists(
                    TA_SIGNATURE_SECTION.into(),
                ));
            }
            continue;
        }

        if name == SHSTRTAB_SECTION || name == SYMTAB_SECTION || name == STRTAB_SECTION {
            continue;
        }

        if sh.sh_type == SHT_NULL {
            continue;
        }

        if sh.sh_type == SHT_NOBITS {
            continue;
        }

        let off = sh.sh_offset as usize;
        let sz = sh.sh_size as usize;
        let end = off.checked_add(sz).ok_or(PlainError::InvalidShstrtab)?;
        if end > bytes.len() {
            return Err(PlainError::SectionOutOfBounds {
                name: name.to_string(),
                offset: sh.sh_offset,
                size: sh.sh_size,
                file_len: bytes.len(),
            });
        }
        extend_plain(&mut out, &bytes[off..end], limits)?;
    }

    Ok(out)
}