mirror of
https://github.com/rustdesk/rustdesk.git
synced 2026-09-15 00:41:01 +03:00
port_forward_mux: fix RecvWindow counter overflow on long transfers
Replace cumulative accounting (granted/received) with remaining credit tracking to prevent u32 overflow after 4 GiB of data on a single channel. Wire behavior is identical, but the fix allows large file transfers without mid-stream channel closure. Add regression test for 8 GiB transfer to verify fix. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EZ49AbZJYfm8NTp5yDPMab
This commit is contained in:
@@ -28,46 +28,43 @@ pub fn charge(len: usize) -> u32 {
|
|||||||
u32::try_from(len).unwrap_or(u32::MAX).max(MIN_FRAME_CHARGE)
|
u32::try_from(len).unwrap_or(u32::MAX).max(MIN_FRAME_CHARGE)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Receiver-side accounting: what we granted, what arrived, what we have
|
/// Receiver-side accounting: the credit the peer still has, and what we have
|
||||||
/// drained to the local socket since the last `window_update`.
|
/// drained to the local socket since the last `window_update`. Both are
|
||||||
|
/// bounded — a cumulative counter would wear out on a long transfer.
|
||||||
pub struct RecvWindow {
|
pub struct RecvWindow {
|
||||||
granted: u32,
|
remaining: u32,
|
||||||
received: u32,
|
|
||||||
drained_since_update: u32,
|
drained_since_update: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl RecvWindow {
|
impl RecvWindow {
|
||||||
pub fn new(granted: u32) -> Self {
|
pub fn new(granted: u32) -> Self {
|
||||||
Self {
|
Self {
|
||||||
granted: effective_window(granted),
|
remaining: effective_window(granted),
|
||||||
received: 0,
|
|
||||||
drained_since_update: 0,
|
drained_since_update: 0,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// False means the peer overran the window: a protocol violation.
|
/// False means the peer overran the window: a protocol violation.
|
||||||
pub fn accept(&mut self, len: usize) -> bool {
|
pub fn accept(&mut self, len: usize) -> bool {
|
||||||
let len = charge(len);
|
match self.remaining.checked_sub(charge(len)) {
|
||||||
match self.received.checked_add(len) {
|
Some(left) => {
|
||||||
Some(total) if total <= self.granted => {
|
self.remaining = left;
|
||||||
self.received = total;
|
|
||||||
true
|
true
|
||||||
}
|
}
|
||||||
_ => false,
|
None => false,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns the amount to advertise in a `window_update` once enough has
|
/// Returns the amount to advertise in a `window_update` once enough has
|
||||||
/// been drained; the grant is extended by the same amount.
|
/// been drained; the same amount is credited back.
|
||||||
pub fn drained(&mut self, n: usize) -> Option<u32> {
|
pub fn drained(&mut self, n: usize) -> Option<u32> {
|
||||||
let n = charge(n);
|
self.drained_since_update = self.drained_since_update.saturating_add(charge(n));
|
||||||
self.drained_since_update = self.drained_since_update.saturating_add(n);
|
|
||||||
if self.drained_since_update < UPDATE_THRESHOLD {
|
if self.drained_since_update < UPDATE_THRESHOLD {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
let add = self.drained_since_update;
|
let add = self.drained_since_update;
|
||||||
self.drained_since_update = 0;
|
self.drained_since_update = 0;
|
||||||
self.granted = self.granted.saturating_add(add);
|
self.remaining = self.remaining.saturating_add(add);
|
||||||
Some(add)
|
Some(add)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -227,6 +224,19 @@ mod tests {
|
|||||||
assert!(!w.accept(1));
|
assert!(!w.accept(1));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn accounting_survives_a_transfer_far_larger_than_the_window() {
|
||||||
|
// Cumulative counters used to overflow around 4 GiB on one channel and
|
||||||
|
// read as a protocol violation mid-transfer.
|
||||||
|
let mut w = RecvWindow::new(CHANNEL_WINDOW);
|
||||||
|
let mut moved: u64 = 0;
|
||||||
|
while moved < 8 * 1024 * 1024 * 1024 {
|
||||||
|
assert!(w.accept(MAX_FRAME));
|
||||||
|
w.drained(MAX_FRAME);
|
||||||
|
moved += MAX_FRAME as u64;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn send_credit_blocks_at_zero_and_resumes_on_add() {
|
fn send_credit_blocks_at_zero_and_resumes_on_add() {
|
||||||
rt().block_on(async {
|
rt().block_on(async {
|
||||||
|
|||||||
Reference in New Issue
Block a user