diff --git a/Cargo.lock b/Cargo.lock index 3bdf475..a5a31c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -228,6 +228,7 @@ dependencies = [ "num", "num-derive", "num-traits", + "rand", "serde", "serde_json", "serde_repr", diff --git a/Cargo.toml b/Cargo.toml index 35b31f0..f2a8e94 100644 --- a/Cargo.toml +++ b/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" \ No newline at end of file +futures-util = "0.3.21" +rand = "0.8.5" \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 74989b4..b6bce46 100644 --- a/src/main.rs +++ b/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?; } } } diff --git a/src/opcodes.rs b/src/opcodes.rs index a3951cb..34c8377 100644 --- a/src/opcodes.rs +++ b/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 = 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(()) }