Move handshake and ping into Connection trait

Cargo.toml

@@ -19,6 +19,7 @@ serde_derive = "^1.0"
 serde_json   = "^1.0"
 byteorder    = "^1.0"
 log          = "~0.4"
+bytes        = "^0.4"
 
 [target.'cfg(windows)'.dependencies]
 named_pipe   = "0.3.0"

src/connection/base.rs

@@ -1,17 +1,37 @@
 use std::{
-    io::{Write, Read},
+    io::{Write, Read, ErrorKind},
     marker::Sized,
     path::PathBuf,
+    thread,
+    time,
 };
 
-use models::message::Message;
-use error::Result;
+use bytes::BytesMut;
+
+use utils;
+use models::message::{Message, OpCode};
+use error::{Error, Result};
+
+
+/// Wait for a non-blocking connection until it's complete.
+macro_rules! try_until_done {
+    [ $e:expr ] => {
+        loop {
+            match $e {
+                Ok(_) => break,
+                Err(Error::IoError(ref err)) if err.kind() == ErrorKind::WouldBlock => (),
+                Err(why) => return Err(why),
+            }
+
+            thread::sleep(time::Duration::from_millis(500));
+        }
+    }
+}
 
 
 pub trait Connection: Sized {
     type Socket: Write + Read;
 
-
     fn socket(&mut self) -> &mut Self::Socket;
 
     fn ipc_path() -> PathBuf;
@@ -24,6 +44,26 @@ pub trait Connection: Sized {
         Self::ipc_path().join(format!("discord-ipc-{}", n))
     }
 
+    fn handshake(&mut self, client_id: u64) -> Result<()> {
+        let hs = json![{
+            "client_id": client_id.to_string(),
+            "v": 1,
+            "nonce": utils::nonce()
+        }];
+
+        try_until_done!(self.send(Message::new(OpCode::Handshake, hs.clone())));
+        try_until_done!(self.recv());
+
+        Ok(())
+    }
+
+    fn ping(&mut self) -> Result<OpCode> {
+        let message = Message::new(OpCode::Ping, json![{}]);
+        self.send(message)?;
+        let response = self.recv()?;
+        Ok(response.opcode)
+    }
+
     fn send(&mut self, message: Message) -> Result<()> {
         match message.encode() {
             Err(why) => error!("{:?}", why),
@@ -36,11 +76,19 @@ pub trait Connection: Sized {
     }
 
     fn recv(&mut self) -> Result<Message> {
-        let mut buf: Vec<u8> = vec![0; 1024];
-        let n = self.socket().read(buf.as_mut_slice())?;
-        buf.resize(n, 0);
+        let mut buf = BytesMut::new();
+        buf.resize(1024, 0);
+        let n = self.socket().read(&mut buf)?;
+        debug!("Received {} bytes", n);
+
+        if n == 0 {
+            return Err(Error::ConnectionClosed);
+        }
+
+        buf = buf.split_to(n);
         let message = Message::decode(&buf)?;
         debug!("<- {:?}", message);
+
         Ok(message)
     }
 }

src/error.rs

@@ -19,6 +19,7 @@ pub enum Error {
     Timeout(ChannelTimeout),
     Conversion,
     SubscriptionFailed,
+    ConnectionClosed,
 }
 
 impl Display for Error {
@@ -32,6 +33,7 @@ impl StdError for Error {
         match *self {
             Error::Conversion => "Failed to convert values",
             Error::SubscriptionFailed => "Failed to subscribe to event",
+            Error::ConnectionClosed => "Connection closed",
             Error::IoError(ref err) => err.description(),
             Error::JsonError(ref err) => err.description(),
             Error::Timeout(ref err) => err.description(),

src/lib.rs

@@ -7,6 +7,7 @@ extern crate serde;
 extern crate serde_json;
 extern crate byteorder;
 extern crate uuid;
+extern crate bytes;
 #[cfg(windows)]
 extern crate named_pipe;