From ffa21e41c01b5bcb916828d4141fc1c6624c7b5c Mon Sep 17 00:00:00 2001 From: Evie Viau Date: Thu, 17 Feb 2022 17:32:13 -0500 Subject: [PATCH] Add check for opcode --- Cargo.lock | 32 +++++++++++++++++++++ Cargo.toml | 4 +++ src/main.rs | 78 ++++++++++++++++++++++++++++++++++++++++++++++++-- src/opcodes.rs | 21 ++++++++++++-- 4 files changed, 131 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c8423ce..3bdf475 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -65,6 +65,12 @@ dependencies = [ "generic-array", ] +[[package]] +name = "dotenv" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77c90badedccf4105eca100756a0b1289e191f6fcbdadd3cee1d2f614f97da8f" + [[package]] name = "fnv" version = "1.0.7" @@ -87,6 +93,17 @@ version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c09fd04b7e4073ac7156a9539b57a484a8ea920f79c7c675d05d289ab6110d3" +[[package]] +name = "futures-macro" +version = "0.3.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33c1e13800337f4d4d7a316bf45a567dbcb6ffe087f16424852d97e97a91f512" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.21" @@ -106,6 +123,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8b7abd5d659d9b90c8cba917f6ec750a74e2dc23902ef9cd4cc8c8b22e6036a" dependencies = [ "futures-core", + "futures-macro", "futures-sink", "futures-task", "pin-project-lite", @@ -205,11 +223,14 @@ dependencies = [ name = "lvsp-server" version = "0.1.0" dependencies = [ + "dotenv", + "futures-util", "num", "num-derive", "num-traits", "serde", "serde_json", + "serde_repr", "tokio", "tokio-tungstenite", ] @@ -513,6 +534,17 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_repr" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98d0516900518c29efa217c298fa1f4e6c6ffc85ae29fd7f4ee48f176e1a9ed5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "sha-1" version = "0.9.8" diff --git a/Cargo.toml b/Cargo.toml index 686a350..35b31f0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,5 +14,9 @@ num-traits = "0.2.14" serde = { version = "1.0.136", features = ["derive"]} serde_json = "1.0.79" +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 diff --git a/src/main.rs b/src/main.rs index e448f1f..74989b4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,82 @@ #[macro_use] extern crate num_derive; +use std::io::Error; + +use dotenv::dotenv; +use std::env; +use std::net::SocketAddr; +use std::time::Duration; +use tokio::net::{TcpListener, TcpStream}; + +use futures_util::{future, SinkExt, StreamExt, TryStreamExt}; +use tokio_tungstenite::tungstenite::Message; +use crate::opcodes::check_if_opcode; + mod opcodes; -fn main() { - println!("the funny"); +#[tokio::main] +async fn main() -> Result<(), Error> { + dotenv().ok(); + + let addr = env::var("LISTEN_ADDR").unwrap_or("0.0.0.0:3621".to_string()); + + let socket = TcpListener::bind(&addr).await.expect("Failed to bind to address!"); + println!("Listening on {}!", &addr); + + while let Ok((stream, _)) = socket.accept().await { + let peer = stream.peer_addr().expect("Failed to connect to peer, missing address?"); + println!("Connecting to peer {}...", &peer); + + tokio::spawn(accept_conn(peer, stream)); + } + + Ok(()) } + +async fn accept_conn(peer: SocketAddr, stream: TcpStream) { + if let Err(e) = handle_conn(peer, stream).await { + match e { + tokio_tungstenite::tungstenite::Error::ConnectionClosed | tokio_tungstenite::tungstenite::Error::Protocol(_) | tokio_tungstenite::tungstenite::Error::Utf8 => (), + err => eprintln!("Error accepting connection from {}!", &peer), + } + } +} + +async fn handle_conn(peer: SocketAddr, stream: TcpStream) -> tokio_tungstenite::tungstenite::Result<()> { + let ws_stream = tokio_tungstenite::accept_async(stream) + .await + .expect("Failed to complete the websocket handshake!"); + 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 + + loop { + tokio::select! { + msg = ws_receiver.next() => { + match msg { + Some(msg) => { + let msg = msg?; + + if msg.is_text() { + if check_if_opcode(msg.clone()).is_ok() { + println!("valid") + } else { + ws_sender.send(Message::Text((opcodes::ErrorCode::DECODE as i32).to_string())).await?; + } + } else if msg.is_close() { + break; + } + }, + None => break, + } + }, + _ = heartbeat.tick() => { + ws_sender.send(Message::Text("deez".to_owned())).await?; + } + } + } + + Ok(()) +} \ No newline at end of file diff --git a/src/opcodes.rs b/src/opcodes.rs index fd06f7d..a3951cb 100644 --- a/src/opcodes.rs +++ b/src/opcodes.rs @@ -10,10 +10,14 @@ //! snowflake type: A string encoding a Discord Snowflake. //! //! [Source](https://gitlab.com/litecord/litecord/-/blob/master/docs/lvsp.md) +use std::any::Any; use serde::{Serialize, Deserialize}; +use serde_repr::{Serialize_repr, Deserialize_repr}; +use tokio_tungstenite::tungstenite::Message; /// Op codes sent/received by Litecord -#[derive(FromPrimitive, Deserialize, Serialize)] +#[derive(FromPrimitive, Serialize_repr, Deserialize_repr)] +#[repr(u8)] pub enum OpCode { /// Sent by the server when a connection is established. HELLO = 0, @@ -56,6 +60,7 @@ pub enum ErrorCode { /// Message data for the socket #[derive(Deserialize, Serialize)] +#[serde(untagged)] pub enum MessageData { /// Sent by the server when a connection is established. HELLO { @@ -108,7 +113,8 @@ pub enum MessageData { } /// Info message types -#[derive(Deserialize, Serialize)] +#[derive(Serialize_repr, Deserialize_repr)] +#[repr(u8)] pub enum InfoType { /// Request a channel to be created inside the voice server. CHANNEL_REQ = 0, @@ -136,6 +142,7 @@ pub enum InfoType { /// Info message data #[derive(Deserialize, Serialize)] +#[serde(untagged)] pub enum InfoData { /// Request a channel to be created inside the voice server. /// @@ -217,4 +224,14 @@ struct SocketMessage { /// Message data d: MessageData +} + +pub fn check_if_opcode(msg: Message) -> Result<(), ()> { + 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)) + } else { + Err(()) + } } \ No newline at end of file