Implement HELLO and HEARTBEAT

also prevents server only ops being sent to the server from the client

Cargo.lock

@@ -228,6 +228,7 @@ dependencies = [
  "num",
  "num-derive",
  "num-traits",
+ "rand",
  "serde",
  "serde_json",
  "serde_repr",

Cargo.toml

@@ -19,4 +19,5 @@ serde_repr = "0.1.7"
 tokio-tungstenite = "0.16.1"
 
 dotenv = "0.15.0"
-futures-util = "0.3.21"
+futures-util = "0.3.21"
+rand = "0.8.5"

src/main.rs

@@ -11,7 +11,13 @@ use tokio::net::{TcpListener, TcpStream};
 
 use futures_util::{future, SinkExt, StreamExt, TryStreamExt};
 use tokio_tungstenite::tungstenite::Message;
-use crate::opcodes::check_if_opcode;
+use crate::OpCode::{HEARTBEAT_ACK, HELLO};
+use crate::opcodes::{check_if_opcode, MessageData, OpCode, SocketMessage};
+
+use rand::prelude::*;
+use rand::distributions::Alphanumeric;
+
+use serde_json::Value::Array;
 
 mod opcodes;
 
@@ -50,7 +56,25 @@ async fn handle_conn(peer: SocketAddr, stream: TcpStream) -> tokio_tungstenite::
     println!("Connected to peer: {}!", &peer);
 
     let (mut ws_sender, mut ws_receiver) = ws_stream.split();
-    let mut heartbeat = tokio::time::interval(Duration::from_millis(1000)); // We need to get this from the HELLO op
+    let mut heartbeat = tokio::time::interval(Duration::from_millis(1000));
+
+    let mut nonce: String = rand::thread_rng()
+        .sample_iter(&Alphanumeric)
+        .take(10)
+        .map(char::from)
+        .collect();
+
+    ws_sender.send(Message::Text(
+        serde_json::to_string(
+            &SocketMessage {
+                op: HELLO,
+                d: MessageData::HELLO {
+                    heartbeat_interval: 1000,
+                    nonce
+                }
+            }
+        ).unwrap().to_owned()
+    )).await?;
 
     loop {
         tokio::select! {
@@ -60,8 +84,38 @@ async fn handle_conn(peer: SocketAddr, stream: TcpStream) -> tokio_tungstenite::
                         let msg = msg?;
 
                         if msg.is_text() {
-                            if check_if_opcode(msg.clone()).is_ok() {
-                                println!("valid")
+                            let op = check_if_opcode(msg.clone());
+                            if op.is_ok() {
+                                match op.unwrap().0 {
+                                    OpCode::IDENTIFY => {
+                                        unimplemented!()
+                                    }
+
+                                    OpCode::RESUME => {
+                                        unimplemented!()
+                                    }
+
+                                    OpCode::HEARTBEAT => {
+                                        ws_sender.send(Message::Text(
+                                            serde_json::to_string(
+                                                &SocketMessage {
+                                                    op: HEARTBEAT_ACK,
+                                                    d: MessageData::HEARTBEAT_ACK {
+                                                        health: 1.0 // trust
+                                                    }
+                                                }
+                                            ).unwrap().to_owned()
+                                        )).await?;
+                                    }
+
+                                    OpCode::INFO => {
+                                        unimplemented!()
+                                    },
+
+                                    _ => {
+                                        ws_sender.send(Message::Text((opcodes::ErrorCode::DECODE as i32).to_string())).await?;
+                                    }
+                                }
                             } else {
                                  ws_sender.send(Message::Text((opcodes::ErrorCode::DECODE as i32).to_string())).await?;
                             }
@@ -73,7 +127,7 @@ async fn handle_conn(peer: SocketAddr, stream: TcpStream) -> tokio_tungstenite::
                 }
             },
             _ = heartbeat.tick() => {
-                ws_sender.send(Message::Text("deez".to_owned())).await?;
+                //ws_sender.send(Message::Text("deez".to_owned())).await?;
             }
         }
     }

src/opcodes.rs

@@ -86,7 +86,7 @@ pub enum MessageData {
     ///
     /// The server MUST reply with a HEARTBEAT_ACK message back in a reasonable
     /// time period.
-    HEARTBEAT,
+    HEARTBEAT {},
 
     /// Sent by the server in reply to a HEARTBEAT message coming from the client.
     ///
@@ -218,19 +218,21 @@ pub enum InfoData {
 /// **Note:** the snowflake type follows the same rules as the Discord Gateway's
 /// snowflake type: A string encoding a Discord Snowflake.
 #[derive(Deserialize, Serialize)]
-struct SocketMessage {
+pub struct SocketMessage {
     /// Operator code
-    op: OpCode,
+    pub op: OpCode,
 
     /// Message data
-    d: MessageData
+    pub d: MessageData
 }
 
-pub fn check_if_opcode(msg: Message) -> Result<(), ()> {
+pub fn check_if_opcode(msg: Message) -> Result<(OpCode, MessageData), ()> {
     let message_json: Result<SocketMessage, serde_json::Error> = serde_json::from_str(msg.to_text().expect("Failed to convert message to str!"));
 
     if message_json.is_ok() {
-        Ok(println!("{}", message_json.unwrap().op as u8))
+        let output = message_json.unwrap();
+
+        Ok((output.op, output.d))
     } else {
         Err(())
     }