#[macro_use]
extern crate log;
use byteorder::BigEndian;
use byteorder::ReadBytesExt;
use byteorder::WriteBytesExt;
use std::io;
use std::io::Cursor;
const HEADER_INIT_LEN: usize = 7;
const HEADER_CONT_LEN: usize = 5;
pub const MAX_PAYLOAD_LEN: usize = 64 - HEADER_INIT_LEN + 128 * (64 - HEADER_CONT_LEN);
pub const MAX_LEN: usize = 129 * 64;
pub trait U2FFraming {
fn encode(&mut self, message: &[u8], buf: &mut [u8]) -> io::Result<usize>;
fn decode(&mut self, buf: &[u8]) -> io::Result<Option<Vec<u8>>>;
fn set_cmd(&mut self, cmd: u8);
}
pub fn parse_header(buf: &[u8]) -> io::Result<(u32, u8, u16)> {
if buf.len() < HEADER_INIT_LEN {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Buffer to short to contain header (7 bytes)",
));
}
let mut rdr = Cursor::new(buf);
let cid = rdr.read_u32::<BigEndian>()?;
let cmd = rdr.read_u8()?;
let len = rdr.read_u16::<BigEndian>()?;
Ok((cid, cmd, len))
}
pub fn encode_header_init(cid: u32, cmd: u8, len: u16, mut buf: &mut [u8]) -> io::Result<usize> {
if buf.len() < HEADER_INIT_LEN {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Buffer to short to contain header (7 bytes)",
));
}
buf.write_u32::<BigEndian>(cid)?;
buf.write_u8(cmd)?;
buf.write_u16::<BigEndian>(len)?;
Ok(7)
}
pub fn encode_header_cont(cid: u32, seq: u8, mut buf: &mut [u8]) -> io::Result<usize> {
if buf.len() < HEADER_CONT_LEN {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Buffer to short to contain header (5 bytes)",
));
}
buf.write_u32::<BigEndian>(cid)?;
buf.write_u8(seq)?;
Ok(5)
}
pub fn generate_cid() -> u32 {
0xff00ff00
}
pub struct U2FWS {
cid: u32,
cmd: u8,
}
impl U2FWS {
pub fn new(cmd: u8) -> Self {
U2FWS {
cid: generate_cid(),
cmd,
}
}
pub fn with_cid(cid: u32, cmd: u8) -> Self {
U2FWS { cid, cmd }
}
}
impl Default for U2FWS {
fn default() -> Self {
Self::new(0)
}
}
impl U2FFraming for U2FWS {
fn encode(&mut self, message: &[u8], mut buf: &mut [u8]) -> io::Result<usize> {
let len = encode_header_init(self.cid, self.cmd, message.len() as u16, buf)?;
buf = &mut buf[len..];
if buf.len() < message.len() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Message won't fit in buffer",
));
}
let buf_slice = &mut buf[..message.len()];
buf_slice.copy_from_slice(message);
Ok(len + message.len())
}
fn decode(&mut self, buf: &[u8]) -> io::Result<Option<Vec<u8>>> {
let (cid, cmd, len) = parse_header(buf)?;
if cid != self.cid {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"Wrong CID",
));
}
if cmd != self.cmd {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"Wrong CMD",
));
}
if buf.len() < HEADER_INIT_LEN + len as usize {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"Invalid length",
));
}
Ok(Some(Vec::from(
&buf[HEADER_INIT_LEN..HEADER_INIT_LEN + len as usize],
)))
}
fn set_cmd(&mut self, cmd: u8) {
self.cmd = cmd;
}
}
pub struct U2FHID {
cid: u32,
cmd: u8,
}
impl U2FHID {
pub fn new(cmd: u8) -> Self {
U2FHID {
cid: generate_cid(),
cmd,
}
}
pub fn with_cid(cid: u32, cmd: u8) -> Self {
U2FHID { cid, cmd }
}
fn get_encoded_len(len: u16) -> usize {
if len < 57 {
64
} else {
let len = len - 57;
64 + 64 * ((59 + len - 1) / 59) as usize
}
}
}
impl Default for U2FHID {
fn default() -> Self {
Self::new(0)
}
}
impl U2FFraming for U2FHID {
fn encode(&mut self, mut message: &[u8], mut buf: &mut [u8]) -> io::Result<usize> {
let enc_len = Self::get_encoded_len(message.len() as u16);
debug!("Will encode {} in {}", message.len(), enc_len);
if buf.len() < enc_len {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Message won't fit in buffer",
));
}
let len = encode_header_init(self.cid, self.cmd, message.len() as u16, buf)?;
buf = &mut buf[len..];
let len = usize::min(64 - len, message.len());
let buf_slice = &mut buf[..len];
buf_slice.copy_from_slice(&message[..len]);
message = &message[len..];
buf = &mut buf[len..];
let mut seq = 0;
while message.len() > 0 {
let len = encode_header_cont(self.cid, seq as u8, buf)?;
buf = &mut buf[len..];
let len = usize::min(64 - len, message.len());
let buf_slice = &mut buf[..len];
buf_slice.copy_from_slice(&message[..len]);
buf = &mut buf[len..];
message = &message[len..];
seq += 1;
if seq > 127 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"More frames than allowed",
));
}
}
Ok(enc_len)
}
fn decode(&mut self, mut buf: &[u8]) -> io::Result<Option<Vec<u8>>> {
debug!("decode: {}", buf.len());
let (cid, cmd, len) = parse_header(buf)?;
if cid != self.cid {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"Wrong CID",
));
}
if cmd != self.cmd {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"Wrong CMD",
));
}
if buf.len() < Self::get_encoded_len(len) {
println!("{}", Self::get_encoded_len(len));
debug!("need more bytes");
return Ok(None);
}
let mut res = Vec::with_capacity(len as usize);
let mut left = len as usize;
let len = usize::min(57, len as usize);
res.extend_from_slice(&buf[HEADER_INIT_LEN..HEADER_INIT_LEN + len]);
buf = &buf[HEADER_INIT_LEN + len..];
left -= len as usize;
while left > 0 {
let len = usize::min(59, left);
res.extend_from_slice(&buf[HEADER_CONT_LEN..HEADER_CONT_LEN + len]);
buf = &buf[HEADER_CONT_LEN + len..];
left -= len;
}
Ok(Some(res))
}
fn set_cmd(&mut self, cmd: u8) {
self.cmd = cmd;
}
}
#[cfg(test)]
mod tests {
use crate::*;
#[test]
fn test_u2fhid_encode_single() {
let mut codec = U2FHID::with_cid(0xEEEEEEEE, 0x55);
let mut data = [0u8; 8000];
let len = codec.encode(b"\x01\x02\x03\x04", &mut data[..]).unwrap();
assert_eq!(len, 64);
let mut expect = [0u8; 64];
&expect[..11].copy_from_slice(b"\xEE\xEE\xEE\xEE\x55\x00\x04\x01\x02\x03\x04");
assert_eq!(&data[..len], &expect[..]);
}
#[test]
fn test_u2fhid_encode_multi() {
let payload: Vec<u8> = (0..65u8).collect();
let mut codec = U2FHID::with_cid(0xEEEEEEEE, 0x55);
let mut data = [0u8; 8000];
let len = codec.encode(&payload[..], &mut data[..]).unwrap();
assert_eq!(len, 128);
let mut expect = [0u8; 128];
&expect[..7].copy_from_slice(b"\xEE\xEE\xEE\xEE\x55\x00\x41");
&expect[7..64].copy_from_slice(&payload[..57]);
&expect[64..69].copy_from_slice(b"\xEE\xEE\xEE\xEE\x00");
&expect[69..77].copy_from_slice(&payload[57..]);
assert_eq!(&data[..len], &expect[..]);
}
#[test]
fn test_u2fhid_decode_single() {
let mut codec = U2FHID::with_cid(0xEEEEEEEE, 0x55);
let mut raw = [0u8; 64];
&raw[..11].copy_from_slice(b"\xEE\xEE\xEE\xEE\x55\x00\x04\x01\x02\x03\x04");
let data = codec.decode(&raw[..]).unwrap().unwrap();
assert_eq!(&data[..], b"\x01\x02\x03\x04");
}
#[test]
fn test_u2fhid_decode_multi() {
let payload: Vec<u8> = (0..65u8).collect();
let mut codec = U2FHID::with_cid(0xEEEEEEEE, 0x55);
let mut raw = [0u8; 128];
&raw[..7].copy_from_slice(b"\xEE\xEE\xEE\xEE\x55\x00\x41");
&raw[7..64].copy_from_slice(&payload[..57]);
&raw[64..69].copy_from_slice(b"\xEE\xEE\xEE\xEE\x00");
&raw[69..77].copy_from_slice(&payload[57..]);
let data = codec.decode(&raw[..]).unwrap().unwrap();
assert_eq!(&data[..], &payload[..]);
}
#[test]
fn test_u2fws_encode_single() {
let mut codec = U2FWS::with_cid(0xEEEEEEEE, 0x55);
let mut data = [0u8; 8000];
let len = codec.encode(b"\x01\x02\x03\x04", &mut data[..]).unwrap();
assert_eq!(len, 11);
assert_eq!(
&data[..len],
b"\xEE\xEE\xEE\xEE\x55\x00\x04\x01\x02\x03\x04"
);
}
#[test]
fn test_u2fws_encode_multi() {
let payload: Vec<u8> = (0..65u8).collect();
let mut codec = U2FWS::with_cid(0xEEEEEEEE, 0x55);
let mut data = [0u8; 8000];
let len = codec.encode(&payload[..], &mut data[..]).unwrap();
assert_eq!(len, 72);
let mut expect = [0u8; 72];
&expect[..7].copy_from_slice(b"\xEE\xEE\xEE\xEE\x55\x00\x41");
&expect[7..72].copy_from_slice(&payload[..]);
assert_eq!(&data[..len], &expect[..]);
}
#[test]
fn test_u2fws_decode_single() {
let mut codec = U2FWS::with_cid(0xEEEEEEEE, 0x55);
let data = codec
.decode(b"\xEE\xEE\xEE\xEE\x55\x00\x04\x01\x02\x03\x04")
.unwrap()
.unwrap();
assert_eq!(&data[..], b"\x01\x02\x03\x04");
}
#[test]
fn test_u2fws_decode_multi() {
let payload: Vec<u8> = (0..65u8).collect();
let mut codec = U2FWS::with_cid(0xEEEEEEEE, 0x55);
let mut raw = [0u8; 128];
&raw[..7].copy_from_slice(b"\xEE\xEE\xEE\xEE\x55\x00\x41");
&raw[7..72].copy_from_slice(&payload[..]);
let data = codec.decode(&raw[..]).unwrap().unwrap();
assert_eq!(&data[..], &payload[..]);
}
}