cpu-temp 0.1.0

#![cfg(target_os = "windows")]
//! Rust translation of LibreHardwareMonitorLib/PawnIo/PawnIo.cs
//!
//! Provides interaction with the PawnIO driver for hardware monitoring.

use std::mem::size_of;

use thiserror::Error;
use windows::{
    core::{Free, PCWSTR},
    Win32::{
        Foundation::{GENERIC_READ, GENERIC_WRITE, HANDLE, INVALID_HANDLE_VALUE},
        Storage::FileSystem::{
            CreateFileW, FILE_ATTRIBUTE_NORMAL, FILE_SHARE_READ, FILE_SHARE_WRITE, OPEN_EXISTING,
        },
        System::IO::DeviceIoControl,
    },
};

const DEVICE_TYPE: u32 = 41_394u32 << 16;
const FN_NAME_LENGTH: usize = 32;
const IOCTL_PIO_EXECUTE_FN: u32 = 0x841 << 2;
const IOCTL_PIO_LOAD_BINARY: u32 = 0x821 << 2;

/// PawnIO version information.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Version {
    pub major: u32,
    pub minor: u32,
    pub build: u32,
    pub revision: u32,
}

impl Version {
    /// Parse a version string like "1.2.3.4".
    pub fn parse(s: &str) -> Option<Self> {
        let mut iter = s.split('.');
        let major = iter.next()?.parse().ok()?;
        let minor = iter.next()?.parse().ok()?;
        // 使用 map 转换后续部分，并填充默认值 0
        let build = iter.next().and_then(|p| p.parse().ok()).unwrap_or(0);
        let revision = iter.next().and_then(|p| p.parse().ok()).unwrap_or(0);

        Some(Version {
            major,
            minor,
            build,
            revision,
        })
    }
}

impl std::fmt::Display for Version {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{}.{}.{}.{}",
            self.major, self.minor, self.build, self.revision
        )
    }
}

/// Result of a PawnIO operation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IoResult {
    /// The output buffer returned from the device.
    pub output: Vec<i64>,
    /// Whether the operation succeeded.
    pub success: bool,
}

/// Error types for PawnIO operations.
#[derive(Debug, Error)]
pub enum PawnIoError {
    /// PawnIO driver is not installed.
    #[error("PawnIO driver is not installed")]
    NotInstalled,

    /// Device not found or driver not installed.
    #[error("PawnIO device not found: {0}")]
    DeviceNotFound(windows::core::Error),

    /// Failed to load binary module.
    #[error("Failed to load PawnIO module")]
    LoadFailed,

    /// Execute returned an error code.
    #[error("Execute error: {0}")]
    ExecuteError(String),

    /// Input buffer too small for the requested size.
    #[error("Input buffer too small")]
    BufferTooSmall,

    /// Buffer size conversion failed (too large for u32).
    #[error("Buffer too large: {0}")]
    BufferTooLarge(String),
}

pub type PawnResult<T = ()> = Result<T, PawnIoError>;

/// Handles interaction with the PawnIO driver.
pub struct PawnIo {
    handle: Option<HANDLE>,
}

impl PawnIo {
    /// Check if PawnIO is installed on the system.
    pub fn is_installed() -> bool {
        Self::version().is_some()
    }

    /// Get the version of the installed PawnIO driver, if any.
    ///
    /// Reads the `DisplayVersion` value from the Windows registry at:
    /// - `HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\PawnIO`
    /// - Falls back to 64-bit registry view if not found
    pub fn version() -> Option<Version> {
        use windows_registry::LOCAL_MACHINE;

        // Try reading from the standard registry path first.
        if let Ok(key) =
            LOCAL_MACHINE.open(r"SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\PawnIO")
        {
            if let Ok(version_str) = key.get_string("DisplayVersion") {
                if let Some(v) = Version::parse(&version_str) {
                    return Some(v);
                }
            }
        }

        // Fall back to 64-bit registry view (for WOW6432Node on 64-bit Windows).
        if let Ok(key) =
            LOCAL_MACHINE.open(r"SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\PawnIO")
        {
            if let Ok(version_str) = key.get_string("DisplayVersion") {
                if let Some(v) = Version::parse(&version_str) {
                    return Some(v);
                }
            }
        }

        None
    }

    /// Load a PawnIO module by opening the device and loading the binary data.
    ///
    /// # Arguments
    /// * `binary_data` - The binary module data to load (e.g., IntelMSR.bin)
    ///
    /// # Returns
    /// `Ok(PawnIo)` if successful, `Err(PawnIoError)` if the driver is not installed
    /// or the binary failed to load.
    pub fn load_module(binary_data: &[u8]) -> PawnResult<Self> {
        if !Self::is_installed() {
            return Err(PawnIoError::NotInstalled);
        }

        let handle = Self::open_device()?;
        Self::device_io_control_load_binary(handle, binary_data)?;
        Ok(PawnIo {
            handle: Some(handle),
        })
    }

