use prost::Message;
pub use bb_ir::proto::bb_core::{CorrelationKind, SlotFill, WireCorrelation, WireEnvelope};
pub const ENVELOPE_SCHEMA_VERSION: u32 = 1;
pub const SUPPORTED_SCHEMA_VERSIONS: &[u32] = &[ENVELOPE_SCHEMA_VERSION];
#[derive(Clone, Copy, Debug)]
pub struct EnvelopeCaps {
pub max_total_bytes: usize,
pub max_slot_fills: usize,
pub max_per_fill_bytes: usize,
pub max_dest_suffix_bytes: usize,
pub max_src_peer_addresses: usize,
pub max_src_peer_address_bytes: usize,
}
impl Default for EnvelopeCaps {
fn default() -> Self {
Self {
max_total_bytes: 16 * 1024 * 1024,
max_slot_fills: 256,
max_per_fill_bytes: 4 * 1024 * 1024,
max_dest_suffix_bytes: 4 * 1024,
max_src_peer_addresses: 8,
max_src_peer_address_bytes: 256,
}
}
}
impl EnvelopeCaps {
pub fn edge() -> Self {
Self {
max_total_bytes: 256 * 1024,
max_slot_fills: 16,
max_per_fill_bytes: 64 * 1024,
max_dest_suffix_bytes: 512,
max_src_peer_addresses: 4,
max_src_peer_address_bytes: 256,
}
}
}
#[derive(Debug)]
pub enum EnvelopeDecodeError {
Malformed(prost::DecodeError),
OversizeEnvelope {
cap_bytes: usize,
got_bytes: usize,
},
OversizeSlotFill {
which: &'static str,
cap_bytes: usize,
got_bytes: usize,
},
TooManySlotFills {
cap: usize,
got: usize,
},
VersionMismatch {
got: u32,
supported: &'static [u32],
},
TooManySrcPeerAddresses {
cap: usize,
got: usize,
},
OversizeSrcPeerAddress {
cap_bytes: usize,
got_bytes: usize,
},
}
impl std::fmt::Display for EnvelopeDecodeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Malformed(e) => write!(f, "malformed envelope bytes: {e}"),
Self::OversizeEnvelope {
cap_bytes,
got_bytes,
} => write!(
f,
"envelope buffer too large: cap={cap_bytes} got={got_bytes}",
),
Self::OversizeSlotFill {
which,
cap_bytes,
got_bytes,
} => write!(
f,
"slot fill {which} exceeds cap: cap={cap_bytes} got={got_bytes}",
),
Self::TooManySlotFills { cap, got } => write!(
f,
"envelope carries too many slot fills: cap={cap} got={got}",
),
Self::VersionMismatch { got, supported } => write!(
f,
"envelope schema_version mismatch: got={got} supported={supported:?}",
),
Self::TooManySrcPeerAddresses { cap, got } => write!(
f,
"envelope carries too many src_peer_addresses: cap={cap} got={got}",
),
Self::OversizeSrcPeerAddress {
cap_bytes,
got_bytes,
} => write!(
f,
"src_peer_addresses entry too large: cap={cap_bytes} got={got_bytes}",
),
}
}
}
impl std::error::Error for EnvelopeDecodeError {}
pub struct EnvelopeCodec;
impl EnvelopeCodec {
pub fn encode(env: &WireEnvelope) -> Vec<u8> {
env.encode_to_vec()
}
pub fn decode_capped(
bytes: &[u8],
caps: &EnvelopeCaps,
) -> Result<WireEnvelope, EnvelopeDecodeError> {
if bytes.len() > caps.max_total_bytes {
return Err(EnvelopeDecodeError::OversizeEnvelope {
cap_bytes: caps.max_total_bytes,
got_bytes: bytes.len(),
});
}
let env = WireEnvelope::decode(bytes).map_err(EnvelopeDecodeError::Malformed)?;
if !SUPPORTED_SCHEMA_VERSIONS.contains(&env.schema_version) && env.schema_version != 0 {
return Err(EnvelopeDecodeError::VersionMismatch {
got: env.schema_version,
supported: SUPPORTED_SCHEMA_VERSIONS,
});
}
if env.fills.len() > caps.max_slot_fills {
return Err(EnvelopeDecodeError::TooManySlotFills {
cap: caps.max_slot_fills,
got: env.fills.len(),
});
}
for fill in &env.fills {
if fill.payload.len() > caps.max_per_fill_bytes {
return Err(EnvelopeDecodeError::OversizeSlotFill {
which: "payload",
cap_bytes: caps.max_per_fill_bytes,
got_bytes: fill.payload.len(),
});
}
if fill.dest_suffix.len() > caps.max_dest_suffix_bytes {
return Err(EnvelopeDecodeError::OversizeSlotFill {
which: "dest_suffix",
cap_bytes: caps.max_dest_suffix_bytes,
got_bytes: fill.dest_suffix.len(),
});
}
}
if env.src_peer_addresses.len() > caps.max_src_peer_addresses {
return Err(EnvelopeDecodeError::TooManySrcPeerAddresses {
cap: caps.max_src_peer_addresses,
got: env.src_peer_addresses.len(),
});
}
for addr in &env.src_peer_addresses {
if addr.len() > caps.max_src_peer_address_bytes {
return Err(EnvelopeDecodeError::OversizeSrcPeerAddress {
cap_bytes: caps.max_src_peer_address_bytes,
got_bytes: addr.len(),
});
}
}
Ok(env)
}
}