diff options
Diffstat (limited to 'vndbapi/src/conn.rs')
-rw-r--r-- | vndbapi/src/conn.rs | 113 |
1 files changed, 113 insertions, 0 deletions
diff --git a/vndbapi/src/conn.rs b/vndbapi/src/conn.rs new file mode 100644 index 0000000..a642a59 --- /dev/null +++ b/vndbapi/src/conn.rs @@ -0,0 +1,113 @@ +use std::io::Write; +use std::sync::Arc; +use std::error::Error; +use std::net::TcpStream; +use netbuf::Buf; +use rustls::{Session,ClientSession,ClientConfig}; + +use msg::Message; + +const VNDB_HOST: &'static str = "api.vndb.org"; +const VNDB_PORT_RAW: u16 = 19534; +const VNDB_PORT_TLS: u16 = 19535; + + +pub struct Connection { + stream: TcpStream, + tls: Option<ClientSession>, + rbuf: Buf, + wbuf: Buf, // Only used for raw connections. In the case of TLS we directly write to the ClientSession buffer. +} + + +impl Connection { + // Connects to the remote API, but does not yet start the TLS handshake. This is deferred to + // the first recv(), to allow rustls to send the initial command during the handshake. + pub fn connect(hostname: Option<&str>, port: Option<u16>, tls: Option<&Arc<ClientConfig>>) -> Result<Connection,Box<Error>> { + let hostname = hostname.unwrap_or(VNDB_HOST); + let port = port.unwrap_or(if tls.is_some() { VNDB_PORT_TLS } else { VNDB_PORT_RAW }); + + Ok(Connection { + stream: TcpStream::connect((hostname, port))?, + tls: tls.map(|c| ClientSession::new(c, hostname)), + rbuf: Buf::new(), + wbuf: Buf::new(), + }) + } + + // Push a new command to the output buffer. Doesn't actually send anything until a recv() is + // performed. Can be used to pipeline multiple commands. + pub fn send(&mut self, cmd: &Message) { + let w = self.tls.as_mut().map(|x| x as &mut Write).unwrap_or(&mut self.wbuf); + write!(w, "{}\x04", cmd).unwrap(); + } + + // XXX: Both this io_tls and the io_raw code may hang if the OS send buffers are full and the + // server is providing backpressure while waiting for us to read from our recv buffers. + // Properly handling that scenario requires async I/O. In practice, the server will disconnect + // us before that happens. + fn io_tls(&mut self) -> Result<(), Box<Error>> { + let tls = self.tls.as_mut().unwrap(); + while tls.wants_write() { + tls.write_tls(&mut self.stream)?; + } + if tls.wants_read() { + tls.read_tls(&mut self.stream)?; + tls.process_new_packets()?; + self.rbuf.read_from(tls)?; + } + Ok(()) + } + + fn io_raw(&mut self) -> Result<(), Box<Error>> { + while self.wbuf.len() > 0 { + self.wbuf.write_to(&mut self.stream)?; + } + self.rbuf.read_from(&mut self.stream)?; + Ok(()) + } + + // Send any outstanding commands and wait for a reply. Only returns a single response, if you + // expected multiple responses, call this method multiple times. This method blocks until at + // least one full response has been received. + pub fn recv(&mut self) -> Result<Message, Box<Error>> { + let len = { + while !self.rbuf.as_ref().contains(&4) { + if self.tls.is_some() { + self.io_tls()? + } else { + self.io_raw()? + } + } + self.rbuf.as_ref().split(|&c| c == 4).next().unwrap().len() + }; + + let m = Message::parse(::std::str::from_utf8(&self.rbuf[..len])?)?; + self.rbuf.consume(len+1); + Ok(m) + } + + // WARNING: This method should not be called when there are still outstanding commands. + // (Might as well get rid of this convenient method to prevent that scenario?) + pub fn cmd(&mut self, cmd: &Message) -> Result<Message, Box<Error>> { + self.send(cmd); + self.recv() + } +} + + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn connect() { + let mut c = Connection::connect(None, None, Some(&::TLS_CONF)).unwrap(); + c.send(&Message::parse("login {\"protocol\":1,\"client\":\"vndb-rust\",\"clientver\":\"0.1\"}").unwrap()); + c.send(&Message::parse("dbstats").unwrap()); + let r1 = format!("{}", c.recv().unwrap()); + let r2 = format!("{}", c.recv().unwrap()); + assert_eq!(&r1, "ok"); + assert!(r2.starts_with("dbstats {")); + } +} |