[Fix] Harden MTP codec, transport, and SDK security
This commit is contained in:
Alex Emmet
parent
188caf56cc
commit
a7e804c603
73 changed files with 11892 additions and 5756 deletions
|
|
@ -1,12 +1,14 @@
|
|||
use crate::ConnectionHandle;
|
||||
use crate::framing::RetryClassifier;
|
||||
#[cfg(feature = "pipes")]
|
||||
use crate::pipe::PipeReader;
|
||||
use mtp_codec::{CommunicationValue, DecodeLimits, TypeMap};
|
||||
use mtp_codec::{CommunicationValue, DecodeError, DecodeLimits, EncodeLimits, TypeMap};
|
||||
use mtp_common::CommunicationError;
|
||||
use std::ops::Deref;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use tokio::sync::{Mutex, Notify, RwLock, Semaphore, mpsc};
|
||||
use tokio::time::{Duration, sleep, timeout};
|
||||
use tokio::time::{Duration, Instant, sleep, timeout, timeout_at};
|
||||
use tracing::{debug, info, instrument, trace, warn};
|
||||
use wtransport::Connection;
|
||||
|
||||
|
|
@ -19,6 +21,63 @@ pub enum TransportEvent<R = wtransport::RecvStream> {
|
|||
|
||||
const APPLICATION_CLOSE_REASON: &str = "mtp-close";
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum DecodeRejectionClass {
|
||||
Malformed,
|
||||
ResourceLimit,
|
||||
DuplicateField,
|
||||
}
|
||||
|
||||
pub fn classify_decode_error(error: &DecodeError) -> DecodeRejectionClass {
|
||||
match error {
|
||||
DecodeError::MalformedEncoding => DecodeRejectionClass::Malformed,
|
||||
DecodeError::DepthLimit
|
||||
| DecodeError::ValueCountLimit
|
||||
| DecodeError::BlobLimit
|
||||
| DecodeError::AllocationLimit
|
||||
| DecodeError::RecipientLimit => DecodeRejectionClass::ResourceLimit,
|
||||
DecodeError::DuplicateField => DecodeRejectionClass::DuplicateField,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct DecodeRejectionCounters {
|
||||
malformed: AtomicU64,
|
||||
resource_limit: AtomicU64,
|
||||
duplicate_field: AtomicU64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct DecodeRejectionCounts {
|
||||
pub malformed: u64,
|
||||
pub resource_limit: u64,
|
||||
pub duplicate_field: u64,
|
||||
}
|
||||
|
||||
impl DecodeRejectionCounters {
|
||||
pub(crate) fn record(&self, error: &DecodeError) {
|
||||
match classify_decode_error(error) {
|
||||
DecodeRejectionClass::Malformed => {
|
||||
self.malformed.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
DecodeRejectionClass::ResourceLimit => {
|
||||
self.resource_limit.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
DecodeRejectionClass::DuplicateField => {
|
||||
self.duplicate_field.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn snapshot(&self) -> DecodeRejectionCounts {
|
||||
DecodeRejectionCounts {
|
||||
malformed: self.malformed.load(Ordering::Relaxed),
|
||||
resource_limit: self.resource_limit.load(Ordering::Relaxed),
|
||||
duplicate_field: self.duplicate_field.load(Ordering::Relaxed),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum SendMode {
|
||||
PersistentStream,
|
||||
|
|
@ -135,27 +194,62 @@ impl Policy {
|
|||
}
|
||||
}
|
||||
|
||||
/// A validated policy snapshot used after a public [`Policy`] crosses into a
|
||||
/// transport implementation. `Policy` intentionally remains a plain public
|
||||
/// struct for source compatibility, so callers can construct it directly and
|
||||
/// bypass builder methods. Every transport constructor takes this snapshot
|
||||
/// before creating channels or semaphores.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct RuntimePolicy(Policy);
|
||||
|
||||
impl RuntimePolicy {
|
||||
pub(crate) fn from_public(policy: &Policy) -> Self {
|
||||
let mut policy = *policy;
|
||||
policy.receiver_queue_capacity = policy.receiver_queue_capacity.max(1);
|
||||
policy.max_concurrent_stream_tasks = policy.max_concurrent_stream_tasks.max(1);
|
||||
Self(policy)
|
||||
}
|
||||
}
|
||||
|
||||
impl Deref for RuntimePolicy {
|
||||
type Target = Policy;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
enum ReceivedFrame {
|
||||
Message(CommunicationValue),
|
||||
ClosedByPeer,
|
||||
Idle,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum SenderState {
|
||||
Open,
|
||||
Closing,
|
||||
Closed,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Sender {
|
||||
send_guard: Arc<Mutex<()>>,
|
||||
stream_guard: Arc<Mutex<Option<wtransport::SendStream>>>,
|
||||
state: Arc<Mutex<SenderState>>,
|
||||
handle: Arc<ConnectionHandle>,
|
||||
connection: Connection,
|
||||
policy: Arc<Policy>,
|
||||
policy: Arc<RuntimePolicy>,
|
||||
type_map: Arc<RwLock<TypeMap>>,
|
||||
}
|
||||
|
||||
impl Sender {
|
||||
pub fn new(connection: Connection, handle: Arc<ConnectionHandle>, policy: Arc<Policy>) -> Self {
|
||||
let policy = Arc::new(RuntimePolicy::from_public(&policy));
|
||||
Self {
|
||||
send_guard: Arc::new(Mutex::new(())),
|
||||
stream_guard: Arc::new(Mutex::new(None)),
|
||||
state: Arc::new(Mutex::new(SenderState::Open)),
|
||||
handle,
|
||||
connection,
|
||||
policy,
|
||||
|
|
@ -174,7 +268,11 @@ impl Sender {
|
|||
data: &CommunicationValue,
|
||||
policy: &Policy,
|
||||
) -> Result<(), CommunicationError> {
|
||||
let bytes = data.to_bytes().map_err(|_| CommunicationError::Encode)?;
|
||||
let bytes = data
|
||||
.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
|
||||
{
|
||||
|
|
@ -269,11 +367,11 @@ impl Sender {
|
|||
return Ok(());
|
||||
}
|
||||
|
||||
let err = res.err().unwrap_or(CommunicationError::StreamError);
|
||||
if !matches!(
|
||||
err,
|
||||
CommunicationError::StreamError | CommunicationError::StreamClosed
|
||||
) {
|
||||
let err = match res {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(error) => error,
|
||||
};
|
||||
if !RetryClassifier::retry_persistent_stream(&err) {
|
||||
return Err(err);
|
||||
}
|
||||
*stream_opt = None;
|
||||
|
|
@ -356,20 +454,24 @@ impl Sender {
|
|||
|
||||
#[instrument(skip(self, data), level = "trace")]
|
||||
pub async fn send(&self, data: &CommunicationValue) -> Result<(), CommunicationError> {
|
||||
if self.handle.is_closed() {
|
||||
return Err(self
|
||||
.handle
|
||||
.close_reason()
|
||||
.unwrap_or(CommunicationError::UseAfterClosed));
|
||||
}
|
||||
|
||||
let _send_lock = self.send_guard.lock().await;
|
||||
|
||||
{
|
||||
let state = self.state.lock().await;
|
||||
if *state != SenderState::Open {
|
||||
return Err(self
|
||||
.handle
|
||||
.close_reason()
|
||||
.unwrap_or(CommunicationError::StreamClosed));
|
||||
}
|
||||
}
|
||||
|
||||
if self.connection.quic_connection().close_reason().is_some() {
|
||||
let reason = self
|
||||
.handle
|
||||
.close_reason()
|
||||
.unwrap_or(CommunicationError::StreamClosed);
|
||||
*self.state.lock().await = SenderState::Closed;
|
||||
self.handle.close(Some(reason.clone()));
|
||||
return Err(reason);
|
||||
}
|
||||
|
|
@ -402,6 +504,7 @@ impl Sender {
|
|||
if self.connection.quic_connection().close_reason().is_some()
|
||||
|| matches!(normalized, CommunicationError::StreamClosed)
|
||||
{
|
||||
*self.state.lock().await = SenderState::Closed;
|
||||
self.handle.close(Some(normalized.clone()));
|
||||
}
|
||||
|
||||
|
|
@ -413,6 +516,12 @@ impl Sender {
|
|||
#[instrument(skip(self), level = "trace")]
|
||||
pub async fn finish_stream(&self) -> Result<(), CommunicationError> {
|
||||
let _send_lock = self.send_guard.lock().await;
|
||||
if *self.state.lock().await != SenderState::Open {
|
||||
return Err(self
|
||||
.handle
|
||||
.close_reason()
|
||||
.unwrap_or(CommunicationError::StreamClosed));
|
||||
}
|
||||
let mut stream_opt = self.stream_guard.lock().await;
|
||||
if let Some(mut stream) = stream_opt.take() {
|
||||
match timeout(self.policy.write_timeout, stream.finish()).await {
|
||||
|
|
@ -445,11 +554,12 @@ impl Sender {
|
|||
pipe_id: u32,
|
||||
description: &str,
|
||||
) -> Result<crate::pipe::PipeWriter, CommunicationError> {
|
||||
if self.handle.is_closed() {
|
||||
let _send_lock = self.send_guard.lock().await;
|
||||
if *self.state.lock().await != SenderState::Open {
|
||||
return Err(self
|
||||
.handle
|
||||
.close_reason()
|
||||
.unwrap_or(CommunicationError::UseAfterClosed));
|
||||
.unwrap_or(CommunicationError::StreamClosed));
|
||||
}
|
||||
|
||||
if self.connection.quic_connection().close_reason().is_some() {
|
||||
|
|
@ -486,31 +596,47 @@ impl Sender {
|
|||
let handle = self.handle.clone();
|
||||
let policy = self.policy.clone();
|
||||
let stream_guard = self.stream_guard.clone();
|
||||
let send_guard = self.send_guard.clone();
|
||||
let state = self.state.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _send_lock = send_guard.lock().await;
|
||||
{
|
||||
let mut sender_state = state.lock().await;
|
||||
if *sender_state != SenderState::Open {
|
||||
return;
|
||||
}
|
||||
*sender_state = SenderState::Closing;
|
||||
}
|
||||
|
||||
if connection.quic_connection().close_reason().is_some() || handle.is_closed() {
|
||||
*state.lock().await = SenderState::Closed;
|
||||
handle.close(Some(CommunicationError::StreamClosed));
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(mut stream) = stream_guard.lock().await.take() {
|
||||
match timeout(policy.write_timeout, stream.finish()).await {
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(wtransport::error::StreamWriteError::Stopped(code))) => warn!(
|
||||
"[Sender] persistent stream finish failed: peer sent STOP_SENDING (error code {code})"
|
||||
),
|
||||
Ok(Err(e)) => {
|
||||
warn!("[Sender] persistent stream finish failed: {e}")
|
||||
{
|
||||
if let Some(mut stream) = stream_guard.lock().await.take() {
|
||||
match timeout(policy.write_timeout, stream.finish()).await {
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(wtransport::error::StreamWriteError::Stopped(code))) => warn!(
|
||||
"[Sender] persistent stream finish failed: peer sent STOP_SENDING (error code {code})"
|
||||
),
|
||||
Ok(Err(e)) => {
|
||||
warn!("[Sender] persistent stream finish failed: {e}")
|
||||
}
|
||||
Err(_) => warn!("[Sender] persistent stream finish timed out"),
|
||||
}
|
||||
Err(_) => warn!("[Sender] persistent stream finish timed out"),
|
||||
}
|
||||
}
|
||||
|
||||
let _ = Self::send_close_frame(&connection, &policy).await;
|
||||
|
||||
*state.lock().await = SenderState::Closed;
|
||||
handle.close(Some(CommunicationError::StreamClosed));
|
||||
info!(target = "mtp.transport", "connection closed");
|
||||
|
||||
drop(_send_lock);
|
||||
sleep(policy.force_close_delay).await;
|
||||
if connection.quic_connection().close_reason().is_none() {
|
||||
connection.quic_connection().close(
|
||||
|
|
@ -528,35 +654,48 @@ impl Sender {
|
|||
let connection = self.connection.clone();
|
||||
let handle = self.handle.clone();
|
||||
let policy = self.policy.clone();
|
||||
let mut stream_opt = self.stream_guard.lock().await;
|
||||
let _send_lock = self.send_guard.lock().await;
|
||||
{
|
||||
let mut state = self.state.lock().await;
|
||||
if *state != SenderState::Open {
|
||||
return;
|
||||
}
|
||||
*state = SenderState::Closing;
|
||||
}
|
||||
|
||||
if connection.quic_connection().close_reason().is_some() || handle.is_closed() {
|
||||
*self.state.lock().await = SenderState::Closed;
|
||||
handle.close(Some(CommunicationError::StreamClosed));
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(mut stream) = stream_opt.take() {
|
||||
let close_bytes = policy.close_frame_len.to_be_bytes();
|
||||
let close_write = async {
|
||||
stream.write_all(&close_bytes).await?;
|
||||
stream.finish().await
|
||||
};
|
||||
{
|
||||
let mut stream_opt = self.stream_guard.lock().await;
|
||||
if let Some(mut stream) = stream_opt.take() {
|
||||
let close_bytes = policy.close_frame_len.to_be_bytes();
|
||||
let close_write = async {
|
||||
stream.write_all(&close_bytes).await?;
|
||||
stream.finish().await
|
||||
};
|
||||
|
||||
match timeout(policy.write_timeout, close_write).await {
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(wtransport::error::StreamWriteError::Stopped(code))) => {
|
||||
warn!("[Sender] close failed: peer sent STOP_SENDING (error code {code})")
|
||||
match timeout(policy.write_timeout, close_write).await {
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(wtransport::error::StreamWriteError::Stopped(code))) => {
|
||||
warn!("[Sender] close failed: peer sent STOP_SENDING (error code {code})")
|
||||
}
|
||||
Ok(Err(e)) => warn!("[Sender] close failed: {e}"),
|
||||
Err(_) => warn!("[Sender] close timed out"),
|
||||
}
|
||||
Ok(Err(e)) => warn!("[Sender] close failed: {e}"),
|
||||
Err(_) => warn!("[Sender] close timed out"),
|
||||
} else {
|
||||
let _ = Self::send_close_frame(&connection, &policy).await;
|
||||
}
|
||||
} else {
|
||||
let _ = Self::send_close_frame(&connection, &policy).await;
|
||||
}
|
||||
|
||||
*self.state.lock().await = SenderState::Closed;
|
||||
handle.close(Some(CommunicationError::StreamClosed));
|
||||
info!(target = "mtp.transport", "connection closed");
|
||||
|
||||
drop(_send_lock);
|
||||
sleep(policy.force_close_delay).await;
|
||||
if connection.quic_connection().close_reason().is_none() {
|
||||
connection.quic_connection().close(
|
||||
|
|
@ -605,6 +744,7 @@ struct ReceiverInner {
|
|||
queue_notify: Arc<Notify>,
|
||||
max_message_size: Arc<AtomicU64>,
|
||||
type_map: Arc<RwLock<TypeMap>>,
|
||||
decode_rejections: Arc<DecodeRejectionCounters>,
|
||||
}
|
||||
|
||||
impl Clone for Receiver {
|
||||
|
|
@ -626,12 +766,26 @@ impl Drop for Receiver {
|
|||
#[derive(Clone, Default)]
|
||||
struct PingControl {
|
||||
pong_sender: Option<Sender>,
|
||||
pong_observer: Option<mpsc::UnboundedSender<CommunicationValue>>,
|
||||
pong_observer: Option<mpsc::Sender<CommunicationValue>>,
|
||||
expected_pong_id: Option<u32>,
|
||||
}
|
||||
|
||||
impl PingControl {
|
||||
fn accepts_pong(&mut self, id: Option<u32>) -> bool {
|
||||
if self.expected_pong_id == id && id.is_some() {
|
||||
self.expected_pong_id = None;
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Receiver {
|
||||
pub fn new(connection: Connection, handle: Arc<ConnectionHandle>, policy: Arc<Policy>) -> Self {
|
||||
Self::new_with_max_message_size(connection, handle, policy.clone(), policy.max_message_size)
|
||||
let policy = Arc::new(RuntimePolicy::from_public(&policy));
|
||||
let max_message_size = policy.max_message_size;
|
||||
Self::new_with_max_message_size(connection, handle, policy, max_message_size)
|
||||
}
|
||||
|
||||
#[cfg(feature = "host")]
|
||||
|
|
@ -640,6 +794,7 @@ impl Receiver {
|
|||
handle: Arc<ConnectionHandle>,
|
||||
policy: Arc<Policy>,
|
||||
) -> Self {
|
||||
let policy = Arc::new(RuntimePolicy::from_public(&policy));
|
||||
let initial_max = policy
|
||||
.handshake_max_message_size
|
||||
.min(policy.max_message_size);
|
||||
|
|
@ -649,7 +804,7 @@ impl Receiver {
|
|||
fn new_with_max_message_size(
|
||||
connection: Connection,
|
||||
handle: Arc<ConnectionHandle>,
|
||||
policy: Arc<Policy>,
|
||||
policy: Arc<RuntimePolicy>,
|
||||
initial_max_message_size: u64,
|
||||
) -> Self {
|
||||
#[cfg(feature = "pipes")]
|
||||
|
|
@ -674,6 +829,8 @@ impl Receiver {
|
|||
let accept_max_message_size = max_message_size.clone();
|
||||
let type_map = Arc::new(RwLock::new(TypeMap::latest()));
|
||||
let accept_type_map = type_map.clone();
|
||||
let decode_rejections = Arc::new(DecodeRejectionCounters::default());
|
||||
let accept_decode_rejections = decode_rejections.clone();
|
||||
let stream_limit = Arc::new(Semaphore::new(policy.max_concurrent_stream_tasks.max(1)));
|
||||
let accept_stream_limit = stream_limit.clone();
|
||||
debug!(
|
||||
|
|
@ -742,6 +899,7 @@ impl Receiver {
|
|||
let stream_ping_control = accept_ping_control.clone();
|
||||
let stream_max_message_size = accept_max_message_size.clone();
|
||||
let stream_type_map = accept_type_map.clone();
|
||||
let stream_decode_rejections = accept_decode_rejections.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _permit = permit;
|
||||
|
|
@ -761,7 +919,14 @@ impl Receiver {
|
|||
}
|
||||
|
||||
let frame_limit = stream_max_message_size.load(Ordering::Relaxed);
|
||||
match Self::read_one_frame(&mut s, &stream_policy, frame_limit).await {
|
||||
match Self::read_one_frame(
|
||||
&mut s,
|
||||
&stream_policy,
|
||||
frame_limit,
|
||||
&stream_decode_rejections,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(ReceivedFrame::Message(mut msg)) => {
|
||||
let negotiated_type_map =
|
||||
stream_type_map.read().await.clone();
|
||||
|
|
@ -805,24 +970,17 @@ impl Receiver {
|
|||
}
|
||||
}
|
||||
|
||||
let control = {
|
||||
let control = stream_ping_control.read().await;
|
||||
if msg.is_type(mtp_codec::CommunicationType::Ping) {
|
||||
control
|
||||
.pong_sender
|
||||
.clone()
|
||||
.map(|sender| (Some(sender), None))
|
||||
} else if msg.is_type(mtp_codec::CommunicationType::Pong) {
|
||||
control
|
||||
.pong_observer
|
||||
.clone()
|
||||
.map(|observer| (None, Some(observer)))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
let pong_sender = if msg.is_type(mtp_codec::CommunicationType::Ping) {
|
||||
stream_ping_control
|
||||
.read()
|
||||
.await
|
||||
.pong_sender
|
||||
.clone()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
if let Some((Some(sender), _)) = control {
|
||||
if let Some(sender) = pong_sender {
|
||||
let mut pong = CommunicationValue::new_with_type_map(
|
||||
mtp_codec::CommunicationType::Pong,
|
||||
&negotiated_type_map,
|
||||
|
|
@ -844,8 +1002,18 @@ impl Receiver {
|
|||
continue;
|
||||
}
|
||||
|
||||
if let Some((_, Some(observer))) = control {
|
||||
let _ = observer.send(msg);
|
||||
if msg.is_type(mtp_codec::CommunicationType::Pong) {
|
||||
let observer = {
|
||||
let mut control = stream_ping_control.write().await;
|
||||
if control.accepts_pong(msg.id()) {
|
||||
control.pong_observer.clone()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
};
|
||||
if let Some(observer) = observer {
|
||||
let _ = observer.try_send(msg);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
|
|
@ -942,6 +1110,7 @@ impl Receiver {
|
|||
queue_notify,
|
||||
max_message_size,
|
||||
type_map,
|
||||
decode_rejections,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
|
@ -958,6 +1127,14 @@ impl Receiver {
|
|||
*self.inner.type_map.write().await = type_map.clone();
|
||||
}
|
||||
|
||||
/// Return local counts for frames rejected by the structured decoder.
|
||||
///
|
||||
/// These counters are intentionally local-only; peers continue to receive
|
||||
/// the generic protocol parse failure.
|
||||
pub fn decode_rejection_counts(&self) -> DecodeRejectionCounts {
|
||||
self.inner.decode_rejections.snapshot()
|
||||
}
|
||||
|
||||
/* Respond to reserved Ping frames without exposing them to application I/O. */
|
||||
pub fn respond_to_pings(&self, sender: Sender) {
|
||||
if let Ok(mut control) = self.inner.ping_control.try_write() {
|
||||
|
|
@ -967,16 +1144,21 @@ impl Receiver {
|
|||
}
|
||||
}
|
||||
|
||||
/* Route reserved Pong frames to a connection-level observer. */
|
||||
pub async fn observe_pongs(&self, observer: mpsc::UnboundedSender<CommunicationValue>) {
|
||||
/* Route only the currently expected reserved Pong through a bounded observer. */
|
||||
pub async fn observe_pongs_bounded(&self, observer: mpsc::Sender<CommunicationValue>) {
|
||||
self.inner.ping_control.write().await.pong_observer = Some(observer);
|
||||
}
|
||||
|
||||
#[instrument(skip(stream, policy), level = "trace")]
|
||||
pub async fn set_expected_pong_id(&self, expected_pong_id: Option<u32>) {
|
||||
self.inner.ping_control.write().await.expected_pong_id = expected_pong_id;
|
||||
}
|
||||
|
||||
#[instrument(skip(stream, policy, decode_rejections), level = "trace")]
|
||||
async fn read_one_frame(
|
||||
stream: &mut wtransport::RecvStream,
|
||||
policy: &Policy,
|
||||
policy: &RuntimePolicy,
|
||||
max_message_size: u64,
|
||||
decode_rejections: &DecodeRejectionCounters,
|
||||
) -> Result<ReceivedFrame, CommunicationError> {
|
||||
use wtransport::error::{StreamReadError, StreamReadExactError};
|
||||
|
||||
|
|
@ -1011,6 +1193,7 @@ impl Receiver {
|
|||
if len == policy.close_frame_len {
|
||||
return Ok(ReceivedFrame::ClosedByPeer);
|
||||
}
|
||||
let deadline = Instant::now() + policy.read_timeout;
|
||||
|
||||
let body_len = len as usize;
|
||||
let frame_len = body_len
|
||||
|
|
@ -1020,29 +1203,29 @@ impl Receiver {
|
|||
return Err(CommunicationError::MessageTooLarge);
|
||||
}
|
||||
|
||||
// Grow in bounded chunks instead of trusting the peer's length prefix
|
||||
// enough to allocate the complete frame up front.
|
||||
let mut buf = Vec::new();
|
||||
buf.try_reserve(body_len.min(16 * 1024))
|
||||
// The length has already been checked against the admitted frame
|
||||
// limit, so reserve one bounded framing buffer and decode it without a
|
||||
// second prefix-plus-body allocation/copy.
|
||||
let mut frame = Vec::new();
|
||||
frame
|
||||
.try_reserve_exact(frame_len)
|
||||
.map_err(|_| CommunicationError::MessageTooLarge)?;
|
||||
while buf.len() < body_len {
|
||||
let chunk_len = (body_len - buf.len()).min(16 * 1024);
|
||||
let mut chunk = [0u8; 16 * 1024];
|
||||
match timeout(
|
||||
policy.read_timeout,
|
||||
stream.read_exact(&mut chunk[..chunk_len]),
|
||||
frame.extend_from_slice(&len_buf);
|
||||
frame.resize(frame_len, 0);
|
||||
let mut body_offset = 4usize;
|
||||
while body_offset < frame_len {
|
||||
let chunk_len = (frame_len - body_offset).min(16 * 1024);
|
||||
match timeout_at(
|
||||
deadline,
|
||||
stream.read_exact(&mut frame[body_offset..body_offset + chunk_len]),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(())) => {
|
||||
buf.try_reserve(chunk_len)
|
||||
.map_err(|_| CommunicationError::MessageTooLarge)?;
|
||||
buf.extend_from_slice(&chunk[..chunk_len]);
|
||||
}
|
||||
Ok(Ok(())) => body_offset += chunk_len,
|
||||
Ok(Err(StreamReadExactError::FinishedEarly(n))) => {
|
||||
warn!(
|
||||
"[Receiver] body read ended early ({}/{body_len} bytes): stream closed by peer",
|
||||
buf.len() + n
|
||||
body_offset.saturating_sub(4) + n
|
||||
);
|
||||
return Err(CommunicationError::StreamError);
|
||||
}
|
||||
|
|
@ -1063,14 +1246,21 @@ impl Receiver {
|
|||
}
|
||||
}
|
||||
|
||||
let mut frame = Vec::with_capacity(frame_len);
|
||||
frame.extend_from_slice(&len_buf);
|
||||
frame.extend_from_slice(&buf);
|
||||
let message = CommunicationValue::from_bytes_with_limits(
|
||||
let message = match CommunicationValue::try_from_bytes_with_limits(
|
||||
&frame,
|
||||
DecodeLimits::for_transport_message_size(max_message_size),
|
||||
)
|
||||
.map_err(|_| CommunicationError::ParseCommunicationValue)?;
|
||||
) {
|
||||
Ok(message) => message,
|
||||
Err(error) => {
|
||||
decode_rejections.record(&error);
|
||||
warn!(
|
||||
?error,
|
||||
class = ?classify_decode_error(&error),
|
||||
"[Receiver] rejected frame during bounded decode"
|
||||
);
|
||||
return Err(CommunicationError::ParseCommunicationValue);
|
||||
}
|
||||
};
|
||||
|
||||
Ok(ReceivedFrame::Message(message))
|
||||
}
|
||||
|
|
@ -1267,4 +1457,61 @@ mod tests {
|
|||
let debug_str = format!("{:?}", p);
|
||||
assert!(debug_str.contains("Policy"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_policy_normalizes_zero_channel_and_task_limits() {
|
||||
let mut policy = Policy::default();
|
||||
policy.receiver_queue_capacity = 0;
|
||||
policy.max_concurrent_stream_tasks = 0;
|
||||
|
||||
let runtime = RuntimePolicy::from_public(&policy);
|
||||
|
||||
assert_eq!(runtime.receiver_queue_capacity, 1);
|
||||
assert_eq!(runtime.max_concurrent_stream_tasks, 1);
|
||||
assert_eq!(policy.receiver_queue_capacity, 0);
|
||||
assert_eq!(policy.max_concurrent_stream_tasks, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ping_control_accepts_only_the_current_expected_id() {
|
||||
let mut control = PingControl {
|
||||
expected_pong_id: Some(7),
|
||||
..PingControl::default()
|
||||
};
|
||||
|
||||
assert!(!control.accepts_pong(Some(6)));
|
||||
assert_eq!(control.expected_pong_id, Some(7));
|
||||
assert!(control.accepts_pong(Some(7)));
|
||||
assert_eq!(control.expected_pong_id, None);
|
||||
assert!(!control.accepts_pong(Some(7)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_rejection_classes_are_stable_and_counted() {
|
||||
assert_eq!(
|
||||
classify_decode_error(&DecodeError::MalformedEncoding),
|
||||
DecodeRejectionClass::Malformed
|
||||
);
|
||||
assert_eq!(
|
||||
classify_decode_error(&DecodeError::AllocationLimit),
|
||||
DecodeRejectionClass::ResourceLimit
|
||||
);
|
||||
assert_eq!(
|
||||
classify_decode_error(&DecodeError::DuplicateField),
|
||||
DecodeRejectionClass::DuplicateField
|
||||
);
|
||||
|
||||
let counters = DecodeRejectionCounters::default();
|
||||
counters.record(&DecodeError::MalformedEncoding);
|
||||
counters.record(&DecodeError::DepthLimit);
|
||||
counters.record(&DecodeError::DuplicateField);
|
||||
assert_eq!(
|
||||
counters.snapshot(),
|
||||
DecodeRejectionCounts {
|
||||
malformed: 1,
|
||||
resource_limit: 1,
|
||||
duplicate_field: 1,
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,22 +2,26 @@ use mtp_common::CommunicationError;
|
|||
use std::net::SocketAddr;
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
atomic::{AtomicBool, AtomicU64, Ordering},
|
||||
};
|
||||
use tokio::sync::watch;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ConnectionHandle {
|
||||
connection_id: u64,
|
||||
closed: AtomicBool,
|
||||
close_tx: watch::Sender<Option<CommunicationError>>,
|
||||
close_rx: watch::Receiver<Option<CommunicationError>>,
|
||||
remote_addr: Option<SocketAddr>,
|
||||
}
|
||||
|
||||
static NEXT_CONNECTION_ID: AtomicU64 = AtomicU64::new(1);
|
||||
|
||||
impl ConnectionHandle {
|
||||
pub fn new() -> Self {
|
||||
let (close_tx, close_rx) = watch::channel(None);
|
||||
Self {
|
||||
connection_id: NEXT_CONNECTION_ID.fetch_add(1, Ordering::Relaxed).max(1),
|
||||
closed: AtomicBool::new(false),
|
||||
close_tx,
|
||||
close_rx,
|
||||
|
|
@ -35,6 +39,11 @@ impl ConnectionHandle {
|
|||
self.remote_addr
|
||||
}
|
||||
|
||||
/// Stable process-local identifier for authentication-rate-limit scopes.
|
||||
pub fn connection_id(&self) -> u64 {
|
||||
self.connection_id
|
||||
}
|
||||
|
||||
pub fn is_open(&self) -> bool {
|
||||
!self.closed.load(Ordering::SeqCst)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,7 +1,21 @@
|
|||
use crate::{Policy, TransportSendStream};
|
||||
use mtp_codec::CommunicationValue;
|
||||
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
|
||||
|
|
@ -12,7 +26,11 @@ pub(crate) async fn write_frame<S: TransportSendStream>(
|
|||
value: &CommunicationValue,
|
||||
policy: &Policy,
|
||||
) -> Result<(), CommunicationError> {
|
||||
let bytes = value.to_bytes().map_err(|_| CommunicationError::Encode)?;
|
||||
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
|
||||
{
|
||||
|
|
|
|||
|
|
@ -5,21 +5,23 @@
|
|||
//! wrappers while the framing implementation below is shared by adapters.
|
||||
|
||||
use crate::{
|
||||
Policy, TransportConnection, TransportRecvStream, TransportSendStream, framing::write_frame,
|
||||
Policy, TransportConnection, TransportRecvStream, TransportSendStream,
|
||||
connection::{DecodeRejectionCounters, RuntimePolicy, classify_decode_error},
|
||||
framing::{RetryClassifier, write_frame},
|
||||
};
|
||||
use mtp_codec::{CommunicationValue, DecodeLimits, TypeMap};
|
||||
use mtp_common::CommunicationError;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use tokio::sync::{Mutex, RwLock, Semaphore, mpsc};
|
||||
use tokio::time::timeout;
|
||||
use tokio::sync::{Mutex, Notify, RwLock, Semaphore, mpsc};
|
||||
use tokio::time::{Instant, timeout, timeout_at};
|
||||
|
||||
#[cfg(feature = "pipes")]
|
||||
use crate::pipe::{PipeReader, PipeWriter};
|
||||
|
||||
pub struct GenericSender<C: TransportConnection> {
|
||||
connection: C,
|
||||
policy: Arc<Policy>,
|
||||
policy: Arc<RuntimePolicy>,
|
||||
persistent: Arc<Mutex<Option<C::SendStream>>>,
|
||||
send_lock: Arc<Mutex<()>>,
|
||||
type_map: Arc<RwLock<TypeMap>>,
|
||||
|
|
@ -39,6 +41,7 @@ impl<C: TransportConnection> Clone for GenericSender<C> {
|
|||
|
||||
impl<C: TransportConnection> GenericSender<C> {
|
||||
pub fn new(connection: C, policy: Arc<Policy>) -> Self {
|
||||
let policy = Arc::new(RuntimePolicy::from_public(&policy));
|
||||
Self {
|
||||
connection,
|
||||
policy,
|
||||
|
|
@ -84,20 +87,30 @@ impl<C: TransportConnection> GenericSender<C> {
|
|||
if stream.is_none() {
|
||||
*stream = Some(self.open().await?);
|
||||
}
|
||||
let result = timeout(
|
||||
self.policy.write_timeout,
|
||||
write_frame(stream.as_mut().unwrap(), value, &self.policy),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| CommunicationError::StreamError)
|
||||
.and_then(|r| r);
|
||||
let result = match stream.as_mut() {
|
||||
Some(stream) => timeout(
|
||||
self.policy.write_timeout,
|
||||
write_frame(stream, value, &self.policy),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| CommunicationError::StreamError)
|
||||
.and_then(|result| result),
|
||||
None => Err(CommunicationError::StreamError),
|
||||
};
|
||||
if result.is_ok() {
|
||||
return result;
|
||||
return Ok(());
|
||||
}
|
||||
let error = match result {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(error) => error,
|
||||
};
|
||||
if !RetryClassifier::retry_persistent_stream(&error) {
|
||||
return Err(error);
|
||||
}
|
||||
*stream = None;
|
||||
attempts += 1;
|
||||
if attempts > self.policy.persistent_stream_max_retries {
|
||||
return result;
|
||||
return Err(error);
|
||||
}
|
||||
tokio::time::sleep(
|
||||
self.policy.persistent_stream_retry_backoff * attempts as u32,
|
||||
|
|
@ -114,6 +127,7 @@ impl<C: TransportConnection> GenericSender<C> {
|
|||
pipe_id: u32,
|
||||
description: &str,
|
||||
) -> Result<PipeWriter<C::SendStream>, CommunicationError> {
|
||||
let _send_lock = self.send_lock.lock().await;
|
||||
if self.connection.close_reason().is_some() {
|
||||
return Err(CommunicationError::StreamClosed);
|
||||
}
|
||||
|
|
@ -179,6 +193,8 @@ pub struct GenericReceiver<C: TransportConnection> {
|
|||
ping_sender: Arc<RwLock<Option<GenericSender<C>>>>,
|
||||
max_message_size: Arc<AtomicU64>,
|
||||
type_map: Arc<RwLock<TypeMap>>,
|
||||
queue_notify: Arc<Notify>,
|
||||
decode_rejections: Arc<DecodeRejectionCounters>,
|
||||
_accept_task: Arc<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
|
|
@ -192,6 +208,8 @@ impl<C: TransportConnection> Clone for GenericReceiver<C> {
|
|||
ping_sender: self.ping_sender.clone(),
|
||||
max_message_size: self.max_message_size.clone(),
|
||||
type_map: self.type_map.clone(),
|
||||
queue_notify: self.queue_notify.clone(),
|
||||
decode_rejections: self.decode_rejections.clone(),
|
||||
_accept_task: self._accept_task.clone(),
|
||||
}
|
||||
}
|
||||
|
|
@ -207,6 +225,7 @@ impl<C: TransportConnection> Drop for GenericReceiver<C> {
|
|||
|
||||
impl<C: TransportConnection> GenericReceiver<C> {
|
||||
pub fn new(connection: C, policy: Arc<Policy>) -> Self {
|
||||
let policy = Arc::new(RuntimePolicy::from_public(&policy));
|
||||
let (tx, rx) = mpsc::channel(policy.receiver_queue_capacity);
|
||||
#[cfg(feature = "pipes")]
|
||||
let (pipe_tx, pipe_rx) = mpsc::channel(policy.receiver_queue_capacity);
|
||||
|
|
@ -222,6 +241,10 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
let task_max_message_size = max_message_size.clone();
|
||||
let type_map = Arc::new(RwLock::new(TypeMap::latest()));
|
||||
let task_type_map = type_map.clone();
|
||||
let queue_notify = Arc::new(Notify::new());
|
||||
let task_queue_notify = queue_notify.clone();
|
||||
let decode_rejections = Arc::new(DecodeRejectionCounters::default());
|
||||
let task_decode_rejections = decode_rejections.clone();
|
||||
let task_accept_task_tx = tx.clone();
|
||||
#[cfg(feature = "pipes")]
|
||||
let task_accept_task_pipe_tx = pipe_tx.clone();
|
||||
|
|
@ -237,8 +260,10 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
#[cfg(not(feature = "pipes"))]
|
||||
let cap_full = task_accept_task_tx.capacity() == 0;
|
||||
|
||||
let notified = task_queue_notify.notified();
|
||||
tokio::pin!(notified);
|
||||
if cap_full {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(1)).await;
|
||||
notified.await;
|
||||
continue;
|
||||
}
|
||||
|
||||
|
|
@ -277,11 +302,12 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
let ping_sender = task_ping_sender.clone();
|
||||
let connection = task_connection.clone();
|
||||
let type_map = task_type_map.clone();
|
||||
let decode_rejections = task_decode_rejections.clone();
|
||||
tokio::spawn(async move {
|
||||
let _permit = permit;
|
||||
let mut stream = stream;
|
||||
let mut frames = 0usize;
|
||||
loop {
|
||||
'stream: loop {
|
||||
if policy
|
||||
.max_frames_per_stream
|
||||
.is_some_and(|max| frames >= max)
|
||||
|
|
@ -304,7 +330,7 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
policy.application_close_code,
|
||||
b"frame header read error",
|
||||
);
|
||||
break;
|
||||
break 'stream;
|
||||
}
|
||||
Err(_) => {
|
||||
break;
|
||||
|
|
@ -314,6 +340,7 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
if len == policy.close_frame_len {
|
||||
break;
|
||||
}
|
||||
let deadline = Instant::now() + policy.read_timeout;
|
||||
let frame_limit = max_message_size.load(Ordering::Relaxed);
|
||||
let body_len = len as usize;
|
||||
let frame_len = match body_len.checked_add(4) {
|
||||
|
|
@ -330,29 +357,25 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
connection.close(policy.application_close_code, b"frame too large");
|
||||
break;
|
||||
}
|
||||
let target_len = body_len;
|
||||
let mut body = Vec::new();
|
||||
if body.try_reserve(target_len.min(16 * 1024)).is_err() {
|
||||
tracing::warn!(
|
||||
target_len,
|
||||
"MTP receive stream could not reserve frame body"
|
||||
);
|
||||
let mut frame = Vec::new();
|
||||
if frame.try_reserve_exact(frame_len).is_err() {
|
||||
tracing::warn!(frame_len, "MTP receive stream could not reserve frame");
|
||||
let _ = tx.send(Err(CommunicationError::MessageTooLarge)).await;
|
||||
connection
|
||||
.close(policy.application_close_code, b"frame allocation failed");
|
||||
break;
|
||||
}
|
||||
while body.len() < target_len {
|
||||
let chunk_len = (target_len - body.len()).min(16 * 1024);
|
||||
let mut chunk = [0u8; 16 * 1024];
|
||||
let body_read = tokio::time::timeout(
|
||||
policy.read_timeout,
|
||||
stream.read_exact(&mut chunk[..chunk_len]),
|
||||
frame.extend_from_slice(&len.to_be_bytes());
|
||||
frame.resize(frame_len, 0);
|
||||
let mut body_offset = 4usize;
|
||||
while body_offset < frame_len {
|
||||
let chunk_len = (frame_len - body_offset).min(16 * 1024);
|
||||
let body_read = timeout_at(
|
||||
deadline,
|
||||
stream.read_exact(&mut frame[body_offset..body_offset + chunk_len]),
|
||||
)
|
||||
.await;
|
||||
if !matches!(&body_read, Ok(Ok(())))
|
||||
|| body.try_reserve(chunk_len).is_err()
|
||||
{
|
||||
if !matches!(&body_read, Ok(Ok(()))) {
|
||||
tracing::warn!(
|
||||
pipe_chunk_len = chunk_len,
|
||||
?body_read,
|
||||
|
|
@ -361,24 +384,23 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
let _ = tx.send(Err(CommunicationError::StreamError)).await;
|
||||
connection
|
||||
.close(policy.application_close_code, b"frame body read error");
|
||||
break;
|
||||
break 'stream;
|
||||
}
|
||||
body.extend_from_slice(&chunk[..chunk_len]);
|
||||
}
|
||||
if body.len() != target_len {
|
||||
break;
|
||||
body_offset += chunk_len;
|
||||
}
|
||||
frames += 1;
|
||||
let mut frame = Vec::with_capacity(frame_len);
|
||||
frame.extend_from_slice(&len.to_be_bytes());
|
||||
frame.extend_from_slice(&body);
|
||||
let mut message = match CommunicationValue::from_bytes_with_limits(
|
||||
let mut message = match CommunicationValue::try_from_bytes_with_limits(
|
||||
&frame,
|
||||
DecodeLimits::for_transport_message_size(frame_limit),
|
||||
) {
|
||||
Ok(message) => message,
|
||||
Err(_) => {
|
||||
tracing::warn!("MTP receive stream contained an invalid frame");
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
?error,
|
||||
class = ?classify_decode_error(&error),
|
||||
"MTP receive stream rejected by bounded decode"
|
||||
);
|
||||
decode_rejections.record(&error);
|
||||
let _ = tx
|
||||
.send(Err(CommunicationError::ParseCommunicationValue))
|
||||
.await;
|
||||
|
|
@ -463,6 +485,8 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
ping_sender,
|
||||
max_message_size,
|
||||
type_map,
|
||||
queue_notify,
|
||||
decode_rejections,
|
||||
_accept_task: Arc::new(accept_task),
|
||||
}
|
||||
}
|
||||
|
|
@ -480,13 +504,23 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
pub async fn set_type_map(&self, type_map: &TypeMap) {
|
||||
*self.type_map.write().await = type_map.clone();
|
||||
}
|
||||
|
||||
/// Return local counts for frames rejected by the structured decoder.
|
||||
pub fn decode_rejection_counts(&self) -> crate::DecodeRejectionCounts {
|
||||
self.decode_rejections.snapshot()
|
||||
}
|
||||
pub async fn receive(&self) -> Result<CommunicationValue, CommunicationError> {
|
||||
self.incoming
|
||||
let result = self
|
||||
.incoming
|
||||
.lock()
|
||||
.await
|
||||
.recv()
|
||||
.await
|
||||
.unwrap_or(Err(CommunicationError::StreamClosed))
|
||||
.unwrap_or(Err(CommunicationError::StreamClosed));
|
||||
if result.is_ok() {
|
||||
self.queue_notify.notify_one();
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
#[cfg(feature = "pipes")]
|
||||
|
|
@ -498,7 +532,10 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
tokio::select! {
|
||||
msg = incoming.recv() => {
|
||||
match msg {
|
||||
Some(Ok(val)) => Ok(crate::TransportEvent::Message(val)),
|
||||
Some(Ok(val)) => {
|
||||
self.queue_notify.notify_one();
|
||||
Ok(crate::TransportEvent::Message(val))
|
||||
}
|
||||
Some(Err(e)) => Err(e),
|
||||
None => Err(self
|
||||
.connection
|
||||
|
|
@ -508,7 +545,10 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
}
|
||||
pipe = pipes.recv() => {
|
||||
match pipe {
|
||||
Some(reader) => Ok(crate::TransportEvent::Pipe(reader)),
|
||||
Some(reader) => {
|
||||
self.queue_notify.notify_one();
|
||||
Ok(crate::TransportEvent::Pipe(reader))
|
||||
}
|
||||
None => Err(self
|
||||
.connection
|
||||
.close_reason()
|
||||
|
|
@ -520,12 +560,17 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
|
||||
#[cfg(feature = "pipes")]
|
||||
pub async fn receive_pipe(&self) -> Result<PipeReader<C::RecvStream>, CommunicationError> {
|
||||
self.pipes
|
||||
let result = self
|
||||
.pipes
|
||||
.lock()
|
||||
.await
|
||||
.recv()
|
||||
.await
|
||||
.ok_or(CommunicationError::StreamClosed)
|
||||
.ok_or(CommunicationError::StreamClosed);
|
||||
if result.is_ok() {
|
||||
self.queue_notify.notify_one();
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
#[cfg(feature = "pipes")]
|
||||
|
|
@ -534,7 +579,10 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
|||
) -> Result<Option<PipeReader<C::RecvStream>>, CommunicationError> {
|
||||
match self.pipes.try_lock() {
|
||||
Ok(mut rx) => match rx.try_recv() {
|
||||
Ok(reader) => Ok(Some(reader)),
|
||||
Ok(reader) => {
|
||||
self.queue_notify.notify_one();
|
||||
Ok(Some(reader))
|
||||
}
|
||||
Err(mpsc::error::TryRecvError::Empty) => Ok(None),
|
||||
Err(mpsc::error::TryRecvError::Disconnected) => {
|
||||
Err(CommunicationError::StreamClosed)
|
||||
|
|
|
|||
|
|
@ -11,7 +11,10 @@ pub mod encrypted_pipe;
|
|||
#[cfg(feature = "pipes")]
|
||||
pub mod pipe;
|
||||
|
||||
pub use connection::{Policy, Receiver, SendMode, Sender};
|
||||
pub use connection::{
|
||||
DecodeRejectionClass, DecodeRejectionCounters, DecodeRejectionCounts, Policy, Receiver,
|
||||
SendMode, Sender, classify_decode_error,
|
||||
};
|
||||
pub use generic_connection::{GenericReceiver, GenericSender};
|
||||
|
||||
#[cfg(feature = "pipes")]
|
||||
|
|
|
|||
Loading…
Reference in a new issue