ssh/model/
flow_control.rs

1use crate::constant::size::LOCAL_WINDOW_SIZE;
2
3use crate::constant::size;
4
5pub(crate) struct FlowControl {
6    local_window: u32,
7    remote_window: u32,
8}
9
10impl FlowControl {
11    pub fn new(remote: u32) -> Self {
12        FlowControl {
13            local_window: LOCAL_WINDOW_SIZE,
14            remote_window: remote,
15        }
16    }
17
18    pub fn tune_on_recv(&mut self, buf: &mut Vec<u8>) {
19        let recv_len = buf.len() as u32;
20
21        if self.local_window >= recv_len {
22            self.local_window -= recv_len;
23        } else {
24            let drop_len = recv_len - self.local_window;
25            tracing::debug!("Recv more than expected, drop len {}", drop_len);
26            buf.truncate(self.local_window as usize);
27            self.local_window = 0;
28        }
29    }
30
31    pub fn tune_on_send(&mut self, buf: &mut Vec<u8>) -> Vec<u8> {
32        let want_send = buf.len();
33
34        let can_send = {
35            let mut can_send = want_send;
36
37            if can_send > self.remote_window as usize {
38                can_send = self.remote_window as usize
39            }
40
41            if can_send > size::BUF_SIZE {
42                can_send = size::BUF_SIZE
43            }
44            can_send
45        };
46
47        self.remote_window -= can_send as u32;
48
49        buf.split_off(can_send)
50    }
51
52    pub fn on_recv(&mut self, size: u32) {
53        self.remote_window += size
54    }
55
56    pub fn on_send(&mut self, size: u32) {
57        self.local_window += size
58    }
59
60    pub fn can_send(&self) -> bool {
61        self.remote_window > 0
62    }
63}