Cargo.lock

@@ -5,6 +5,8 @@ name = "avarice"
 version = "0.1.0"
 dependencies = [
  "custom_error 1.9.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "serde 1.0.126 (registry+https://github.com/rust-lang/crates.io-index)",
+ "serde_json 1.0.66 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
 [[package]]
@@ -12,5 +14,34 @@ name = "custom_error"
 version = "1.9.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 
+[[package]]
+name = "itoa"
+version = "0.4.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+
+[[package]]
+name = "ryu"
+version = "1.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+
+[[package]]
+name = "serde"
+version = "1.0.126"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+
+[[package]]
+name = "serde_json"
+version = "1.0.66"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "itoa 0.4.7 (registry+https://github.com/rust-lang/crates.io-index)",
+ "ryu 1.0.5 (registry+https://github.com/rust-lang/crates.io-index)",
+ "serde 1.0.126 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
 [metadata]
 "checksum custom_error 1.9.2 (registry+https://github.com/rust-lang/crates.io-index)" = "4f8a51dd197fa6ba5b4dc98a990a43cc13693c23eb0089ebb0fcc1f04152bca6"
+"checksum itoa 0.4.7 (registry+https://github.com/rust-lang/crates.io-index)" = "dd25036021b0de88a0aff6b850051563c6516d0bf53f8638938edbb9de732736"
+"checksum ryu 1.0.5 (registry+https://github.com/rust-lang/crates.io-index)" = "71d301d4193d031abdd79ff7e3dd721168a9572ef3fe51a1517aba235bd8f86e"
+"checksum serde 1.0.126 (registry+https://github.com/rust-lang/crates.io-index)" = "ec7505abeacaec74ae4778d9d9328fe5a5d04253220a85c4ee022239fc996d03"
+"checksum serde_json 1.0.66 (registry+https://github.com/rust-lang/crates.io-index)" = "336b10da19a12ad094b59d870ebde26a45402e5b470add4b5fd03c5048a32127"

Cargo.toml

@@ -7,4 +7,6 @@ edition = "2018"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-custom_error="1.9.2"
+custom_error="1.*"
+serde ="1.*"
+serde_json="1.*"

src/link/message.rs

@@ -0,0 +1,38 @@
+//! Implements Avarice message: target service, message type, json parameters
+use serde_json;
+use serde_json::json;
+use std::fmt;
+
+pub struct AvariceMessage {
+    pub service: String,
+    pub message_type: String,
+    pub parameters: serde_json::Value,
+}
+
+impl AvariceMessage {
+    /// Parses JSON form of a message into `AvariceMessage` struct
+    pub fn from(message_str: &str) -> Option<AvariceMessage> {
+        let mut message_json: serde_json::Value = serde_json::from_str(message_str).unwrap();
+        let message_json = message_json.as_object_mut()?;
+        Some(AvariceMessage {
+            service: message_json.remove("s")?.as_str()?.to_owned(),
+            message_type: message_json.remove("t")?.as_str()?.to_owned(),
+            parameters: message_json.remove("p")?,
+        })
+    }
+}
+
+impl fmt::Display for AvariceMessage {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "{}",
+            json!({
+                "s": self.service,
+                "t": self.message_type,
+                "p": self.parameters,
+            })
+            .to_string()
+        )
+    }
+}

src/link/mod.rs

@@ -1,62 +1,111 @@
-//! Implements reader and writer to use when talking to UE2 game server.
-
+//! This module provides a simple interface to exchange messages (`message::AvariceMessage`) between
+//! Avarice and ue-server. The model is simple - we create an mpsc-channel of `Sender` and
+//! `Receiver`, then pass `Sender` to the `network` sub-module and wait for its messages about
+//! ue-servers' connections using `Receiver`.
+use std::collections::HashMap;
 use std::error::Error;
-use std::io::{Read, Write};
-use std::net::{SocketAddr, TcpListener};
-use std::sync::Arc;
-use std::{str, thread};
-
-mod reader;
-mod writer;
-pub use reader::MessageReader;
+use std::io::Write;
+use std::net::{SocketAddr, TcpStream};
+use std::sync::mpsc::{channel, Receiver};
 pub use writer::MessageWriter;
 
