use crate::*;
#[derive(Default, Debug, Copy, Clone, PartialEq, PartialOrd, Serialize, Deserialize)]
pub struct Header {
pub msgtype: u8,
pub psize: u32,
pub pubkey: [u8; 32],
pub nonce: [u8; 24],
pub encrypted: bool,
}
impl Header {
pub fn new() -> Header {
Header::default()
}
pub fn set_msgtype<T: Into<u8>>(mut self, msgtype: T) -> Header {
self.msgtype = msgtype.into();
self
}
pub fn set_psize<T: Into<u32>>(mut self, psize: T) -> Header {
self.psize = psize.into();
self
}
pub fn set_pubkey<T: Into<[u8; 32]>>(mut self, pubkey: T) -> Header {
self.pubkey = pubkey.into();
self
}
pub fn set_nonce<T: RngCore + CryptoRng>(mut self, csprng: &mut T) -> Result<Header> {
let mut nonce: [u8; 24] = [0u8; 24];
csprng.try_fill_bytes(&mut nonce)?;
self.nonce = nonce;
Ok(self)
}
pub fn set_encrypted(mut self, encrypted: bool) -> Header {
self.encrypted = encrypted;
self
}
pub fn to_bytes(self) -> [u8; 64] {
let mut header_bytes = [0u8; 64];
let reserved = [0u8; 2];
[self.msgtype]
.iter()
.chain(self.psize.to_be_bytes().iter())
.chain(self.pubkey.iter())
.chain(self.nonce.iter())
.chain([self.encrypted as u8].iter())
.chain(reserved.iter())
.enumerate()
.for_each(|(i, x)| header_bytes[i] = *x);
header_bytes
}
}
#[cfg(test)]
mod tests {
use super::*;
use rmp_serde as rmps;
enum MsgType {
Hello,
Unknown,
}
impl From<u8> for MsgType {
fn from(t: u8) -> MsgType {
match t {
0 => MsgType::Hello,
_ => MsgType::Unknown,
}
}
}
impl From<MsgType> for u8 {
fn from(t: MsgType) -> u8 {
match t {
MsgType::Hello => 0,
MsgType::Unknown => 255,
}
}
}
#[derive(Default, Debug, Clone, PartialEq, PartialOrd, Serialize, Deserialize)]
pub struct Hello {
name: String,
age: u8,
}
impl Job for Hello {
fn encode(&self) -> Vec<u8> {
rmps::to_vec(&self).unwrap()
}
fn decode(input: &[u8]) -> Hello {
let hello: Hello = rmps::from_read(input).unwrap();
hello
}
fn ack(&self) -> Vec<u8> {
let name = &self.name;
let age = &self.age;
let ack_string = format!("Hello from {name}, aged {age}");
Vec::from(ack_string)
}
fn run(&self) -> std::result::Result<(), Box<dyn std::error::Error>> {
self.ack();
Ok(())
}
}
#[test]
fn test_header_new() {
let header = Header::new();
assert_eq!(header.msgtype, 0u8);
assert_eq!(header.psize, 0u32);
assert_eq!(header.pubkey, [0u8; 32]);
assert_eq!(header.nonce, [0u8; 24]);
assert!(!header.encrypted);
}
#[test]
fn test_header_msgtype() {
let header = Header::new().set_msgtype(MsgType::Hello);
assert_eq!(header.msgtype, 0u8)
}
#[test]
fn test_header_psize() {
let header = Header::new().set_psize(10u32);
assert_eq!(header.psize, 10);
}
#[test]
fn test_header_pubkey() {
let mut rng = rand::thread_rng();
let secret_key = crypto_box::SecretKey::generate(&mut rng);
let pubkey = secret_key.public_key().as_bytes().to_owned();
let header = Header::new().set_pubkey(pubkey);
assert_eq!(header.pubkey, pubkey);
}
#[test]
fn test_header_nonce() {
let mut rng = rand::thread_rng();
let header = Header::new().set_nonce(&mut rng).unwrap();
assert_ne!(header.nonce, [0u8; 24]);
}
#[test]
fn test_header_encrypted() {
let header = Header::new().set_encrypted(true);
assert!(header.encrypted)
}
#[test]
fn test_header_bytes() {
let mut rng = rand::thread_rng();
let secret_key = crypto_box::SecretKey::generate(&mut rng);
let pubkey = secret_key.public_key().as_bytes().to_owned();
let header = Header::new()
.set_msgtype(MsgType::Hello)
.set_psize(10u32)
.set_pubkey(pubkey)
.set_nonce(&mut rng)
.unwrap();
let header_bytes = header.to_bytes();
assert_eq!(header_bytes[..5], [0, 0, 0, 0, 10]);
assert_eq!(header_bytes[5..37], pubkey);
assert_ne!(header_bytes[37..61], [0u8; 24]);
assert_eq!(header_bytes[61], 0u8);
assert_eq!(header_bytes[62..], [0u8; 2]);
}
}