xref: /aosp_15_r20/external/rust/beto-rust/nearby/presence/np_adv_dynamic/src/extended.rs

// Copyright 2023 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crypto_provider::CryptoProvider;
use np_adv::extended::salt::MultiSalt;
use np_adv::{extended::data_elements::*, extended::serialize::*, shared_data::*};
use sink::Sink;
use std::fmt::{Display, Formatter};

/// An advertisement builder for V1 advertisements where the
/// presence/absence of salt is determined at run-time instead of compile-time.
pub struct BoxedAdvBuilder {
    adv_builder: Box<AdvBuilder>,
}

impl From<AdvBuilder> for BoxedAdvBuilder {
    fn from(adv_builder: AdvBuilder) -> Self {
        let adv_builder = Box::new(adv_builder);
        BoxedAdvBuilder { adv_builder }
    }
}

/// Error possibly generated when attempting to add a section to
/// a BoxedAdvBuilder.
#[derive(Debug)]
pub enum BoxedAddSectionError {
    /// An error which was generated by the underlying AdvBuilder wrapped by the BoxedAdvBuilder
    Underlying(AddSectionError),
    /// An error generated when the boxed advertisement builder is unsalted, but the section
    /// identity requires salt.
    IdentityRequiresSaltError,
}

impl Display for BoxedAddSectionError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            BoxedAddSectionError::Underlying(u) => {
                write!(f, "{0}", u)
            }
            BoxedAddSectionError::IdentityRequiresSaltError => {
                write!(f, "Error generated when the BoxedAdvBuilder is unsalted, but the section identity requires salt.")
            }
        }
    }
}

impl std::error::Error for BoxedAddSectionError {}

impl From<AddSectionError> for BoxedAddSectionError {
    fn from(wrapped: AddSectionError) -> Self {
        BoxedAddSectionError::Underlying(wrapped)
    }
}

fn wrap_owning_section_builder<S: Into<BoxedSectionBuilder<AdvBuilder>>>(
    maybe_section_builder: Result<S, (AdvBuilder, AddSectionError)>,
) -> Result<BoxedSectionBuilder<AdvBuilder>, (BoxedAdvBuilder, BoxedAddSectionError)> {
    match maybe_section_builder {
        Ok(section_builder) => Ok(section_builder.into()),
        Err((adv_builder, err)) => Err((adv_builder.into(), err.into())),
    }
}

fn wrap_mut_ref_section_builder<'a, S: Into<BoxedSectionBuilder<&'a mut AdvBuilder>>>(
    maybe_section_builder: Result<S, AddSectionError>,
) -> Result<BoxedSectionBuilder<&'a mut AdvBuilder>, BoxedAddSectionError> {
    let section_builder = maybe_section_builder?;
    Ok(section_builder.into())
}

impl BoxedAdvBuilder {
    /// Gets the current count of sections which have been added
    /// to this advertisement builder. Does not count currently-
    /// outstanding section builders.
    pub fn section_count(&self) -> usize {
        self.adv_builder.section_count()
    }

    /// Create a section builder using the given identity,
    /// taking ownership of this advertisement builder.
    ///
    /// Unlike `BoxedAdvBuilder#section_builder`, the returned
    /// section builder will take ownership of this advertisement
    /// builder, if the operation was successful. Otherwise,
    /// this advertisement builder will be returned back to the
    /// caller unaltered as part of the `Err` arm.
    pub fn into_section_builder(
        self,
        identity: BoxedEncoder,
    ) -> Result<BoxedSectionBuilder<AdvBuilder>, (Self, BoxedAddSectionError)> {
        match identity {
            BoxedEncoder::Unencrypted => wrap_owning_section_builder(
                self.adv_builder.into_section_builder(UnencryptedSectionEncoder),
            ),
            BoxedEncoder::MicEncrypted(ident) => {
                wrap_owning_section_builder(self.adv_builder.into_section_builder(ident))
            }
            BoxedEncoder::SignedEncrypted(ident) => {
                wrap_owning_section_builder(self.adv_builder.into_section_builder(ident))
            }
        }
    }

    /// Create a section builder using the given encoder.
    ///
    /// Returns `Err` if the underlying advertisement builder
    /// yields an error when attempting to add a new section
    /// (typically because there's no more available adv space),
    /// or if the requested identity requires salt, and the
    /// advertisement builder is salt-less.
    pub fn section_builder(
        &mut self,
        encoder: BoxedEncoder,
    ) -> Result<BoxedSectionBuilder<&mut AdvBuilder>, BoxedAddSectionError> {
        match encoder {
            BoxedEncoder::Unencrypted => wrap_mut_ref_section_builder(
                self.adv_builder.section_builder(UnencryptedSectionEncoder),
            ),
            BoxedEncoder::MicEncrypted(ident) => {
                wrap_mut_ref_section_builder(self.adv_builder.section_builder(ident))
            }
            BoxedEncoder::SignedEncrypted(ident) => {
                wrap_mut_ref_section_builder(self.adv_builder.section_builder(ident))
            }
        }
    }

