wolfcose 0.1.0

#![allow(missing_docs)]
#![no_std]

extern crate std;

use core::ptr;
use wolfcose::{raw, Algorithm, CborDecoder, CborEncoder, Error, Result};

const DRONE_ID: &str = "drone-17";
const PROTOCOL: &str = "drone-protocol/v1";
const TELEMETRY_KEY: [u8; 16] = [0x71; 16];
const CONTROL_KEY: [u8; 32] = [0x91; 32];
const CONTROL_KID: &[u8] = b"drone-17-control";

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct TelemetryFrame {
    sequence: u64,
    unix_ms: u64,
    lat_e7: i32,
    lon_e7: i32,
    alt_cm: i32,
    ground_speed_cm_s: u32,
    battery_percent: u8,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct ControlFrame {
    sequence: u64,
    lat_e7: i32,
    lon_e7: i32,
    alt_cm: i32,
}

fn main() -> Result<()> {
    // Host runner for no-std-style protocol code. The protocol path below uses
    // fixed buffers, borrowed slices, and raw wolfCOSE handles; `std` is linked
    // only so this example can print results on normal desktop targets.
    let telemetry = TelemetryFrame {
        sequence: 42,
        unix_ms: 1_725_000_000_123,
        lat_e7: 377_749_000,
        lon_e7: -1_224_194_000,
        alt_cm: 12_340,
        ground_speed_cm_s: 850,
        battery_percent: 88,
    };
    let control = ControlFrame {
        sequence: 43,
        lat_e7: 377_750_000,
        lon_e7: -1_224_190_000,
        alt_cm: 15_000,
    };

    let mut telemetry_payload = [0; 192];
    let telemetry_payload_len = encode_telemetry(&telemetry, &mut telemetry_payload)?;
    let mut telemetry_aad = [0; 96];
    let telemetry_aad_len = encode_aad("telemetry", telemetry.sequence, &mut telemetry_aad)?;
    let telemetry_iv = iv_from_sequence(*b"TELM", telemetry.sequence);
    let mut encrypted_telemetry = [0; 384];
    let mut scratch = [0; 1024];
    let encrypted_telemetry_len = encrypt_telemetry(
        &telemetry_payload[..telemetry_payload_len],
        &telemetry_aad[..telemetry_aad_len],
        &telemetry_iv,
        &mut encrypted_telemetry,
        &mut scratch,
    )?;

    let mut decrypted = [0; 192];
    let decoded_telemetry = decrypt_telemetry(
        &encrypted_telemetry[..encrypted_telemetry_len],
        &telemetry_aad[..telemetry_aad_len],
        &mut decrypted,
        &mut scratch,
    )?;
    assert_eq!(decoded_telemetry, telemetry);

    let mut control_payload = [0; 192];
    let control_payload_len = encode_control(&control, &mut control_payload)?;
    let mut control_aad = [0; 96];
    let control_aad_len = encode_aad("control", control.sequence, &mut control_aad)?;
    let mut signed_control = [0; 384];
    let signed_control_len = mac_control(
        &control_payload[..control_payload_len],
        &control_aad[..control_aad_len],
        &mut signed_control,
        &mut scratch,
    )?;

    let decoded_control = verify_control(
        &signed_control[..signed_control_len],
        &control_aad[..control_aad_len],
        &mut scratch,
    )?;
    assert_eq!(decoded_control, control);

    std::println!(
        "no_std drone protocol: telemetry_payload={} encrypted_telemetry={} control_payload={} signed_control={}",
        telemetry_payload_len,
        encrypted_telemetry_len,
        control_payload_len,
        signed_control_len
    );
    Ok(())
}

fn encode_telemetry(frame: &TelemetryFrame, out: &mut [u8]) -> Result<usize> {
    let mut encoder = CborEncoder::new(out);
    encoder.encode_array_start(6)?;
    encoder.encode_tstr(DRONE_ID)?;
    encoder.encode_u64(frame.sequence)?;
    encoder.encode_u64(frame.unix_ms)?;
    encoder.encode_array_start(4)?;
    encoder.encode_i64(frame.lat_e7 as i64)?;
    encoder.encode_i64(frame.lon_e7 as i64)?;
    encoder.encode_i64(frame.alt_cm as i64)?;
    encoder.encode_u64(frame.ground_speed_cm_s as u64)?;
    encoder.encode_u64(frame.battery_percent as u64)?;
    encoder.encode_tstr("mission")?;
    Ok(encoder.len())
}

fn decode_telemetry(input: &[u8]) -> Result<TelemetryFrame> {
    let mut decoder = CborDecoder::new(input);
    expect_len(decoder.decode_array_start()?, 6)?;
    expect_bytes(decoder.decode_tstr_bytes()?, DRONE_ID.as_bytes())?;
    let sequence = decoder.decode_u64()?;
    let unix_ms = decoder.decode_u64()?;
    expect_len(decoder.decode_array_start()?, 4)?;
    let lat_e7 = decode_i32(&mut decoder)?;
    let lon_e7 = decode_i32(&mut decoder)?;
    let alt_cm = decode_i32(&mut decoder)?;
    let ground_speed_cm_s = decode_u32(&mut decoder)?;
    let battery_percent = decode_u8(&mut decoder)?;
    expect_bytes(decoder.decode_tstr_bytes()?, b"mission")?;
    if !decoder.is_finished() {
        return Err(Error::CborMalformed);
    }
    Ok(TelemetryFrame {
        sequence,
        unix_ms,
        lat_e7,
        lon_e7,
        alt_cm,
        ground_speed_cm_s,
        battery_percent,
    })
}

fn encode_control(frame: &ControlFrame, out: &mut [u8]) -> Result<usize> {
    let mut encoder = CborEncoder::new(out);
    encoder.encode_array_start(3)?;
    encoder.encode_tstr(DRONE_ID)?;
    encoder.encode_u64(frame.sequence)?;
    encoder.encode_array_start(4)?;
    encoder.encode_tstr("set_waypoint")?;
    encoder.encode_i64(frame.lat_e7 as i64)?;
    encoder.encode_i64(frame.lon_e7 as i64)?;
    encoder.encode_i64(frame.alt_cm as i64)?;
    Ok(encoder.len())
}

fn decode_control(input: &[u8]) -> Result<ControlFrame> {
    let mut decoder = CborDecoder::new(input);
    expect_len(decoder.decode_array_start()?, 3)?;
    expect_bytes(decoder.decode_tstr_bytes()?, DRONE_ID.as_bytes())?;
    let sequence = decoder.decode_u64()?;
    expect_len(decoder.decode_array_start()?, 4)?;
    expect_bytes(decoder.decode_tstr_bytes()?, b"set_waypoint")?;
    let lat_e7 = decode_i32(&mut decoder)?;
    let lon_e7 = decode_i32(&mut decoder)?;
    let alt_cm = decode_i32(&mut decoder)?;
    if !decoder.is_finished() {
        return Err(Error::CborMalformed);
    }
    Ok(ControlFrame {
        sequence,
        lat_e7,
        lon_e7,
        alt_cm,
    })
}

fn encode_aad(channel: &str, sequence: u64, out: &mut [u8]) -> Result<usize> {
    let mut encoder = CborEncoder::new(out);
    encoder.encode_array_start(4)?;
    encoder.encode_tstr(PROTOCOL)?;
    encoder.encode_tstr(channel)?;
    encoder.encode_tstr(DRONE_ID)?;
    encoder.encode_u64(sequence)?;
    Ok(encoder.len())
}

fn encrypt_telemetry(
    plaintext: &[u8],
    aad: &[u8],
    iv: &[u8],
    out: &mut [u8],
    scratch: &mut [u8],
) -> Result<usize> {
    let mut key = RawKey::new(&TELEMETRY_KEY)?;
    let mut out_len = 0;
    let mut detached_len = 0;
    // SAFETY: slice pointers are valid for their lengths, null detached output
    // selects attached ciphertext mode, and output lengths bound C writes.
    Error::from_code(unsafe {
        raw::wc_CoseEncrypt0_Encrypt(
            key.as_mut_ptr(),
            Algorithm::A128GCM.id(),
            ptr_or_null(iv),
            iv.len(),
            ptr_or_null(plaintext),
            plaintext.len(),
            ptr::null_mut(),
            0,
            &mut detached_len,
            ptr_or_null(aad),
            aad.len(),
            scratch.as_mut_ptr(),
            scratch.len(),
            out.as_mut_ptr(),
            out.len(),
            &mut out_len,
        )
    })?;
    Ok(out_len)
}

fn decrypt_telemetry(
    message: &[u8],
    aad: &[u8],
    plaintext: &mut [u8],
    scratch: &mut [u8],
) -> Result<TelemetryFrame> {
    let mut key = RawKey::new(&TELEMETRY_KEY)?;
    let mut header = raw::WOLFCOSE_HDR::default();
    let mut plaintext_len = 0;
    // SAFETY: pointers and output locations are valid for the supplied lengths.
    // Null detached ciphertext selects attached ciphertext mode.
    Error::from_code(unsafe {
        raw::wc_CoseEncrypt0_Decrypt(
            key.as_mut_ptr(),
            ptr_or_null(message),
            message.len(),
            ptr::null(),
            0,
            ptr_or_null(aad),
            aad.len(),
            scratch.as_mut_ptr(),
            scratch.len(),
            &mut header,
            plaintext.as_mut_ptr(),
            plaintext.len(),
            &mut plaintext_len,
        )
    })?;
    decode_telemetry(&plaintext[..plaintext_len])
}

fn mac_control(payload: &[u8], aad: &[u8], out: &mut [u8], scratch: &mut [u8]) -> Result<usize> {
    let key = RawKey::new(&CONTROL_KEY)?;
    let mut out_len = 0;
    // SAFETY: slice pointers are valid for their lengths, null detached payload
    // selects attached mode, and output lengths bound C writes.
    Error::from_code(unsafe {
        raw::wc_CoseMac0_Create(
            key.as_ptr(),
            Algorithm::HMAC256.id(),
            CONTROL_KID.as_ptr(),
            CONTROL_KID.len(),
            ptr_or_null(payload),
            payload.len(),
            ptr::null(),
            0,
            ptr_or_null(aad),
            aad.len(),
            scratch.as_mut_ptr(),
            scratch.len(),
            out.as_mut_ptr(),
            out.len(),
            &mut out_len,
        )
    })?;
    Ok(out_len)
}

fn verify_control(message: &[u8], aad: &[u8], scratch: &mut [u8]) -> Result<ControlFrame> {
    let key = RawKey::new(&CONTROL_KEY)?;
    let mut header = raw::WOLFCOSE_HDR::default();
    let mut payload = ptr::null();
    let mut payload_len = 0;
    // SAFETY: inputs and output pointers are valid for the supplied sizes. Null
    // detached payload selects attached mode.
    Error::from_code(unsafe {
        raw::wc_CoseMac0_Verify(
            key.as_ptr(),
            ptr_or_null(message),
            message.len(),
            ptr::null(),
            0,
            ptr_or_null(aad),
            aad.len(),
            scratch.as_mut_ptr(),
            scratch.len(),
            &mut header,
            &mut payload,
            &mut payload_len,
        )
    })?;
    if payload.is_null() {
        return Err(Error::DetachedPayload);
    }
    // SAFETY: wolfCOSE returns an attached payload pointer into `message`, valid
    // for `payload_len` while `message` is alive.
    let payload = unsafe { core::slice::from_raw_parts(payload, payload_len) };
    decode_control(payload)
}

struct RawKey(raw::WOLFCOSE_KEY);

impl RawKey {
    fn new(material: &[u8]) -> Result<Self> {
        let mut key = raw::WOLFCOSE_KEY::default();
        // SAFETY: `key` points to valid writable storage.
        Error::from_code(unsafe { raw::wc_CoseKey_Init(&mut key) })?;
        // SAFETY: `material` points to caller-owned key bytes and remains alive
        // for this RawKey's operation-scoped lifetime.
        if let Err(error) = Error::from_code(unsafe {
            raw::wc_CoseKey_SetSymmetric(&mut key, ptr_or_null(material), material.len())
        }) {
            // SAFETY: initialized key is being released after setup failure.
            unsafe { raw::wc_CoseKey_Free(&mut key) };
            return Err(error);
        }
        Ok(Self(key))
    }

    fn as_ptr(&self) -> *const raw::WOLFCOSE_KEY {
        &self.0
    }

    fn as_mut_ptr(&mut self) -> *mut raw::WOLFCOSE_KEY {
        &mut self.0
    }
}

impl Drop for RawKey {
    fn drop(&mut self) {
        // SAFETY: RawKey always contains a key initialized by wc_CoseKey_Init.
        unsafe { raw::wc_CoseKey_Free(&mut self.0) };
    }
}

fn decode_i32(decoder: &mut CborDecoder<'_>) -> Result<i32> {
    let value = decoder.decode_i64()?;
    i32::try_from(value).map_err(|_| Error::CborOverflow)
}

fn decode_u32(decoder: &mut CborDecoder<'_>) -> Result<u32> {
    let value = decoder.decode_u64()?;
    u32::try_from(value).map_err(|_| Error::CborOverflow)
}

fn decode_u8(decoder: &mut CborDecoder<'_>) -> Result<u8> {
    let value = decoder.decode_u64()?;
    u8::try_from(value).map_err(|_| Error::CborOverflow)
}

fn expect_len(actual: usize, expected: usize) -> Result<()> {
    if actual == expected {
        Ok(())
    } else {
        Err(Error::CborMalformed)
    }
}

fn expect_bytes(actual: &[u8], expected: &[u8]) -> Result<()> {
    if actual == expected {
        Ok(())
    } else {
        Err(Error::CborType)
    }
}

fn ptr_or_null(data: &[u8]) -> *const u8 {
    if data.is_empty() {
        ptr::null()
    } else {
        data.as_ptr()
    }
}

fn iv_from_sequence(prefix: [u8; 4], sequence: u64) -> [u8; 12] {
    let mut iv = [0; 12];
    iv[..4].copy_from_slice(&prefix);
    iv[4..].copy_from_slice(&sequence.to_be_bytes());
    iv
}