Add check for opcode

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"

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"

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(())
+}

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<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))
+    } else {
+        Err(())
+    }
 }