    /// Convert the builder into an encoded advertisement.
    pub fn into_advertisement(self) -> EncodedAdvertisement {
        self.adv_builder.into_advertisement()
    }
}

/// A wrapped v1 encoder whose type is given at run-time.
pub enum BoxedEncoder {
    /// Unencrypted encoder.
    Unencrypted,
    /// An encrypted encoder leveraging MIC for verification.
    MicEncrypted(MicEncryptedSectionEncoder<MultiSalt>),
    /// An encrypted encoder leveraging signatures for verification.
    SignedEncrypted(SignedEncryptedSectionEncoder),
}

/// A `SectionBuilder` whose corresponding Identity
/// and salted-ness is given at run-time instead of
/// at compile-time.
pub enum BoxedSectionBuilder<R: AsMut<AdvBuilder>> {
    /// A builder for a public section.
    Public(Box<SectionBuilder<R, UnencryptedSectionEncoder>>),
    /// A builder for a MIC-verified section.
    MicEncrypted(Box<SectionBuilder<R, MicEncryptedSectionEncoder<MultiSalt>>>),
    /// A builder for a signature-verified section.
    SignedEncrypted(Box<SectionBuilder<R, SignedEncryptedSectionEncoder>>),
}

impl BoxedSectionBuilder<AdvBuilder> {
    /// Gets the 0-based index of the section currently under construction
    /// in the context of the containing advertisement.
    pub fn section_index(&self) -> usize {
        match self {
            BoxedSectionBuilder::Public(x) => x.section_index(),
            BoxedSectionBuilder::MicEncrypted(x) => x.section_index(),
            BoxedSectionBuilder::SignedEncrypted(x) => x.section_index(),
        }
    }
    /// Add this builder to the advertisement that created it,
    /// returning the containing advertisement builder.
    pub fn add_to_advertisement<C: CryptoProvider>(self) -> BoxedAdvBuilder {
        let adv_builder = match self {
            BoxedSectionBuilder::Public(x) => x.add_to_advertisement::<C>(),
            BoxedSectionBuilder::MicEncrypted(x) => x.add_to_advertisement::<C>(),
            BoxedSectionBuilder::SignedEncrypted(x) => x.add_to_advertisement::<C>(),
        };
        BoxedAdvBuilder::from(adv_builder)
    }
}

impl<'a> BoxedSectionBuilder<&'a mut AdvBuilder> {
    /// Add this builder to the advertisement that created it.
    pub fn add_to_advertisement<C: CryptoProvider>(self) {
        match self {
            BoxedSectionBuilder::Public(x) => x.add_to_advertisement::<C>(),
            BoxedSectionBuilder::MicEncrypted(x) => x.add_to_advertisement::<C>(),
            BoxedSectionBuilder::SignedEncrypted(x) => x.add_to_advertisement::<C>(),
        }
    }
}

impl<R: AsMut<AdvBuilder>> BoxedSectionBuilder<R> {
    /// Returns true if this wraps a section builder which
    /// leverages some encrypted identity.
    pub fn is_encrypted(&self) -> bool {
        match self {
            BoxedSectionBuilder::Public(_) => false,
            BoxedSectionBuilder::MicEncrypted(_) => true,
            BoxedSectionBuilder::SignedEncrypted(_) => true,
        }
    }
    /// Gets the derived salt of the next DE to be added to the section,
    /// if this section-builder corresponds to an encrypted section that can
    /// provide per-DE salts.
    /// Otherwise, returns nothing.
    ///
    /// Suitable for scenarios (like FFI) where a closure would be inappropriate
    /// for DE construction, and interaction with the client is preferred.
    pub fn next_de_salt(&self) -> Option<DeSalt> {
        match self {
            BoxedSectionBuilder::Public(_) => None,
            BoxedSectionBuilder::MicEncrypted(x) => x.next_de_salt(),
            BoxedSectionBuilder::SignedEncrypted(x) => Some(x.next_de_salt()),
        }
    }