    /// Check if the underlying handle is currently valid and open.
    pub fn is_loaded(&self) -> bool {
        self.handle.is_some_and(|h| !Self::is_invalid_handle(h))
    }

    /// Execute a function on the PawnIO device.
    ///
    /// # Arguments
    /// * `name` - Function name to execute (ASCII, max 31 chars)
    /// * `input` - Input buffer of i64 values
    /// * `out_length` - Expected number of i64 values in the output
    ///
    /// # Returns
    /// `IoResult` containing the output buffer and success flag
    pub fn execute(&self, name: &str, input: &[u64], out_length: usize) -> IoResult {
        if !self.is_loaded() {
            return IoResult {
                output: vec![0i64; out_length],
                success: false,
            };
        }

        let handle = self.handle.unwrap();
        let name_bytes = Self::prepare_name_bytes(name);
        let total_input_size = FN_NAME_LENGTH + std::mem::size_of_val(input);
        let mut total_input = vec![0u8; total_input_size];

        // Copy function name (ASCII, null-padded)
        let name_copy_len = name_bytes.len().min(FN_NAME_LENGTH - 1);
        total_input[..name_copy_len].copy_from_slice(&name_bytes[..name_copy_len]);

        // Copy input buffer as little-endian bytes
        let input_bytes: Vec<u8> = input.iter().flat_map(|v| v.to_le_bytes()).collect();
        let input_byte_len = std::mem::size_of_val(input);
        total_input[FN_NAME_LENGTH..FN_NAME_LENGTH + input_byte_len]
            .copy_from_slice(&input_bytes[..input_byte_len]);

        let mut output = vec![0u8; out_length * size_of::<i64>()];
        let mut bytes_returned: u32 = 0;

        let res = Self::device_io_control(
            handle,
            ControlCode::Execute,
            &total_input,
            &mut output,
            &mut bytes_returned,
        );

        if res.is_ok() {
            let actual_length = (bytes_returned as usize / size_of::<i64>()) as usize;
            let mut outp = vec![0i64; actual_length.min(out_length)];
            for i in 0..actual_length.min(out_length) {
                let bytes = &output[i * 8..(i + 1) * 8];
                outp[i] = i64::from_le_bytes(bytes.try_into().unwrap());
            }
            return IoResult {
                output: outp,
                success: true,
            };
        }

        IoResult {
            output: vec![0i64; out_length],
            success: false,
        }
    }

    /// Execute a function with detailed error reporting.
    ///
    /// # Arguments
    /// * `name` - Function name to execute
    /// * `in_buffer` - Input buffer
    /// * `in_size` - Number of i64 values to read from input
    /// * `out_buffer` - Output buffer to write results to
    /// * `out_size` - Number of i64 values to read from output
    ///
    /// # Returns
    /// `Ok(u32)` with actual number of i64 values returned, or `Err(PawnIoError)`
    pub fn execute_hr(
        &self,
        name: &str,
        in_buffer: &[i64],
        in_size: u32,
        out_buffer: &mut [i64],
        out_size: u32,
    ) -> Result<u32, PawnIoError> {
        if (in_buffer.len() as u32) < in_size {
            return Err(PawnIoError::BufferTooSmall);
        }

        if (out_buffer.len() as u32) < out_size {
            // In the original code it throws ArgumentOutOfRangeException
            return Err(PawnIoError::BufferTooSmall);
        }

        if !self.is_loaded() {
            return Ok(0);
        }

        let handle = self.handle.unwrap();
        let name_bytes = Self::prepare_name_bytes(name);
        let total_input_size = FN_NAME_LENGTH + (in_size as usize) * size_of::<i64>();
        let mut total_input = vec![0u8; total_input_size];

        let name_copy_len = name_bytes.len().min(FN_NAME_LENGTH - 1);
        total_input[..name_copy_len].copy_from_slice(&name_bytes[..name_copy_len]);

        let input_bytes: Vec<u8> = in_buffer[..in_size as usize]
            .iter()
            .flat_map(|v| v.to_le_bytes())
            .collect();
        total_input[FN_NAME_LENGTH..FN_NAME_LENGTH + (in_size as usize) * size_of::<i64>()]
            .copy_from_slice(&input_bytes);

        let mut output = vec![0u8; (out_size as usize) * size_of::<i64>()];
        let mut bytes_returned: u32 = 0;

        let res = Self::device_io_control(
            handle,
            ControlCode::Execute,
            &total_input,
            &mut output,
            &mut bytes_returned,
        );

        match res {
            Ok(()) => {
                let returned_count = bytes_returned / size_of::<i64>() as u32;
                let copy_len = (returned_count as usize).min(out_buffer.len());
                for i in 0..copy_len {
                    out_buffer[i] = i64::from_le_bytes([
                        output[i * 8],
                        output[i * 8 + 1],
                        output[i * 8 + 2],
                        output[i * 8 + 3],
                        output[i * 8 + 4],
                        output[i * 8 + 5],
                        output[i * 8 + 6],
                        output[i * 8 + 7],
                    ]);
                }
                Ok(returned_count)
            }
            Err(e) => Err(PawnIoError::ExecuteError(format!("{e}"))),
        }
    }

