#![no_main]
#![no_std]

// CHANGE ME!
// new brand string must be EXACTLY 48 characters (pad with spaces)
static NEW_BRAND_STRING: &'static str = "Lappy 486                                       ";
// 48 character reference:               ------------------------------------------------

use uefi::print;
use uefi::prelude::*;
use x86::cpuid::CpuId;
use core::arch::asm;
use core::convert::TryInto;

unsafe fn wrmsr_amd(msr: u32, value: u64) {
    let low = value as u32;
    let high = (value >> 32) as u32;
    // wrmsr with the AMD "secret" key - not necessary for these MSRs
    let secret = 0x9c5a203a as u32;
    unsafe {
        asm!("wrmsr", in("ecx") msr, in("eax") low, in("edx") high, in("edi") secret);
    }
}

pub fn brand_string_to_int64(brand_str: &str) -> Option<[u64; 6]> {
    // make sure we have exactly 48 characters for the brand string
    if brand_str.len() != 48 {
        return None
    }
    let str_bytes = brand_str.as_bytes();
    let mut out = [0u64; 6];
    // convert it to an array of 6 LE integers
    // does rust seriously not support for loops
    let mut i = 0;
    while i < 6 {
        let start = i * 8;
        let chunk: [u8; 8] = str_bytes[start..start + 8].try_into().unwrap();
        out[i] = u64::from_le_bytes(chunk);
        i += 1;
    }
    Some(out)
}

#[entry]
fn main() -> Status {
    uefi::helpers::init().unwrap();
    // get the vendor (we only work on AMD)
    let cpu_id = CpuId::new();
    let vendor_name = cpu_id.get_vendor_info().unwrap();
    print!("CPU vendor: {}\n", vendor_name.as_str());
    if vendor_name.as_str() == "AuthenticAMD" {
        // get the current CPU brand string
        let old_brand_string = cpu_id.get_processor_brand_string().unwrap();
        print!("Current brand: \"{}\"\n", old_brand_string.as_str());
        // convert to integers to set in the MSRs
        match brand_string_to_int64(NEW_BRAND_STRING) {
            None => print!("Invalid format for brand string\n\"{}\".\n", NEW_BRAND_STRING),
            Some(ints) => {
                // set the Core::X86::Msr::ProcNameString registers to our new string
                let mut i = 0;
                while i < 6 {
                    unsafe {
                        wrmsr_amd(0xC0010030 + i, ints[i as usize]);
                    }
                    i += 1;
                }
                // verify that the string has actually changed
                let verify_brand_string = cpu_id.get_processor_brand_string().unwrap();
                print!("New brand:     \"{}\"\n", verify_brand_string.as_str());
            }
        }
    } else {
        print!("This program only supports AuthenticAMD.\n");
    }
    // stall for 7 seconds to let you read the console output
    boot::stall(7_000_000);
    Status::SUCCESS
}