+mod message;
+mod network;
+mod reader;
+mod writer;
+pub use message::AvariceMessage;
+pub use network::{run_server, NetworkMessage};
+
+/// For collecting messages from all connected ue-servers and providing a way to reply back.
+pub struct AvariceServer {
+    connected_links: HashMap<SocketAddr, Link>,
+    receiver: Receiver<NetworkMessage>,
+}
+
+/// For representing a link to one of the connected ue-servers, can be used to send messages back.
+/// To receive messages use `AvariceServer`'s `next()` method instead.
 pub struct Link {
-    reader: MessageReader,
+    ue_server_address: SocketAddr,
     writer: MessageWriter,
+    writing_stream: TcpStream,
+}
+
+impl AvariceServer {
+    /// Blocks until a new message arrives from one of the ue-servers. Returns a pair of `Link`,
+    /// corresponding to the ue-server that sent next message and `AvariceMessage`,
+    /// representing that message.
+    pub fn next(&mut self) -> Option<(&mut Link, AvariceMessage)> {
+        loop {
+            match self.receiver.recv() {
+                Ok(NetworkMessage::ConnectionEstablished(ue_server_address, writing_stream)) => {
+                    //  If `ue_server_address` was already present in `self.connected_links`
+                    //  hash map, then it means we have failed to clean it up after it
+                    //  has disconnected. We can just throw away the old value here.
+                    self.connected_links.insert(
+                        ue_server_address,
+                        Link {
+                            ue_server_address,
+                            writing_stream,
+                            writer: MessageWriter::new(),
+                        },
+                    );
+                    continue;
+                }
+                Ok(NetworkMessage::ConnectionLost(ue_server_address)) => {
+                    self.connected_links.remove(&ue_server_address);
+                    continue;
+                }
+                Ok(NetworkMessage::InvalidDataReceived(ue_server_address)) => {
+                    self.connected_links.remove(&ue_server_address);
+                    continue;
+                }
+                Ok(NetworkMessage::UEReceivedUpdate(ue_server_address, ue_received_bytes)) => {
+                    if let Some(link) = self.connected_links.get_mut(&ue_server_address) {
+                        link.update_ue_received_bytes(ue_received_bytes)
+                    }
+                    continue;
+                }
+                Ok(NetworkMessage::MessageReceived(ue_server_address, message)) => {
+                    //  Not having a link with key `ue_server_address` should be impossible here
+                    return self
+                        .connected_links
+                        .get_mut(&ue_server_address)
+                        .and_then(|x| Some((x, message)));
+                }
+                _ => return None,
+            }
+        }
+    }
 }
 
 impl Link {
-    pub fn run<F>(port: u16, handler: F) -> Result<(), Box<dyn Error>>
-    where
-        F: Fn(&mut Link, &str) -> () + Send + Sync + 'static,
-    {
-        let address = SocketAddr::from(([0, 0, 0, 0], port));
-        let listener = TcpListener::bind(address)?;
-        let handler = Arc::new(handler);
-        loop {
-            //  Listen to new (multiple) connection
-            let mut reading_stream = listener.accept()?.0;
-            let mut writing_stream = reading_stream.try_clone()?;
-            let mut avarice_link = Link {
-                reader: MessageReader::new(),
-                writer: MessageWriter::new(),
-            };
-            let handler_clone = handler.clone();
-            //  On connection - spawn a new thread
-            thread::spawn(move || loop {
-                let mut buffer = [0; 1024];
-                //  Reading cycle
-                match reading_stream.read(&mut buffer) {
-                    Ok(n) => avarice_link.reader.push(&buffer[..n]).unwrap(),
-                    _ => panic!("Connection issue!"),
-                };
-                //  Handling cycle
-                while let Some(message) = avarice_link.reader.pop() {
-                    handler_clone(&mut avarice_link, &message);
-                }
-                //  Writing
-                avarice_link
-                    .writer
-                    .update_ue_received_bytes(avarice_link.reader.ue_received_bytes());
-                if let Some(bytes) = avarice_link.writer.try_pop() {
-                    writing_stream.write_all(&bytes).unwrap();
-                }
-            });
-        }
+    pub fn send(&mut self, message: AvariceMessage) {
+        self.writer.push(&message.to_string());
+        self.flush();
     }
 
-    pub fn write(&mut self, message: &str) {
-        self.writer.push(message);
+    pub fn socket_address(&self) -> SocketAddr {
+        self.ue_server_address
+    }
+
+    fn update_ue_received_bytes(&mut self, ue_received_bytes: u64) {
+        self.writer.update_ue_received_bytes(ue_received_bytes);
+        self.flush();
+    }
+
+    fn flush(&mut self) {
+        if let Some(bytes) = self.writer.try_pop() {
+            self.writing_stream.write_all(&bytes).unwrap();
+        }
     }
 }
+
+/// Creates a new `AvariceServer` that will listen for ue-server connections on the specified port.
+pub fn start_avarice(port: u16) -> Result<AvariceServer, Box<dyn Error>> {
+    let (sender, receiver) = channel();
+    run_server(port, sender)?;
+    Ok(AvariceServer {
+        connected_links: HashMap::new(),
+        receiver,
+    })
+}

src/link/network.rs

@@ -0,0 +1,106 @@
+//!     Implements a network model where messages from all the ue-servers are collected in a single
+//! main thread. For that we spawn a new thread that listens for new connections, which in turn
+//! spawns a new thread for every connected ue-server to handle reading data from it.
+//!     Since all reading is handled in ue-servers' own threads, to collect messages they have
+//! received in the main thread we use `std::sync::mpsc::Sender`. Conversely, all the writing to
+//! ue-server is handled in the main thread itself. Writing `TcpStream` is sent to the main thread
+//! by the same `std::sync::mpsc::Sender` object.
+use super::message::AvariceMessage;
+pub use super::reader::MessageReader;
+use std::error::Error;
+use std::io::Read;
+use std::net::{SocketAddr, TcpListener, TcpStream};
+use std::sync::mpsc::Sender;
+use std::thread;
+
+pub struct UEConnection {
+    pub address: SocketAddr,
+    pub reader: MessageReader,
+    pub reading_stream: TcpStream,
+    pub message_sender: Sender<NetworkMessage>,
+}
+
+/// Possible messages to the main thread
+pub enum NetworkMessage {
+    ConnectionEstablished(SocketAddr, TcpStream),
+    InvalidDataReceived(SocketAddr),
+    ConnectionLost(SocketAddr),
+    MessageReceived(SocketAddr, AvariceMessage),
+    UEReceivedUpdate(SocketAddr, u64),
+}
+
+pub fn run_server(port: u16, message_sender: Sender<NetworkMessage>) -> Result<(), Box<dyn Error>> {
+    let address = SocketAddr::from(([0, 0, 0, 0], port));
+    let listener = TcpListener::bind(address)?;
+    thread::spawn(move || loop {
+        //  Listen to new (multiple) connection
+        let (reading_stream, address) = listener.accept().unwrap();
+        let writing_stream = reading_stream.try_clone().unwrap();
+        message_sender
+            .send(NetworkMessage::ConnectionEstablished(
+                address,
+                writing_stream,
+            ))
+            .unwrap();
+        //  On connection - spawn a new thread
+        let sender_clone = message_sender.clone();
+        thread::spawn(move || {
+            manage_connection(UEConnection {
+                reader: MessageReader::new(),
+                message_sender: sender_clone,
+                reading_stream,
+                address,
+            })
+        });
+    });
+    Ok(())
+}
+
+fn manage_connection(mut connection: UEConnection) {
+    let mut buffer = [0; 1024];
+    loop {
+        //  Reading cycle
+        match connection.reading_stream.read(&mut buffer) {
+            Ok(n) => connection.reader.push(&buffer[..n]).unwrap(),
+            _ => {
+                connection
+                    .message_sender
+                    .send(NetworkMessage::ConnectionLost(connection.address))
+                    .unwrap();
+                return;
+            }
+        };
+        if connection.reader.is_broken() {
+            connection
+                .message_sender
+                .send(NetworkMessage::InvalidDataReceived(connection.address))
+                .unwrap();
+            return;
+        }
+        //  Decoding cycle
+        while let Some(text_message) = connection.reader.pop() {
+            if let Some(avarice_message) = AvariceMessage::from(&text_message) {
+                connection
+                    .message_sender
+                    .send(NetworkMessage::MessageReceived(
+                        connection.address,
+                        avarice_message,
+                    ))
+                    .unwrap();
+            } else {
+                connection
+                    .message_sender
+                    .send(NetworkMessage::InvalidDataReceived(connection.address))
+                    .unwrap();
+                return;
+            }
+        }
+        connection
+            .message_sender
+            .send(NetworkMessage::UEReceivedUpdate(
+                connection.address,
+                connection.reader.ue_received_bytes(),
+            ))
+            .unwrap();
+    }
+}

src/main.rs

@@ -1,9 +1,14 @@
+// ! This utility provides a way to communicate with Unreal Engine 2 servers running Acedia mutator.
 mod link;
-use link::Link;
+use link::start_avarice;
 
 fn main() {
-    match Link::run(1234, |link, message| { link.write(message);}) {
-        Ok(_) => print!("Connect!"),
-        _ => print!("Connection error!"),
-    };
+    let mut server = start_avarice(1234).unwrap();
+    while let Some((link, message)) = server.next() {
+        println!("{}: {}", link.socket_address(), message.to_string());
+        if message.message_type != "end" {
+            link.send(message);
+        }
+    }
+    println!("Avarice has shut down!");
 }