use std::io;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::sync::Notify;
use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
pub type PeerStream = Compat<yamux::Stream>;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct AvailabilityItem {
pub store_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub root: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub retrieval_key: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct AvailabilityRequest {
pub items: Vec<AvailabilityItem>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct AvailabilityAnswer {
pub available: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub roots: Option<Vec<String>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub total_length: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub chunk_count: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub complete: Option<bool>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct AvailabilityResponse {
pub items: Vec<AvailabilityAnswer>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct RangeRequest {
pub store_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub retrieval_key: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub root: Option<String>,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub capsule: bool,
#[serde(default)]
pub offset: u64,
pub length: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub skip_layout: Option<bool>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct RangeFrame {
pub offset: u64,
pub length: u64,
#[serde(with = "base64_bytes")]
pub bytes: Vec<u8>,
pub complete: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub total_length: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub chunk_lens: Option<Vec<u64>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub chunk_index: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub inclusion_proof: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub root: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub chunk_count: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub chunk_lens_offset: Option<u64>,
}
mod base64_bytes {
use base64::Engine as _;
use serde::de::{SeqAccess, Visitor};
use serde::{Deserializer, Serializer};
use std::fmt;
pub fn serialize<S: Serializer>(bytes: &[u8], s: S) -> Result<S::Ok, S::Error> {
s.serialize_str(&base64::engine::general_purpose::STANDARD.encode(bytes))
}
pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<u8>, D::Error> {
d.deserialize_any(Base64OrArray)
}
struct Base64OrArray;
impl<'de> Visitor<'de> for Base64OrArray {
type Value = Vec<u8>;
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("base64-encoded ciphertext (or a legacy byte array)")
}
fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<Self::Value, E> {
base64::engine::general_purpose::STANDARD
.decode(v)
.map_err(|e| E::custom(format!("range frame bytes are not valid base64: {e}")))
}
fn visit_bytes<E: serde::de::Error>(self, v: &[u8]) -> Result<Self::Value, E> {
Ok(v.to_vec())
}
fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value, A::Error> {
let mut out = Vec::with_capacity(seq.size_hint().unwrap_or_default());
while let Some(b) = seq.next_element::<u8>()? {
out.push(b);
}
Ok(out)
}
}
}
impl AvailabilityItem {
pub fn store(store_id: impl Into<String>) -> Self {
AvailabilityItem {
store_id: store_id.into(),
root: None,
retrieval_key: None,
}
}
pub fn with_root(mut self, root: impl Into<String>) -> Self {
self.root = Some(root.into());
self
}
pub fn with_retrieval_key(mut self, retrieval_key: impl Into<String>) -> Self {
self.retrieval_key = Some(retrieval_key.into());
self
}
}
impl AvailabilityRequest {
pub fn new(items: Vec<AvailabilityItem>) -> Self {
AvailabilityRequest { items }
}
pub fn encode(&self) -> io::Result<Vec<u8>> {
encode_framed(self)
}
pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
decode_framed(r).await
}
}
impl AvailabilityAnswer {
pub fn available() -> Self {
AvailabilityAnswer::with_availability(true)
}
pub fn unavailable() -> Self {
AvailabilityAnswer::with_availability(false)
}
fn with_availability(available: bool) -> Self {
AvailabilityAnswer {
available,
roots: None,
total_length: None,
chunk_count: None,
complete: None,
}
}
pub fn with_roots(mut self, roots: Vec<String>) -> Self {
self.roots = Some(roots);
self
}
pub fn with_total_length(mut self, total_length: u64) -> Self {
self.total_length = Some(total_length);
self
}
pub fn with_chunk_count(mut self, chunk_count: u64) -> Self {
self.chunk_count = Some(chunk_count);
self
}
pub fn with_complete(mut self, complete: bool) -> Self {
self.complete = Some(complete);
self
}
}
impl AvailabilityResponse {
pub fn new(items: Vec<AvailabilityAnswer>) -> Self {
AvailabilityResponse { items }
}
pub fn encode(&self) -> io::Result<Vec<u8>> {
encode_framed(self)
}
pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
decode_framed(r).await
}
}
impl RangeRequest {
pub fn resource(
store_id: impl Into<String>,
retrieval_key: impl Into<String>,
offset: u64,
length: u64,
) -> Self {
RangeRequest::resource_or_capsule(
store_id,
Some(retrieval_key.into()),
false,
offset,
length,
)
}
pub fn capsule(store_id: impl Into<String>, offset: u64, length: u64) -> Self {
RangeRequest::resource_or_capsule(store_id, None, true, offset, length)
}
fn resource_or_capsule(
store_id: impl Into<String>,
retrieval_key: Option<String>,
capsule: bool,
offset: u64,
length: u64,
) -> Self {
RangeRequest {
store_id: store_id.into(),
retrieval_key,
root: None,
capsule,
offset,
length,
skip_layout: None,
}
}
pub fn with_root(mut self, root: impl Into<String>) -> Self {
self.root = Some(root.into());
self
}
pub fn with_skip_layout(mut self, skip_layout: bool) -> Self {
self.skip_layout = Some(skip_layout);
self
}
pub fn encode(&self) -> io::Result<Vec<u8>> {
encode_framed(self)
}
pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Self> {
decode_framed(r).await
}
}
impl RangeFrame {
pub fn data(offset: u64, bytes: Vec<u8>) -> Self {
RangeFrame {
offset,
length: bytes.len() as u64,
bytes,
complete: false,
total_length: None,
chunk_lens: None,
chunk_index: None,
inclusion_proof: None,
root: None,
chunk_count: None,
chunk_lens_offset: None,
}
}
pub fn with_complete(mut self, complete: bool) -> Self {
self.complete = complete;
self
}
pub fn with_identity(
mut self,
root: impl Into<String>,
total_length: u64,
chunk_count: u64,
) -> Self {
self.root = Some(root.into());
self.total_length = Some(total_length);
self.chunk_count = Some(chunk_count);
self
}
pub fn with_chunk_lens_page(mut self, chunk_lens_offset: u64, chunk_lens: Vec<u64>) -> Self {
self.chunk_lens_offset = Some(chunk_lens_offset);
self.chunk_lens = Some(chunk_lens);
self
}
pub fn split_chunk_lens_pages(chunk_lens: &[u64]) -> Vec<(u64, Vec<u64>)> {
chunk_lens
.chunks(MAX_CHUNK_LENS_PER_FRAME)
.enumerate()
.map(|(page, entries)| ((page * MAX_CHUNK_LENS_PER_FRAME) as u64, entries.to_vec()))
.collect()
}
pub fn with_chunk_index(mut self, chunk_index: u64) -> Self {
self.chunk_index = Some(chunk_index);
self
}
pub fn with_inclusion_proof(mut self, inclusion_proof: impl Into<String>) -> Self {
self.inclusion_proof = Some(inclusion_proof.into());
self
}
pub fn with_declared_length(mut self, length: u64) -> Self {
self.length = length;
self
}
pub fn encode(&self) -> io::Result<Vec<u8>> {
if self.bytes.len() > MAX_RANGE_FRAME_PAYLOAD {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"RangeFrame payload {} exceeds MAX_RANGE_FRAME_PAYLOAD {MAX_RANGE_FRAME_PAYLOAD}; \
split the range into ceiling-sized frames",
self.bytes.len()
),
));
}
if let Some(proof) = &self.inclusion_proof {
if proof.len() > MAX_INCLUSION_PROOF_B64 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"RangeFrame inclusion_proof {} exceeds MAX_INCLUSION_PROOF_B64 \
{MAX_INCLUSION_PROOF_B64}; the resource has no conforming range stream, so \
answer RANGE_METADATA_UNREPRESENTABLE rather than streaming frames",
proof.len()
),
));
}
}
encode_framed(self)
}
pub async fn decode<R: AsyncRead + Unpin>(r: &mut R) -> io::Result<Option<Self>> {
decode_framed_opt(r).await
}
}
pub const MAX_FRAMED_BODY: usize = 64 * 1024;
pub const MAX_RANGE_FRAME_PAYLOAD: usize = 32 * 1024;
pub const MAX_INCLUSION_PROOF_B64: usize = 4096;
pub const MAX_CHUNK_LENS_PER_FRAME: usize = 2048;
pub const MAX_FIRST_FRAME_CHUNK_LENS: usize = 2_486;
pub const MAX_RESOURCE_CHUNK_COUNT: usize = 1_048_576;
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum ChunkLensError {
#[error(
"declared chunk_count {chunk_count} exceeds MAX_RESOURCE_CHUNK_COUNT {MAX_RESOURCE_CHUNK_COUNT}"
)]
ChunkCountTooLarge {
chunk_count: usize,
},
#[error("cannot reserve a {chunk_count}-entry chunk_lens array ({bytes} bytes)")]
AllocationFailed {
chunk_count: usize,
bytes: usize,
},
#[error("chunk_lens page at offset {offset} is empty")]
EmptyPage {
offset: u64,
},
#[error(
"chunk_lens page of {entries} entries exceeds MAX_CHUNK_LENS_PER_FRAME {MAX_CHUNK_LENS_PER_FRAME}"
)]
PageTooLarge {
entries: usize,
},
#[error(
"chunk_lens_offset {offset} is not a multiple of MAX_CHUNK_LENS_PER_FRAME {MAX_CHUNK_LENS_PER_FRAME}"
)]
MisalignedOffset {
offset: u64,
},
#[error("chunk_lens_offset {offset} is beyond the declared chunk_count {chunk_count}")]
OffsetOutOfRange {
offset: u64,
chunk_count: usize,
},
#[error(
"chunk_lens page of {entries} entries at offset {offset} extends past the declared chunk_count \
{chunk_count}"
)]
PageExtendsPastEnd {
offset: u64,
entries: usize,
chunk_count: usize,
},
#[error(
"chunk_lens page at offset {offset} has {entries} entries; this page must have exactly {expected}"
)]
UnexpectedPageLength {
offset: u64,
entries: usize,
expected: usize,
},
#[error("a chunk_lens page at offset {offset} was already accepted")]
DuplicatePage {
offset: u64,
},
#[error("incomplete chunk_lens prologue: {have} of {want} entries")]
Incomplete {
have: usize,
want: usize,
},
}
#[derive(Debug, Clone)]
pub struct ChunkLensAssembler {
lens: Vec<u64>,
filled: Vec<bool>,
filled_pages: usize,
}
impl ChunkLensAssembler {
pub fn new(chunk_count: usize) -> Result<Self, ChunkLensError> {
if chunk_count > MAX_RESOURCE_CHUNK_COUNT {
return Err(ChunkLensError::ChunkCountTooLarge { chunk_count });
}
let mut lens = Vec::new();
lens.try_reserve_exact(chunk_count)
.map_err(|_| ChunkLensError::AllocationFailed {
chunk_count,
bytes: chunk_count * std::mem::size_of::<u64>(),
})?;
lens.resize(chunk_count, 0);
let page_count = chunk_count.div_ceil(MAX_CHUNK_LENS_PER_FRAME);
let mut filled = Vec::new();
filled
.try_reserve_exact(page_count)
.map_err(|_| ChunkLensError::AllocationFailed {
chunk_count,
bytes: page_count,
})?;
filled.resize(page_count, false);
Ok(ChunkLensAssembler {
lens,
filled,
filled_pages: 0,
})
}
pub fn accept_page(&mut self, offset: u64, page: &[u64]) -> Result<(), ChunkLensError> {
let chunk_count = self.lens.len();
let page_size = MAX_CHUNK_LENS_PER_FRAME as u64;
if page.is_empty() {
return Err(ChunkLensError::EmptyPage { offset });
}
if page.len() > MAX_CHUNK_LENS_PER_FRAME {
return Err(ChunkLensError::PageTooLarge {
entries: page.len(),
});
}
if offset % page_size != 0 {
return Err(ChunkLensError::MisalignedOffset { offset });
}
if offset >= chunk_count as u64 {
return Err(ChunkLensError::OffsetOutOfRange {
offset,
chunk_count,
});
}
let start = offset as usize;
let end = start + page.len();
if end > chunk_count {
return Err(ChunkLensError::PageExtendsPastEnd {
offset,
entries: page.len(),
chunk_count,
});
}
let expected = MAX_CHUNK_LENS_PER_FRAME.min(chunk_count - start);
if page.len() != expected {
return Err(ChunkLensError::UnexpectedPageLength {
offset,
entries: page.len(),
expected,
});
}
let slot = start / MAX_CHUNK_LENS_PER_FRAME;
if self.filled[slot] {
return Err(ChunkLensError::DuplicatePage { offset });
}
self.lens[start..end].copy_from_slice(page);
self.filled[slot] = true;
self.filled_pages += 1;
Ok(())
}
pub fn is_complete(&self) -> bool {
self.filled_pages == self.filled.len()
}
pub fn into_chunk_lens(self) -> Result<Vec<u64>, ChunkLensError> {
if !self.is_complete() {
let have = self
.filled
.iter()
.enumerate()
.filter(|(_, filled)| **filled)
.map(|(slot, _)| {
let start = slot * MAX_CHUNK_LENS_PER_FRAME;
MAX_CHUNK_LENS_PER_FRAME.min(self.lens.len() - start)
})
.sum();
return Err(ChunkLensError::Incomplete {
have,
want: self.lens.len(),
});
}
Ok(self.lens)
}
}
fn encode_framed<T: Serialize>(value: &T) -> io::Result<Vec<u8>> {
let body =
serde_json::to_vec(value).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
if body.len() > MAX_FRAMED_BODY {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"framed body {} exceeds MAX_FRAMED_BODY {MAX_FRAMED_BODY}; split the payload on \
MAX_RANGE_FRAME_PAYLOAD ({MAX_RANGE_FRAME_PAYLOAD})",
body.len()
),
));
}
let mut out = Vec::with_capacity(4 + body.len());
out.extend_from_slice(&(body.len() as u32).to_be_bytes());
out.extend_from_slice(&body);
Ok(out)
}
async fn decode_framed<T: for<'de> Deserialize<'de>, R: AsyncRead + Unpin>(
r: &mut R,
) -> io::Result<T> {
let mut len_buf = [0u8; 4];
r.read_exact(&mut len_buf).await?;
let len = u32::from_be_bytes(len_buf) as usize;
if len > MAX_FRAMED_BODY {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"control message too large",
));
}
let mut body = vec![0u8; len];
r.read_exact(&mut body).await?;
serde_json::from_slice(&body).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
}
async fn decode_framed_opt<T: for<'de> Deserialize<'de>, R: AsyncRead + Unpin>(
r: &mut R,
) -> io::Result<Option<T>> {
let mut len_buf = [0u8; 4];
match r.read_exact(&mut len_buf).await {
Ok(_) => {}
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
Err(e) => return Err(e),
}
let len = u32::from_be_bytes(len_buf) as usize;
if len > MAX_FRAMED_BODY {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"control message too large",
));
}
let mut body = vec![0u8; len];
r.read_exact(&mut body).await?;
serde_json::from_slice(&body)
.map(Some)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
}
enum MuxCommand {
OpenOutbound(tokio::sync::oneshot::Sender<Result<yamux::Stream, String>>),
}
pub struct PeerSession {
cmd_tx: tokio::sync::mpsc::Sender<MuxCommand>,
inbound_rx: tokio::sync::mpsc::Receiver<PeerStream>,
closed_flag: Arc<AtomicBool>,
closed_notify: Arc<Notify>,
}
#[derive(Clone)]
pub struct ClosedHandle {
flag: Arc<AtomicBool>,
notify: Arc<Notify>,
}
impl ClosedHandle {
pub fn is_closed(&self) -> bool {
self.flag.load(Ordering::Acquire)
}
pub async fn closed(&self) {
loop {
if self.flag.load(Ordering::Acquire) {
return;
}
let notified = self.notify.notified();
if self.flag.load(Ordering::Acquire) {
return;
}
notified.await;
}
}
}
impl std::fmt::Debug for PeerSession {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PeerSession").finish_non_exhaustive()
}
}
impl PeerSession {
pub fn client<S>(io: S) -> Self
where
S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
{
Self::new(io, yamux::Mode::Client)
}
pub fn server<S>(io: S) -> Self
where
S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
{
Self::new(io, yamux::Mode::Server)
}
fn new<S>(io: S, mode: yamux::Mode) -> Self
where
S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
{
let (cmd_tx, cmd_rx) = tokio::sync::mpsc::channel::<MuxCommand>(64);
let (inbound_tx, inbound_rx) = tokio::sync::mpsc::channel::<PeerStream>(64);
let conn = yamux::Connection::new(io.compat(), yamux::Config::default(), mode);
let closed_flag = Arc::new(AtomicBool::new(false));
let closed_notify = Arc::new(Notify::new());
tokio::spawn(drive_connection(
conn,
cmd_rx,
inbound_tx,
Arc::clone(&closed_flag),
Arc::clone(&closed_notify),
));
PeerSession {
cmd_tx,
inbound_rx,
closed_flag,
closed_notify,
}
}
pub fn closed_handle(&self) -> ClosedHandle {
ClosedHandle {
flag: Arc::clone(&self.closed_flag),
notify: Arc::clone(&self.closed_notify),
}
}
pub async fn open_stream(&mut self) -> io::Result<PeerStream> {
let (tx, rx) = tokio::sync::oneshot::channel();
self.cmd_tx
.send(MuxCommand::OpenOutbound(tx))
.await
.map_err(|_| io::Error::other("mux driver closed"))?;
let stream = rx
.await
.map_err(|_| io::Error::other("mux driver dropped request"))?
.map_err(io::Error::other)?;
Ok(stream.compat())
}
pub async fn accept_stream(&mut self) -> Option<PeerStream> {
self.inbound_rx.recv().await
}
pub async fn open_range_stream(&mut self, req: &RangeRequest) -> io::Result<PeerStream> {
let mut stream = self.open_stream().await?;
stream.write_all(&req.encode()?).await?;
stream.flush().await?;
Ok(stream)
}
pub async fn query_availability(
&mut self,
items: Vec<AvailabilityItem>,
) -> io::Result<AvailabilityResponse> {
let req = AvailabilityRequest { items };
let mut stream = self.open_stream().await?;
stream.write_all(&req.encode()?).await?;
stream.flush().await?;
AvailabilityResponse::decode(&mut stream).await
}
}
async fn drive_connection<T>(
mut conn: yamux::Connection<T>,
mut cmd_rx: tokio::sync::mpsc::Receiver<MuxCommand>,
inbound_tx: tokio::sync::mpsc::Sender<PeerStream>,
closed_flag: Arc<AtomicBool>,
closed_notify: Arc<Notify>,
) where
T: futures::AsyncRead + futures::AsyncWrite + Send + Unpin + 'static,
{
use std::future::poll_fn;
loop {
tokio::select! {
cmd = cmd_rx.recv() => {
match cmd {
Some(MuxCommand::OpenOutbound(reply)) => {
let res = poll_fn(|cx| conn.poll_new_outbound(cx)).await;
let _ = reply.send(res.map_err(|e| e.to_string()));
}
None => {
let _ = poll_fn(|cx| conn.poll_close(cx)).await;
break;
}
}
}
inbound = poll_fn(|cx| conn.poll_next_inbound(cx)) => {
match inbound {
Some(Ok(stream)) => {
let _ = inbound_tx.try_send(stream.compat());
}
Some(Err(_)) | None => {
break;
}
}
}
}
}
closed_flag.store(true, Ordering::Release);
closed_notify.notify_waiters();
}