use crate::{
envelope,
flags::{has, FLAG_CLOSE, FLAG_OPEN, FLAG_REPLY},
layout::{echo_packet_len, PacketLayout},
params::Params,
ProtoError, Result, SEQ_SIZE, TOKEN_SIZE,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RequestToEncode<'a> {
Open { params: &'a Params, no_test: bool },
Close { token: u64 },
Echo {
token: u64,
sequence: u32,
params: &'a Params,
payload: &'a [u8],
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DecodedRequest<'a> {
pub hmac_present: bool,
pub kind: DecodedRequestKind<'a>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DecodedRequestKind<'a> {
Open { no_test: bool, params: &'a [u8] },
Close { token: u64 },
Echo {
token: u64,
sequence: u32,
tail: &'a [u8],
},
}
pub fn encode_request(request: RequestToEncode<'_>, hmac_key: Option<&[u8]>) -> Result<Vec<u8>> {
match request {
RequestToEncode::Open { params, no_test } => encode_open(params, no_test, hmac_key),
RequestToEncode::Close { token } => encode_close(token, hmac_key),
RequestToEncode::Echo {
token,
sequence,
params,
payload,
} => encode_echo(token, sequence, params, payload, hmac_key),
}
}
pub fn decode_request(packet: &[u8]) -> Result<DecodedRequest<'_>> {
let envelope = envelope::decode_structural(packet)?;
if has(envelope.flags, FLAG_REPLY) {
return Err(ProtoError::UnexpectedFlag(FLAG_REPLY));
}
let body = envelope.body_offset;
let kind = if has(envelope.flags, FLAG_OPEN) {
require_len(packet, body)?;
DecodedRequestKind::Open {
no_test: has(envelope.flags, FLAG_CLOSE),
params: &packet[body..],
}
} else if has(envelope.flags, FLAG_CLOSE) {
require_len(packet, body + TOKEN_SIZE)?;
DecodedRequestKind::Close {
token: read_token(packet, body),
}
} else {
let tail_offset = body + TOKEN_SIZE + SEQ_SIZE;
require_len(packet, tail_offset)?;
DecodedRequestKind::Echo {
token: read_token(packet, body),
sequence: read_sequence(packet, body + TOKEN_SIZE),
tail: &packet[tail_offset..],
}
};
Ok(DecodedRequest {
hmac_present: envelope.hmac_present,
kind,
})
}
fn encode_open(params: &Params, no_test: bool, hmac_key: Option<&[u8]>) -> Result<Vec<u8>> {
let flags = if no_test {
FLAG_OPEN | FLAG_CLOSE
} else {
FLAG_OPEN
};
let encoded = params.encode();
let mut out = envelope::begin(
flags,
hmac_key,
PacketLayout::open_request(hmac_key.is_some()).header_len() + encoded.len(),
)?;
out.extend_from_slice(&encoded);
envelope::finish(out, hmac_key)
}
fn encode_close(token: u64, hmac_key: Option<&[u8]>) -> Result<Vec<u8>> {
let mut out = envelope::begin(
FLAG_CLOSE,
hmac_key,
PacketLayout::close_request(hmac_key.is_some()).header_len(),
)?;
out.extend_from_slice(&token.to_le_bytes());
envelope::finish(out, hmac_key)
}
fn encode_echo(
token: u64,
sequence: u32,
params: &Params,
payload: &[u8],
hmac_key: Option<&[u8]>,
) -> Result<Vec<u8>> {
let layout = PacketLayout::echo(hmac_key.is_some(), params);
let len = echo_packet_len(hmac_key.is_some(), params)?;
let payload_offset = layout.header_len();
let available_payload_len = len.saturating_sub(payload_offset);
if payload.len() > available_payload_len {
return Err(ProtoError::PayloadTooLarge {
available: available_payload_len,
provided: payload.len(),
});
}
let mut out = envelope::begin(0, hmac_key, len)?;
out.extend_from_slice(&token.to_le_bytes());
out.extend_from_slice(&sequence.to_le_bytes());
out.resize(len, 0);
out[payload_offset..payload_offset + payload.len()].copy_from_slice(payload);
envelope::finish(out, hmac_key)
}
fn require_len(packet: &[u8], needed: usize) -> Result<()> {
if packet.len() < needed {
return Err(ProtoError::PacketTooShort {
needed,
actual: packet.len(),
});
}
Ok(())
}
fn read_token(packet: &[u8], offset: usize) -> u64 {
u64::from_le_bytes(packet[offset..offset + TOKEN_SIZE].try_into().unwrap())
}
fn read_sequence(packet: &[u8], offset: usize) -> u32 {
u32::from_le_bytes(packet[offset..offset + SEQ_SIZE].try_into().unwrap())
}