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

use crate::common::protocol::messages::request::HandShake;
use crate::common::protocol::messages::{Request, Response};
use crate::sync::protocol::WiredProtocol;
use crate::OrientResult;
use byteorder::{BigEndian, ReadBytesExt};
use std::io::Write;
use std::net::Shutdown;
use std::net::{SocketAddr, TcpStream};

pub struct Connection {
    stream: TcpStream,
    protocol: WiredProtocol,
}

impl Connection {
    pub fn connect(addr: &SocketAddr) -> OrientResult<Self> {
        let mut stream = TcpStream::connect(addr)?;
        let p = stream.read_i16::<BigEndian>()?;
        let protocol = WiredProtocol::from_version(p)?;
        let conn = Connection { stream, protocol };
        conn.handshake()
    }

    fn handshake(mut self) -> OrientResult<Connection> {
        let handshake = HandShake {
            p_version: self.protocol.version,
            name: String::from("Rust Driver"),
            version: String::from("0.1"),
        };
        self.send_and_forget(handshake.into())?;
        Ok(self)
    }

    pub fn close(&mut self) -> OrientResult<()> {
        self.stream.shutdown(Shutdown::Both)?;
        Ok(())
    }

    pub fn send_and_forget(&mut self, request: Request) -> OrientResult<()> {
        let buf = self.protocol.encode(request)?;
        self.stream.write_all(buf.as_slice())?;
        Ok(())
    }
    pub fn send(&mut self, request: Request) -> OrientResult<Response> {
        self.send_and_forget(request)?;
        self.protocol.decode(&mut self.stream)
    }
}