scaproust 0.3.2

Nanomsg scalability protocols implementation in rust. Various messaging patterns over pluggable transports
Documentation 
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

// Copyright (c) 2015-2017 Contributors as noted in the AUTHORS file.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0>
// or the MIT license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your option.
// This file may not be copied, modified, or distributed except according to those terms.

use std::io;

use byteorder::{ BigEndian, ByteOrder };

use core::Message;
use transport::async::stub::*;
use io_error::*;

pub struct RecvOperation {
    step: Option<RecvOperationStep>
}

impl RecvOperation {
    pub fn new(recv_max_size: u64) -> RecvOperation {
        RecvOperation {
            step: Some(RecvOperationStep::Header([0; 9], 0, recv_max_size))
        }
    }

    pub fn run<T:io::Read>(&mut self, stream: &mut T) -> io::Result<Option<Message>> {
        if let Some(step) = self.step.take() {
            self.resume_at(stream, step)
        } else {
            Err(other_io_error("Cannot resume already finished recv operation"))
        }
    }

    fn resume_at<T:io::Read>(&mut self, stream: &mut T, step: RecvOperationStep) -> io::Result<Option<Message>> {
        let mut cur_step = step;

        loop {
            let (passed, next_step) = try!(cur_step.advance(stream));

            if !passed {
                self.step = Some(next_step);
                return Ok(None);
            }

            match next_step {
                RecvOperationStep::Terminal(msg) => return Ok(Some(msg)),
                other => cur_step = other
            }
        }
    }
}

enum RecvOperationStep {
    Header([u8; 9], usize, u64),
    Payload(Vec<u8>, usize),
    Terminal(Message)
}

impl RecvOperationStep {
    fn advance<T:io::Read>(self, stream: &mut T) -> io::Result<(bool, RecvOperationStep)> {
        match self {
            RecvOperationStep::Header(buffer, read, max_size) => read_header(stream, buffer, read, max_size),
            RecvOperationStep::Payload(buffer, read) => read_payload(stream, buffer, read),
            RecvOperationStep::Terminal(_) => Err(other_io_error("Cannot advance terminal step of recv operation"))
        }
    }
}

fn read_header<T:io::Read>(stream: &mut T, mut buffer: [u8; 9], mut read: usize, max_size: u64) -> io::Result<(bool, RecvOperationStep)> {
    read += try!(stream.read_buffer(&mut buffer[read..]));

    if read == 9 {
        let msg_len = BigEndian::read_u64(&buffer[1..]);
        if msg_len > max_size {
            Err(invalid_data_io_error("message is too long"))
        } else {
            let payload = vec![0u8; msg_len as usize];

            Ok((true, RecvOperationStep::Payload(payload, 0)))
        }
    } else {
        Ok((false, RecvOperationStep::Header(buffer, read, max_size)))
    }
}

fn read_payload<T:io::Read>(stream: &mut T, mut buffer: Vec<u8>, mut read: usize) -> io::Result<(bool, RecvOperationStep)> {
    read += try!(stream.read_buffer(&mut buffer[read..]));

    if read == buffer.capacity() {
        Ok((true, RecvOperationStep::Terminal(Message::from_body(buffer))))
    } else {
        Ok((false, RecvOperationStep::Payload(buffer, read)))
    }
}

#[cfg(test)]
mod tests {
    use std::io;

    use super::*;

    #[test]
    fn recv_in_one_run() {
        let buffer = vec![1, 0, 0, 0, 0, 0, 0, 0, 8, 1, 4, 3, 2, 65, 66, 67, 69];
        let mut stream = io::Cursor::new(buffer);
        let mut operation = RecvOperation::new(1024);
        let msg = operation.run(&mut stream).
            expect("recv should have succeeded").
            expect("recv should be done");
        let expected_bytes = [1, 4, 3, 2, 65, 66, 67, 69];

        assert_eq!(&expected_bytes, msg.get_body());
    }
}