rustkey 0.5.0

// rusTkey, a rust crate/library that provides a development API for the tillitis TKey.
// Copyright (C) 2024  Danny van Heumen
// SPDX-License-Identifier: BSD-2-Clause

#![no_std]
#![no_main]

use core::{arch::global_asm, panic::PanicInfo};

use rustkey::{frame, led, timer, trng};

const BENCH_TRNG: u8 = 0;
const BENCH_BLAKE2S: u8 = 1;
const BENCH_RANDOM: u8 = 2;
const GENERATE_RANDOMNESS: u8 = 3;

fn bench_trng(secs: u32) -> u64 {
    timer::set_prescaler(timer::PRESCALE_SECONDS);
    timer::initialize(secs);
    timer::start();
    let mut count = 0u64;
    while timer::running() {
        if trng::ready() {
            _ = trng::read();
            count += 1;
        }
    }
    count
}

/// `bench_blake2s` runs repeated `blake2s` calls for 32-byte digest, given a key and a buffer of
/// given length. (The loop consists of buffer-filling then calling `blake2s`.)
fn bench_blake2s<const N: usize>(secs: u32, key: &[u8]) -> u64 {
    timer::set_prescaler(timer::PRESCALE_SECONDS);
    timer::initialize(secs);
    timer::start();
    let mut count = 0u64;
    let mut buffer = [0u8; N];
    let mut out = [0u8; 32];
    while timer::running() {
        buffer.fill(count as u8);
        rustkey::blake2s(&mut out, key, &buffer).unwrap();
        count += 1
    }
    count
}

/// `bench_random` benchmarks `rustkey::random`. Check current implementation for details on how
/// `random` operates exactly. Implementation, at time of writing, takes one sample from TRNG for
/// every 32-byte range of the buffer. (We ignore seed, as this is a one-time operation before
/// starting randomness-loop.)
fn bench_random<const N: usize>(secs: u32) -> u64 {
    timer::set_prescaler(timer::PRESCALE_SECONDS);
    timer::initialize(secs);
    timer::start();
    let mut count = 0u64;
    let mut buffer = [0u8; N];
    while timer::running() {
        rustkey::random(&mut buffer, &[]);
        count += 1
    }
    count
}

/// `randomness` produces `count` 128-byte frames of random data and sends these to the client.
/// First response is 1-length frame with benchmark ID for confirmation, followed by `count`
/// 128-byte frames filled with random data.
fn randomness(request: &frame::Header, count: u32) -> u64 {
    let mut response = *request;
    response.error = false;
    response.length = frame::CommandLength::Length1;
    frame::write(&response, &[GENERATE_RANDOMNESS]);
    response.length = frame::CommandLength::Length128;
    let mut buffer = [0u8; 128];
    for _ in 0..count {
        rustkey::random(&mut buffer, &[]);
        frame::write(&response, &buffer);
    }
    count as u64
}

#[no_mangle]
extern "C" fn main() -> ! {
    led::set(led::LED_OFF);
    let mut header = frame::Header {
        id: 0,
        endpoint: frame::Endpoint::Software,
        error: false,
        length: frame::CommandLength::Length4,
    };
    let mut buffer = [0u8; frame::LENGTH_MAX];
    // Read frame: first byte is command-byte, followed by 3-byte BE value for the number of
    // seconds to query the TRNG.
    frame::read_into(&mut header, &mut buffer).unwrap();
    let cmd = buffer[0];
    let count = match cmd {
        BENCH_TRNG => bench_trng(u32::from_be_bytes([0, buffer[1], buffer[2], buffer[3]])),
        BENCH_BLAKE2S => bench_blake2s::<64>(
            u32::from_be_bytes([0, buffer[1], buffer[2], buffer[3]]),
            "some-key".as_bytes(),
        ),
        BENCH_RANDOM => {
            bench_random::<32>(u32::from_be_bytes([0, buffer[1], buffer[2], buffer[3]]))
        }
        GENERATE_RANDOMNESS => randomness(
            &header,
            u32::from_be_bytes([0, buffer[1], buffer[2], buffer[3]]),
        ),
        _ => panic!(),
    };
    // Send count, i.e. number of reads from TRNG, back to client.
    buffer.fill(0);
    header.endpoint = frame::Endpoint::Software;
    header.error = false;
    header.length = frame::CommandLength::Length32;
    buffer[0] = cmd;
    buffer[1..9].copy_from_slice(&count.to_be_bytes());
    frame::write(&header, &buffer);
    led::set(led::LED_GREEN);
    rustkey::done()
}

#[panic_handler]
fn panic(_: &PanicInfo) -> ! {
    rustkey::abort()
}

global_asm!(
    ".section \".text.init\"",
    ".global _start",
    "_start:",
    "li x1, 0",
    "li x2, 0",
    "li x3, 0",
    "li x4, 0",
    "li x5, 0",
    "li x6, 0",
    "li x7, 0",
    "li x8, 0",
    "li x9, 0",
    "li x10,0",
    "li x11,0",
    "li x12,0",
    "li x13,0",
    "li x14,0",
    "li x15,0",
    "li x16,0",
    "li x17,0",
    "li x18,0",
    "li x19,0",
    "li x20,0",
    "li x21,0",
    "li x22,0",
    "li x23,0",
    "li x24,0",
    "li x25,0",
    "li x26,0",
    "li x27,0",
    "li x28,0",
    "li x29,0",
    "li x30,0",
    "li x31,0",
    /* init stack below 0x40020000 (TK1_RAM_BASE+TK1_RAM_SIZE) */
    "la sp, _stack_start",
    /* zero-init bss section */
    "la a0, _sbss",
    "la a1, _ebss",
    "bge a0, a1, end_init_bss",
    "loop_init_bss:",
    "sw zero, 0(a0)",
    "addi a0, a0, 4",
    "blt a0, a1, loop_init_bss",
    "end_init_bss:",
    "call main",
    options(raw)
);