    /// Close the device handle.
    pub fn close(&mut self) {
        if self.is_loaded() {
            if let Some(mut handle) = self.handle.take() {
                unsafe { handle.free() };
            }
        }
    }

    /// Check if a handle value is invalid.
    fn is_invalid_handle(handle: HANDLE) -> bool {
        handle == INVALID_HANDLE_VALUE || handle.0.is_null()
    }

    /// Open the PawnIO device handle.
    fn open_device() -> PawnResult<HANDLE> {
        let path = r"\\?\GLOBALROOT\Device\PawnIO";
        let path_wide: Vec<u16> = path.encode_utf16().chain(Some(0)).collect();

        let res = unsafe {
            CreateFileW(
                PCWSTR(path_wide.as_ptr()),
                (GENERIC_READ | GENERIC_WRITE).0,
                FILE_SHARE_READ | FILE_SHARE_WRITE,
                None,
                OPEN_EXISTING,
                FILE_ATTRIBUTE_NORMAL,
                None,
            )
        };

        match res {
            Ok(handle) => Ok(handle),
            Err(e) => Err(PawnIoError::DeviceNotFound(e)),
        }
    }

    /// Load the PawnIO binary module via DeviceIoControl.
    fn device_io_control_load_binary(handle: HANDLE, data: &[u8]) -> PawnResult<()> {
        let mut bytes_returned: u32 = 0;
        Self::device_io_control(
            handle,
            ControlCode::LoadBinary,
            data,
            &mut [],
            &mut bytes_returned,
        )?;
        Ok(())
    }

    /// Generic DeviceIoControl helper.
    fn device_io_control(
        handle: HANDLE,
        control_code: ControlCode,
        in_buffer: &[u8],
        out_buffer: &mut [u8],
        bytes_returned: &mut u32,
    ) -> PawnResult<()> {
        let in_size = u32::try_from(in_buffer.len())
            .map_err(|e| PawnIoError::BufferTooLarge(format!("Input buffer too large: {e}")))?;

        let out_size = u32::try_from(out_buffer.len())
            .map_err(|e| PawnIoError::BufferTooLarge(format!("Output buffer too large: {e}")))?;

        let in_ptr = if in_buffer.is_empty() {
            None
        } else {
            Some(in_buffer.as_ptr() as _)
        };

        let out_ptr = if out_buffer.is_empty() {
            None
        } else {
            Some(out_buffer.as_mut_ptr() as _)
        };

        unsafe {
            DeviceIoControl(
                handle,
                control_code as u32,
                in_ptr,
                in_size,
                out_ptr,
                out_size,
                Some(bytes_returned),
                None,
            )
        }
        .map_err(|e| {
            PawnIoError::ExecuteError(format!(
                "Device IO control failed (0x{:08X}): {}",
                e.code().0,
                e.message()
            ))
        })
    }

    /// Prepare function name bytes (ASCII, null-padded to FN_NAME_LENGTH - 1).
    fn prepare_name_bytes(name: &str) -> Vec<u8> {
        let mut bytes = name.bytes().collect::<Vec<u8>>();
        bytes.resize(FN_NAME_LENGTH - 1, 0);
        bytes
    }
}

impl Drop for PawnIo {
    fn drop(&mut self) {
        self.close();
    }
}

#[derive(Debug, Clone, Copy)]
#[repr(u32)]
enum ControlCode {
    LoadBinary = DEVICE_TYPE | IOCTL_PIO_LOAD_BINARY,
    Execute = DEVICE_TYPE | IOCTL_PIO_EXECUTE_FN,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_version_parse() {
        assert_eq!(
            Version::parse("1.2.3.4"),
            Some(Version {
                major: 1,
                minor: 2,
                build: 3,
                revision: 4
            })
        );
        assert_eq!(
            Version::parse("1.0"),
            Some(Version {
                major: 1,
                minor: 0,
                build: 0,
                revision: 0
            })
        );
        assert_eq!(Version::parse("invalid"), None);
    }

    #[test]
    fn test_prepare_name_bytes() {
        let bytes = PawnIo::prepare_name_bytes("Test");
        assert_eq!(bytes.len(), FN_NAME_LENGTH - 1);
        assert_eq!(&bytes[..4], b"Test");
    }

    #[test]
    fn test_get_version() {
        let version = PawnIo::version().unwrap();
        println!("pawnio version: {}", version);
    }
}