    /// Add a data element to the section with a closure that returns a `Result`.
    ///
    /// The provided `build_de` closure will be invoked with the derived salt for this DE,
    /// if any salt has been specified for the surrounding advertisement.
    pub fn add_de_res<E>(
        &mut self,
        build_de: impl FnOnce(Option<DeSalt>) -> Result<BoxedWriteDataElement, E>,
    ) -> Result<(), AddDataElementError<E>> {
        match self {
            BoxedSectionBuilder::Public(x) => {
                let build_de_modified = |()| build_de(None);
                x.add_de_res(build_de_modified)
            }
            BoxedSectionBuilder::MicEncrypted(x) => x.add_de_res(build_de),
            BoxedSectionBuilder::SignedEncrypted(x) => {
                let build_de_modified = |de_salt: DeSalt| build_de(Some(de_salt));
                x.add_de_res(build_de_modified)
            }
        }
    }
    /// Like add_de_res, but for infalliable closures
    pub fn add_de(
        &mut self,
        build_de: impl FnOnce(Option<DeSalt>) -> BoxedWriteDataElement,
    ) -> Result<(), AddDataElementError<()>> {
        self.add_de_res(|derived_salt| Ok::<_, ()>(build_de(derived_salt)))
    }
}

impl<R: AsMut<AdvBuilder>> From<SectionBuilder<R, UnencryptedSectionEncoder>>
    for BoxedSectionBuilder<R>
{
    fn from(section_builder: SectionBuilder<R, UnencryptedSectionEncoder>) -> Self {
        BoxedSectionBuilder::Public(Box::new(section_builder))
    }
}

impl<R: AsMut<AdvBuilder>> From<SectionBuilder<R, MicEncryptedSectionEncoder<MultiSalt>>>
    for BoxedSectionBuilder<R>
{
    fn from(section_builder: SectionBuilder<R, MicEncryptedSectionEncoder<MultiSalt>>) -> Self {
        BoxedSectionBuilder::MicEncrypted(Box::new(section_builder))
    }
}

impl<R: AsMut<AdvBuilder>> From<SectionBuilder<R, SignedEncryptedSectionEncoder>>
    for BoxedSectionBuilder<R>
{
    fn from(section_builder: SectionBuilder<R, SignedEncryptedSectionEncoder>) -> Self {
        BoxedSectionBuilder::SignedEncrypted(Box::new(section_builder))
    }
}

/// Mutable trait object reference to a `Sink<u8>`
pub struct DynSink<'a> {
    wrapped: &'a mut dyn Sink<u8>,
}

impl<'a> Sink<u8> for DynSink<'a> {
    fn try_extend_from_slice(&mut self, items: &[u8]) -> Option<()> {
        self.wrapped.try_extend_from_slice(items)
    }
    fn try_push(&mut self, item: u8) -> Option<()> {
        self.wrapped.try_push(item)
    }
}

impl<'a> From<&'a mut dyn Sink<u8>> for DynSink<'a> {
    fn from(wrapped: &'a mut dyn Sink<u8>) -> Self {
        DynSink { wrapped }
    }
}

/// A version of the WriteDataElement trait which is object-safe
pub trait DynWriteDataElement {
    /// Gets the data-element header for the data element
    fn de_header(&self) -> DeHeader;
    /// Writes the contents of the DE payload to the given DynSink.
    /// Returns Some(()) if the write operation was successful,
    /// and None if it was unsuccessful
    fn write_de_contents(&self, sink: DynSink) -> Option<()>;
}

impl<T: WriteDataElement> DynWriteDataElement for T {
    fn de_header(&self) -> DeHeader {
        WriteDataElement::de_header(self)
    }
    fn write_de_contents(&self, mut sink: DynSink) -> Option<()> {
        WriteDataElement::write_de_contents(self, &mut sink)
    }
}

/// Trait object wrapper for DynWriteDataElement instances
pub struct BoxedWriteDataElement {
    wrapped: Box<dyn DynWriteDataElement>,
}

impl BoxedWriteDataElement {
    /// Constructs a new `BoxedWriteDataElement` from a `WriteDataElement`
    /// whose trait impl is valid for a `'static` lifetime.
    pub fn new<D: WriteDataElement + 'static>(wrapped: D) -> Self {
        let wrapped = Box::new(wrapped);
        Self { wrapped }
    }
}

impl WriteDataElement for BoxedWriteDataElement {
    fn de_header(&self) -> DeHeader {
        self.wrapped.de_header()
    }
    fn write_de_contents<S: Sink<u8>>(&self, sink: &mut S) -> Option<()> {
        let sink: &mut dyn Sink<u8> = sink;
        let dyn_sink = DynSink::from(sink);
        self.wrapped.write_de_contents(dyn_sink)
    }
}

impl From<TxPower> for BoxedWriteDataElement {
    fn from(tx_power: TxPower) -> Self {
        BoxedWriteDataElement::new::<TxPowerDataElement>(tx_power.into())
    }
}

impl From<ContextSyncSeqNum> for BoxedWriteDataElement {
    fn from(context_sync_sequence_num: ContextSyncSeqNum) -> Self {
        BoxedWriteDataElement::new::<ContextSyncSeqNumDataElement>(context_sync_sequence_num.into())
    }
}

impl From<ActionsDataElement> for BoxedWriteDataElement {
    fn from(data: ActionsDataElement) -> Self {
        BoxedWriteDataElement::new(data)
    }
}