1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
use std;

#[derive(Debug, Clone, Copy)]
pub enum DebtHandling {
    Silent,
    Warn,
    DropMessage,
}

pub enum ZeroMessagesHandling {
    Drop,
    Deliver,
}

pub enum ProcessMessageResult {
    Return(std::result::Result<usize, std::io::Error>),
    Recurse,
}

/// A `Read` utility to deal with partial reads
pub struct ReadDebt(pub Option<Vec<u8>>, pub DebtHandling, pub ZeroMessagesHandling);
impl ReadDebt {
    pub fn process_message(&mut self, buf: &mut [u8], buf_in: &[u8]) -> ProcessMessageResult {
        assert_eq!(self.0, None);
        let mut l = buf_in.len();
        if l > buf.len() {
            match self.1 {
                DebtHandling::Silent => (),
                DebtHandling::Warn => {
                    warn!("Incoming message too long ({} > {}): splitting it to parts.\nUse -B option to increase buffer size or -S option to drop messages instead of splitting.", l, buf.len());
                }
                DebtHandling::DropMessage => {
                    error!("Dropping too large message ({} > {}). Use -B option to increase buffer size.", l, buf.len());
                    return ProcessMessageResult::Recurse;
                }
            }
            l = buf.len();
        }
        buf[..l].copy_from_slice(&buf_in[..l]);

        if l < buf_in.len() {
            self.0 = Some(buf_in[l..].to_vec());
        }

        debug!("Fullfulling the debt of {} bytes", l);
        if l == 0 {
            match self.2 {
                ZeroMessagesHandling::Deliver => (),
                ZeroMessagesHandling::Drop => {
                    info!("Dropping incoming zero-length message");
                    return ProcessMessageResult::Recurse;
                }
            }
        }
        ProcessMessageResult::Return(Ok(l))
    }
    pub fn check_debt(
        &mut self,
        buf: &mut [u8],
    ) -> Option<std::result::Result<usize, std::io::Error>> {
        if let Some(debt) = self.0.take() {
            match self.process_message(buf, debt.as_slice()) {
                ProcessMessageResult::Return(x) => Some(x),
                ProcessMessageResult::Recurse => unreachable!(),
            }
        } else {
            None
        }
    }
}