105 lines
3.4 KiB
Rust
105 lines
3.4 KiB
Rust
use crate::{Policy, TransportSendStream};
|
|
use mtp_codec::{CommunicationValue, EncodeLimits};
|
|
use mtp_common::CommunicationError;
|
|
|
|
/// Classifies failures that may be recovered by replacing a persistent
|
|
/// application stream. Encoding and frame-size failures are deterministic and
|
|
/// must reach the caller without opening more streams.
|
|
pub(crate) struct RetryClassifier;
|
|
|
|
impl RetryClassifier {
|
|
pub(crate) fn retry_persistent_stream(error: &CommunicationError) -> bool {
|
|
matches!(
|
|
error,
|
|
CommunicationError::StreamError | CommunicationError::StreamClosed
|
|
)
|
|
}
|
|
}
|
|
|
|
/// Writes the canonical self-framed MTP value used by every transport.
|
|
///
|
|
/// `CommunicationValue` already begins with the four-byte body length. The
|
|
/// transport writes that representation directly so a frame does not carry a
|
|
/// redundant outer length prefix.
|
|
pub(crate) async fn write_frame<S: TransportSendStream>(
|
|
stream: &mut S,
|
|
value: &CommunicationValue,
|
|
policy: &Policy,
|
|
) -> Result<(), CommunicationError> {
|
|
let bytes = value
|
|
.to_bytes_with_limits(EncodeLimits::for_transport_message_size(
|
|
policy.max_message_size,
|
|
))
|
|
.map_err(|_| CommunicationError::Encode)?;
|
|
if bytes.len() as u64 > policy.max_message_size
|
|
|| bytes.len() as u64 >= policy.close_frame_len as u64
|
|
{
|
|
return Err(CommunicationError::MessageTooLarge);
|
|
}
|
|
stream.write_all(&bytes).await
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use async_trait::async_trait;
|
|
use mtp_codec::{CommunicationType, DataValue};
|
|
use std::pin::Pin;
|
|
use std::task::{Context, Poll};
|
|
use tokio::io::AsyncWrite;
|
|
|
|
#[derive(Default)]
|
|
struct BufferStream {
|
|
bytes: Vec<u8>,
|
|
}
|
|
|
|
impl AsyncWrite for BufferStream {
|
|
fn poll_write(
|
|
mut self: Pin<&mut Self>,
|
|
_cx: &mut Context<'_>,
|
|
buf: &[u8],
|
|
) -> Poll<std::io::Result<usize>> {
|
|
self.bytes.extend_from_slice(buf);
|
|
Poll::Ready(Ok(buf.len()))
|
|
}
|
|
|
|
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
|
|
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl TransportSendStream for BufferStream {
|
|
async fn write_all(&mut self, buf: &[u8]) -> Result<(), CommunicationError> {
|
|
self.bytes.extend_from_slice(buf);
|
|
Ok(())
|
|
}
|
|
|
|
async fn finish(&mut self) -> Result<(), CommunicationError> {
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn framing_preserves_a_generic_payload() {
|
|
let payload = DataValue::Array(vec![
|
|
DataValue::Str("payload".into()),
|
|
DataValue::Bytes(vec![1, 2, 3]),
|
|
]);
|
|
let frame = CommunicationValue::new(CommunicationType::Pong).with_payload(payload.clone());
|
|
let mut stream = BufferStream::default();
|
|
|
|
write_frame(&mut stream, &frame, &Policy::default())
|
|
.await
|
|
.unwrap();
|
|
|
|
let body_len = u32::from_be_bytes(stream.bytes[..4].try_into().unwrap()) as usize;
|
|
assert_eq!(body_len, stream.bytes.len() - 4);
|
|
let decoded = CommunicationValue::from_bytes(&stream.bytes).unwrap();
|
|
assert_eq!(decoded.into_payload(), payload);
|
|
}
|
|
}
|