amqp 0.0.10

AMQP/RabbitMQ protocol client
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
115
116
117

use std::net::TcpStream;
use std::io::Write;
use std::cmp;

#[cfg(feature = "tls")]
use openssl::ssl::{SslContext, SslMethod, SslStream};

#[cfg(not(feature = "tls"))]
use url;

#[cfg(not(feature = "tls"))]
use amqp_error::AMQPError;

use amqp_error::AMQPResult;
use framing::{Frame, FrameType};

enum AMQPStream {
	Cleartext(TcpStream),
#[cfg(feature = "tls")]
	Tls(SslStream<TcpStream>)
}

pub struct Connection {
    socket: AMQPStream,
    pub frame_max_limit: u32
}

impl Clone for Connection {
    fn clone(&self) -> Connection {
        match self.socket {
            AMQPStream::Cleartext(ref stream) => Connection {socket: AMQPStream::Cleartext(stream.try_clone().unwrap()), frame_max_limit: self.frame_max_limit},
#[cfg(feature = "tls")]
            AMQPStream::Tls(ref stream) => Connection {socket: AMQPStream::Tls(stream.try_clone().unwrap()), frame_max_limit: self.frame_max_limit},
        }
    }
}

impl Connection {
#[cfg(feature = "tls")]
    pub fn open(host: &str, port: u16, use_tls: bool) -> AMQPResult<Connection> {
        let connection;
        let mut socket = try!(TcpStream::connect((host, port)));
        if use_tls {
            let ctx = try!(SslContext::new(SslMethod::Sslv23));
            let mut tls_socket = try!(SslStream::connect(&ctx, socket));
            try!(tls_socket.write_all(&[b'A', b'M', b'Q', b'P', 0, 0, 9, 1]));
            connection = Connection { socket: AMQPStream::Tls(tls_socket), frame_max_limit: 131072 };
        } else {
            try!(socket.write_all(&[b'A', b'M', b'Q', b'P', 0, 0, 9, 1]));
            connection = Connection { socket: AMQPStream::Cleartext(socket), frame_max_limit: 131072 };
        }
        Ok(connection)
    }

#[cfg(not(feature = "tls"))]
    pub fn open(host: &str, port: u16, use_tls: bool) -> AMQPResult<Connection> {
        if use_tls {
            return Err(AMQPError::UrlParseError(url::ParseError::InvalidScheme))
        }
        let mut socket = try!(TcpStream::connect((host, port)));
        try!(socket.write_all(&[b'A', b'M', b'Q', b'P', 0, 0, 9, 1]));
        let connection = Connection { socket: AMQPStream::Cleartext(socket), frame_max_limit: 131072 };
        Ok(connection)
    }


    pub fn write(&mut self, frame: Frame) -> AMQPResult<()>{
        match frame.frame_type {
            FrameType::BODY => {
                //TODO: Check if need to include frame header + end octet into calculation. (9 bytes extra)
                for content_frame in split_content_into_frames(frame.payload, self.frame_max_limit).into_iter() {
                    try!(self.write_frame(Frame { frame_type: frame.frame_type, channel: frame.channel, payload: content_frame}))
                }
                Ok(())
            },
            _ => self.write_frame(frame)
        }
    }

    pub fn read(&mut self) -> AMQPResult<Frame> {
        match self.socket {
            AMQPStream::Cleartext(ref mut stream) => Frame::decode(stream),
#[cfg(feature = "tls")]
            AMQPStream::Tls(ref mut stream) => Frame::decode(stream),
        }
    }

    fn write_frame(&mut self, frame: Frame) -> AMQPResult<()>{
        match self.socket {
            AMQPStream::Cleartext(ref mut stream) => Ok(try!(stream.write_all(&frame.encode()))),
#[cfg(feature = "tls")]
            AMQPStream::Tls(ref mut stream) => Ok(try!(stream.write_all(&frame.encode()))),
        }
    }

}


fn split_content_into_frames(content: Vec<u8>, frame_limit: u32) -> Vec<Vec<u8>> {
    assert!(frame_limit > 0, "Can't have frame_max_limit == 0");
    let mut content_frames = vec!();
    let mut current_pos = 0;
    while current_pos < content.len() {
        let new_pos = current_pos + cmp::min(content.len() - current_pos, frame_limit as usize);
        content_frames.push(content[current_pos .. new_pos].to_vec());
        current_pos = new_pos;
    }
    content_frames
}


#[test]
fn test_split_content_into_frames() {
    let content = vec!(1,2,3,4,5,6,7,8,9,10);
    let frames = split_content_into_frames(content, 3);
    assert_eq!(frames, vec!(vec!(1, 2, 3), vec!(4, 5, 6), vec!(7, 8, 9), vec!(10)));
}