diff --git a/Cargo.lock b/Cargo.lock index aea48a0..e0d9c92 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1312,7 +1312,6 @@ dependencies = [ name = "ql-btp" version = "0.1.0" dependencies = [ - "bytes", "ql-codec", ] @@ -1321,6 +1320,17 @@ name = "ql-codec" version = "0.1.0" dependencies = [ "bytes", + "ql-codec-macro", +] + +[[package]] +name = "ql-codec-macro" +version = "0.1.0" +dependencies = [ + "proc-macro2", + "ql-codec", + "quote", + "syn", ] [[package]] @@ -1347,7 +1357,6 @@ dependencies = [ name = "ql-keyos" version = "0.1.0" dependencies = [ - "bytes", "ql-codec", "ql-rpc", "serde", diff --git a/Cargo.toml b/Cargo.toml index 3325265..6c7bb73 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,7 @@ resolver = "2" members = [ "btp", "codec", + "codec/macro", "common", "keyos", "router", @@ -21,7 +22,8 @@ license = "MIT OR Apache-2.0" bytes = "1.11.1" # workspace crates -ql-codec = { path = "codec" } +ql-codec = { path = "codec", features = ["derive"] } +ql-codec-macro = { path = "codec/macro" } ql-common = { path = "common" } ql-fsm = { path = "fsm" } ql-keyos = { path = "keyos" } diff --git a/btp/Cargo.toml b/btp/Cargo.toml index d5dac0e..6e14bbe 100644 --- a/btp/Cargo.toml +++ b/btp/Cargo.toml @@ -6,5 +6,4 @@ homepage.workspace = true license.workspace = true [dependencies] -bytes = { workspace = true } ql-codec = { workspace = true } diff --git a/btp/src/chunk.rs b/btp/src/chunk.rs index 86e3099..c6d957d 100644 --- a/btp/src/chunk.rs +++ b/btp/src/chunk.rs @@ -1,7 +1,6 @@ use std::sync::atomic::{AtomicU32, Ordering}; -use bytes::BufMut; -use ql_codec::Encode; +use ql_codec::{Encode, Writer}; use crate::{Header, APP_MTU, CHUNK_DATA_SIZE, MAX_RECORD_SIZE, V2_TAG_BITS}; @@ -28,7 +27,7 @@ impl Iterator for Chunker<'_> { tagged_record_len: V2_TAG_BITS | self.data.len() as u32, }; let mut chunk = [0; APP_MTU]; - let mut output = &mut chunk[..]; + let mut output = Writer::new(&mut chunk[..]); header.encode(&mut output); output.put_slice(data); self.index += 1; diff --git a/btp/src/lib.rs b/btp/src/lib.rs index f9276b1..3e8b5e7 100644 --- a/btp/src/lib.rs +++ b/btp/src/lib.rs @@ -37,11 +37,10 @@ pub fn packet_version(data: &[u8]) -> Option { } } -ql_codec::codec! { - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - struct Header { - sequence: u16, - index: u16, - tagged_record_len: u32, - } +#[derive(Debug, Clone, Copy, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +struct Header { + sequence: u16, + index: u16, + tagged_record_len: u32, } diff --git a/codec/Cargo.toml b/codec/Cargo.toml index 86a559f..39b7c1e 100644 --- a/codec/Cargo.toml +++ b/codec/Cargo.toml @@ -7,3 +7,7 @@ license.workspace = true [dependencies] bytes = { workspace = true } +ql-codec-macro = { workspace = true, optional = true } + +[features] +derive = ["dep:ql-codec-macro"] diff --git a/codec/macro/Cargo.toml b/codec/macro/Cargo.toml new file mode 100644 index 0000000..be3597d --- /dev/null +++ b/codec/macro/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "ql-codec-macro" +version = "0.1.0" +edition = "2021" +description = "QuantumLink codec derive" +license.workspace = true + +[lib] +proc-macro = true + +[dependencies] +proc-macro2 = "1.0.101" +quote = "1.0.40" +syn = { version = "2.0.106", default-features = false, features = [ + "clone-impls", + "derive", + "parsing", + "printing", + "proc-macro", +] } + +[dev-dependencies] +ql-codec = { workspace = true, features = ["derive"] } diff --git a/codec/macro/src/lib.rs b/codec/macro/src/lib.rs new file mode 100644 index 0000000..5b116f3 --- /dev/null +++ b/codec/macro/src/lib.rs @@ -0,0 +1,691 @@ +use proc_macro2::TokenStream; +use quote::{format_ident, quote}; +use syn::{parse_macro_input, parse_quote}; + +/// implements the codec traits and generates `Ref` unless the type has no fields +/// +/// See the `ql_codec` crate docs for the encoding and the `#[codec(...)]` attributes. +#[proc_macro_derive(Codec, attributes(codec))] +pub fn derive_codec_entry(item: proc_macro::TokenStream) -> proc_macro::TokenStream { + derive_codec(parse_macro_input!(item as syn::DeriveInput)) + .unwrap_or_else(|e| e.to_compile_error()) + .into() +} + +fn derive_codec(item: syn::DeriveInput) -> Result { + // The borrowed form already borrows its input, and borrowing from two places isn't worth it. + if let Some(lifetime) = item.generics.lifetimes().next() { + return Err(syn::Error::new_spanned( + lifetime, + "Lifetime parameters are not supported", + )); + } + let options = Options::parse(&item.attrs)?; + // Named apart from everything in the item, or it would capture a type or parameter called `B`. + let mut input = format_ident!("B"); + while mentions(quote!(#item), &input) { + input = format_ident!("{input}_"); + } + let ref_generics = ref_generics(&item, &input); + let expansion = match &item.data { + syn::Data::Struct(_) if options.discriminants.is_some() => { + return Err(syn::Error::new_spanned( + &item.ident, + "Only enums have discriminants", + )) + } + syn::Data::Struct(data) => derive_struct( + &item, + &ref_generics, + Fields::new(&data.fields, options.frozen, &input)?, + ), + syn::Data::Enum(data) => derive_enum(&item, &ref_generics, data, &options, &input)?, + syn::Data::Union(_) => { + return Err(syn::Error::new_spanned( + &item.ident, + "Unions are not supported", + )) + } + }; + Ok(expand( + &item, + &ref_generics, + &options.derives, + &input, + expansion, + )) +} + +#[derive(Default)] +struct Options { + frozen: bool, + derives: Vec, + discriminants: Option, +} + +impl Options { + fn parse(attrs: &[syn::Attribute]) -> Result { + let mut options = Self::default(); + for attr in attrs.iter().filter(|attr| attr.path().is_ident("codec")) { + attr.parse_nested_meta(|meta| { + if meta.path.is_ident("frozen") { + options.frozen = true; + Ok(()) + } else if meta.path.is_ident("derive") { + meta.parse_nested_meta(|derive| { + options.derives.push(derive.path); + Ok(()) + }) + } else if meta.path.is_ident("discriminants") { + options.discriminants = Some(meta.value()?.parse()?); + Ok(()) + } else { + Err(meta.error("Expected `frozen`, `derive(...)` or `discriminants = Name`")) + } + })?; + } + Ok(options) + } +} + +/// The parts that differ between structs and enums +struct Expansion { + /// declaration of the borrowed form when the type has fields + ref_declaration: Option, + decode_ref: TokenStream, + from_ref: TokenStream, + decode: TokenStream, + missing: TokenStream, + missing_owned: TokenStream, + encoded_len: TokenStream, + encode: TokenStream, + /// The discriminants enum and its accessor, if requested + discriminants: TokenStream, +} + +/// The form a decode builds +#[derive(Clone, Copy)] +enum Form { + Owned, + Borrowed, +} + +impl Form { + /// The form's type, as named inside the `Decode` impl + fn path(self) -> TokenStream { + match self { + Self::Owned => quote!(Self), + Self::Borrowed => quote!(Self::Ref), + } + } +} + +fn expand( + item: &syn::DeriveInput, + ref_generics: &syn::Generics, + derives: &[syn::Path], + input: &syn::Ident, + expansion: Expansion, +) -> TokenStream { + let Expansion { + ref_declaration, + decode_ref, + from_ref, + decode, + missing, + missing_owned, + encoded_len, + encode, + discriminants, + } = expansion; + let ident = &item.ident; + let ref_ident = ref_ident(ident); + let (_, ty_generics, _) = item.generics.split_for_impl(); + let (ref_impl_generics, ref_ty_generics, ref_bounds) = ref_generics.split_for_impl(); + + let encode_impl = |generics: syn::Generics, ty: TokenStream| { + let (impl_generics, _, where_clause) = generics.split_for_impl(); + quote!( + #[automatically_derived] + impl #impl_generics ::ql_codec::Encode for #ty #where_clause { + #[inline] + fn encoded_len(&self) -> usize { + #encoded_len + } + + fn encode(&self, out: &mut ::ql_codec::Writer) { + #encode + } + } + + #[automatically_derived] + impl #impl_generics ::ql_codec::Element for #ty #where_clause {} + ) + }; + let owned_encode = encode_impl( + with_field_bounds( + &item.generics, + item, + |ty| parse_quote!(#ty: ::ql_codec::Encode), + ), + quote!(#ident #ty_generics), + ); + let (ref_items, ref_type) = match ref_declaration { + Some(ref_declaration) => { + let doc = format!("[`{ident}`] with fields borrowed from its input container"); + let derives = (!derives.is_empty()).then(|| quote!(#[derive(#(#derives),*)])); + let ref_encode = encode_impl( + with_field_bounds( + ref_generics, + item, + |ty| parse_quote!(<#ty as ::ql_codec::Decode<#input>>::Ref: ::ql_codec::Encode), + ), + quote!(#ref_ident #ref_ty_generics), + ); + let ref_items = quote!( + #[doc = #doc] + #derives + #ref_declaration + + impl #ref_impl_generics #ref_ident #ref_ty_generics #ref_bounds { + /// copies borrowed data into the owned type + #[allow(dead_code)] + #[inline] + pub fn into_owned(self) -> #ident #ty_generics { + self.into() + } + } + + #[automatically_derived] + impl #ref_impl_generics ::core::convert::From<#ref_ident #ref_ty_generics> + for #ident #ty_generics #ref_bounds + { + #[inline] + fn from(value: #ref_ident #ref_ty_generics) -> Self { + >::from_ref(value) + } + } + + #ref_encode + ); + (ref_items, quote!(#ref_ident #ref_ty_generics)) + } + None => (quote!(), quote!(Self)), + }; + + quote!( + #ref_items + #discriminants + + #[automatically_derived] + impl #ref_impl_generics ::ql_codec::Decode<#input> for #ident #ty_generics #ref_bounds { + type Ref = #ref_type; + + #[inline] + fn decode(reader: &mut ::ql_codec::Reader<#input>) + -> ::core::result::Result + { + #decode + } + + #[inline] + fn decode_ref(reader: &mut ::ql_codec::Reader<#input>) + -> ::core::result::Result + { + #decode_ref + } + + #[inline] + fn from_ref(value: Self::Ref) -> Self { + #from_ref + } + + fn missing() -> ::core::option::Option { + #missing_owned + } + + #[inline] + fn missing_ref() -> ::core::option::Option { + #missing + } + } + + #owned_encode + ) +} + +fn derive_struct( + item: &syn::DeriveInput, + ref_generics: &syn::Generics, + fields: Fields, +) -> Expansion { + let vis = &item.vis; + let ref_ident = ref_ident(&item.ident); + + let pattern = fields.pattern(); + let decode_prelude = decode_prelude(fields.frozen); + let decode = |form: Form| { + let value = fields.decode(&form.path(), form); + quote!(#decode_prelude ::core::result::Result::Ok(#value)) + }; + let encoded_len = fields.encoded_len(); + let encode = fields.encode(); + let shape = fields.ref_shape(); + let where_clause = &ref_generics.where_clause; + let ref_declaration = match fields.shape { + syn::Fields::Unnamed(_) => { + quote!(#vis struct #ref_ident #ref_generics #shape #where_clause;) + } + _ => quote!(#vis struct #ref_ident #ref_generics #where_clause #shape), + }; + let owned = fields.owned("e!(Self)); + + Expansion { + ref_declaration: (!fields.shape.is_empty()).then_some(ref_declaration), + decode_ref: decode(Form::Borrowed), + from_ref: quote!(let Self::Ref #pattern = value; #owned), + decode: decode(Form::Owned), + missing: fields.missing(Form::Borrowed), + missing_owned: fields.missing(Form::Owned), + encoded_len: quote!(let Self #pattern = self; #encoded_len), + encode: quote!(let Self #pattern = self; #encode), + discriminants: quote!(), + } +} + +fn derive_enum( + item: &syn::DeriveInput, + ref_generics: &syn::Generics, + data: &syn::DataEnum, + options: &Options, + input: &syn::Ident, +) -> Result { + let frozen = options.frozen; + let vis = &item.vis; + let ref_ident = ref_ident(&item.ident); + + let mut unknown = None; + let mut next_enumerator = Some(0u32); + let mut borrows = false; + let mut ref_variants = TokenStream::new(); + let mut encoded_len_arms = TokenStream::new(); + let mut encode_arms = TokenStream::new(); + let mut decode_arms = TokenStream::new(); + let mut decode_ref_arms = TokenStream::new(); + let mut from_ref_arms = TokenStream::new(); + let mut discriminants = TokenStream::new(); + let mut discriminant_arms = TokenStream::new(); + for variant in &data.variants { + let enumerator = match &variant.discriminant { + Some((_, discriminant)) => parse_enumerator(discriminant)?, + None => next_enumerator + .ok_or_else(|| syn::Error::new_spanned(variant, "Enumerator overflows u32"))?, + }; + next_enumerator = enumerator.checked_add(1); + + if has_flag(&variant.attrs, "unknown")? { + if frozen { + return Err(syn::Error::new_spanned( + variant, + "A frozen enum has no unknown variant", + )); + } + if !matches!(variant.fields, syn::Fields::Unit) { + return Err(syn::Error::new_spanned( + variant, + "The unknown variant must be a unit variant", + )); + } + if unknown.is_some() { + return Err(syn::Error::new_spanned( + variant, + "Only one variant can be unknown", + )); + } + unknown = Some(&variant.ident); + } + + let ident = &variant.ident; + let fields = Fields::new(&variant.fields, frozen, input)?; + let cfg: Vec<_> = cfg_attrs(&variant.attrs).collect(); + let attrs = ref_attrs(&variant.attrs); + let shape = fields.ref_shape(); + let pattern = fields.pattern(); + let encoded_len = fields.encoded_len(); + let encode = fields.encode(); + let decode = fields.decode("e!(Self::#ident), Form::Owned); + let decode_ref = fields.decode("e!(Self::Ref::#ident), Form::Borrowed); + let owned = fields.owned("e!(Self::#ident)); + borrows |= !variant.fields.is_empty(); + ref_variants.extend(quote!(#(#attrs)* #ident #shape,)); + encoded_len_arms.extend(quote!(#(#cfg)* Self::#ident #pattern => ::ql_codec::varint::encoded_len(#enumerator) + #encoded_len,)); + encode_arms.extend(quote!(#(#cfg)* Self::#ident #pattern => { + ::ql_codec::varint::encode(#enumerator, out); + #encode + })); + decode_arms.extend(quote!(#(#cfg)* #enumerator => #decode,)); + decode_ref_arms.extend(quote!(#(#cfg)* #enumerator => #decode_ref,)); + from_ref_arms.extend(quote!(#(#cfg)* Self::Ref::#ident #pattern => #owned,)); + if let Some(name) = &options.discriminants { + let enumerator = proc_macro2::Literal::u32_unsuffixed(enumerator); + discriminants.extend(quote!(#(#cfg)* #ident = #enumerator,)); + discriminant_arms.extend(quote!(#(#cfg)* Self::#ident { .. } => #name::#ident,)); + } + } + // `match self {}` is rejected for a reference to an empty enum. + let scrutinee = if data.variants.is_empty() { + quote!(*self) + } else { + quote!(self) + }; + + let decode_prelude = decode_prelude(frozen); + let decode = |form: Form, arms: TokenStream| { + let path = form.path(); + let unknown_arm = match unknown { + Some(unknown) => quote!(_ => #path::#unknown {},), + None => quote!( + _ => return ::core::result::Result::Err(::ql_codec::Error::InvalidDiscriminant), + ), + }; + quote!( + let enumerator = reader.decode_varint::()?; + #decode_prelude + ::core::result::Result::Ok(match enumerator { + #arms + #unknown_arm + }) + ) + }; + + let discriminants = options.discriminants.as_ref().map(|name| { + let ident = &item.ident; + let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl(); + let doc = format!("[`{ident}`]'s variants without their fields, numbered by enumerator"); + quote!( + #[doc = #doc] + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + #vis enum #name { #discriminants } + + impl #impl_generics #ident #ty_generics #where_clause { + /// The variant without its fields + #[allow(dead_code)] + #[inline] + #vis fn discriminant(&self) -> #name { + match #scrutinee { #discriminant_arms } + } + } + ) + }); + + let where_clause = &ref_generics.where_clause; + Ok(Expansion { + ref_declaration: borrows + .then(|| quote!(#vis enum #ref_ident #ref_generics #where_clause { #ref_variants })), + decode_ref: decode(Form::Borrowed, decode_ref_arms), + from_ref: quote!(match value { #from_ref_arms }), + decode: decode(Form::Owned, decode_arms), + missing: match unknown { + Some(unknown) => quote!(::core::option::Option::Some(Self::Ref::#unknown {})), + None => quote!(::core::option::Option::None), + }, + missing_owned: match unknown { + Some(unknown) => quote!(::core::option::Option::Some(Self::#unknown {})), + None => quote!(::core::option::Option::None), + }, + encoded_len: quote!(match #scrutinee { #encoded_len_arms }), + encode: quote!(match #scrutinee { #encode_arms }), + discriminants: discriminants.unwrap_or_default(), + }) +} + +struct Fields<'a> { + shape: &'a syn::Fields, + frozen: bool, + input: &'a syn::Ident, +} + +impl<'a> Fields<'a> { + fn new( + shape: &'a syn::Fields, + frozen: bool, + input: &'a syn::Ident, + ) -> Result { + if let Some(attr) = shape + .iter() + .flat_map(|field| &field.attrs) + .find(|attr| attr.path().is_ident("codec")) + { + return Err(syn::Error::new_spanned( + attr, + "Fields take no codec attributes", + )); + } + Ok(Self { + shape, + frozen, + input, + }) + } + + /// `{ a: field_0, 1: field_1 }` also matches tuple and unit shapes + fn pattern(&self) -> TokenStream { + self.construct("e!(), |_, binding| quote!(#binding)) + } + + /// length of the fields bound by [`Self::pattern`] + fn encoded_len(&self) -> TokenStream { + let fields_len = self.fields_len(); + if self.frozen { + fields_len + } else { + quote!(::ql_codec::encoded_len_prefixed(#fields_len)) + } + } + + /// encodes the fields bound by [`Self::pattern`] into `out` + fn encode(&self) -> TokenStream { + let fields = self.shape.iter().enumerate().map(|(index, field)| { + let binding = format_ident!("field_{index}"); + let cfg = cfg_attrs(&field.attrs); + quote!(#(#cfg)* ::ql_codec::Encode::encode(#binding, out);) + }); + let len = (!self.frozen).then(|| { + let fields_len = self.fields_len(); + quote!(::ql_codec::encode_len_prefix(#fields_len, out);) + }); + quote!(#len #(#fields)*) + } + + fn fields_len(&self) -> TokenStream { + let fields = self.shape.iter().enumerate().map(|(index, field)| { + let binding = format_ident!("field_{index}"); + let cfg = cfg_attrs(&field.attrs); + quote!(#(#cfg)* { len += ::ql_codec::Encode::encoded_len(#binding); }) + }); + quote!({ + #[allow(unused_mut)] + let mut len = 0usize; + #(#fields)* + len + }) + } + + fn decode(&self, path: &TokenStream, form: Form) -> TokenStream { + let input = self.input; + self.construct(path, |field, _| { + let ty = &field.ty; + match (self.frozen, form) { + (true, Form::Owned) => quote!(<#ty as ::ql_codec::Decode<#input>>::decode(reader)?), + (true, Form::Borrowed) => { + quote!(<#ty as ::ql_codec::Decode<#input>>::decode_ref(reader)?) + } + (false, Form::Owned) => quote!(::ql_codec::decode_field::<#ty, _>(&mut fields)?), + (false, Form::Borrowed) => { + quote!(::ql_codec::decode_field_ref::<#ty, _>(&mut fields)?) + } + } + }) + } + + fn missing(&self, form: Form) -> TokenStream { + let input = self.input; + let value = self.construct(&form.path(), |field, _| { + let ty = &field.ty; + match form { + Form::Owned => quote!(<#ty as ::ql_codec::Decode<#input>>::missing()?), + Form::Borrowed => quote!(<#ty as ::ql_codec::Decode<#input>>::missing_ref()?), + } + }); + quote!(::core::option::Option::Some(#value)) + } + + /// builds the owned value from the borrowed fields bound by [`Self::pattern`] + fn owned(&self, path: &TokenStream) -> TokenStream { + let input = self.input; + self.construct(path, |field, binding| { + let ty = &field.ty; + quote!(<#ty as ::ql_codec::Decode<#input>>::from_ref(#binding)) + }) + } + + /// `path { .. }` with each field set to `value(field, binding)`, `binding` being its name in + /// [`Self::pattern`] + fn construct( + &self, + path: &TokenStream, + value: impl Fn(&syn::Field, &syn::Ident) -> TokenStream, + ) -> TokenStream { + let fields = self.shape.iter().enumerate().map(|(index, field)| { + let member = field + .ident + .clone() + .map_or_else(|| syn::Member::Unnamed(index.into()), syn::Member::Named); + let cfg = cfg_attrs(&field.attrs); + let value = value(field, &format_ident!("field_{index}")); + quote!(#(#cfg)* #member: #value,) + }); + quote!(#path { #(#fields)* }) + } + + /// the fields of the borrowed form in the declared shape + fn ref_shape(&self) -> TokenStream { + let input = self.input; + let fields = self.shape.iter().map(|field| { + let vis = &field.vis; + let ty = &field.ty; + let attrs = ref_attrs(&field.attrs); + let name = field.ident.as_ref().map(|name| quote!(#name:)); + quote!(#(#attrs)* #vis #name <#ty as ::ql_codec::Decode<#input>>::Ref,) + }); + match self.shape { + syn::Fields::Named(_) => quote!({ #(#fields)* }), + syn::Fields::Unnamed(_) => quote!(( #(#fields)* )), + syn::Fields::Unit => quote!(), + } + } +} + +/// Takes the length-prefixed fields out of `reader` as `fields`, for [`Fields::decode`] +fn decode_prelude(frozen: bool) -> TokenStream { + if frozen { + return quote!(); + } + quote!( + // Unused without fields, but the body still has to be skipped. + #[allow(unused_mut, unused_variables)] + let mut fields = reader.take_len_prefixed_reader()?; + ) +} + +/// Whether `attrs` hold `#[codec(name)]`, the only codec attribute allowed there +fn has_flag(attrs: &[syn::Attribute], name: &str) -> Result { + let mut found = false; + for attr in attrs.iter().filter(|attr| attr.path().is_ident("codec")) { + attr.parse_nested_meta(|meta| { + if !meta.path.is_ident(name) { + return Err(meta.error(format!("Expected `{name}`"))); + } + found = true; + Ok(()) + })?; + } + Ok(found) +} + +fn parse_enumerator(discriminant: &syn::Expr) -> Result { + let syn::Expr::Lit(syn::ExprLit { + lit: syn::Lit::Int(int), + .. + }) = discriminant + else { + return Err(syn::Error::new_spanned( + discriminant, + "Discriminant must be an integer literal", + )); + }; + int.base10_parse() +} + +fn ref_ident(ident: &syn::Ident) -> syn::Ident { + format_ident!("{ident}Ref") +} + +/// Whether `ident` occurs anywhere in `tokens` +fn mentions(tokens: TokenStream, ident: &syn::Ident) -> bool { + tokens.into_iter().any(|tree| match tree { + proc_macro2::TokenTree::Ident(other) => other == *ident, + proc_macro2::TokenTree::Group(group) => mentions(group.stream(), ident), + _ => false, + }) +} + +/// the input container followed by the type's own generics +fn ref_generics(item: &syn::DeriveInput, input: &syn::Ident) -> syn::Generics { + let mut generics = item.generics.clone(); + generics + .params + .insert(0, parse_quote!(#input: ::ql_codec::ByteSlice)); + with_field_bounds( + &generics, + item, + |ty| parse_quote!(#ty: ::ql_codec::Decode<#input>), + ) +} + +/// `generics` with `bound` on each field type, if the item is generic +fn with_field_bounds( + generics: &syn::Generics, + item: &syn::DeriveInput, + bound: impl Fn(&syn::Type) -> syn::WherePredicate, +) -> syn::Generics { + let mut generics = generics.clone(); + if item.generics.type_params().next().is_none() { + return generics; + } + let predicates = &mut generics.make_where_clause().predicates; + match &item.data { + syn::Data::Struct(data) => { + predicates.extend(data.fields.iter().map(|field| bound(&field.ty))) + } + syn::Data::Enum(data) => predicates.extend( + data.variants + .iter() + .flat_map(|variant| &variant.fields) + .map(|field| bound(&field.ty)), + ), + syn::Data::Union(_) => {} + } + generics +} + +fn cfg_attrs(attrs: &[syn::Attribute]) -> impl Iterator { + attrs.iter().filter(|attr| attr.path().is_ident("cfg")) +} + +/// Attributes carried over to the borrowed form +fn ref_attrs(attrs: &[syn::Attribute]) -> impl Iterator { + attrs + .iter() + .filter(|attr| attr.path().is_ident("cfg") || attr.path().is_ident("doc")) +} diff --git a/codec/macro/tests/compatibility.rs b/codec/macro/tests/compatibility.rs new file mode 100644 index 0000000..b6de2ac --- /dev/null +++ b/codec/macro/tests/compatibility.rs @@ -0,0 +1,248 @@ +use std::{fmt::Debug, num::NonZeroU8}; + +use ql_codec::{Codec, Decode, Encode, Error}; + +mod v1 { + use ql_codec::Codec; + + #[derive(Debug, PartialEq, Eq, Codec)] + pub struct DeviceInfo { + pub name: String, + pub battery: u8, + } + + #[derive(Debug, PartialEq, Eq, Codec)] + pub enum UsbEvent { + Connected, + Disconnected, + #[codec(unknown)] + Unknown, + } + + #[derive(Debug, PartialEq, Eq, Codec)] + pub struct Status { + pub device: DeviceInfo, + pub events: Vec, + pub uptime: u32, + } +} + +/// `v1` grown in the compatible ways: fields and variants appended, and a unit variant gaining +/// fields +mod v2 { + use ql_codec::Codec; + + #[derive(Debug, PartialEq, Eq, Codec)] + pub struct DeviceInfo { + pub name: String, + pub battery: u8, + pub serial: Option, + pub charging: bool, + } + + #[derive(Debug, PartialEq, Eq, Codec)] + pub enum UsbEvent { + Connected { + speed: u32, + }, + Disconnected, + #[codec(unknown)] + Unknown, + Suspended, + } + + #[derive(Debug, PartialEq, Eq, Codec)] + pub struct Status { + pub device: DeviceInfo, + pub events: Vec, + pub uptime: u32, + pub tags: Vec, + } +} + +#[derive(Debug, PartialEq, Eq, Codec)] +#[codec(frozen)] +pub struct ImageSize { + pub width: u16, + pub height: u16, +} + +#[derive(Debug, PartialEq, Eq, Codec)] +#[codec(frozen)] +pub enum PixelFormat { + Rgb = 1, + Rgba, + AlphaMap = 7, +} + +#[derive(Debug, PartialEq, Eq, Codec)] +#[codec(frozen)] +pub enum Shape { + Point, + Image(ImageSize, PixelFormat), +} + +#[derive(Debug, PartialEq, Eq, Codec)] +pub enum ChargeState { + Idle, + Charging, +} + +#[derive(Debug, PartialEq, Eq, Codec)] +pub struct Process { + pub pid: NonZeroU8, +} + +#[rustfmt::skip] +const V1_STATUS: &[u8] = &[ + 14, // length + 4, 2, b'k', b'b', 80, // device + 2, 0, 0, 1, 0, // events + 0x02, 0x01, 0x00, 0x00, // uptime +]; + +#[rustfmt::skip] +const V2_STATUS: &[u8] = &[ + 28, // length + 9, 2, b'k', b'b', 80, 1, 2, b'A', b'1', 1, // device + 3, 0, 4, 0xe0, 0x01, 0x00, 0x00, 3, 0, 1, 0, // events + 0x02, 0x01, 0x00, 0x00, // uptime + 1, 1, b'x', // tags +]; + +const IMAGE_SHAPE: &[u8] = &[1, 0x40, 0x01, 0xf0, 0x00, 7]; + +fn v1_status() -> v1::Status { + v1::Status { + device: v1::DeviceInfo { + name: "kb".into(), + battery: 80, + }, + events: vec![v1::UsbEvent::Connected, v1::UsbEvent::Disconnected], + uptime: 0x0102, + } +} + +fn v2_status() -> v2::Status { + v2::Status { + device: v2::DeviceInfo { + name: "kb".into(), + battery: 80, + serial: Some("A1".into()), + charging: true, + }, + events: vec![ + v2::UsbEvent::Connected { speed: 480 }, + v2::UsbEvent::Suspended, + v2::UsbEvent::Disconnected, + ], + uptime: 0x0102, + tags: vec!["x".into()], + } +} + +fn image_shape() -> Shape { + Shape::Image( + ImageSize { + width: 320, + height: 240, + }, + PixelFormat::AlphaMap, + ) +} + +/// decodes both forms +fn decode<'a, T: Decode<&'a [u8]> + Debug + PartialEq>(bytes: &'a [u8]) -> T { + let owned = T::decode_bytes(bytes).unwrap(); + assert_eq!(T::from_ref(T::decode_bytes_ref(bytes).unwrap()), owned); + owned +} + +#[test] +fn encodings_match_fixed_bytes() { + assert_eq!(v1_status().encode_vec(), V1_STATUS); + assert_eq!(v2_status().encode_vec(), V2_STATUS); + assert_eq!(image_shape().encode_vec(), IMAGE_SHAPE); + assert_eq!(PixelFormat::Rgba.encode_vec(), [2]); + + assert_eq!(decode::(V1_STATUS), v1_status()); + assert_eq!(decode::(V2_STATUS), v2_status()); + assert_eq!(decode::(IMAGE_SHAPE), image_shape()); +} + +#[test] +fn integer_lists_are_aligned_to_their_size() { + assert_eq!((1u8, vec![0x0302u16]).encode_vec(), [1, 1, 2, 3, 0]); + assert_eq!((1u8, 2u8, vec![0x0302u16]).encode_vec(), [1, 2, 1, 0, 2, 3]); + assert_eq!( + (1u8, [0x0504_0302u32]).encode_vec(), + [1, 0, 0, 0, 2, 3, 4, 5] + ); +} + +#[test] +fn new_decoder_fills_missing_fields() { + let status: v2::Status = decode(V1_STATUS); + assert_eq!( + status, + v2::Status { + device: v2::DeviceInfo { + name: "kb".into(), + battery: 80, + serial: None, + charging: false, + }, + events: vec![ + v2::UsbEvent::Connected { speed: 0 }, + v2::UsbEvent::Disconnected + ], + uptime: 0x0102, + tags: vec![], + } + ); +} + +#[test] +fn old_decoder_skips_new_fields_and_variants() { + let status: v1::Status = decode(V2_STATUS); + assert_eq!( + status, + v1::Status { + device: v1::DeviceInfo { + name: "kb".into(), + battery: 80 + }, + events: vec![ + v1::UsbEvent::Connected, + v1::UsbEvent::Unknown, + v1::UsbEvent::Disconnected, + ], + uptime: 0x0102, + } + ); +} + +#[test] +fn undecodable_encodings() { + // An enumerator past the known ones needs an unknown variant to land in. + assert_eq!( + ChargeState::decode_bytes(&[2, 0][..]), + Err(Error::InvalidDiscriminant) + ); + assert_eq!( + PixelFormat::decode_bytes(&[9][..]), + Err(Error::InvalidDiscriminant) + ); + // Only fields with a missing value can be left out. + assert_eq!(Process::decode_bytes(&[0][..]), Err(Error::UnexpectedEof)); + // A struct may end between fields, not inside one, even if the bytes after it would complete + // the field. + assert_eq!( + v1::DeviceInfo::decode_bytes(&[2, 5, b'k', b'b', b'x', b'y', b'z', 80][..]), + Err(Error::UnexpectedEof) + ); + assert_eq!( + Shape::decode_bytes(&IMAGE_SHAPE[..5]), + Err(Error::UnexpectedEof) + ); +} diff --git a/codec/macro/tests/malformed.rs b/codec/macro/tests/malformed.rs new file mode 100644 index 0000000..c7739a7 --- /dev/null +++ b/codec/macro/tests/malformed.rs @@ -0,0 +1,94 @@ +use std::{ + alloc::{GlobalAlloc, Layout, System}, + num::NonZeroU8, + time::{Duration, SystemTime}, +}; + +use ql_codec::{Codec, Decode, Error}; + +#[derive(Debug, PartialEq, Eq, Codec)] +pub enum Command { + Reboot, + SetBrightness(u8), +} + +/// Refuses allocations over 64 MiB, so reserving memory for a corrupt length aborts the test +/// instead of quietly succeeding under overcommit +struct CappedAllocator; + +// SAFETY: There's no safe way to limit allocations. This delegates to `System` and only ever +// refuses one, which `GlobalAlloc` permits. +unsafe impl GlobalAlloc for CappedAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + if layout.size() > 64 << 20 { + return std::ptr::null_mut(); + } + // SAFETY: The caller upholds `alloc`'s contract, which is the same for `System`. + unsafe { System.alloc(layout) } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + // SAFETY: `ptr` came from `System.alloc` with this `layout`. + unsafe { System.dealloc(ptr, layout) } + } +} + +#[global_allocator] +static ALLOCATOR: CappedAllocator = CappedAllocator; + +#[test] +fn corrupt_lengths() { + let huge_length = [0xff, 0xff, 0xff, 0xff, 0x0f]; + let input = [huge_length.as_slice(), &[0; 1 << 20]].concat(); + // Reserving room for `u32::MAX` items, or for an item per byte left, would hit the + // allocator's cap. + assert_eq!( + Vec::<[u8; 1024]>::decode_bytes(input.as_slice()), + Err(Error::UnexpectedEof) + ); + assert_eq!( + Vec::::decode_bytes(input.as_slice()), + Err(Error::UnexpectedEof) + ); +} + +#[test] +fn invalid_values() { + assert_eq!( + Command::decode_bytes(&[2, 1, 50][..]), + Err(Error::InvalidDiscriminant) + ); + assert_eq!( + bool::decode_bytes(&[2][..]), + Err(Error::InvalidDiscriminant) + ); + assert_eq!( + Option::::decode_bytes(&[2, 0][..]), + Err(Error::InvalidDiscriminant) + ); + assert_eq!( + Result::::decode_bytes(&[2, 0][..]), + Err(Error::InvalidDiscriminant) + ); + assert_eq!(NonZeroU8::decode_bytes(&[0][..]), Err(Error::InvalidData)); + assert_eq!( + String::decode_bytes(&[2, 0xc3, 0x28][..]), + Err(Error::InvalidUtf8) + ); + assert_eq!( + Duration::decode_bytes( + [[0; 8].as_slice(), &1_000_000_000u32.to_le_bytes()] + .concat() + .as_slice() + ), + Err(Error::InvalidRange) + ); + assert_eq!( + SystemTime::decode_bytes( + [u64::MAX.to_le_bytes().as_slice(), &[0; 4]] + .concat() + .as_slice() + ), + Err(Error::InvalidRange) + ); +} diff --git a/codec/macro/tests/mutable.rs b/codec/macro/tests/mutable.rs new file mode 100644 index 0000000..fa4fa44 --- /dev/null +++ b/codec/macro/tests/mutable.rs @@ -0,0 +1,205 @@ +use std::{ + collections::{hash_map::RandomState, BTreeMap}, + hash::BuildHasher, +}; + +use ql_codec::{Codec, Decode, Encode, Error, Reader, Ref, Target}; + +#[derive(Debug, PartialEq, Codec)] +enum Item { + Text(B), + #[codec(unknown)] + Unknown, +} + +#[derive(Debug, PartialEq, Codec)] +struct Message { + items: Vec>>, + outcome: Result, String>, + attributes: BTreeMap, + pair: (Box<[u8]>, String), + count: u32, +} + +#[test] +fn nested_views_mutate_the_encoding_and_outlive_the_reader() { + let mut expected = Message { + items: vec![Some(Item::Text("hello".into())), None], + outcome: Ok(Box::new([1, 2])), + attributes: BTreeMap::from([("key".into(), "value".into())]), + pair: (Box::from(&b"bytes"[..]), "tail".into()), + count: 42, + }; + let mut bytes = expected.encode_vec(); + bytes.push(9); + let (mut view, tail) = { + let mut reader = Reader::new(bytes.as_mut_slice()); + let view: MessageRef<&mut [u8]> = reader.decode_ref::().unwrap(); + let tail = reader.take_all(); + (view, tail) + }; + let Some(ItemRef::Text(text)) = &mut view.items[0] else { + panic!("expected text"); + }; + text.make_ascii_uppercase(); + let fixed = view.outcome.as_mut().unwrap(); + AsMut::<[u8]>::as_mut(fixed)[0] = 3; + assert_eq!(AsRef::<[u8]>::as_ref(fixed), &[3, 2]); + let (mut key, mut value) = view.attributes.pop_first().unwrap(); + key.make_ascii_uppercase(); + value.make_ascii_uppercase(); + view.attributes.insert(key, value); + view.pair.0[0] = b'B'; + view.pair.1.make_ascii_uppercase(); + tail[0] = 10; + + expected.items[0] = Some(Item::Text("HELLO".into())); + expected.outcome = Ok(Box::new([3, 2])); + expected.attributes = BTreeMap::from([("KEY".into(), "VALUE".into())]); + expected.pair = (Box::from(&b"Bytes"[..]), "TAIL".into()); + assert_eq!(view.encode_vec(), expected.encode_vec()); + assert_eq!(view.into_owned(), expected); + assert_eq!(bytes.last(), Some(&10)); + let shared: MessageRef<&[u8]> = Reader::new(bytes.as_slice()) + .decode_ref::() + .unwrap(); + assert_eq!(shared.into_owned(), expected); +} + +#[test] +fn missing_mutable_fields_only_default_without_backing_storage() { + #[derive(Codec)] + struct Defaults { + text: String, + bytes: Box<[u8]>, + items: Vec, + fixed: Option<[u32; 2]>, + empty: Box<[u8; 0]>, + } + + let mut bytes = [0]; + let view = Reader::new(bytes.as_mut_slice()) + .decode_ref::() + .unwrap(); + assert_eq!(view.text, ""); + assert!(view.bytes.is_empty()); + assert!(view.items.is_empty()); + assert!(view.fixed.is_none()); + assert!(view.empty.is_empty()); + + #[derive(Codec)] + struct Fixed { + bytes: Box<[u8; 2]>, + } + + assert_eq!(Fixed::decode_bytes_ref(&[0][..]).unwrap().bytes, &[0, 0]); + assert!(matches!( + Reader::new(bytes.as_mut_slice()).decode_ref::(), + Err(Error::UnexpectedEof) + )); + assert_eq!( + *Fixed::decode_bytes(bytes.as_mut_slice()).unwrap().bytes, + [0, 0] + ); + assert!(<[u32; 2] as Decode<&mut [u8]>>::missing_ref().is_none()); + assert!(<[u32; 0] as Decode<&mut [u8]>>::missing_ref() + .unwrap() + .is_empty()); + + #[derive(Codec)] + struct Integers { + words: [u32; 2], + samples: Vec, + } + + let integers = Integers::decode_bytes_ref(&[0][..]).unwrap(); + assert_eq!(*integers.words, [0, 0]); + assert!(integers.samples.is_empty()); +} + +#[test] +fn mutable_strings_validate_utf8_and_length() { + let mut invalid = [1, 0xff]; + assert!(matches!( + String::decode_bytes_ref(invalid.as_mut_slice()), + Err(Error::InvalidUtf8) + )); + let mut truncated = [2, b'a']; + assert!(matches!( + String::decode_bytes(truncated.as_mut_slice()), + Err(Error::UnexpectedEof) + )); +} + +#[test] +fn shared_views_derive_value_traits_and_fixed_refs_validate_length() { + #[derive(Codec)] + #[codec(frozen, derive(Debug, Clone, Copy, PartialEq, Eq, Hash))] + struct Header { + bytes: [u8; 2], + text: String, + } + + let bytes = Header { + bytes: [1, 2], + text: "hello".into(), + } + .encode_vec(); + let shared = Header::decode_bytes_ref(bytes.as_slice()).unwrap(); + let copied = shared; + assert_eq!(copied, shared); + let fixed: &[u8] = copied.bytes.as_ref(); + assert_eq!(fixed, &[1, 2]); + assert_eq!(fixed.as_ptr(), bytes.as_ptr()); + let hasher = RandomState::new(); + assert_eq!( + hasher.hash_one(shared), + hasher.hash_one(([1u8, 2], "hello")) + ); + assert_eq!(shared.into_owned().bytes, [1, 2]); + assert!(matches!( + Ref::<_, [u8; 2]>::new(&bytes[..1]), + Err(Error::InvalidData) + )); +} + +#[test] +fn custom_targets_validate_elements() { + #[repr(u8)] + enum Switch { + Off, + On, + } + + // safety: validation accepts exactly the valid discriminants and the layout has no padding + unsafe impl Target for Switch { + fn validate(bytes: &[u8]) -> Result<(), Error> { + ::validate(bytes)?; + match bytes[0] { + 0 | 1 => Ok(()), + _ => Err(Error::InvalidData), + } + } + + unsafe fn from_bytes(bytes: &[u8]) -> &Self { + // safety: the caller checked the layout and discriminant + unsafe { &*bytes.as_ptr().cast() } + } + + unsafe fn from_bytes_mut(bytes: &mut [u8]) -> &mut Self { + // safety: the caller checked the layout and discriminant and provides exclusive access + unsafe { &mut *bytes.as_mut_ptr().cast() } + } + } + + assert!(matches!( + Ref::<_, [[Switch; 2]; 1]>::new(&[0, 2][..]), + Err(Error::InvalidData) + )); + + let mut bytes = [0, 1]; + let mut view = Ref::<_, [Switch]>::new(&mut bytes[..]).unwrap(); + view[0] = Switch::On; + view[1] = Switch::Off; + assert_eq!(view.into_bytes(), &[1, 0]); +} diff --git a/codec/macro/tests/roundtrip.rs b/codec/macro/tests/roundtrip.rs new file mode 100644 index 0000000..b7d597e --- /dev/null +++ b/codec/macro/tests/roundtrip.rs @@ -0,0 +1,356 @@ +use std::{ + collections::BTreeMap, + fmt::Debug, + num::NonZeroU8, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use ql_codec::{Codec, Decode, Encode, Reader, Ref}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Codec)] +#[codec(frozen)] +pub struct AppId(pub [u8; 16]); + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Codec)] +pub enum Location { + System, + AppData, + Usb, +} + +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +#[codec(derive(Debug, PartialEq))] +pub enum Permission { + Camera, + Files { location: Location, read_only: bool }, + Custom(String), +} + +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +pub struct InstalledApp { + pub app_id: AppId, + pub pid: Option, + pub name: String, + pub icon: Vec, + pub permissions: Vec, + pub signatures: Vec>, + pub keywords: Vec, + pub nine_slice: Option<[u16; 4]>, + pub size: usize, + pub installed_at: SystemTime, +} + +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +pub enum AppManagerError { + NotInstalled(AppId), + PermissionDenied { + location: Location, + }, + Io(String), + #[codec(unknown)] + InternalError, +} + +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +pub struct WriteFile { + pub path: String, + pub data: Vec, + pub location: Location, +} + +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +pub struct ListApps; + +/// Named like the borrowed form's input parameter, which it must not capture +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +pub struct B(pub u32); + +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +pub struct Holder { + pub b: B, +} + +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +pub struct Transceive(pub (Vec, Vec, Duration)); + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Codec)] +#[codec(discriminants = SeekFromKind)] +pub enum SeekFrom { + Start(u64), + End(i64), + Current(i64), +} + +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +pub struct IconSet(pub BTreeMap>); + +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +#[codec(frozen)] +pub struct Envelope { + pub cid: u32, + pub message: B, +} + +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +pub enum Reply { + Done(T), + Failed(Vec), + #[codec(unknown)] + Unknown, +} + +/// A module with its own `Result`, which the derive's output must not pick up +#[allow(dead_code)] +mod own_result { + use ql_codec::Codec; + + pub type Result = std::result::Result; + + #[derive(Codec)] + pub struct Request { + pub id: u32, + } + + #[derive(Codec)] + pub enum Response { + Done(u32), + #[codec(unknown)] + Unknown, + } +} + +fn installed_app() -> InstalledApp { + InstalledApp { + app_id: AppId(*b"wallet-app-id-16"), + pid: NonZeroU8::new(42), + name: "Wallet".into(), + icon: (0..=255).collect(), + permissions: vec![ + Permission::Camera, + Permission::Files { + location: Location::AppData, + read_only: true, + }, + Permission::Custom("dev".into()), + ], + signatures: vec![vec![0xaa; 64], vec![]], + keywords: vec!["bitcoin".into(), "".into()], + nine_slice: Some([4, 4, 8, 8]), + size: 3 << 20, + installed_at: UNIX_EPOCH + Duration::new(1_700_000_000, 999_999_999), + } +} + +fn roundtrip(value: &T) +where + T: Encode + Debug + PartialEq, + for<'a> T: Decode<&'a [u8]> + Decode<&'a mut [u8]>, +{ + let mut bytes = value.encode_vec(); + assert_eq!(T::decode_bytes(bytes.as_slice()).unwrap(), *value); + let mut reader = Reader::new(bytes.as_mut_slice()); + assert_eq!(T::decode(&mut reader).unwrap(), *value); + assert!(reader.is_empty()); +} + +#[test] +fn borrowed_form_points_into_the_encoding() { + let app = installed_app(); + let bytes = app.encode_vec(); + let decoded = InstalledApp::decode_bytes_ref(bytes.as_slice()).unwrap(); + + let within = |slice: &[u8]| bytes.as_ptr_range().contains(&slice.as_ptr()); + let app_id: &[u8; 16] = &decoded.app_id.0; + let icon: &[u8] = &decoded.icon; + let name: &str = &decoded.name; + let signatures: &[Ref<&[u8], [u8]>] = &decoded.signatures; + let keywords: &[Ref<&[u8], str>] = &decoded.keywords; + assert!(within(app_id)); + assert!(within(icon)); + assert!(within(name.as_bytes())); + assert!(within(&signatures[0])); + assert!(within(keywords[0].as_bytes())); + assert_eq!( + decoded.permissions[2], + PermissionRef::Custom(Ref::new(&b"dev"[..]).unwrap()) + ); + + assert_eq!(decoded.encode_vec(), bytes); + assert_eq!(InstalledApp::from(decoded), app); +} + +#[test] +fn borrowed_form_encodes_like_the_owned_one() { + let data = [0x5a; 300]; + let write = WriteFileRef { + path: Ref::new(&b"notes.txt"[..]).unwrap(), + data: Ref::new(data.as_slice()).unwrap(), + location: Location::AppData, + }; + let bytes = write.encode_vec(); + + let owned = WriteFile { + path: "notes.txt".into(), + data: data.to_vec(), + location: Location::AppData, + }; + assert_eq!(bytes, owned.encode_vec()); + assert_eq!(WriteFile::decode_bytes(bytes.as_slice()).unwrap(), owned); +} + +#[test] +fn responses_roundtrip() { + let sparse = InstalledApp { + pid: None, + icon: vec![], + permissions: vec![], + nine_slice: None, + ..installed_app() + }; + let responses: [Result, AppManagerError>; 6] = [ + Ok(vec![]), + Ok(vec![installed_app(), sparse]), + Err(AppManagerError::NotInstalled(AppId([7; 16]))), + Err(AppManagerError::PermissionDenied { + location: Location::Usb, + }), + Err(AppManagerError::Io("disk full".into())), + Err(AppManagerError::InternalError), + ]; + for response in &responses { + roundtrip(response); + } +} + +#[test] +fn maps_roundtrip() { + let icons = IconSet(BTreeMap::from([ + ("wallet".to_string(), vec![1, 2, 3]), + ("camera".to_string(), vec![]), + ])); + roundtrip(&icons); + roundtrip(&IconSet(BTreeMap::new())); + + let bytes = icons.encode_vec(); + let borrowed = IconSet::decode_bytes_ref(bytes.as_slice()).unwrap(); + assert_eq!(*borrowed.0["wallet"], [1, 2, 3]); + assert_eq!(borrowed.encode_vec(), bytes); +} + +#[test] +fn generic_types_roundtrip() { + let envelope = Envelope { + cid: 7, + message: WriteFile { + path: "notes.txt".into(), + data: vec![1, 2, 3], + location: Location::Usb, + }, + }; + roundtrip(&envelope); + let bytes = envelope.encode_vec(); + let borrowed = Envelope::::decode_bytes_ref(bytes.as_slice()).unwrap(); + assert_eq!(borrowed.message.path, "notes.txt"); + assert_eq!(borrowed.encode_vec(), bytes); + assert_eq!(borrowed.into_owned(), envelope); + + for reply in [ + Reply::Done(3u64), + Reply::Failed(vec!["disk full".to_string()]), + Reply::Unknown, + ] { + roundtrip(&reply); + } +} + +#[test] +fn messages_roundtrip() { + roundtrip(&ListApps); + roundtrip(&Holder { b: B(7) }); + roundtrip(&Transceive(( + vec![0x00, 0xa4, 0x04, 0x00], + vec![0x90, 0x00], + Duration::from_millis(250), + ))); + for seek in [ + SeekFrom::Start(u64::MAX), + SeekFrom::End(-17), + SeekFrom::Current(i64::MIN), + ] { + roundtrip(&seek); + assert_eq!(seek.encode_vec()[0], seek.discriminant() as u8); + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Codec)] +pub struct Samples { + pub name: String, + pub values: Vec, + pub sid: [u64; 2], +} + +fn samples(name: &str) -> Samples { + Samples { + name: name.into(), + values: vec![1, u32::MAX, 3], + sid: [u64::MAX, 5], + } +} + +#[repr(align(8))] +struct AlignedBuffer([u8; 256]); + +/// `bytes` copied into a buffer aligned for any integer, `offset` bytes in +fn copy_aligned(bytes: &[u8], offset: usize) -> AlignedBuffer { + let mut buffer = AlignedBuffer([0; 256]); + buffer.0[offset..][..bytes.len()].copy_from_slice(bytes); + buffer +} + +#[test] +fn integer_lists_are_lent_aligned() { + // Each name shifts the integers after it differently. + let list = Vec::from(["", "a", "ab", "abc"].map(samples)); + let bytes = list.encode_vec(); + let buffer = copy_aligned(&bytes, 0); + let encoding = &buffer.0[..bytes.len()]; + let decoded = Vec::::decode_bytes_ref(encoding).unwrap(); + + for (decoded, samples) in decoded.iter().zip(&list) { + let values: &[u32] = &decoded.values; + let sid: &[u64; 2] = &decoded.sid; + assert_eq!(values, samples.values); + assert_eq!(sid, &samples.sid); + assert!(encoding.as_ptr_range().contains(&values.as_ptr().cast())); + assert!(encoding.as_ptr_range().contains(&sid.as_ptr().cast())); + } + assert_eq!(decoded.encode_vec(), bytes); +} + +#[test] +fn mutable_integer_lists_change_the_encoding_in_place() { + let mut expected = samples("ab"); + let bytes = expected.encode_vec(); + let mut buffer = copy_aligned(&bytes, 0); + let encoding = &mut buffer.0[..bytes.len()]; + let mut view = Samples::decode_bytes_ref(&mut *encoding).unwrap(); + view.values.reverse(); + view.sid[1] = 9; + + expected.values.reverse(); + expected.sid[1] = 9; + assert_eq!(Samples::decode_bytes(&*encoding).unwrap(), expected); +} + +#[test] +fn misaligned_integer_lists_are_only_copied() { + let samples = samples("ab"); + let bytes = samples.encode_vec(); + let buffer = copy_aligned(&bytes, 1); + let encoding = &buffer.0[1..][..bytes.len()]; + assert!(matches!( + Samples::decode_bytes_ref(encoding), + Err(ql_codec::Error::InvalidAlignment) + )); + assert_eq!(Samples::decode_bytes(encoding).unwrap(), samples); +} diff --git a/codec/src/aligned.rs b/codec/src/aligned.rs new file mode 100644 index 0000000..43d6637 --- /dev/null +++ b/codec/src/aligned.rs @@ -0,0 +1,178 @@ +use crate::{ + codec::len_prefix, varint, ByteSlice, Decode, Element, Encode, Error, Output, Reader, Ref, + Target, Writer, +}; + +#[cfg(target_endian = "big")] +compile_error!("Integer lists are lent as encoded, little-endian"); + +/// Zeros needed at `position` to align an integer of `size` bytes +#[inline] +fn padding(position: usize, size: usize) -> usize { + position.next_multiple_of(size) - position +} + +/// Encoded length of `len` integers of type `T`, without a length prefix +/// +/// `size_of::() - 1` zeros surround them, enough before to align them and the rest after, so +/// the length doesn't depend on where they're encoded. +#[inline] +fn encoded_len_ints(len: usize) -> usize { + size_of::() - 1 + len * size_of::() +} + +fn encode_ints(items: &[T], out: &mut Writer) { + let padding = padding(out.position(), size_of::()); + out.put_zeros(padding); + out.put_slice(as_bytes(items)); + out.put_zeros(size_of::() - 1 - padding); +} + +/// Takes the bytes of `len` integers of type `T`, skipping the zeros around them +fn take_ints(reader: &mut Reader, len: usize) -> Result { + let size = size_of::(); + let padding = padding(reader.position(), size); + reader.take_n(padding)?; + let bytes = reader.take_n(len.checked_mul(size).ok_or(Error::LengthOverflow)?)?; + reader.take_n(size - 1 - padding)?; + Ok(bytes) +} + +fn take_list(reader: &mut Reader) -> Result { + let len = reader.decode_varint::()?; + take_ints::(reader, len as usize) +} + +/// Copies integers out of `bytes`, which needn't be aligned +fn copy_ints(bytes: &[u8]) -> Vec { + const { assert!(T::ANY_BIT_PATTERN) }; + // safety: the flag guarantees that zero and the copied bytes are valid + let mut items = vec![unsafe { core::mem::zeroed() }; bytes.len() / size_of::()]; + as_bytes_mut(&mut items).copy_from_slice(bytes); + items +} + +/// See [`copy_ints`] +fn copy_int_array(bytes: &[u8]) -> [T; N] { + const { assert!(T::ANY_BIT_PATTERN) }; + // safety: the flag guarantees that zero and the copied bytes are valid + let mut items = [unsafe { core::mem::zeroed() }; N]; + as_bytes_mut(&mut items).copy_from_slice(bytes); + items +} + +/// Implements the codec traits for lists and arrays of a fixed-width integer +/// +/// Separate impls for each integer, because generic ones would overlap with those for an +/// [`Element`]. +macro_rules! impl_int { + ($($ty:ty),* $(,)?) => { + $( + impl Encode for [$ty] { + #[inline] + fn encoded_len(&self) -> usize { + varint::encoded_len(len_prefix(self.len())) + + encoded_len_ints::<$ty>(self.len()) + } + + fn encode(&self, out: &mut Writer) { + varint::encode(len_prefix(self.len()), out); + encode_ints(self, out); + } + } + + impl Encode for Vec<$ty> { + #[inline] + fn encoded_len(&self) -> usize { + self.as_slice().encoded_len() + } + + fn encode(&self, out: &mut Writer) { + self.as_slice().encode(out); + } + } + + impl Decode for Vec<$ty> { + type Ref = Ref; + + #[inline] + fn decode(reader: &mut Reader) -> Result { + Ok(copy_ints(&take_list::<$ty, B>(reader)?)) + } + + #[inline] + fn decode_ref(reader: &mut Reader) -> Result { + Ref::new(take_list::<$ty, B>(reader)?) + } + + #[inline] + fn from_ref(value: Self::Ref) -> Self { + value.to_vec() + } + + #[inline] + fn missing_ref() -> Option { + Ref::new(B::from_static(&[])?).ok() + } + } + + impl Element for Vec<$ty> {} + + impl Encode for [$ty; N] { + #[inline] + fn encoded_len(&self) -> usize { + encoded_len_ints::<$ty>(N) + } + + fn encode(&self, out: &mut Writer) { + encode_ints(self, out); + } + } + + impl Decode for [$ty; N] { + type Ref = Ref; + + #[inline] + fn decode(reader: &mut Reader) -> Result { + Ok(copy_int_array(&take_ints::<$ty, B>(reader, N)?)) + } + + #[inline] + fn decode_ref(reader: &mut Reader) -> Result { + Ref::new(take_ints::<$ty, B>(reader, N)?) + } + + #[inline] + fn from_ref(value: Self::Ref) -> Self { + *value + } + + #[inline] + fn missing() -> Option { + Some([0; N]) + } + + #[inline] + fn missing_ref() -> Option { + let zeros = as_bytes::<$ty>(&const { [0; N] }); + Ref::new(B::from_static(zeros)?).ok() + } + } + + impl Element for [$ty; N] {} + )* + }; +} + +impl_int!(u8, u16, u32, u64, i8, i16, i32, i64); + +fn as_bytes(items: &[T]) -> &[u8] { + // safety: Target has no padding and the slice covers exactly its initialized bytes + unsafe { core::slice::from_raw_parts(items.as_ptr().cast(), size_of_val(items)) } +} + +fn as_bytes_mut(items: &mut [T]) -> &mut [u8] { + const { assert!(T::ANY_BIT_PATTERN) }; + // safety: the flag permits arbitrary bytes and the slice is exclusively borrowed + unsafe { core::slice::from_raw_parts_mut(items.as_mut_ptr().cast(), size_of_val(items)) } +} diff --git a/codec/src/buf_view.rs b/codec/src/buf_view.rs index 5f112ec..c148c5c 100644 --- a/codec/src/buf_view.rs +++ b/codec/src/buf_view.rs @@ -41,6 +41,7 @@ impl BufView for [u8] { where Self: 'a; + #[inline] fn buf(&self) -> Self::Buf<'_> { self } @@ -63,6 +64,7 @@ impl BufView for Vec { where Self: 'a; + #[inline] fn buf(&self) -> Self::Buf<'_> { self.as_slice() } @@ -74,6 +76,7 @@ impl BufView for Box<[u8]> { where Self: 'a; + #[inline] fn buf(&self) -> Self::Buf<'_> { self.as_ref() } @@ -85,6 +88,7 @@ impl BufView for Bytes { where Self: 'a; + #[inline] fn buf(&self) -> Self::Buf<'_> { self.as_ref() } diff --git a/codec/src/codec.rs b/codec/src/codec.rs index 2a11608..9e7ad13 100644 --- a/codec/src/codec.rs +++ b/codec/src/codec.rs @@ -1,183 +1,485 @@ -use bytes::{Buf, BufMut}; +use std::{ + collections::BTreeMap, + num::NonZeroU8, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; -use crate::{varint, BufView, ByteSlice, Decode, Encode, Error, Reader}; +use bytes::{Buf, Bytes}; -impl Decode for [u8; N] { - fn decode(reader: &mut Reader) -> Result { - let bytes = reader.take_n(N)?; - let mut out = [0u8; N]; - out.copy_from_slice(&bytes); - Ok(out) - } -} +use crate::{ + varint, BufView, ByteSlice, Decode, DecodeValue, Element, Encode, Error, Output, Reader, Ref, + Writer, +}; -impl Encode for [u8; N] { - fn encoded_len(&self) -> usize { - N - } - - fn encode(&self, out: &mut W) { - out.put_slice(self); - } -} - -impl Decode for Box<[u8; N]> { - fn decode(reader: &mut Reader) -> Result { - let bytes = reader.take_n(N)?; - let mut out = Self::new_uninit(); - let src = bytes.as_ptr(); - let dst = out.as_mut_ptr().cast::(); - // SAFETY: `take_bytes(N)` guarantees the source has exactly `N` bytes. - unsafe { - std::ptr::copy_nonoverlapping(src, dst, N); - Ok(out.assume_init()) - } - } -} +macro_rules! impl_codec { + (fixed_integer: $($ty:ty),* $(,)?) => { + $( + impl Encode for $ty { + #[inline] + fn encoded_len(&self) -> usize { + size_of::() + } -impl Encode for Box<[u8; N]> { - fn encoded_len(&self) -> usize { - N - } + fn encode(&self, out: &mut Writer) { + out.put_slice(&self.to_le_bytes()); + } + } - fn encode(&self, out: &mut W) { - out.put_slice(self.as_ref()); - } -} + impl DecodeValue for $ty { + #[inline] + fn decode_value(reader: &mut Reader) -> Result { + let bytes = reader.take_n(size_of::())?; + Ok(Self::from_le_bytes((&*bytes).try_into().unwrap())) + } -macro_rules! impl_codec { - (byte_encode: $($ty:ty),* $(,)?) => { + #[inline] + fn missing() -> Option { + Some(0) + } + } + )* + }; + (byte_buffer: $($ty:ty => $from_ref:expr),* $(,)?) => { $( impl Encode for $ty { + #[inline] fn encoded_len(&self) -> usize { encoded_len_bytes(self) } - fn encode(&self, out: &mut W) { + + fn encode(&self, out: &mut Writer) { encode_bytes(self, out); } } - )* - }; - (owned_byte_decode: $($ty:ty),* $(,)?) => { - $( - impl Decode for $ty { - fn decode(reader: &mut Reader) -> Result { - Ok(<$ty>::from(&*reader.take_len_prefixed()?)) + + impl Decode for $ty { + type Ref = Ref; + + #[inline] + fn decode_ref(reader: &mut Reader) -> Result { + Ref::new(reader.take_len_prefixed()?) + } + + #[inline] + fn from_ref(value: Self::Ref) -> Self { + $from_ref(&*value) + } + + #[inline] + fn missing_ref() -> Option { + Ref::new(B::from_static(&[])?).ok() } } )* }; - (fixed_integer: $($ty:ty),* $(,)?) => { + (tuple: $(($($name:ident $index:tt),+)),* $(,)?) => { $( - impl Decode for $ty { - fn decode(reader: &mut Reader) -> Result { - Ok(Self::from_le_bytes(reader.decode()?)) + impl<$($name: Encode),+> Encode for ($($name,)+) { + fn encoded_len(&self) -> usize { + 0 $(+ self.$index.encoded_len())+ + } + + fn encode(&self, out: &mut Writer) { + $(self.$index.encode(out);)+ } } - impl Encode for $ty { - fn encoded_len(&self) -> usize { - size_of::() + + impl ),+> Decode for ($($name,)+) { + type Ref = ($($name::Ref,)+); + + fn decode(reader: &mut Reader) -> Result { + Ok(($($name::decode(reader)?,)+)) } - fn encode(&self, out: &mut W) { - out.put_slice(&self.to_le_bytes()); + + fn decode_ref(reader: &mut Reader) -> Result { + Ok(($($name::decode_ref(reader)?,)+)) + } + + fn from_ref(value: Self::Ref) -> Self { + ($($name::from_ref(value.$index),)+) + } + + fn missing() -> Option { + Some(($($name::missing()?,)+)) + } + + fn missing_ref() -> Option { + Some(($($name::missing_ref()?,)+)) } } + + impl<$($name),+> Element for ($($name,)+) {} )* }; + (element: $($ty:ty),* $(,)?) => { + $(impl Element for $ty {})* + }; } -impl_codec!(byte_encode: [u8], Vec, Box<[u8]>, bytes::Bytes); +impl_codec!(fixed_integer: u8, u16, u32, u64, i8, i16, i32, i64); +impl_codec!( + byte_buffer: + Box<[u8]> => Box::from, + Bytes => Bytes::copy_from_slice, +); +impl_codec!(tuple: (T0 0, T1 1), (T0 0, T1 1, T2 2)); +// Not the fixed-width integers, so their lists can borrow as `Ref`. +impl_codec!( + element: usize, bool, NonZeroU8, (), Box<[u8]>, Bytes, String, Duration, SystemTime, +); + +impl Encode for Box<[u8; N]> { + fn encoded_len(&self) -> usize { + N + } + + fn encode(&self, out: &mut Writer) { + (**self).encode(out); + } +} + +impl Decode for Box<[u8; N]> { + type Ref = Ref; + + fn decode_ref(reader: &mut Reader) -> Result { + <[u8; N]>::decode_ref(reader) + } + + fn from_ref(value: Self::Ref) -> Self { + // copy through a slice to avoid putting a large array on the stack + Box::<[u8]>::from(value.as_slice()) + .try_into() + .expect("slice of an N byte array") + } + + fn missing() -> Option { + Some( + Box::<[u8]>::from(&const { [0; N] }[..]) + .try_into() + .expect("slice of an N byte array"), + ) + } -impl<'a> Decode<&'a [u8]> for &'a [u8] { - fn decode(reader: &mut Reader<&'a [u8]>) -> Result { - reader.take_len_prefixed() + fn missing_ref() -> Option { + <[u8; N]>::missing_ref() } } -impl<'a> Decode<&'a mut [u8]> for &'a mut [u8] { - fn decode(reader: &mut Reader<&'a mut [u8]>) -> Result { - reader.take_len_prefixed() +impl Element for Box<[u8; N]> {} + +impl Encode for [T; N] { + fn encoded_len(&self) -> usize { + self.iter().map(Encode::encoded_len).sum() + } + + fn encode(&self, out: &mut Writer) { + for item in self { + item.encode(out); + } } } -impl Decode for bytes::Bytes { - fn decode(reader: &mut Reader) -> Result { - reader.take_len_prefixed() +impl + Element, const N: usize> Decode for [T; N] { + type Ref = [T::Ref; N]; + + fn decode(reader: &mut Reader) -> Result { + decode_array(|| T::decode(reader)) + } + + fn decode_ref(reader: &mut Reader) -> Result { + decode_array(|| T::decode_ref(reader)) + } + + fn from_ref(value: Self::Ref) -> Self { + value.map(T::from_ref) + } + + fn missing() -> Option { + decode_array(|| T::missing().ok_or(())).ok() + } + + fn missing_ref() -> Option { + decode_array(|| T::missing_ref().ok_or(())).ok() } } -impl_codec!(owned_byte_decode: Vec, Box<[u8]>); +impl Element for [T; N] {} -impl Encode for str { +fn decode_array( + mut decode_item: impl FnMut() -> Result, +) -> Result<[T; N], E> { + let mut error = None; + let items: [Option; N] = std::array::from_fn(|_| { + if error.is_some() { + return None; + } + match decode_item() { + Ok(item) => Some(item), + Err(err) => { + error = Some(err); + None + } + } + }); + match error { + Some(error) => Err(error), + None => Ok(items.map(|item| item.expect("decoded every array item"))), + } +} + +impl Encode for [T] { fn encoded_len(&self) -> usize { - self.as_bytes().encoded_len() + varint::encoded_len(len_prefix(self.len())) + + self.iter().map(Encode::encoded_len).sum::() } - fn encode(&self, out: &mut W) { - self.as_bytes().encode(out); + fn encode(&self, out: &mut Writer) { + varint::encode(len_prefix(self.len()), out); + for item in self { + item.encode(out); + } + } +} + +impl Encode for Vec { + fn encoded_len(&self) -> usize { + self.as_slice().encoded_len() + } + + fn encode(&self, out: &mut Writer) { + self.as_slice().encode(out); + } +} + +impl + Element> Decode for Vec { + type Ref = Vec; + + fn decode(reader: &mut Reader) -> Result { + decode_list(reader, T::decode) + } + + fn decode_ref(reader: &mut Reader) -> Result { + decode_list(reader, T::decode_ref) + } + + fn from_ref(value: Self::Ref) -> Self { + value.into_iter().map(T::from_ref).collect() + } + + fn missing_ref() -> Option { + Some(Vec::new()) + } +} + +impl Element for Vec {} + +fn decode_list( + reader: &mut Reader, + decode_item: impl Fn(&mut Reader) -> Result, +) -> Result, Error> { + let len = reader.decode_varint::()? as usize; + // At most as much memory as the input, or a corrupt length could reserve gigabytes. + let mut items = Vec::with_capacity(len.min(reader.remaining_len() / size_of::().max(1))); + for _ in 0..len { + items.push(decode_item(reader)?); + } + Ok(items) +} + +/// Encodes like a list of key and value pairs +impl Encode for BTreeMap { + fn encoded_len(&self) -> usize { + varint::encoded_len(len_prefix(self.len())) + + self + .iter() + .map(|(key, value)| key.encoded_len() + value.encoded_len()) + .sum::() + } + + fn encode(&self, out: &mut Writer) { + varint::encode(len_prefix(self.len()), out); + for (key, value) in self { + key.encode(out); + value.encode(out); + } + } +} + +/// a key that occurs twice keeps its last value +impl Decode for BTreeMap +where + B: ByteSlice, + K: Decode + Ord, + V: Decode, + K::Ref: Ord, +{ + type Ref = BTreeMap; + + fn decode(reader: &mut Reader) -> Result { + decode_map(reader, K::decode, V::decode) + } + + fn decode_ref(reader: &mut Reader) -> Result { + decode_map(reader, K::decode_ref, V::decode_ref) + } + + fn from_ref(value: Self::Ref) -> Self { + value + .into_iter() + .map(|(key, value)| (K::from_ref(key), V::from_ref(value))) + .collect() + } + + fn missing_ref() -> Option { + Some(BTreeMap::new()) + } +} + +impl Element for BTreeMap {} + +fn decode_map( + reader: &mut Reader, + decode_key: impl Fn(&mut Reader) -> Result, + decode_value: impl Fn(&mut Reader) -> Result, +) -> Result, Error> { + let len = reader.decode_varint::()? as usize; + let mut entries = BTreeMap::new(); + for _ in 0..len { + entries.insert(decode_key(reader)?, decode_value(reader)?); } + Ok(entries) } -impl<'a> Decode<&'a [u8]> for &'a str { - fn decode(reader: &mut Reader<&'a [u8]>) -> Result { - std::str::from_utf8(reader.take_len_prefixed()?).map_err(|_| Error::InvalidUtf8) +impl Encode for str { + #[inline] + fn encoded_len(&self) -> usize { + self.as_bytes().encoded_len() + } + + fn encode(&self, out: &mut Writer) { + self.as_bytes().encode(out); } } impl Encode for String { + #[inline] fn encoded_len(&self) -> usize { self.as_str().encoded_len() } - fn encode(&self, out: &mut W) { + fn encode(&self, out: &mut Writer) { self.as_str().encode(out); } } impl Decode for String { - fn decode(reader: &mut Reader) -> Result { - let bytes = reader.take_len_prefixed()?; - std::str::from_utf8(&bytes) - .map(str::to_owned) - .map_err(|_| Error::InvalidUtf8) + type Ref = Ref; + + #[inline] + fn decode_ref(reader: &mut Reader) -> Result { + Ref::new(reader.take_len_prefixed()?) + } + + #[inline] + fn from_ref(value: Self::Ref) -> Self { + value.to_string() + } + + #[inline] + fn missing_ref() -> Option { + Ref::new(B::from_static(&[])?).ok() } } -impl Decode for u8 { - fn decode(reader: &mut Reader) -> Result { - reader.take_u8() +impl Encode for NonZeroU8 { + #[inline] + fn encoded_len(&self) -> usize { + size_of::() + } + + fn encode(&self, out: &mut Writer) { + out.put_u8(self.get()); } } -impl Encode for u8 { +impl DecodeValue for NonZeroU8 { + #[inline] + fn decode_value(reader: &mut Reader) -> Result { + Self::new(reader.take_u8()?).ok_or(Error::InvalidData) + } +} + +/// `usize` is a varint, because its width differs between targets +/// +/// A target too narrow for the value fails to decode it with [`Error::InvalidRange`]. +impl Encode for usize { + #[inline] fn encoded_len(&self) -> usize { - size_of::() + varint::encoded_len(*self as u64) } - fn encode(&self, out: &mut W) { - out.put_u8(*self); + fn encode(&self, out: &mut Writer) { + varint::encode(*self as u64, out); } } -impl_codec!(fixed_integer: u16, u32, u64); +impl DecodeValue for usize { + #[inline] + fn decode_value(reader: &mut Reader) -> Result { + Self::try_from(reader.decode_varint::()?).map_err(|_| Error::InvalidRange) + } -impl Decode for bool { - fn decode(reader: &mut Reader) -> Result { - match reader.decode::()? { + #[inline] + fn missing() -> Option { + Some(0) + } +} + +impl Encode for bool { + #[inline] + fn encoded_len(&self) -> usize { + size_of::() + } + + fn encode(&self, out: &mut Writer) { + out.put_u8(u8::from(*self)); + } +} + +impl DecodeValue for bool { + #[inline] + fn decode_value(reader: &mut Reader) -> Result { + match reader.take_u8()? { 0 => Ok(false), 1 => Ok(true), _ => Err(Error::InvalidDiscriminant), } } + + #[inline] + fn missing() -> Option { + Some(false) + } } -impl Encode for bool { +impl Encode for () { + #[inline] fn encoded_len(&self) -> usize { - size_of::() + 0 } - fn encode(&self, out: &mut W) { - out.put_u8(u8::from(*self)); + fn encode(&self, _out: &mut Writer) {} +} + +impl DecodeValue for () { + #[inline] + fn decode_value(_reader: &mut Reader) -> Result { + Ok(()) + } + + #[inline] + fn missing() -> Option { + Some(()) } } @@ -186,7 +488,7 @@ impl Encode for Option { 1 + self.as_ref().map_or(0, Encode::encoded_len) } - fn encode(&self, out: &mut W) { + fn encode(&self, out: &mut Writer) { match self { None => out.put_u8(0), Some(inner) => { @@ -198,43 +500,172 @@ impl Encode for Option { } impl> Decode for Option { + type Ref = Option; + fn decode(reader: &mut Reader) -> Result { - match reader.decode::()? { - 0 => Ok(None), - 1 => Ok(Some(reader.decode::()?)), - _ => Err(Error::InvalidDiscriminant), + decode_option(reader, T::decode) + } + + fn decode_ref(reader: &mut Reader) -> Result { + decode_option(reader, T::decode_ref) + } + + fn from_ref(value: Self::Ref) -> Self { + value.map(T::from_ref) + } + + fn missing_ref() -> Option { + Some(None) + } +} + +impl Element for Option {} + +fn decode_option( + reader: &mut Reader, + decode_value: impl Fn(&mut Reader) -> Result, +) -> Result, Error> { + match reader.take_u8()? { + 0 => Ok(None), + 1 => Ok(Some(decode_value(reader)?)), + _ => Err(Error::InvalidDiscriminant), + } +} + +impl Encode for Result { + fn encoded_len(&self) -> usize { + 1 + match self { + Ok(value) => value.encoded_len(), + Err(error) => error.encoded_len(), + } + } + + fn encode(&self, out: &mut Writer) { + match self { + Ok(value) => { + out.put_u8(0); + value.encode(out); + } + Err(error) => { + out.put_u8(1); + error.encode(out); + } + } + } +} + +impl, E: Decode> Decode for Result { + type Ref = Result; + + fn decode(reader: &mut Reader) -> Result { + decode_result(reader, T::decode, E::decode) + } + + fn decode_ref(reader: &mut Reader) -> Result { + decode_result(reader, T::decode_ref, E::decode_ref) + } + + fn from_ref(value: Self::Ref) -> Self { + value.map(T::from_ref).map_err(E::from_ref) + } +} + +impl Element for Result {} + +fn decode_result( + reader: &mut Reader, + decode_value: impl Fn(&mut Reader) -> Result, + decode_error: impl Fn(&mut Reader) -> Result, +) -> Result, Error> { + match reader.take_u8()? { + 0 => Ok(Ok(decode_value(reader)?)), + 1 => Ok(Err(decode_error(reader)?)), + _ => Err(Error::InvalidDiscriminant), + } +} + +impl Encode for Duration { + #[inline] + fn encoded_len(&self) -> usize { + size_of::() + size_of::() + } + + fn encode(&self, out: &mut Writer) { + self.as_secs().encode(out); + self.subsec_nanos().encode(out); + } +} + +impl DecodeValue for Duration { + #[inline] + fn decode_value(reader: &mut Reader) -> Result { + let secs: u64 = reader.decode()?; + let nanos: u32 = reader.decode()?; + if nanos >= 1_000_000_000 { + return Err(Error::InvalidRange); } + Ok(Self::new(secs, nanos)) + } + + #[inline] + fn missing() -> Option { + Some(Self::ZERO) + } +} + +/// Encodes as the [`Duration`] since the Unix epoch, which earlier times clamp to. +impl Encode for SystemTime { + #[inline] + fn encoded_len(&self) -> usize { + Duration::ZERO.encoded_len() + } + + fn encode(&self, out: &mut Writer) { + self.duration_since(UNIX_EPOCH) + .unwrap_or_default() + .encode(out); + } +} + +impl DecodeValue for SystemTime { + #[inline] + fn decode_value(reader: &mut Reader) -> Result { + UNIX_EPOCH + .checked_add(reader.decode()?) + .ok_or(Error::InvalidRange) } } /// Length prefixes are `u32`, so a 64-bit host cannot emit one a 32-bit peer would reject: a /// `usize` varint is capped at ten bytes on one target and five on the other. #[track_caller] -fn len_prefix(len: usize) -> u32 { - u32::try_from(len).expect("byte field longer than u32::MAX") +#[inline] +pub(crate) fn len_prefix(len: usize) -> u32 { + u32::try_from(len).expect("length prefix above u32::MAX") } -pub fn encoded_len_bytes(bytes: &B) -> usize { - let len = bytes.buf().remaining(); +/// Encoded length of `len` bytes behind their length prefix +#[inline] +pub fn encoded_len_prefixed(len: usize) -> usize { varint::encoded_len(len_prefix(len)) + len } -pub fn encode_bytes(bytes: &B, out: &mut W) -where - B: BufView + ?Sized, - W: BufMut + ?Sized, -{ - varint::encode(len_prefix(bytes.buf().remaining()), out); +/// Writes the length prefix for `len` bytes +pub fn encode_len_prefix(len: usize, out: &mut Writer) { + varint::encode(len_prefix(len), out); +} + +pub fn encoded_len_bytes(bytes: &B) -> usize { + encoded_len_prefixed(bytes.buf().remaining()) +} + +pub fn encode_bytes(bytes: &B, out: &mut Writer) { + encode_len_prefix(bytes.buf().remaining(), out); encode_bytes_raw(bytes, out); } /// Writes the bytes without a length prefix, for a field that runs to the end of the message. -pub fn encode_bytes_raw(bytes: &B, out: &mut W) -where - B: BufView + ?Sized, - W: BufMut + ?Sized, -{ - // `BufMut::put` needs a sized writer, so feed it a chunk at a time. +pub fn encode_bytes_raw(bytes: &B, out: &mut Writer) { let mut bytes = bytes.buf(); while bytes.has_remaining() { let chunk = bytes.chunk(); @@ -242,3 +673,32 @@ where bytes.advance(chunk.len()); } } + +/// Decodes the next field of a struct that isn't frozen, or its missing value once the encoding +/// ran out +/// +/// A field cut short is still an error. +pub fn decode_field(fields: &mut Reader) -> Result +where + T: Decode, + B: ByteSlice, +{ + if fields.is_empty() { + T::missing().ok_or(Error::UnexpectedEof) + } else { + T::decode(fields) + } +} + +/// Like [`decode_field`], but decodes the borrowed form +pub fn decode_field_ref(fields: &mut Reader) -> Result +where + T: Decode, + B: ByteSlice, +{ + if fields.is_empty() { + T::missing_ref().ok_or(Error::UnexpectedEof) + } else { + T::decode_ref(fields) + } +} diff --git a/codec/src/error.rs b/codec/src/error.rs index 861c3c0..99b408c 100644 --- a/codec/src/error.rs +++ b/codec/src/error.rs @@ -2,6 +2,7 @@ use core::fmt; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Error { + InvalidAlignment, InvalidData, UnexpectedEof, InvalidDiscriminant, @@ -14,6 +15,7 @@ pub enum Error { impl fmt::Display for Error { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let message = match self { + Self::InvalidAlignment => "invalid alignment", Self::InvalidData => "invalid data", Self::UnexpectedEof => "unexpected end of input", Self::InvalidDiscriminant => "invalid discriminant", diff --git a/codec/src/lib.rs b/codec/src/lib.rs index 662c144..2292e2b 100644 --- a/codec/src/lib.rs +++ b/codec/src/lib.rs @@ -1,63 +1,224 @@ //! binary codec primitives shared by QuantumLink crates //! //! [`Encode`] and [`Decode`] define the wire encoding for protocol types -//! [`Reader`] and [`ByteSlice`] support decoding from borrowed or owned bytes +//! [`Decode::decode_bytes`] reads a value from the start of a byte container, and [`Reader`] +//! several in a row. Their `_ref` forms borrow from the encoding instead. +//! +//! `#[derive(Codec)]` (feature `derive`) implements both for structs and enums, generic ones +//! included. By default their encoding is backward and forward compatible: a struct, and the +//! fields of an enum variant, are preceded by their length. A decoder fills fields that a shorter +//! encoding lacks with [`Decode::missing`], and skips the ones it doesn't know. +//! `#[codec(frozen)]` drops the lengths for types that won't change. Enumerators are varints +//! either way, so only those below 128 take a single byte. Unless frozen, an enum decodes +//! enumerators it doesn't know to its `#[codec(unknown)]` variant, or fails without one. The +//! derive generates `Ref` as the borrowed form, which encodes the same as the owned one, +//! and `#[codec(derive(...))]` derives traits on it for non-generic types. +//! `#[codec(discriminants = Name)]` on an enum also generates `Name`, its variants without their +//! fields and numbered by enumerator, and a `discriminant()` method returning it. +//! +//! Borrowed forms decode from `&[u8]`, `&mut [u8]`, `Bytes` or `BytesMut`, without copying. +//! Strings, byte buffers, and vectors and arrays of fixed-width integers are lent as [`Ref`]s, +//! which dereference to the value, mutably to change the encoding in place if the container is +//! mutable. Other vectors and arrays decode their elements into `Vec` and `[T::Ref; N]`. +//! +//! The integers are laid out like in memory: little-endian, starting a multiple of +//! `size_of::()` from the start of the encoding, with `size_of::() - 1` zeros around them. +//! [`Writer`] and [`Reader`] count from where they start, so an encoding must be decoded from its +//! start. To borrow the integers, it must also start aligned for `T` in memory, or borrowed +//! decoding fails with [`Error::InvalidAlignment`]. The owned form decodes either way. +//! +//! ``` +//! # #[cfg(feature = "derive")] +//! # mod example { +//! use ql_codec::{Codec, Decode, Encode}; +//! +//! #[derive(Codec)] +//! #[codec(frozen)] +//! struct Message { +//! name: String, +//! bytes: Vec, +//! } +//! +//! # fn example() -> Result<(), ql_codec::Error> { +//! let message = Message { name: "hello".into(), bytes: vec![1, 2] }; +//! let mut bytes = message.encode_vec(); +//! let shared: MessageRef<&[u8]> = Message::decode_bytes_ref(bytes.as_slice())?; +//! assert_eq!(shared.name, "hello"); +//! +//! let mut mutable: MessageRef<&mut [u8]> = Message::decode_bytes_ref(bytes.as_mut_slice())?; +//! mutable.name.make_ascii_uppercase(); +//! mutable.bytes[0] = 3; +//! # Ok(()) +//! # } +//! # } +//! ``` +//! +//! Decoding recurses once per level of nesting, so a type decoded from untrusted input must not +//! be recursive, or a deep enough encoding overflows the stack. +mod aligned; mod buf_view; mod codec; mod error; -mod macros; mod reader; +mod reference; mod slice; pub mod varint; +mod writer; pub use buf_view::BufView; -pub use codec::{encode_bytes, encode_bytes_raw, encoded_len_bytes}; +pub use codec::{ + decode_field, decode_field_ref, encode_bytes, encode_bytes_raw, encode_len_prefix, + encoded_len_bytes, encoded_len_prefixed, +}; pub use error::Error; +#[cfg(feature = "derive")] +pub use ql_codec_macro::Codec; pub use reader::Reader; +pub use reference::{Ref, Target}; pub use slice::ByteSlice; pub use varint::Varint; +pub use writer::{Output, Writer}; pub trait Encode { fn encoded_len(&self) -> usize; - fn encode(&self, out: &mut W); + fn encode(&self, out: &mut Writer); fn encode_vec(&self) -> Vec { let len = self.encoded_len(); - let mut out = Vec::with_capacity(len); + let mut out = Writer::new(Vec::with_capacity(len)); self.encode(&mut out); - assert_eq!(out.len(), len); - out + assert_eq!(out.position(), len); + out.into_inner() } } +/// owned and borrowed decoding from an input container +/// +/// implement once over `B` to support shared and mutable input. pub trait Decode: Sized { - fn decode(reader: &mut Reader) -> Result; + type Ref; + + /// decodes the owned form, through the borrowed form by default + /// + /// a type holding other values should decode them directly, or decoding allocates twice and + /// misaligned integer lists fail to borrow. + fn decode(reader: &mut Reader) -> Result { + Self::decode_ref(reader).map(Self::from_ref) + } + + fn decode_ref(reader: &mut Reader) -> Result; + fn from_ref(value: Self::Ref) -> Self; + /// decodes from the start of `bytes`, ignoring whatever follows fn decode_bytes(bytes: B) -> Result { - let mut reader = Reader::new(bytes); - Self::decode(&mut reader) + Self::decode(&mut Reader::new(bytes)) + } + + /// decodes the borrowed form, preserving the input container's mutability + fn decode_bytes_ref(bytes: B) -> Result { + Self::decode_ref(&mut Reader::new(bytes)) + } + + /// the value of a field omitted by an older encoder + fn missing() -> Option { + Self::missing_ref().map(Self::from_ref) + } + + /// borrowed default for a field omitted by an older encoder + /// + /// a mutable container can't lend a non-empty array, so an array appended to a struct should + /// be an `Option`. + fn missing_ref() -> Option { + None } } +/// implements owned and borrowed decoding for values that borrow nothing +pub trait DecodeValue: Sized { + fn decode_value(reader: &mut Reader) -> Result; + + /// see [`Decode::missing`] + fn missing() -> Option { + None + } +} + +impl Decode for T { + type Ref = Self; + + fn decode_ref(reader: &mut Reader) -> Result { + ::decode_value(reader) + } + + fn from_ref(value: Self) -> Self { + value + } + + fn missing_ref() -> Option { + ::missing() + } +} + +/// Marks the types whose lists and arrays decode item by item +/// +/// That's every encodable type except the fixed-width integers, whose lists and arrays are lent +/// as a whole, as `Ref`. The two kinds need separate `Vec` and `[T; N]` impls, and +/// Rust only keeps those apart if the item-by-item ones require a marker the integers lack. +/// +/// The derive and `varint_wrapper!` implement it. A hand-written codec impl needs an +/// `impl Element for T {}`, or `Vec` and `[T; N]` fail with an unsatisfied `T: Element`. +pub trait Element {} + +impl Encode for &T { + fn encoded_len(&self) -> usize { + (**self).encoded_len() + } + + fn encode(&self, out: &mut Writer) { + (**self).encode(out); + } +} + +impl Element for &T {} + +impl Encode for &mut T { + fn encoded_len(&self) -> usize { + (**self).encoded_len() + } + + fn encode(&self, out: &mut Writer) { + (**self).encode(out); + } +} + +impl Element for &mut T {} + +/// Implements the codec traits for a newtype over an integer, carried as a varint #[macro_export] macro_rules! varint_wrapper { ($name:ty, $inner:ty) => { - impl ql_codec::Encode for $name { + impl $crate::Encode for $name { + #[inline] fn encoded_len(&self) -> usize { - ql_codec::varint::encoded_len::<$inner>(self.0) + $crate::varint::encoded_len::<$inner>(self.0) } - fn encode(&self, out: &mut W) { - ql_codec::varint::encode::<$inner, W>(self.0, out); + fn encode(&self, out: &mut $crate::Writer) { + $crate::varint::encode::<$inner>(self.0, out); } } - impl ql_codec::Decode for $name { - fn decode(reader: &mut ql_codec::Reader) -> Result { + impl $crate::DecodeValue for $name { + #[inline] + fn decode_value( + reader: &mut $crate::Reader, + ) -> Result { Ok(Self(reader.decode_varint::<$inner>()?)) } } + + impl $crate::Element for $name {} }; } diff --git a/codec/src/macros.rs b/codec/src/macros.rs deleted file mode 100644 index 50cee54..0000000 --- a/codec/src/macros.rs +++ /dev/null @@ -1,215 +0,0 @@ -/// generates `Encode` and `Decode` for newtypes, structs, and enums. -/// -/// newtypes encode as their wrapped value, structs in field order, and enums as a `u8` discriminant -/// followed by an optional payload; the lone generic parameter is treated as the reader's byte -/// container. -/// -/// `enum Frame as FrameKind` also generates `FrameKind` and `Frame::kind()`: -/// ``` -/// use ql_codec::Encode; -/// -/// ql_codec::codec! { -/// #[derive(Debug, PartialEq)] -/// pub enum Frame as FrameKind { -/// Ping = 1, -/// Close(u16) = 2, -/// } -/// } -/// -/// let frame = Frame::Close(7); -/// assert_eq!(frame.kind(), FrameKind::Close); -/// assert_eq!(frame.encode_vec(), [2, 7, 0]); -/// ``` -#[macro_export] -macro_rules! codec { - // newtype - ( - $(#[$meta:meta])* - $vis:vis struct $name:ident($field_vis:vis $inner:ty); - ) => { - $(#[$meta])* - $vis struct $name($field_vis $inner); - - impl $crate::Encode for $name { - fn encoded_len(&self) -> usize { - $crate::Encode::encoded_len(&self.0) - } - - fn encode(&self, out: &mut W) { - $crate::Encode::encode(&self.0, out); - } - } - - impl $crate::Decode for $name { - fn decode(reader: &mut $crate::Reader) -> Result { - Ok(Self(reader.decode()?)) - } - } - }; - - // struct with an optional byte container - ( - $(#[$meta:meta])* - $vis:vis struct $name:ident $(<$bytes:ident>)? { - $($(#[$field_meta:meta])* $field_vis:vis $field:ident: $ty:ty),* $(,)? - } - ) => { - $(#[$meta])* - $vis struct $name$(<$bytes>)? { - $($(#[$field_meta])* $field_vis $field: $ty,)* - } - - impl$(<$bytes: $crate::BufView>)? $crate::Encode for $name$(<$bytes>)? { - fn encoded_len(&self) -> usize { - $($crate::Encode::encoded_len(&self.$field) +)* 0 - } - - fn encode(&self, out: &mut W) { - $($crate::Encode::encode(&self.$field, out);)* - } - } - - $crate::codec!(@struct_decode $name$(<$bytes>)?, $($field),*); - }; - - // payload enum with a separate discriminant enum - ( - $(#[$meta:meta])* - $vis:vis enum $name:ident $(<$bytes:ident>)? as $kind:ident { - $($(#[$variant_meta:meta])* $variant:ident $(($payload:ty))? = $value:literal),* $(,)? - } - ) => { - $crate::codec! { - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - $vis enum $kind { - $($variant = $value,)* - } - } - - $(#[$meta])* - $vis enum $name$(<$bytes>)? { - $($(#[$variant_meta])* $variant $(($payload))?,)* - } - - impl$(<$bytes>)? $name$(<$bytes>)? { - $vis fn kind(&self) -> $kind { - match self { - $(Self::$variant { .. } => $kind::$variant,)* - } - } - } - - impl$(<$bytes: $crate::BufView>)? $crate::Encode for $name$(<$bytes>)? { - #[allow(unreachable_patterns)] - fn encoded_len(&self) -> usize { - $crate::Encode::encoded_len(&self.kind()) - + match self { - $($(Self::$variant(payload) => - <$payload as $crate::Encode>::encoded_len(payload),)?)* - _ => 0, - } - } - - #[allow(unreachable_patterns)] - fn encode(&self, out: &mut W) { - $crate::Encode::encode(&self.kind(), out); - match self { - $($(Self::$variant(payload) => - <$payload as $crate::Encode>::encode(payload, out),)?)* - _ => {} - } - } - } - - $crate::codec! { - @enum_decode $name$(<$bytes>)?, $kind, $($variant $(($payload))? = $value),* - } - }; - - // u8 discriminant enum - ( - $(#[$meta:meta])* - $vis:vis enum $name:ident { - $($(#[$variant_meta:meta])* $variant:ident = $value:literal),* $(,)? - } - ) => { - $(#[$meta])* - #[repr(u8)] - $vis enum $name { - $($(#[$variant_meta])* $variant = $value,)* - } - - impl TryFrom for $name { - type Error = $crate::Error; - - fn try_from(value: u8) -> Result { - match value { - $($value => Ok(Self::$variant),)* - _ => Err($crate::Error::InvalidDiscriminant), - } - } - } - - impl $crate::Encode for $name { - fn encoded_len(&self) -> usize { - size_of::() - } - - fn encode(&self, out: &mut W) { - ::bytes::BufMut::put_u8(out, *self as u8); - } - } - - impl $crate::Decode for $name { - fn decode(reader: &mut $crate::Reader) -> Result { - reader.decode::()?.try_into() - } - } - }; - - // struct decoding with its byte container - (@struct_decode $name:ident<$bytes:ident>, $($field:ident),* $(,)?) => { - impl<$bytes: $crate::ByteSlice> $crate::Decode<$bytes> for $name<$bytes> { - fn decode(reader: &mut $crate::Reader<$bytes>) -> Result { - Ok(Self { $($field: reader.decode()?,)* }) - } - } - }; - - // struct decoding with a fresh byte container - (@struct_decode $name:ident, $($field:ident),* $(,)?) => { - impl $crate::Decode for $name { - fn decode(reader: &mut $crate::Reader) -> Result { - Ok(Self { $($field: reader.decode()?,)* }) - } - } - }; - - // payload enum decoding with its byte container - ( - @enum_decode $name:ident<$bytes:ident>, $kind:ident, - $($variant:ident $(($payload:ty))? = $value:literal),* $(,)? - ) => { - impl<$bytes: $crate::ByteSlice> $crate::Decode<$bytes> for $name<$bytes> { - fn decode(reader: &mut $crate::Reader<$bytes>) -> Result { - Ok(match reader.decode::<$kind>()? { - $($kind::$variant => Self::$variant $((reader.decode::<$payload>()?))?,)* - }) - } - } - }; - - // payload enum decoding with a fresh byte container - ( - @enum_decode $name:ident, $kind:ident, - $($variant:ident $(($payload:ty))? = $value:literal),* $(,)? - ) => { - impl $crate::Decode for $name { - fn decode(reader: &mut $crate::Reader) -> Result { - Ok(match reader.decode::<$kind>()? { - $($kind::$variant => Self::$variant $((reader.decode::<$payload>()?))?,)* - }) - } - } - }; -} diff --git a/codec/src/reader.rs b/codec/src/reader.rs index 35ed703..ed2269d 100644 --- a/codec/src/reader.rs +++ b/codec/src/reader.rs @@ -3,12 +3,19 @@ use crate::{varint, ByteSlice, Decode, Error}; #[derive(Clone)] pub struct Reader { remaining: B, + /// Position of the end of `remaining` + end: usize, } impl Reader { + /// Reads an encoding from its start #[inline] pub fn new(bytes: B) -> Self { - Self { remaining: bytes } + let end = bytes.len(); + Self { + remaining: bytes, + end, + } } #[inline] @@ -21,6 +28,13 @@ impl Reader { self.remaining.len() } + /// Bytes read since the encoding started + #[inline] + pub fn position(&self) -> usize { + self.end - self.remaining.len() + } + + #[inline] pub fn take_n(&mut self, len: usize) -> Result { if len > self.remaining.len() { return Err(Error::UnexpectedEof); @@ -33,23 +47,38 @@ impl Reader { self.remaining.take_u8().ok_or(Error::UnexpectedEof) } + #[inline] pub fn take_all(&mut self) -> B { self.remaining.split_off_front(self.remaining.len()) } + #[inline] pub fn take_len_prefixed(&mut self) -> Result { let len = self.decode_varint::()?; self.take_n(len as usize) } + /// Like [`Reader::take_len_prefixed`], but reads the bytes as part of this encoding, so + /// integer lists in them stay aligned #[inline] - pub fn decode(&mut self) -> Result - where - T: Decode, - { + pub fn take_len_prefixed_reader(&mut self) -> Result { + let remaining = self.take_len_prefixed()?; + Ok(Self { + remaining, + end: self.position(), + }) + } + + #[inline] + pub fn decode>(&mut self) -> Result { T::decode(self) } + #[inline] + pub fn decode_ref>(&mut self) -> Result { + T::decode_ref(self) + } + #[inline] pub fn decode_varint(&mut self) -> Result where diff --git a/codec/src/reference.rs b/codec/src/reference.rs new file mode 100644 index 0000000..7d4dae6 --- /dev/null +++ b/codec/src/reference.rs @@ -0,0 +1,347 @@ +use core::{ + borrow::{Borrow, BorrowMut}, + fmt, + hash::{Hash, Hasher}, + marker::PhantomData, + ops::{Deref, DerefMut}, +}; + +use crate::{BufView, ByteSlice, Element, Encode, Error, Output, Writer}; + +/// a validated typed reference with the mutability of its byte container +/// +/// supports strings and other layouts implementing Target. mutable access preserves +/// the length and validity of the encoding. +/// +/// shared storage cannot be mutated through the typed reference +/// +/// ```compile_fail +/// use ql_codec::Ref; +/// +/// let mut text = Ref::<_, str>::new(&b"hello"[..]).unwrap(); +/// text.make_ascii_uppercase(); +/// ``` +pub struct Ref { + bytes: B, + target: PhantomData, +} + +impl Ref { + pub fn new(bytes: B) -> Result { + T::validate(&bytes)?; + Ok(Self { + bytes, + target: PhantomData, + }) + } +} + +impl Ref { + pub fn into_bytes(self) -> B { + self.bytes + } +} + +impl Deref for Ref { + type Target = T; + + fn deref(&self) -> &T { + // safety: construction validated these stable bytes and typed mutations preserve validity + unsafe { T::from_bytes(&self.bytes) } + } +} + +impl DerefMut for Ref { + fn deref_mut(&mut self) -> &mut T { + // safety: mutable dereferencing provides the same validated bytes exclusively + unsafe { T::from_bytes_mut(&mut self.bytes) } + } +} + +impl Clone for Ref { + fn clone(&self) -> Self { + Self::new(self.bytes.clone()).expect("cloned bytes retain a valid layout") + } +} + +impl Copy for Ref {} + +impl AsRef for Ref { + fn as_ref(&self) -> &T { + self + } +} + +impl AsRef<[T]> for Ref +where + [T; N]: Target, +{ + fn as_ref(&self) -> &[T] { + self.as_slice() + } +} + +impl AsMut for Ref { + fn as_mut(&mut self) -> &mut T { + self + } +} + +impl AsMut<[T]> for Ref +where + [T; N]: Target, +{ + fn as_mut(&mut self) -> &mut [T] { + self.as_mut_slice() + } +} + +impl Borrow for Ref { + fn borrow(&self) -> &T { + self + } +} + +impl BorrowMut for Ref { + fn borrow_mut(&mut self) -> &mut T { + self + } +} + +impl Encode for Ref { + fn encoded_len(&self) -> usize { + (**self).encoded_len() + } + + fn encode(&self, out: &mut Writer) { + (**self).encode(out); + } +} + +impl Element for Ref {} + +impl BufView for Ref { + type Buf<'a> + = T::Buf<'a> + where + Self: 'a; + + fn buf(&self) -> Self::Buf<'_> { + (**self).buf() + } +} + +impl fmt::Debug for Ref { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + (**self).fmt(f) + } +} + +impl fmt::Display for Ref { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + (**self).fmt(f) + } +} + +impl PartialEq> + for Ref +{ + fn eq(&self, other: &Ref) -> bool { + **self == **other + } +} + +impl Eq for Ref {} + +impl PartialEq<&T> for Ref { + fn eq(&self, other: &&T) -> bool { + **self == **other + } +} + +impl PartialOrd for Ref { + fn partial_cmp(&self, other: &Self) -> Option { + (**self).partial_cmp(&**other) + } +} + +impl Ord for Ref { + fn cmp(&self, other: &Self) -> core::cmp::Ordering { + (**self).cmp(&**other) + } +} + +impl Hash for Ref { + fn hash(&self, state: &mut H) { + (**self).hash(state); + } +} + +/// layouts whose validated bytes remain valid after typed mutation +/// +/// # Safety +/// +/// the type must have no padding and shared access must not modify its bytes. casts must +/// borrow the input bytes in the type's native layout. validation must check size, alignment +/// and contents so both casts are sound. typed writes must initialize every byte and preserve +/// validity. the flag must only be true if every bit pattern is valid. +/// callers of the casts must provide validated bytes. mutable casts additionally require +/// exclusive access. +pub unsafe trait Target { + /// allows containers to skip element validation after checking layout + const ANY_BIT_PATTERN: bool = false; + + /// checks the native size, alignment and contents + fn validate(bytes: &[u8]) -> Result<(), Error>; + + /// # Safety + /// + /// `bytes` must have passed [`Self::validate`] + unsafe fn from_bytes(bytes: &[u8]) -> &Self; + + /// # Safety + /// + /// `bytes` must have passed [`Self::validate`] and be exclusively borrowed + unsafe fn from_bytes_mut(bytes: &mut [u8]) -> &mut Self; +} + +// safety: safe string operations preserve utf8 and initialize every byte +unsafe impl Target for str { + fn validate(bytes: &[u8]) -> Result<(), Error> { + std::str::from_utf8(bytes) + .map(|_| ()) + .map_err(|_| Error::InvalidUtf8) + } + + unsafe fn from_bytes(bytes: &[u8]) -> &Self { + // safety: the caller provides validated utf8 + unsafe { std::str::from_utf8_unchecked(bytes) } + } + + unsafe fn from_bytes_mut(bytes: &mut [u8]) -> &mut Self { + // safety: the caller provides exclusive validated utf8 + unsafe { std::str::from_utf8_unchecked_mut(bytes) } + } +} + +// safety: validation checks the layout and each element's contents when required +unsafe impl Target for [T] { + const ANY_BIT_PATTERN: bool = T::ANY_BIT_PATTERN; + + fn validate(bytes: &[u8]) -> Result<(), Error> { + if bytes.is_empty() { + return Ok(()); + } + if size_of::() == 0 { + return Err(Error::InvalidData); + } + if !bytes.as_ptr().cast::().is_aligned() { + return Err(Error::InvalidAlignment); + } + if bytes.len() % size_of::() != 0 { + return Err(Error::InvalidData); + } + if !T::ANY_BIT_PATTERN { + for bytes in bytes.chunks_exact(size_of::()) { + T::validate(bytes)?; + } + } + Ok(()) + } + + unsafe fn from_bytes(bytes: &[u8]) -> &Self { + if bytes.is_empty() { + return &[]; + } + // safety: the caller validated length and alignment and each element's contents + unsafe { core::slice::from_raw_parts(bytes.as_ptr().cast(), bytes.len() / size_of::()) } + } + + unsafe fn from_bytes_mut(bytes: &mut [u8]) -> &mut Self { + if bytes.is_empty() { + return &mut []; + } + // safety: the caller provides exclusive bytes with validated length and alignment + unsafe { + core::slice::from_raw_parts_mut(bytes.as_mut_ptr().cast(), bytes.len() / size_of::()) + } + } +} + +// safety: see `[T]`, and validation additionally checks the fixed array length +unsafe impl Target for [T; N] { + const ANY_BIT_PATTERN: bool = N == 0 || T::ANY_BIT_PATTERN; + + fn validate(bytes: &[u8]) -> Result<(), Error> { + if bytes.len() != size_of::() { + return Err(Error::InvalidData); + } + if N != 0 && size_of::() == 0 && !T::ANY_BIT_PATTERN { + return T::validate(bytes); + } + <[T] as Target>::validate(bytes) + } + + unsafe fn from_bytes(bytes: &[u8]) -> &Self { + let ptr = if bytes.is_empty() { + core::ptr::NonNull::::dangling().as_ptr() + } else { + bytes.as_ptr().cast::() + }; + // safety: validation ensures an aligned array layout and empty bytes mean zero size + unsafe { &*ptr } + } + + unsafe fn from_bytes_mut(bytes: &mut [u8]) -> &mut Self { + let ptr = if bytes.is_empty() { + core::ptr::NonNull::::dangling().as_ptr() + } else { + bytes.as_mut_ptr().cast::() + }; + // safety: the array layout is validated and exclusively borrowed or has zero size + unsafe { &mut *ptr } + } +} + +macro_rules! impl_target { + ($($ty:ty),* $(,)?) => { + $( + // safety: these primitive layouts have no padding and accept every bit pattern + unsafe impl Target for $ty { + const ANY_BIT_PATTERN: bool = true; + + fn validate(bytes: &[u8]) -> Result<(), Error> { + <[Self; 1] as Target>::validate(bytes) + } + + unsafe fn from_bytes(bytes: &[u8]) -> &Self { + // safety: the caller provides the validated primitive layout + unsafe { &*bytes.as_ptr().cast() } + } + + unsafe fn from_bytes_mut(bytes: &mut [u8]) -> &mut Self { + // safety: the caller provides the validated primitive layout exclusively + unsafe { &mut *bytes.as_mut_ptr().cast() } + } + } + )* + }; +} + +impl_target!( + u8, + u16, + u32, + u64, + u128, + usize, + i8, + i16, + i32, + i64, + i128, + isize, + f32, + f64, + () +); diff --git a/codec/src/slice.rs b/codec/src/slice.rs index 9fafbdb..d877d19 100644 --- a/codec/src/slice.rs +++ b/codec/src/slice.rs @@ -1,20 +1,31 @@ use core::{mem, ops::Deref}; -use bytes::{Buf, Bytes}; +use bytes::{Buf, Bytes, BytesMut}; -/// A byte slice owner used by the codec reader -pub trait ByteSlice: Deref + Sized { - /// splits `self[..mid]` off the front, leaving `self[mid..]` behind +/// bytes that retain their container type when split +/// +/// # Safety +/// +/// dereferencing must return the same address and length until the container is split. +/// mutable dereferencing must refer to those same bytes. moving the container must preserve +/// the allocation and shared access must not modify its bytes. +pub unsafe trait ByteSlice: Deref + Sized { + /// splits off the first `mid` bytes and retains the rest /// - /// # Panics - /// - /// Panics if `mid` exceeds the slice length. + /// panics if `mid` exceeds the slice length fn split_off_front(&mut self, mid: usize) -> Self; - fn take_u8(&mut self) -> Option; + /// backing storage for a missing borrowed field, which a mutable container only has if empty + fn from_static(bytes: &'static [u8]) -> Option; + + /// takes the first byte, if there is one + fn take_u8(&mut self) -> Option { + (!self.is_empty()).then(|| self.split_off_front(1)[0]) + } } -impl ByteSlice for &[u8] { +// safety: slice references retain their allocation and mutable slices are exclusive +unsafe impl ByteSlice for &[u8] { #[inline] fn split_off_front(&mut self, mid: usize) -> Self { let (head, tail) = self.split_at(mid); @@ -22,39 +33,60 @@ impl ByteSlice for &[u8] { head } - #[inline] - fn take_u8(&mut self) -> Option { - let (&byte, remaining) = self.split_first()?; - *self = remaining; - Some(byte) + fn from_static(bytes: &'static [u8]) -> Option { + Some(bytes) } } -impl ByteSlice for &mut [u8] { +// safety: slice references retain their allocation and mutable slices are exclusive +unsafe impl ByteSlice for &mut [u8] { #[inline] fn split_off_front(&mut self, mid: usize) -> Self { + assert!(mid <= self.len()); let (head, tail) = mem::take(self).split_at_mut(mid); *self = tail; head } + fn from_static(bytes: &'static [u8]) -> Option { + bytes.is_empty().then_some(&mut []) + } +} + +// safety: the buffer lives on the heap, so moving its handle keeps the address, and splitting +// leaves both halves' bytes in place. `Bytes` never modifies shared bytes, and `BytesMut` owns +// its range exclusively. +unsafe impl ByteSlice for Bytes { + #[inline] + fn split_off_front(&mut self, mid: usize) -> Self { + self.split_to(mid) + } + + fn from_static(bytes: &'static [u8]) -> Option { + Some(Self::from_static(bytes)) + } + + // not through `split_off_front`, which would count a reference for every byte #[inline] fn take_u8(&mut self) -> Option { - let (byte, remaining) = mem::take(self).split_first_mut()?; - let byte = *byte; - *self = remaining; - Some(byte) + self.try_get_u8().ok() } } -impl ByteSlice for Bytes { +// safety: see `Bytes` +unsafe impl ByteSlice for BytesMut { #[inline] fn split_off_front(&mut self, mid: usize) -> Self { self.split_to(mid) } + fn from_static(bytes: &'static [u8]) -> Option { + bytes.is_empty().then(Self::new) + } + + // see `Bytes` #[inline] fn take_u8(&mut self) -> Option { - Buf::try_get_u8(self).ok() + self.try_get_u8().ok() } } diff --git a/codec/src/varint.rs b/codec/src/varint.rs index 0e5e1cb..ae93fbb 100644 --- a/codec/src/varint.rs +++ b/codec/src/varint.rs @@ -1,8 +1,6 @@ use core::{fmt, ops::Deref}; -use bytes::BufMut; - -use crate::{ByteSlice, Decode, Encode, Error, Reader}; +use crate::{ByteSlice, DecodeValue, Element, Encode, Error, Output, Reader, Writer}; /// An integer field carried as a varint /// @@ -14,12 +12,13 @@ pub struct Varint(pub T); pub trait Primitive: Copy { const MAX_ENCODED_LEN: usize; + const LAST_PAYLOAD_MAX: u8; fn from_u8(value: u8) -> Self; fn low_7_bits(self) -> u8; fn shr_7(self) -> Self; fn needs_more(self) -> bool; - fn checked_add_payload(self, payload: u8, shift: usize) -> Option; + fn add_payload(self, payload: u8, shift: usize) -> Self; } impl Varint { @@ -46,17 +45,23 @@ impl Encode for Varint { encoded_len(self.0) } - fn encode(&self, out: &mut W) { + fn encode(&self, out: &mut Writer) { encode(self.0, out); } } -impl Decode for Varint { - fn decode(reader: &mut Reader) -> Result { +impl DecodeValue for Varint { + fn decode_value(reader: &mut Reader) -> Result { self::decode(reader).map(Self) } + + fn missing() -> Option { + Some(Self(T::from_u8(0))) + } } +impl Element for Varint {} + pub fn encoded_len(mut value: T) -> usize { let mut len = 1; while value.needs_more() { @@ -66,11 +71,7 @@ pub fn encoded_len(mut value: T) -> usize { len } -pub fn encode(mut value: T, out: &mut W) -where - T: Primitive, - W: BufMut + ?Sized, -{ +pub fn encode(mut value: T, out: &mut Writer) { while value.needs_more() { out.put_u8(value.low_7_bits() | 0x80); value = value.shr_7(); @@ -78,30 +79,31 @@ where out.put_u8(value.low_7_bits()); } +#[inline] pub fn decode(reader: &mut Reader) -> Result where T: Primitive, B: ByteSlice, { - let mut value = T::from_u8(0); - - for index in 0..T::MAX_ENCODED_LEN { - let byte = reader.decode::()?; - let payload = byte & 0x7f; - - value = value - .checked_add_payload(payload, index * 7) - .ok_or(Error::InvalidVarint)?; + let first = reader.take_u8()?; + if first < 0x80 { + return Ok(T::from_u8(first)); + } - if byte & 0x80 == 0 { - if index > 0 && payload == 0 { - return Err(Error::InvalidVarint); - } + let mut value = T::from_u8(first & 0x7f); + for index in 1..T::MAX_ENCODED_LEN - 1 { + let byte = reader.take_u8()?; + value = value.add_payload(byte & 0x7f, index * 7); + if byte < 0x80 { return Ok(value); } } - Err(Error::InvalidVarint) + let last = reader.take_u8()?; + if last > T::LAST_PAYLOAD_MAX { + return Err(Error::InvalidVarint); + } + Ok(value.add_payload(last, (T::MAX_ENCODED_LEN - 1) * 7)) } macro_rules! impl_varint { @@ -109,6 +111,8 @@ macro_rules! impl_varint { $( impl Primitive for $ty { const MAX_ENCODED_LEN: usize = (size_of::() * 8).div_ceil(7); + #[allow(clippy::cast_possible_truncation)] + const LAST_PAYLOAD_MAX: u8 = (Self::MAX >> ((Self::MAX_ENCODED_LEN - 1) * 7)) as u8; #[inline] fn from_u8(value: u8) -> Self { @@ -128,17 +132,8 @@ macro_rules! impl_varint { self >= 0x80 } #[inline] - fn checked_add_payload(self, payload: u8, shift: usize) -> Option { - let bit_width = size_of::() * 8; - if shift >= bit_width { - return (payload == 0).then_some(self); - } - let max_payload = Self::MAX >> shift; - let payload = Self::from(payload); - if payload > max_payload { - return None; - } - Some(self | (payload << shift)) + fn add_payload(self, payload: u8, shift: usize) -> Self { + self | (Self::from(payload) << shift) } } )* @@ -157,8 +152,9 @@ mod tests { where T: Primitive + Debug + PartialEq, { - let mut out = Vec::new(); + let mut out = Writer::new(Vec::new()); encode(value, &mut out); + let out = out.into_inner(); assert_eq!(out.len(), encoded_len(value)); let mut reader = Reader::new(out.as_slice()); assert_eq!(decode(&mut reader), Ok(value)); @@ -189,14 +185,14 @@ mod tests { } assert_decodes(max); - let mut encoded = Vec::new(); + let mut encoded = Writer::new(Vec::new()); encode(T::from_u8(127), &mut encoded); - encoded.push(0x55); + encoded.put_u8(0x55); + let encoded = encoded.into_inner(); let mut reader = Reader::new(encoded.as_slice()); assert_eq!(decode(&mut reader), Ok(T::from_u8(127))); assert_eq!(reader.take_all(), &[0x55]); - assert_error::(&[0x80, 0x00], Error::InvalidVarint); assert_error::(&[0x80], Error::UnexpectedEof); assert_error::(&vec![0xff; T::MAX_ENCODED_LEN], Error::InvalidVarint); diff --git a/codec/src/writer.rs b/codec/src/writer.rs new file mode 100644 index 0000000..492164b --- /dev/null +++ b/codec/src/writer.rs @@ -0,0 +1,92 @@ +use core::mem; + +use bytes::BytesMut; + +/// A destination for encoded bytes +pub trait Output { + fn put_slice(&mut self, bytes: &[u8]); +} + +impl Output for Vec { + #[inline] + fn put_slice(&mut self, bytes: &[u8]) { + self.extend_from_slice(bytes); + } +} + +impl Output for BytesMut { + #[inline] + fn put_slice(&mut self, bytes: &[u8]) { + self.extend_from_slice(bytes); + } +} + +/// Fills the slice from the front, leaving the rest unwritten +/// +/// # Panics +/// +/// Panics if the bytes don't fit. +impl Output for &mut [u8] { + #[inline] + fn put_slice(&mut self, bytes: &[u8]) { + let (head, tail) = mem::take(self).split_at_mut(bytes.len()); + head.copy_from_slice(bytes); + *self = tail; + } +} + +impl Output for &mut O { + #[inline] + fn put_slice(&mut self, bytes: &[u8]) { + (**self).put_slice(bytes); + } +} + +/// Writes an encoding to an [`Output`], counting its bytes to align lists of integers +pub struct Writer { + out: O, + position: usize, +} + +impl Writer { + /// Starts an encoding, which a decoder must read from its start + #[inline] + pub fn new(out: O) -> Self { + Self::with_position(out, 0) + } + + /// Continues an encoding of which `position` bytes are already written + #[inline] + pub fn with_position(out: O, position: usize) -> Self { + Self { out, position } + } + + /// Bytes written since the encoding started + #[inline] + pub fn position(&self) -> usize { + self.position + } + + #[inline] + pub fn put_slice(&mut self, bytes: &[u8]) { + self.out.put_slice(bytes); + self.position += bytes.len(); + } + + #[inline] + pub fn put_u8(&mut self, byte: u8) { + self.put_slice(&[byte]); + } + + #[inline] + pub fn put_zeros(&mut self, len: usize) { + for _ in 0..len { + self.put_u8(0); + } + } + + #[inline] + pub fn into_inner(self) -> O { + self.out + } +} diff --git a/common/src/lib.rs b/common/src/lib.rs index c59d4b0..ba2fcf6 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -50,11 +50,10 @@ impl std::fmt::Display for ResetCode { ql_codec::varint_wrapper!(ResetCode, u64); -ql_codec::codec! { - #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] - #[repr(transparent)] - pub struct QID(pub [u8; Self::SIZE]); -} +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, ql_codec::Codec)] +#[codec(frozen)] +#[repr(transparent)] +pub struct QID(pub [u8; QID::SIZE]); impl QID { pub const SIZE: usize = 16; diff --git a/fsm/src/fsm.rs b/fsm/src/fsm.rs index 573e7db..6e5cc84 100644 --- a/fsm/src/fsm.rs +++ b/fsm/src/fsm.rs @@ -1,7 +1,7 @@ use std::{collections::VecDeque, time::Instant}; -use bytes::Bytes; -use ql_codec::{Decode, Reader}; +use bytes::BytesMut; +use ql_codec::Reader; use ql_common::StreamId; use ql_wire::{self as wire, QlCrypto, SessionCloseCode}; @@ -91,16 +91,14 @@ pub fn handle_connect_kk( handshake::handle_connect_kk(fsm, crypto) } -pub fn receive( +pub fn receive( fsm: &mut QlFsm, - mut bytes: B, + bytes: BytesMut, crypto: &impl QlCrypto, -) -> Result -where - B: AsMut<[u8]> + Into, -{ - let mut reader = Reader::new(bytes.as_mut()); - let header = wire::RecordHeader::decode(&mut reader) +) -> Result { + let mut reader = Reader::new(bytes); + let header: wire::RecordHeader = reader + .decode() .map_err(|error| ReceiveError::wire(ReceiveStage::RecordHeader, error))?; if header.version != wire::QL_WIRE_VERSION { @@ -112,7 +110,8 @@ where match header.record_type { wire::RecordType::Handshake => { - let record = wire::QlHandshakeRecord::decode(&mut reader) + let record = reader + .decode() .map_err(|error| ReceiveError::wire(ReceiveStage::HandshakeRecord, error))?; handshake::handle_handshake_record(fsm, crypto, header.route, &record) } @@ -123,13 +122,14 @@ where if header.route.sender != conn.transport.remote_qid { return Err(ReceiveError::InvalidQid); } - let (decrypt_len, seq) = { - let record = wire::QlSessionRecord::decode(&mut reader) + let (plaintext, seq) = { + let record = reader + .decode_ref::>>() .map_err(|error| ReceiveError::wire(ReceiveStage::SessionRecord, error))?; if conn.session.is_replay(record.header.seq) { return Ok(false); } - let payload = wire::decrypt_record( + let plaintext = wire::decrypt_record( crypto, &header, &record.header, @@ -137,12 +137,10 @@ where &conn.transport.rx_key, ) .map_err(|error| ReceiveError::wire(ReceiveStage::SessionPayload, error))?; - (payload.len(), record.header.seq) + (plaintext.freeze(), record.header.seq) }; - let bytes: Bytes = bytes.into(); - let plaintext = bytes.slice(bytes.len() - decrypt_len..); - let frames = wire::parse_session_frames(plaintext); + let frames = wire::parse_session_frames(&plaintext); let mut emit = EventSink::new(events); conn.session.receive(state.now, seq, frames, &mut emit); diff --git a/fsm/src/lib.rs b/fsm/src/lib.rs index 11b5a41..8ba8037 100644 --- a/fsm/src/lib.rs +++ b/fsm/src/lib.rs @@ -37,6 +37,7 @@ use std::{ }; pub use bytes::Bytes; +use bytes::BytesMut; pub use error::*; pub use pairing::PairingInvite; use ql_common::{ResetCode, StreamId, QID}; @@ -268,17 +269,17 @@ impl QlFsm { } /// handles one inbound wire message and returns whether it was processed rather than ignored - pub fn receive( + /// + /// `bytes` is decrypted in place and backs the stream data it carries, so nothing is copied + /// unless something else still references its buffer + pub fn receive( &mut self, now: Instant, - bytes: B, + bytes: impl Into, crypto: &impl QlCrypto, - ) -> Result - where - B: AsMut<[u8]> + Into, - { + ) -> Result { self.state.now = now; - fsm::receive(self, bytes, crypto) + fsm::receive(self, BytesMut::from(bytes.into()), crypto) } /// returns the next queued event, if any diff --git a/fsm/src/pairing.rs b/fsm/src/pairing.rs index 12eec3e..c82b47e 100644 --- a/fsm/src/pairing.rs +++ b/fsm/src/pairing.rs @@ -1,14 +1,13 @@ use ql_common::QID; use ql_wire::PairingToken; -ql_codec::codec! { - /// Out-of-band invite consumed by the initiator of an XX pairing - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - pub struct PairingInvite { - pub version: u8, - pub qid: QID, - pub token: PairingToken, - } +/// Out-of-band invite consumed by the initiator of an XX pairing +#[derive(Debug, Clone, Copy, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct PairingInvite { + pub version: u8, + pub qid: QID, + pub token: PairingToken, } impl PairingInvite { diff --git a/fsm/src/session/tests.rs b/fsm/src/session/tests.rs index c11a24d..6a7422b 100644 --- a/fsm/src/session/tests.rs +++ b/fsm/src/session/tests.rs @@ -138,7 +138,7 @@ fn receive_events( assert!(builder.push_frame(frame)); } let bytes = Bytes::from(builder.bytes().to_vec()); - let frames = parse_session_frames(bytes); + let frames = parse_session_frames(&bytes); let mut events = Vec::new(); let mut emit = |event| events.push(event); fsm.receive(now, seq, frames, &mut emit); @@ -428,7 +428,6 @@ fn ack_reopens_write_capacity() { let mut stream = fsm.stream(stream_id).unwrap(); assert_eq!(stream.io().writer().active().unwrap().write(&mut bytes), 4); stream.metadata_mut().pending_write = Bytes::from_static(b"z"); - drop(stream); let (record_seq, _record) = next_outbound(&mut fsm, now).unwrap(); let mut events = Vec::new(); @@ -464,7 +463,6 @@ fn ack_of_fin_notifies_metadata_once() { let mut stream = fsm.stream(stream_id).unwrap(); assert_eq!(stream.io().writer().active().unwrap().write(&mut bytes), 4); stream.io().writer().active().unwrap().finish(); - drop(stream); let (record_seq, record) = next_outbound(&mut fsm, now).unwrap(); assert!(matches!( @@ -692,7 +690,6 @@ fn inbound_stream_data_queues_opened_and_notifies_metadata() { assert_eq!(bytes, b"hello"); reader.commit_read(bytes.len()).unwrap(); } - drop(stream); let mut stream = fsm.stream(stream_id).unwrap(); assert_eq!(stream.metadata().inbound_finished, 1); assert!(stream.io().reader().active().is_none()); @@ -1008,7 +1005,6 @@ fn draining_the_last_bytes_reaps_a_terminal_stream_on_next_poll() { assert_eq!(bytes, b"hello"); reader.commit_read(bytes.len()).unwrap(); } - drop(stream); assert!(fsm.state.streams.contains_key(&stream_id)); fsm.take_next_write(now); diff --git a/fsm/src/session/tracked.rs b/fsm/src/session/tracked.rs index e03716f..92c6be4 100644 --- a/fsm/src/session/tracked.rs +++ b/fsm/src/session/tracked.rs @@ -16,6 +16,7 @@ pub struct TrackedRecord { } #[derive(Debug, Clone)] +#[allow(clippy::enum_variant_names)] pub enum TrackedFrame { StreamData(TrackedStreamData), StreamReset(StreamReset), diff --git a/fsm/src/tests/handshake.rs b/fsm/src/tests/handshake.rs index 195e48b..cde754b 100644 --- a/fsm/src/tests/handshake.rs +++ b/fsm/src/tests/handshake.rs @@ -355,7 +355,8 @@ fn handshake_timeout_drops_single_ik_attempt_without_resend() { harness.connect_ik(Side::A).unwrap(); harness.drain_events(Side::A); let first = harness.next_outbound(Side::A).unwrap(); - let (_, first) = ql_wire::decode_record::(first.as_slice()).unwrap(); + let (_, first) = ql_wire::decode_record::(first.as_slice()).unwrap(); + let first = first.into_owned(); assert!(matches!(first, ql_wire::QlHandshakeRecord::Ik1(_))); assert!(harness.next_outbound(Side::A).is_none()); @@ -427,7 +428,8 @@ fn simultaneous_ik_and_kk_connect_prefers_ik() { } fn handshake_id(record: &[u8]) -> ql_wire::HandshakeId { - let (_, record) = ql_wire::decode_record(record).unwrap(); + let (_, record) = ql_wire::decode_record::(record).unwrap(); + let record = record.into_owned(); match record { ql_wire::QlHandshakeRecord::Ik1(message) => message.handshake_id, ql_wire::QlHandshakeRecord::Ik2(message) => message.handshake_id, diff --git a/fsm/src/tests/mod.rs b/fsm/src/tests/mod.rs index 995ef36..f7b0bb0 100644 --- a/fsm/src/tests/mod.rs +++ b/fsm/src/tests/mod.rs @@ -336,7 +336,7 @@ fn decrypt_record( session_key: &SessionKey, ) -> (ql_wire::SessionHeader, Vec>>) { let (header, record) = - ql_wire::decode_record::, _>(record).unwrap(); + ql_wire::decode_record::>>(record).unwrap(); let plaintext = ql_wire::decrypt_record( crypto, &header, diff --git a/keyos/Cargo.toml b/keyos/Cargo.toml index 79b4f3a..27ed8ee 100644 --- a/keyos/Cargo.toml +++ b/keyos/Cargo.toml @@ -6,7 +6,6 @@ description = "KeyOS QuantumLink addressing primitives" license.workspace = true [dependencies] -bytes = { workspace = true } ql-codec = { workspace = true } ql-rpc = { workspace = true } serde = { version = "1.0", features = ["derive"], optional = true } diff --git a/keyos/src/lib.rs b/keyos/src/lib.rs index 923bd40..5365533 100644 --- a/keyos/src/lib.rs +++ b/keyos/src/lib.rs @@ -4,12 +4,13 @@ //! route keys pair these identifiers with [`RouteId`] for RPC dispatch //! [`PeerPermissions`] lists the apps a peer may access -use ql_codec::{ByteSlice, Decode, Encode, Reader}; +use ql_codec::{Encode, Reader, Writer}; -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, ql_codec::Codec)] #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[codec(frozen)] #[repr(transparent)] -pub struct AppId(pub [u8; Self::SIZE]); +pub struct AppId(pub [u8; AppId::SIZE]); impl AppId { pub const SIZE: usize = 16; @@ -39,22 +40,6 @@ impl AppId { } } -impl Encode for AppId { - fn encoded_len(&self) -> usize { - Self::SIZE - } - - fn encode(&self, out: &mut W) { - self.0.encode(out); - } -} - -impl Decode for AppId { - fn decode(reader: &mut Reader) -> Result { - Ok(Self(reader.decode()?)) - } -} - // todo: version #[derive(Debug, Clone, Default, PartialEq, Eq)] #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] @@ -104,9 +89,8 @@ impl ql_rpc::RpcRouteKey for ServiceRouteKey { self.service_id.encoded_len() + self.route_id.encoded_len() } - fn encode(&self, out: &mut W) { - self.service_id.encode(out); - self.route_id.encode(out); + fn encode(&self, out: &mut Vec) { + (self.service_id, self.route_id).encode(&mut Writer::new(out)); } fn decode(bytes: &[u8]) -> Option { @@ -130,9 +114,8 @@ impl ql_rpc::RpcRouteKey for AppRouteKey { self.app_id.encoded_len() + self.route_id.encoded_len() } - fn encode(&self, out: &mut W) { - self.app_id.encode(out); - self.route_id.encode(out); + fn encode(&self, out: &mut Vec) { + (self.app_id, self.route_id).encode(&mut Writer::new(out)); } fn decode(bytes: &[u8]) -> Option { diff --git a/router/src/protocol.rs b/router/src/protocol.rs index e647cf4..acbce97 100644 --- a/router/src/protocol.rs +++ b/router/src/protocol.rs @@ -2,7 +2,7 @@ use std::{fmt, io}; -use ql_codec::{codec, Decode, Encode, Reader}; +use ql_codec::{Decode, Encode, Reader, Writer}; use ql_common::QID; use ql_wire::{ Ik1, Ik2, Nonce, QlAead, QlHandshakeRecord, QlHash, QlRandom, RecordHeader, SessionKey, @@ -71,64 +71,59 @@ impl From for io::Error { } } -codec! { - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - pub enum PacketKind { - Confirm = 1, - Attach = 2, - Record = 3, - } +#[derive(Debug, Clone, Copy, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub enum PacketKind { + Confirm = 1, + Attach = 2, + Record = 3, } -codec! { - pub struct TransportResponse { - pub header: RecordHeader, - pub handshake: QlHandshakeRecord, - } +#[derive(ql_codec::Codec)] +#[codec(frozen)] +pub struct TransportResponse { + pub header: RecordHeader, + pub handshake: QlHandshakeRecord, } -codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub enum AttachRecord as AttachKind { - Initiate(Ik1) = 1, - Challenge(AttachChallenge) = 2, - Confirm(AttachConfirm) = 3, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +#[repr(u8)] +pub enum AttachRecord { + Initiate(Ik1) = 1, + Challenge(AttachChallenge) = 2, + Confirm(AttachConfirm) = 3, } -codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct AttachChallenge { - pub handshake: Ik2, - pub cookie: AttachCookie, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct AttachChallenge { + pub handshake: Ik2, + pub cookie: AttachCookie, } -codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct AttachConfirm { - pub cookie: AttachCookie, - pub tag: [u8; ENCRYPTED_MESSAGE_AUTH_SIZE], - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct AttachConfirm { + pub cookie: AttachCookie, + pub tag: [u8; ENCRYPTED_MESSAGE_AUTH_SIZE], } -codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct AttachCookie { - pub nonce: Nonce, - pub ciphertext: [u8; COOKIE_STATE_SIZE], - pub tag: [u8; ENCRYPTED_MESSAGE_AUTH_SIZE], - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct AttachCookie { + pub nonce: Nonce, + pub ciphertext: [u8; COOKIE_STATE_SIZE], + pub tag: [u8; ENCRYPTED_MESSAGE_AUTH_SIZE], } -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - struct CookieState { - qid: QID, - key: SessionKey, - connection_id: u64, - expires_at: u64, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +struct CookieState { + qid: QID, + key: SessionKey, + connection_id: u64, + expires_at: u64, } pub fn seal_cookie( @@ -148,7 +143,7 @@ pub fn seal_cookie( connection_id, expires_at, } - .encode(&mut encoded.as_mut_slice()); + .encode(&mut Writer::new(encoded.as_mut_slice())); let tag = crypto.aes256_gcm_encrypt(key, &nonce, COOKIE_DOMAIN, &mut encoded); AttachCookie { nonce, @@ -275,8 +270,7 @@ pub fn seal_packet_into( .ok_or(Error::FrameTooLarge)?; packet.clear(); packet.reserve(FRAME_HEADER_SIZE + length); - (length as u32).encode(packet); - kind.encode(packet); + (length as u32, kind).encode(&mut Writer::new(&mut *packet)); packet.extend_from_slice(payload); let tag = crypto.aes256_gcm_encrypt(key, &Nonce::from_counter(nonce), packet, &mut []); packet.extend_from_slice(&tag); @@ -354,7 +348,7 @@ pub async fn write_frame(writer: &mut (impl AsyncWrite + Unpin), payload: &[u8]) return Err(Error::FrameTooLarge.into()); } let mut header = [0; FRAME_HEADER_SIZE]; - (payload.len() as u32).encode(&mut header.as_mut_slice()); + (payload.len() as u32).encode(&mut Writer::new(header.as_mut_slice())); writer.write_all(&header).await?; writer.write_all(payload).await } diff --git a/rpc/src/rpc/mod.rs b/rpc/src/rpc/mod.rs index dd3572f..cd0124f 100644 --- a/rpc/src/rpc/mod.rs +++ b/rpc/src/rpc/mod.rs @@ -5,8 +5,6 @@ //! route dispatch uses caller-defined route keys and the submodules provide the //! matching client and server helpers for encoding, decoding, and handler glue -use bytes::BufMut; - pub mod download; pub mod duplex; pub mod notification; @@ -22,7 +20,7 @@ mod utils; pub trait RpcRouteKey: Sized + std::fmt::Debug + Ord + Clone + 'static { fn encoded_len(&self) -> usize; - fn encode(&self, out: &mut W); + fn encode(&self, out: &mut Vec); fn decode(bytes: &[u8]) -> Option; } diff --git a/rpc/src/rpc/tests.rs b/rpc/src/rpc/tests.rs index 43b1e13..0c8df81 100644 --- a/rpc/src/rpc/tests.rs +++ b/rpc/src/rpc/tests.rs @@ -7,7 +7,7 @@ use std::{ task::{Context as TaskContext, Poll, Waker}, }; -use bytes::{BufMut, Bytes}; +use bytes::Bytes; use ql_common::{ResetCode, StreamId, QID}; use super::{download, duplex, notification, progress, request, subscription, upload}; @@ -190,8 +190,8 @@ impl RpcRouteKey for TestKey { 1 } - fn encode(&self, out: &mut W) { - out.put_u8(self.0); + fn encode(&self, out: &mut Vec) { + out.push(self.0); } fn decode(bytes: &[u8]) -> Option { diff --git a/runtime/src/driver/mod.rs b/runtime/src/driver/mod.rs index 28ee5ff..b3953cb 100644 --- a/runtime/src/driver/mod.rs +++ b/runtime/src/driver/mod.rs @@ -292,7 +292,6 @@ impl DriverState { let (reader, writer, reader_io, writer_io) = io::new_stream(stream_id, self.runtime_tx.clone()); metadata.initialize(writer_io, reader_io, stream_io); - drop(stream); log::info!("delivering inbound stream to platform: stream_id={stream_id}"); platform.handle_inbound( diff --git a/runtime/src/tests/rpc.rs b/runtime/src/tests/rpc.rs index 1e8f335..97011d7 100644 --- a/runtime/src/tests/rpc.rs +++ b/runtime/src/tests/rpc.rs @@ -7,8 +7,7 @@ use std::{ time::Duration, }; -use bytes::BufMut; -use ql_codec::Encode; +use ql_codec::{Encode, Writer}; use ql_rpc::{ progress::{ProgressHandlerLocal, ProgressResponder}, request::{RequestHandler, Response}, @@ -26,8 +25,8 @@ impl ql_rpc::RpcRouteKey for TestRouteKey { ql_codec::Varint(self.0).encoded_len() } - fn encode(&self, out: &mut W) { - ql_codec::Varint(self.0).encode(out); + fn encode(&self, out: &mut Vec) { + ql_codec::Varint(self.0).encode(&mut Writer::new(out)); } fn decode(bytes: &[u8]) -> Option { diff --git a/wire/src/crypto.rs b/wire/src/crypto.rs index 380e98d..48e931a 100644 --- a/wire/src/crypto.rs +++ b/wire/src/crypto.rs @@ -3,11 +3,10 @@ use crate::{ ENCRYPTED_MESSAGE_AUTH_SIZE, }; -ql_codec::codec! { - #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] - #[repr(transparent)] - pub struct Nonce(pub [u8; Self::SIZE]); -} +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, ql_codec::Codec)] +#[codec(frozen)] +#[repr(transparent)] +pub struct Nonce(pub [u8; Nonce::SIZE]); impl Nonce { pub const SIZE: usize = 12; diff --git a/wire/src/encrypted/ack.rs b/wire/src/encrypted/ack.rs index 8be99b9..abe59bd 100644 --- a/wire/src/encrypted/ack.rs +++ b/wire/src/encrypted/ack.rs @@ -1,6 +1,6 @@ use std::{fmt, ops::RangeInclusive}; -use ql_codec::{ByteSlice, Encode, Error, Varint}; +use ql_codec::{Encode, Error, Output, Varint, Writer}; use crate::RecordSeq; @@ -135,7 +135,7 @@ impl Encode for RecordAck { .sum::() } - fn encode(&self, out: &mut W) { + fn encode(&self, out: &mut Writer) { self.largest_acked.encode(out); Varint(Self::block_count(self.blocks.len())).encode(out); self.first_range_len.encode(out); @@ -146,12 +146,17 @@ impl Encode for RecordAck { } } -impl ql_codec::Decode for RecordAck { - fn decode(reader: &mut ql_codec::Reader) -> Result { +impl ql_codec::DecodeValue for RecordAck { + fn decode_value( + reader: &mut ql_codec::Reader, + ) -> Result { let largest_acked = reader.decode()?; let block_count = *reader.decode::>()? as usize; let first_range_len = reader.decode()?; - let mut blocks = Vec::with_capacity(block_count); + // At most as much memory as the input, or a corrupt count could reserve gigabytes. + let mut blocks = Vec::with_capacity( + block_count.min(reader.remaining_len() / size_of::()), + ); for _ in 0..block_count { blocks.push(RecordAckBlock { gap: reader.decode()?, diff --git a/wire/src/encrypted/builder.rs b/wire/src/encrypted/builder.rs index 82ccef0..470a02d 100644 --- a/wire/src/encrypted/builder.rs +++ b/wire/src/encrypted/builder.rs @@ -1,5 +1,4 @@ -use bytes::BufMut; -use ql_codec::{BufView, Encode, Varint}; +use ql_codec::{BufView, Encode, Varint, Writer}; use super::{RecordAck, SessionClose, SessionFrame, StreamData, StreamReset, StreamWindow}; use crate::{ @@ -118,21 +117,29 @@ impl SessionRecordBuilder { &mut self.bytes[self.prefix_len..], ); - let mut prefix = &mut self.bytes[..self.prefix_len]; + let mut prefix = Writer::new(&mut self.bytes[..self.prefix_len]); record_header.encode(&mut prefix); header.encode(&mut prefix); auth.encode(&mut prefix); - debug_assert!(prefix.is_empty()); + debug_assert!(prefix.into_inner().is_empty()); self.bytes } - fn push_wire_size(&mut self, wire_size: usize, encode: impl FnOnce(&mut Vec)) -> bool { + fn push_wire_size( + &mut self, + wire_size: usize, + encode: impl FnOnce(&mut Writer<&mut Vec>), + ) -> bool { if !self.can_push_len(wire_size) { return false; } self.ensure_prefix_capacity(wire_size); let start = self.bytes.len(); - encode(&mut self.bytes); + // The frames are decoded from the start of the plaintext, after the prefix. + encode(&mut Writer::with_position( + &mut self.bytes, + start - self.prefix_len, + )); debug_assert_eq!(self.bytes.len(), start + wire_size); true } diff --git a/wire/src/encrypted/close.rs b/wire/src/encrypted/close.rs index 1b76f3a..2b45f58 100644 --- a/wire/src/encrypted/close.rs +++ b/wire/src/encrypted/close.rs @@ -1,9 +1,8 @@ -ql_codec::codec! { - /// closes the whole session immediately with a reset code. - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct SessionClose { - pub code: SessionCloseCode, - } +/// closes the whole session immediately with a reset code. +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct SessionClose { + pub code: SessionCloseCode, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] diff --git a/wire/src/encrypted/mod.rs b/wire/src/encrypted/mod.rs index efd1c91..c89a09f 100644 --- a/wire/src/encrypted/mod.rs +++ b/wire/src/encrypted/mod.rs @@ -1,4 +1,5 @@ -use ql_codec::{ByteSlice, Reader}; +use bytes::Bytes; +use ql_codec::Reader; use ql_common::StreamId; use crate::{ @@ -19,60 +20,64 @@ pub use stream_data::*; pub use stream_reset::*; pub use stream_window::*; -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub enum SessionFrame as SessionFrameKind { - // todo: do we need ping as explicit frame? - Ping = 1, - Ack(RecordAck) = 2, - StreamData(StreamData) = 3, - StreamWindow(StreamWindow) = 4, - StreamReset(StreamReset) = 5, - Close(SessionClose) = 6, - Unpair = 7, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen, discriminants = SessionFrameKind)] +#[repr(u8)] +pub enum SessionFrame { + // todo: do we need ping as explicit frame? + Ping = 1, + Ack(RecordAck) = 2, + StreamData(StreamData) = 3, + StreamWindow(StreamWindow) = 4, + StreamReset(StreamReset) = 5, + Close(SessionClose) = 6, + Unpair = 7, } -impl SessionFrame { - pub fn into_owned(self) -> SessionFrame> { +impl SessionFrameRef> { + fn into_frame(self) -> SessionFrame { match self { Self::Ping => SessionFrame::Ping, Self::Unpair => SessionFrame::Unpair, Self::Ack(frame) => SessionFrame::Ack(frame), - Self::StreamData(frame) => SessionFrame::StreamData(frame.into_owned()), - Self::StreamWindow(frame) => SessionFrame::StreamWindow(frame), - Self::StreamReset(frame) => SessionFrame::StreamReset(frame), - Self::Close(frame) => SessionFrame::Close(frame), + Self::StreamData(frame) => SessionFrame::StreamData(frame), + Self::StreamWindow(frame) => SessionFrame::StreamWindow(frame.into_owned()), + Self::StreamReset(frame) => SessionFrame::StreamReset(frame.into_owned()), + Self::Close(frame) => SessionFrame::Close(frame.into_owned()), } } } -pub fn parse_session_frames(bytes: B) -> SessionFrameIter { +/// Parses the frames in `bytes`, slicing their stream data out of it without copying +pub fn parse_session_frames(bytes: &Bytes) -> SessionFrameIter { SessionFrameIter { - reader: Reader::new(bytes), + reader: Reader::new(bytes.clone()), } } pub fn decode_session_frames(bytes: &[u8]) -> Result>>, Error> { - parse_session_frames(bytes) - .map(|frame| frame.map(SessionFrame::into_owned)) - .collect() + let mut reader = Reader::new(bytes); + let mut frames = Vec::new(); + while !reader.is_empty() { + frames.push(reader.decode()?); + } + Ok(frames) } #[derive(Clone)] -pub struct SessionFrameIter { - reader: Reader, +pub struct SessionFrameIter { + reader: Reader, } -impl Iterator for SessionFrameIter { - type Item = Result, Error>; +impl Iterator for SessionFrameIter { + type Item = Result, Error>; fn next(&mut self) -> Option { if self.reader.is_empty() { - None - } else { - Some(self.reader.decode::>().map_err(Into::into)) + return None; } + let frame = self.reader.decode_ref::>>(); + Some(frame.map(SessionFrameRef::into_frame).map_err(Into::into)) } } diff --git a/wire/src/encrypted/stream_data.rs b/wire/src/encrypted/stream_data.rs index 5d3a0c5..ba4b4d7 100644 --- a/wire/src/encrypted/stream_data.rs +++ b/wire/src/encrypted/stream_data.rs @@ -1,5 +1,6 @@ use ql_codec::{ - encode_bytes, encoded_len_bytes, BufView, ByteSlice, Decode, Encode, Error, Varint, + encode_bytes, encoded_len_bytes, BufView, ByteSlice, Decode, Encode, Error, Output, Reader, + Varint, Writer, }; use ql_common::StreamId; @@ -37,11 +38,13 @@ impl StreamData { + Varint::::MAX_ENCODED_LEN; // header length } -impl Decode for StreamData { - fn decode(reader: &mut ql_codec::Reader) -> Result { +impl Decode for StreamData> { + type Ref = StreamData; + + fn decode_ref(reader: &mut Reader) -> Result { let stream_id = reader.decode()?; let offset = reader.decode()?; - let flags = reader.decode::()?; + let flags = reader.take_u8()?; let fin = (flags & flag::FIN) != 0; let has_open = (flags & flag::OPEN) != 0; let open = if has_open { @@ -55,7 +58,7 @@ impl Decode for StreamData { }; let bytes = reader.take_len_prefixed()?; - Ok(Self { + Ok(StreamData { stream_id, offset, open, @@ -63,26 +66,34 @@ impl Decode for StreamData { bytes, }) } + + fn from_ref(value: Self::Ref) -> Self { + value.map_bytes(|bytes| bytes.to_vec()) + } } -impl StreamData { - pub fn into_owned(self) -> StreamData> - where - B: ByteSlice, - H: ByteSlice, - { +impl StreamData { + /// Converts the payload and the header with `f` + pub fn map_bytes(self, mut f: impl FnMut(B) -> C) -> StreamData { StreamData { stream_id: self.stream_id, offset: self.offset, open: self.open.map(|open| StreamOpen { - header: open.header.to_vec(), + header: f(open.header), receive_window: open.receive_window, requested_peer_receive_window: open.requested_peer_receive_window, }), fin: self.fin, - bytes: self.bytes.to_vec(), + bytes: f(self.bytes), } } + + pub fn into_owned(self) -> StreamData> + where + B: AsRef<[u8]>, + { + self.map_bytes(|bytes| bytes.as_ref().to_vec()) + } } impl Encode for StreamData { @@ -98,7 +109,7 @@ impl Encode for StreamData { + encoded_len_bytes(&self.bytes) } - fn encode(&self, out: &mut W) { + fn encode(&self, out: &mut Writer) { debug_assert!( *self.offset == 0 || self.open.is_none(), "stream opening is only valid at offset 0" diff --git a/wire/src/encrypted/stream_reset.rs b/wire/src/encrypted/stream_reset.rs index 711d786..c2b6197 100644 --- a/wire/src/encrypted/stream_reset.rs +++ b/wire/src/encrypted/stream_reset.rs @@ -2,29 +2,27 @@ use ql_common::ResetCode; use super::StreamId; -ql_codec::codec! { - /// aborts one or both lanes of a stream with a reset code - /// - /// stream origin is the peer that opened the stream - /// origin lane carries bytes sent by the stream origin - /// return lane carries bytes sent back toward the stream origin - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct StreamReset { - pub stream_id: StreamId, - pub target: ResetTarget, - pub code: ResetCode, - } +/// aborts one or both lanes of a stream with a reset code +/// +/// stream origin is the peer that opened the stream +/// origin lane carries bytes sent by the stream origin +/// return lane carries bytes sent back toward the stream origin +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct StreamReset { + pub stream_id: StreamId, + pub target: ResetTarget, + pub code: ResetCode, } -ql_codec::codec! { - /// selects which stream lane a [`StreamReset`] applies to - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - pub enum ResetTarget { - /// reset the lane sent by the stream origin - Origin = 1, - /// reset the lane sent back toward the stream origin - Return = 2, - /// reset both stream lanes - Both = 3, - } +/// selects which stream lane a [`StreamReset`] applies to +#[derive(Debug, Clone, Copy, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub enum ResetTarget { + /// reset the lane sent by the stream origin + Origin = 1, + /// reset the lane sent back toward the stream origin + Return = 2, + /// reset both stream lanes + Both = 3, } diff --git a/wire/src/encrypted/stream_window.rs b/wire/src/encrypted/stream_window.rs index 49b6134..fef8b20 100644 --- a/wire/src/encrypted/stream_window.rs +++ b/wire/src/encrypted/stream_window.rs @@ -2,11 +2,10 @@ use ql_codec::Varint; use super::StreamId; -ql_codec::codec! { - /// advertises the highest byte offset the peer may send on a stream. - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct StreamWindow { - pub stream_id: StreamId, - pub maximum_offset: Varint, - } +/// advertises the highest byte offset the peer may send on a stream. +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct StreamWindow { + pub stream_id: StreamId, + pub maximum_offset: Varint, } diff --git a/wire/src/encrypted_message.rs b/wire/src/encrypted_message.rs index 0b2ddfe..a18a962 100644 --- a/wire/src/encrypted_message.rs +++ b/wire/src/encrypted_message.rs @@ -1,5 +1,5 @@ use bytes::Buf; -use ql_codec::{encode_bytes_raw, BufView, ByteSlice, Decode, Encode}; +use ql_codec::{encode_bytes_raw, BufView, ByteSlice, Decode, Encode, Output, Reader, Writer}; use crate::ENCRYPTED_MESSAGE_AUTH_SIZE; @@ -9,25 +9,31 @@ pub struct EncryptedMessage { pub ciphertext: B, } -impl EncryptedMessage { - pub fn into_owned(self) -> EncryptedMessage> - where - B: ByteSlice, - { +impl> EncryptedMessage { + pub fn into_owned(self) -> EncryptedMessage> { EncryptedMessage { auth: self.auth, - ciphertext: self.ciphertext.to_vec(), + ciphertext: self.ciphertext.as_ref().to_vec(), } } } -impl Decode for EncryptedMessage { - fn decode(reader: &mut ql_codec::Reader) -> Result { - Ok(Self { +impl Decode for EncryptedMessage> { + type Ref = EncryptedMessage; + + fn decode_ref(reader: &mut Reader) -> Result { + Ok(EncryptedMessage { auth: reader.decode()?, ciphertext: reader.take_all(), }) } + + fn from_ref(value: Self::Ref) -> Self { + Self { + auth: value.auth, + ciphertext: value.ciphertext.to_vec(), + } + } } impl Encode for EncryptedMessage { @@ -35,7 +41,7 @@ impl Encode for EncryptedMessage { self.auth.encoded_len() + self.ciphertext.buf().remaining() } - fn encode(&self, out: &mut W) { + fn encode(&self, out: &mut Writer) { self.auth.encode(out); encode_bytes_raw(&self.ciphertext, out); } diff --git a/wire/src/error.rs b/wire/src/error.rs index 7135923..46e6bc4 100644 --- a/wire/src/error.rs +++ b/wire/src/error.rs @@ -3,6 +3,7 @@ use core::fmt; #[derive(Debug, Clone, PartialEq, Eq)] pub enum Error { // codec errors + InvalidAlignment, UnexpectedEof, InvalidData, InvalidDiscriminant, @@ -28,6 +29,7 @@ pub enum Error { impl fmt::Display for Error { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let message = match self { + Self::InvalidAlignment => "invalid alignment", Self::UnexpectedEof => "unexpected end of input", Self::InvalidData => "invalid data", Self::InvalidDiscriminant => "invalid discriminant", @@ -53,6 +55,7 @@ impl std::error::Error for Error {} impl From for Error { fn from(error: ql_codec::Error) -> Self { match error { + ql_codec::Error::InvalidAlignment => Self::InvalidAlignment, ql_codec::Error::InvalidData => Self::InvalidData, ql_codec::Error::UnexpectedEof => Self::UnexpectedEof, ql_codec::Error::InvalidDiscriminant => Self::InvalidDiscriminant, diff --git a/wire/src/handshake/id.rs b/wire/src/handshake/id.rs index 4f34602..d46c1e7 100644 --- a/wire/src/handshake/id.rs +++ b/wire/src/handshake/id.rs @@ -1,5 +1,4 @@ -ql_codec::codec! { - #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] - #[repr(transparent)] - pub struct HandshakeId(pub u32); -} +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, ql_codec::Codec)] +#[codec(frozen)] +#[repr(transparent)] +pub struct HandshakeId(pub u32); diff --git a/wire/src/handshake/ik.rs b/wire/src/handshake/ik.rs index ec8da42..0bd9bbe 100644 --- a/wire/src/handshake/ik.rs +++ b/wire/src/handshake/ik.rs @@ -1,6 +1,6 @@ use std::borrow::Borrow; -use ql_codec::{ByteSlice, Encode}; +use ql_codec::{Encode, Output, Writer}; use super::{ decrypt_mlkem_ciphertext, decrypt_peer_bundle, encrypt_mlkem_ciphertext, encrypt_peer_bundle, @@ -35,7 +35,7 @@ impl Encode for Ik1 { .map_or(0, EncryptedPeerBundle::encoded_len) } - fn encode(&self, out: &mut W) { + fn encode(&self, out: &mut Writer) { self.handshake_id.encode(out); self.transport_params.encode(out); self.skem_ciphertext.encode(out); @@ -46,8 +46,10 @@ impl Encode for Ik1 { } } -impl ql_codec::Decode for Ik1 { - fn decode(reader: &mut ql_codec::Reader) -> Result { +impl ql_codec::DecodeValue for Ik1 { + fn decode_value( + reader: &mut ql_codec::Reader, + ) -> Result { let handshake_id = reader.decode()?; let transport_params = reader.decode()?; let skem_ciphertext = reader.decode()?; @@ -67,14 +69,13 @@ impl ql_codec::Decode for Ik1 { } } -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct Ik2 { - pub handshake_id: HandshakeId, - pub transport_params: TransportParams, - pub ekem_ciphertext: MlKemCiphertext, - pub skem_ciphertext: EncryptedMlKemCiphertext, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct Ik2 { + pub handshake_id: HandshakeId, + pub transport_params: TransportParams, + pub ekem_ciphertext: MlKemCiphertext, + pub skem_ciphertext: EncryptedMlKemCiphertext, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/wire/src/handshake/mod.rs b/wire/src/handshake/mod.rs index 03910d3..e5fd6ea 100644 --- a/wire/src/handshake/mod.rs +++ b/wire/src/handshake/mod.rs @@ -1,6 +1,6 @@ //! authenticated QuantumLink handshake patterns -use ql_codec::{ByteSlice, Decode, Encode}; +use ql_codec::{Decode, Encode, Output, Writer}; use crate::{ Error, HandshakeKind, MlKemCiphertext, MlKemKeyPair, MlKemPublicKey, Nonce, PeerBundle, @@ -24,17 +24,15 @@ const PROTOCOL_IK: &[u8] = b"ql-wire:pq-ik:v1"; const PROTOCOL_KK: &[u8] = b"ql-wire:pq-kk:v1"; const HANDSHAKE_PREAMBLE_DOMAIN: &[u8] = b"ql-wire:handshake-preamble:v1"; -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct EphemeralPublicKey { - pub mlkem_public_key: MlKemPublicKey, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct EphemeralPublicKey { + pub mlkem_public_key: MlKemPublicKey, } -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct EncryptedMlKemCiphertext(pub Box<[u8; Self::SIZE]>); -} +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct EncryptedMlKemCiphertext(pub Box<[u8; EncryptedMlKemCiphertext::SIZE]>); impl EncryptedMlKemCiphertext { pub const SIZE: usize = MlKemCiphertext::SIZE + ENCRYPTED_MESSAGE_AUTH_SIZE; @@ -62,15 +60,16 @@ impl Encode for EncryptedPeerBundle { self.0.len() } - fn encode(&self, out: &mut W) { + fn encode(&self, out: &mut Writer) { out.put_slice(self.as_bytes()); } } -impl ql_codec::Decode for EncryptedPeerBundle { - fn decode(reader: &mut ql_codec::Reader) -> Result { - let data = reader.take_all(); - Ok(Self(Box::from(&*data))) +impl ql_codec::DecodeValue for EncryptedPeerBundle { + fn decode_value( + reader: &mut ql_codec::Reader, + ) -> Result { + Ok(Self(Box::from(&*reader.take_all()))) } } @@ -293,8 +292,7 @@ fn mix_hash_pairing_handshake( pairing_id: PairingId, transport_params: TransportParams, ) { - let mut preamble = header.encode_vec(); - pairing_id.encode(&mut preamble); + let preamble = (header, pairing_id).encode_vec(); mix_hash_handshake_preamble( symmetric, crypto, diff --git a/wire/src/handshake/pairing.rs b/wire/src/handshake/pairing.rs index 9e3aa94..18b680e 100644 --- a/wire/src/handshake/pairing.rs +++ b/wire/src/handshake/pairing.rs @@ -5,11 +5,10 @@ use crate::QlCrypto; const PAIRING_ID_DOMAIN: &[u8] = b"ql-wire:pairing-id:v1"; const PAIRING_PSK_DOMAIN: &[u8] = b"ql-wire:pairing-psk:v1"; -ql_codec::codec! { - #[derive(Clone, Copy, PartialEq, Eq, Hash)] - #[repr(transparent)] - pub struct PairingToken(pub [u8; Self::SIZE]); -} +#[derive(Clone, Copy, PartialEq, Eq, Hash, ql_codec::Codec)] +#[codec(frozen)] +#[repr(transparent)] +pub struct PairingToken(pub [u8; PairingToken::SIZE]); impl fmt::Debug for PairingToken { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { @@ -41,11 +40,10 @@ impl Display for PairingToken { } } -ql_codec::codec! { - #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] - #[repr(transparent)] - pub struct PairingId(pub [u8; Self::SIZE]); -} +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, ql_codec::Codec)] +#[codec(frozen)] +#[repr(transparent)] +pub struct PairingId(pub [u8; PairingId::SIZE]); impl PairingId { pub const SIZE: usize = 16; diff --git a/wire/src/handshake/transport_params.rs b/wire/src/handshake/transport_params.rs index b0c6cb5..0cc5b35 100644 --- a/wire/src/handshake/transport_params.rs +++ b/wire/src/handshake/transport_params.rs @@ -1,10 +1,9 @@ -ql_codec::codec! { - /// Session parameters advertised in the handshake - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - pub struct TransportParams { - /// Initial per-stream receive credit granted to the remote peer - pub initial_stream_receive_window: u32, - } +/// Session parameters advertised in the handshake +#[derive(Debug, Clone, Copy, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct TransportParams { + /// Initial per-stream receive credit granted to the remote peer + pub initial_stream_receive_window: u32, } impl Default for TransportParams { diff --git a/wire/src/handshake/xx.rs b/wire/src/handshake/xx.rs index 34c256e..4803366 100644 --- a/wire/src/handshake/xx.rs +++ b/wire/src/handshake/xx.rs @@ -14,41 +14,37 @@ use crate::{ const PROTOCOL_XX: &[u8] = b"ql-wire:pq-xx:v1"; -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct Xx1 { - pub handshake_id: HandshakeId, - pub pairing_id: PairingId, - pub transport_params: TransportParams, - pub ephemeral: EphemeralPublicKey, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct Xx1 { + pub handshake_id: HandshakeId, + pub pairing_id: PairingId, + pub transport_params: TransportParams, + pub ephemeral: EphemeralPublicKey, } -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct Xx2 { - pub handshake_id: HandshakeId, - pub transport_params: TransportParams, - pub ekem_ciphertext: MlKemCiphertext, - pub static_bundle: EncryptedPeerBundle, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct Xx2 { + pub handshake_id: HandshakeId, + pub transport_params: TransportParams, + pub ekem_ciphertext: MlKemCiphertext, + pub static_bundle: EncryptedPeerBundle, } -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct Xx3 { - pub handshake_id: HandshakeId, - pub skem_ciphertext: EncryptedMlKemCiphertext, - pub static_bundle: EncryptedPeerBundle, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct Xx3 { + pub handshake_id: HandshakeId, + pub skem_ciphertext: EncryptedMlKemCiphertext, + pub static_bundle: EncryptedPeerBundle, } -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct Xx4 { - pub handshake_id: HandshakeId, - pub skem_ciphertext: EncryptedMlKemCiphertext, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct Xx4 { + pub handshake_id: HandshakeId, + pub skem_ciphertext: EncryptedMlKemCiphertext, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/wire/src/header.rs b/wire/src/header.rs index a40d96f..631553c 100644 --- a/wire/src/header.rs +++ b/wire/src/header.rs @@ -1,22 +1,19 @@ -use ::bytes::BufMut; -use ql_codec::Encode; +use ql_codec::{Encode, Writer}; use ql_common::QID; use crate::QL_WIRE_VERSION; -ql_codec::codec! { - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - pub struct RouteHeader { - pub sender: QID, - pub recipient: QID, - } +#[derive(Debug, Clone, Copy, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct RouteHeader { + pub sender: QID, + pub recipient: QID, } -ql_codec::codec! { - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - pub struct SessionHeader { - pub seq: RecordSeq, - } +#[derive(Debug, Clone, Copy, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct SessionHeader { + pub seq: RecordSeq, } #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] @@ -41,13 +38,13 @@ impl SessionHeader { + size_of::() + route.encoded_len() + self.seq.encoded_len(); - let mut aad = Vec::with_capacity(aad_len); + let mut aad = Writer::new(Vec::with_capacity(aad_len)); aad.put_slice(Self::AAD_DOMAIN); aad.put_u8(QL_WIRE_VERSION); aad.put_u8(Self::AAD_RECORD_KIND_SESSION); route.encode(&mut aad); self.seq.encode(&mut aad); - debug_assert_eq!(aad.len(), aad_len); - aad + debug_assert_eq!(aad.position(), aad_len); + aad.into_inner() } } diff --git a/wire/src/identity.rs b/wire/src/identity.rs index ee39e37..cd76c9f 100644 --- a/wire/src/identity.rs +++ b/wire/src/identity.rs @@ -2,15 +2,14 @@ use ql_common::QID; use crate::{derive_qid, Error, MlKemKeyPair, MlKemPrivateKey, MlKemPublicKey, QlCrypto, QlHash}; -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct PeerBundle { - pub version: u16, - pub qid: QID, - pub capabilities: u32, - pub mlkem_public_key: MlKemPublicKey, - pub name: String, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct PeerBundle { + pub version: u16, + pub qid: QID, + pub capabilities: u32, + pub mlkem_public_key: MlKemPublicKey, + pub name: String, } impl PeerBundle { @@ -24,15 +23,14 @@ impl PeerBundle { } } -ql_codec::codec! { - #[derive(Debug, Clone)] - pub struct QlIdentity { - pub qid: QID, - pub mlkem_private_key: MlKemPrivateKey, - pub mlkem_public_key: MlKemPublicKey, - pub capabilities: u32, - pub name: String, - } +#[derive(Debug, Clone, ql_codec::Codec)] +#[codec(frozen)] +pub struct QlIdentity { + pub qid: QID, + pub mlkem_private_key: MlKemPrivateKey, + pub mlkem_public_key: MlKemPublicKey, + pub capabilities: u32, + pub name: String, } impl QlIdentity { diff --git a/wire/src/pq.rs b/wire/src/pq.rs index 7803f85..ca1e0bb 100644 --- a/wire/src/pq.rs +++ b/wire/src/pq.rs @@ -8,10 +8,9 @@ const ML_KEM_1024_PUBLIC_KEY_SIZE: usize = 1568; const ML_KEM_1024_PRIVATE_KEY_SIZE: usize = 3168; const ML_KEM_1024_CIPHERTEXT_SIZE: usize = 1568; -ql_codec::codec! { - #[derive(Clone, PartialEq, Eq, Hash)] - pub struct SessionKey(pub [u8; Self::SIZE]); -} +#[derive(Clone, PartialEq, Eq, Hash, ql_codec::Codec)] +#[codec(frozen)] +pub struct SessionKey(pub [u8; SessionKey::SIZE]); impl std::fmt::Debug for SessionKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -39,10 +38,9 @@ impl Drop for SessionKey { } } -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq, Hash)] - pub struct MlKemPublicKey(Box<[u8; MlKemPublicKey::SIZE]>); -} +#[derive(Debug, Clone, PartialEq, Eq, Hash, ql_codec::Codec)] +#[codec(frozen)] +pub struct MlKemPublicKey(Box<[u8; MlKemPublicKey::SIZE]>); impl MlKemPublicKey { pub const SIZE: usize = ML_KEM_1024_PUBLIC_KEY_SIZE; @@ -62,10 +60,9 @@ impl Drop for MlKemPublicKey { } } -ql_codec::codec! { - #[derive(Clone, PartialEq, Eq, Hash)] - pub struct MlKemPrivateKey(Box<[u8; MlKemPrivateKey::SIZE]>); -} +#[derive(Clone, PartialEq, Eq, Hash, ql_codec::Codec)] +#[codec(frozen)] +pub struct MlKemPrivateKey(Box<[u8; MlKemPrivateKey::SIZE]>); impl std::fmt::Debug for MlKemPrivateKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -91,10 +88,9 @@ impl Drop for MlKemPrivateKey { } } -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq, Hash)] - pub struct MlKemCiphertext(Box<[u8; MlKemCiphertext::SIZE]>); -} +#[derive(Debug, Clone, PartialEq, Eq, Hash, ql_codec::Codec)] +#[codec(frozen)] +pub struct MlKemCiphertext(Box<[u8; MlKemCiphertext::SIZE]>); impl MlKemCiphertext { pub const SIZE: usize = ML_KEM_1024_CIPHERTEXT_SIZE; diff --git a/wire/src/record.rs b/wire/src/record.rs index 6484b7d..a94ee64 100644 --- a/wire/src/record.rs +++ b/wire/src/record.rs @@ -1,4 +1,4 @@ -use ql_codec::{ByteSlice, Decode, Encode}; +use ql_codec::{BufView, ByteSlice, Decode, Encode, Output, Reader, Writer}; use ql_common::QID; use crate::{ @@ -7,37 +7,36 @@ use crate::{ Error, RouteHeader, SessionHeader, QL_WIRE_VERSION, }; -pub fn encode_record(out: &mut W, header: RecordHeader, body: &T) -where - W: bytes::BufMut + ?Sized, - T: Encode + ?Sized, -{ +pub fn encode_record( + out: &mut Writer, + header: RecordHeader, + body: &T, +) { header.encode(out); body.encode(out); } pub fn encode_record_vec(header: RecordHeader, body: &T) -> Vec { - let mut out = Vec::with_capacity(header.encoded_len() + body.encoded_len()); + let mut out = Writer::new(Vec::with_capacity( + header.encoded_len() + body.encoded_len(), + )); encode_record(&mut out, header, body); - out + out.into_inner() } -pub fn decode_record(bytes: B) -> Result<(RecordHeader, T), Error> -where - T: Decode, - B: ByteSlice, -{ - let mut reader = ql_codec::Reader::new(bytes); - Ok((reader.decode()?, reader.decode()?)) +pub fn decode_record<'a, T: Decode<&'a [u8]>>( + bytes: &'a [u8], +) -> Result<(RecordHeader, T::Ref), Error> { + let mut reader = Reader::new(bytes); + Ok((reader.decode()?, reader.decode_ref::()?)) } -ql_codec::codec! { - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - pub struct RecordHeader { - pub version: u8, - pub route: RouteHeader, - pub record_type: RecordType, - } +#[derive(Debug, Clone, Copy, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub struct RecordHeader { + pub version: u8, + pub route: RouteHeader, + pub record_type: RecordType, } impl RecordHeader { @@ -52,37 +51,34 @@ impl RecordHeader { } } -ql_codec::codec! { - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - pub enum RecordType { - Handshake = 1, - Session = 2, - } +#[derive(Debug, Clone, Copy, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen)] +pub enum RecordType { + Handshake = 1, + Session = 2, } -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub enum QlHandshakeRecord as HandshakeKind { - Ik1(Ik1) = 1, - Ik2(Ik2) = 2, - Kk1(Ik1) = 3, - Kk2(Ik2) = 4, - Xx1(Xx1) = 5, - Xx2(Xx2) = 6, - Xx3(Xx3) = 7, - Xx4(Xx4) = 8, - } +#[derive(Debug, Clone, PartialEq, Eq, ql_codec::Codec)] +#[codec(frozen, discriminants = HandshakeKind)] +#[repr(u8)] +pub enum QlHandshakeRecord { + Ik1(Ik1) = 1, + Ik2(Ik2) = 2, + Kk1(Ik1) = 3, + Kk2(Ik2) = 4, + Xx1(Xx1) = 5, + Xx2(Xx2) = 6, + Xx3(Xx3) = 7, + Xx4(Xx4) = 8, } -ql_codec::codec! { - #[derive(Debug, Clone, PartialEq, Eq)] - pub struct QlSessionRecord { - pub header: SessionHeader, - pub payload: EncryptedMessage, - } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QlSessionRecord { + pub header: SessionHeader, + pub payload: EncryptedMessage, } -impl QlSessionRecord { +impl> QlSessionRecord { pub fn into_owned(self) -> QlSessionRecord> { QlSessionRecord { header: self.header, @@ -90,3 +86,32 @@ impl QlSessionRecord { } } } + +impl Encode for QlSessionRecord { + fn encoded_len(&self) -> usize { + self.header.encoded_len() + self.payload.encoded_len() + } + + fn encode(&self, out: &mut Writer) { + self.header.encode(out); + self.payload.encode(out); + } +} + +impl Decode for QlSessionRecord> { + type Ref = QlSessionRecord; + + fn decode_ref(reader: &mut Reader) -> Result { + Ok(QlSessionRecord { + header: reader.decode()?, + payload: reader.decode_ref::>>()?, + }) + } + + fn from_ref(value: Self::Ref) -> Self { + Self { + header: value.header, + payload: > as Decode>::from_ref(value.payload), + } + } +} diff --git a/wire/src/tests.rs b/wire/src/tests.rs index c0b4652..a55ef2c 100644 --- a/wire/src/tests.rs +++ b/wire/src/tests.rs @@ -4,11 +4,14 @@ use ql_common::{ResetCode, StreamId, QID}; use super::*; fn decode_handshake_record(bytes: &[u8]) -> QlHandshakeRecord { - decode_record(bytes).unwrap().1 + decode_record::(bytes) + .unwrap() + .1 + .into_owned() } fn decode_session_record(bytes: &[u8]) -> QlSessionRecord> { - let (_, record) = decode_record::, _>(bytes).unwrap(); + let (_, record) = decode_record::>>(bytes).unwrap(); record.into_owned() } @@ -677,7 +680,7 @@ fn encrypted_session_record_round_trip_authenticates_header() { let record = encrypt_record(&crypto, session_route, header, &session_key, &body); let record_header = RecordHeader::new(session_route, RecordType::Session); - let bytes = encode_record_vec(record_header, &record); + let mut bytes = encode_record_vec(record_header, &record); assert_eq!( RecordHeader::decode_bytes(bytes.as_slice()).unwrap(), RecordHeader { @@ -725,6 +728,28 @@ fn encrypted_session_record_round_trip_authenticates_header() { ), Err(Error::DecryptFailed) ); + + let payload_start = + record_header.encoded_len() + header.encoded_len() + ENCRYPTED_MESSAGE_AUTH_SIZE; + let ciphertext_ptr = bytes[payload_start..].as_ptr(); + let mut reader = ql_codec::Reader::new(bytes.as_mut_slice()); + let decoded_header = reader.decode::().unwrap(); + let decoded = reader.decode_ref::>>().unwrap(); + assert!(reader.is_empty()); + assert_eq!(decoded.payload.ciphertext.as_ptr(), ciphertext_ptr); + let plaintext = encrypted::decrypt_record( + &crypto, + &decoded_header, + &decoded.header, + decoded.payload, + &session_key, + ) + .unwrap(); + assert_eq!(decode_session_frames(plaintext).unwrap(), body); + assert_eq!( + decode_session_frames(&bytes[payload_start..]).unwrap(), + body + ); } #[test] @@ -951,7 +976,6 @@ fn encoded_len_matches_encoding() { &RecordHeader::new(route(1, 2), RecordType::Handshake), ); assert_encoded_len("RecordType", &RecordType::Session); - assert_encoded_len("HandshakeKind", &HandshakeKind::Xx4); assert_encoded_len("SessionHeader", &SessionHeader { seq: RecordSeq(7) }); assert_encoded_len("ResetTarget", &ResetTarget::Both); assert_encoded_len(