230 lines
8.8 KiB
Rust
230 lines
8.8 KiB
Rust
//! Implements writer that sends data to UE2 game server.
|
|
|
|
use std::cmp::{max, min};
|
|
use std::collections::VecDeque;
|
|
use std::convert::TryFrom;
|
|
use std::iter::Extend;
|
|
|
|
// Defines how many bytes is used to encode "LENGTH" field in the chunk sent to ue-server
|
|
const CHUNK_LENGTH_FIELD: usize = 2;
|
|
// Maximum amount of bytes ue-server is able to receive at once
|
|
const UE_INPUT_BUFFER: usize = 4095;
|
|
// Minimal payload size (in bytes) to send, unless there is not enough data left
|
|
const MIN_PAYLOAD_SIZE: usize = 50;
|
|
|
|
/// For converting byte stream that is expected from the ue-server into actual messages.
|
|
/// Conversion process has two steps:
|
|
/// 1. Every string message is converted into it's utf8 representation and is pre-pended with
|
|
/// it's own length in format:
|
|
///
|
|
/// | Data | Length |
|
|
/// |---------|---------|
|
|
/// | Message `LENGTH` | 4 bytes: u32 BE |
|
|
/// | UTF8-encoded string | `LENGTH` bytes|
|
|
///
|
|
/// Resulting byte sequences from all the messages then concatenated, in order, into
|
|
/// a single data stream.
|
|
///
|
|
/// 2. Resulting data stream is then separated into "chunks" that can be accepted by
|
|
/// the ue-server (each no longer than `UE_INPUT_BUFFER` in total) and are sent in a format:
|
|
///
|
|
/// | Data | Length |
|
|
/// |---------|---------|
|
|
/// | Chunk `LENGTH` | 2 bytes: u16 BE |
|
|
/// | UTF8-encoded string | `LENGTH` bytes|
|
|
///
|
|
/// Use `push()` to input string messages and `try_pop()` to retrieve next chunk, if ue-server
|
|
/// can accept it.
|
|
/// NOTE: `try_pop()` can return `None` even if not all message data has been transferred,
|
|
/// in case ue-server's buffer does not have enough space.
|
|
///
|
|
/// Call `update_ue_received_bytes()` to update `MessageWriter`'s information about
|
|
/// how many bytes ue-server has received so far.
|
|
///
|
|
/// Use `is_empty()` call to check for whether `MessageWriter` has depleted all it's data.
|
|
pub struct MessageWriter {
|
|
sent_bytes: u64,
|
|
ue_received_bytes: u64,
|
|
pending_data: VecDeque<u8>,
|
|
}
|
|
|
|
impl MessageWriter {
|
|
pub fn new() -> MessageWriter {
|
|
MessageWriter {
|
|
sent_bytes: 0,
|
|
ue_received_bytes: 0,
|
|
// This value should be more than enough for typical use:
|
|
// it will take at least one second to send this much data to a 30 tick rate ue-server.
|
|
pending_data: VecDeque::with_capacity(30 * UE_INPUT_BUFFER),
|
|
}
|
|
}
|
|
|
|
pub fn push(&mut self, message: &str) {
|
|
let message_as_utf8 = message.as_bytes();
|
|
let message_as_utf8 = [
|
|
&(message_as_utf8.len() as u32).to_be_bytes(),
|
|
message_as_utf8,
|
|
]
|
|
.concat();
|
|
self.pending_data.extend(message_as_utf8.into_iter());
|
|
}
|
|
|
|
/// This method will always return chunk if all remaining data will fit inside it, otherwise it
|
|
/// will wait until ue-server's buffer has enough space for at least `MIN_PAYLOAD_SIZE` bytes.
|
|
pub fn try_pop(&mut self) -> Option<Vec<u8>> {
|
|
if self.is_empty() {
|
|
return None;
|
|
}
|
|
let required_payload_size = min(self.pending_data.len(), MIN_PAYLOAD_SIZE);
|
|
let available_payload_space = self
|
|
.available_ue_buffer_capacity()
|
|
.checked_sub(CHUNK_LENGTH_FIELD)
|
|
.unwrap_or_default();
|
|
if required_payload_size > available_payload_space {
|
|
return None;
|
|
}
|
|
let payload_size = min(available_payload_space, self.pending_data.len());
|
|
let mut bytes_to_send = Vec::with_capacity(CHUNK_LENGTH_FIELD + payload_size);
|
|
bytes_to_send.extend((payload_size as u16).to_be_bytes().iter());
|
|
for next_byte in self.pending_data.drain(..payload_size) {
|
|
bytes_to_send.push(next_byte);
|
|
}
|
|
self.sent_bytes += bytes_to_send.len() as u64;
|
|
Some(bytes_to_send)
|
|
}
|
|
|
|
pub fn is_empty(&self) -> bool {
|
|
self.pending_data.is_empty()
|
|
}
|
|
|
|
pub fn update_ue_received_bytes(&mut self, ue_received_bytes: u64) {
|
|
self.ue_received_bytes = max(ue_received_bytes, self.ue_received_bytes);
|
|
}
|
|
|
|
fn available_ue_buffer_capacity(&self) -> usize {
|
|
match usize::try_from(self.sent_bytes - self.ue_received_bytes).ok() {
|
|
Some(ue_buffered_bytes) => max(0, UE_INPUT_BUFFER - ue_buffered_bytes),
|
|
_ => 0,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn writer_contents_after_creation() {
|
|
let writer = MessageWriter::new();
|
|
assert_eq!(writer.is_empty(), true);
|
|
assert_eq!(writer.available_ue_buffer_capacity(), UE_INPUT_BUFFER);
|
|
}
|
|
|
|
#[test]
|
|
fn writer_content_after_push() {
|
|
let mut writer = MessageWriter::new();
|
|
writer.push("Hello, world!");
|
|
assert_eq!(writer.is_empty(), false);
|
|
assert_eq!(writer.available_ue_buffer_capacity(), UE_INPUT_BUFFER);
|
|
}
|
|
|
|
#[test]
|
|
fn writing_single_short_message() {
|
|
let mut writer = MessageWriter::new();
|
|
writer.push("Hello, world!");
|
|
let resulting_bytes = writer.try_pop().unwrap();
|
|
let expected_bytes = [
|
|
0, 17, // Bytes in the chunk = message length (4 bytes) + message (13 bytes)
|
|
0, 0, 0, 13, // Bytes in the message
|
|
b'H', b'e', b'l', b'l', b'o', b',', b' ', b'w', b'o', b'r', b'l', b'd', b'!',
|
|
];
|
|
assert_eq!(writer.is_empty(), true);
|
|
assert_eq!(
|
|
writer.available_ue_buffer_capacity(),
|
|
UE_INPUT_BUFFER - "Hello, world!".len() - 2 - 4
|
|
);
|
|
assert_eq!(resulting_bytes, expected_bytes);
|
|
assert_eq!(writer.sent_bytes, expected_bytes.len() as u64);
|
|
}
|
|
|
|
#[test]
|
|
fn writing_first_chunk_of_single_long_message() {
|
|
let mut writer = MessageWriter::new();
|
|
// Because we also have to pass lengths, this will go over the sending limit
|
|
let long_message = "Q".repeat(UE_INPUT_BUFFER);
|
|
writer.push(&long_message);
|
|
let resulting_bytes = writer.try_pop().unwrap();
|
|
assert_eq!(writer.is_empty(), false);
|
|
assert_eq!(resulting_bytes.len(), UE_INPUT_BUFFER);
|
|
assert_eq!(writer.available_ue_buffer_capacity(), 0);
|
|
// Bytes in the chunk = 4095 - 2 = 4093 = 0x0ffd
|
|
assert_eq!(resulting_bytes[0], 0x0f);
|
|
assert_eq!(resulting_bytes[1], 0xfd);
|
|
// Bytes in message = 4095 = 0x0fff
|
|
assert_eq!(resulting_bytes[2], 0);
|
|
assert_eq!(resulting_bytes[3], 0);
|
|
assert_eq!(resulting_bytes[4], 0x0f);
|
|
assert_eq!(resulting_bytes[5], 0xff);
|
|
for &byte in resulting_bytes[6..].iter() {
|
|
assert_eq!(byte, b'Q');
|
|
}
|
|
assert_eq!(writer.try_pop(), None);
|
|
assert_eq!(writer.is_empty(), false);
|
|
}
|
|
|
|
#[test]
|
|
fn writing_second_chunk_of_single_long_message() {
|
|
let mut writer = MessageWriter::new();
|
|
// Because we also have to pass lengths, this will go over the sending limit
|
|
let long_message = "Q".repeat(UE_INPUT_BUFFER);
|
|
writer.push(&long_message);
|
|
// This pops all but 6 bytes of `long_message`, that were required to encode lengths of
|
|
// message and first chunk
|
|
let first_bytes = writer.try_pop().unwrap();
|
|
writer.update_ue_received_bytes(first_bytes.len() as u64);
|
|
let resulting_bytes = writer.try_pop().unwrap();
|
|
assert_eq!(
|
|
writer.available_ue_buffer_capacity(),
|
|
UE_INPUT_BUFFER - resulting_bytes.len()
|
|
);
|
|
assert_eq!(writer.is_empty(), true);
|
|
// Bytes in the chunk = 6
|
|
assert_eq!(resulting_bytes[0], 0);
|
|
assert_eq!(resulting_bytes[1], 6);
|
|
assert_eq!(resulting_bytes[2..], [b'Q', b'Q', b'Q', b'Q', b'Q', b'Q'])
|
|
}
|
|
|
|
#[test]
|
|
fn will_write_small_chunks_if_no_more_data() {
|
|
let mut writer = MessageWriter::new();
|
|
// Because we also have to pass lengths (of chunk `CHUNK_LENGTH_FIELD` amd of message `4`),
|
|
// sending this will leave us with exactly 10 free bytes in the buffer
|
|
let long_message = "Q".repeat(UE_INPUT_BUFFER / 2);
|
|
writer.push(&long_message);
|
|
writer.try_pop();
|
|
let short_message = "Hello, world!";
|
|
writer.push(&short_message);
|
|
let expected_bytes = [
|
|
0, 17, // Bytes in the chunk = message length (4 bytes) + message (13 bytes)
|
|
0, 0, 0, 13, // Bytes in the message
|
|
b'H', b'e', b'l', b'l', b'o', b',', b' ', b'w', b'o', b'r', b'l', b'd', b'!',
|
|
];
|
|
// There should be enough space in the ue-server buffer to send `short_message`
|
|
let resulting_bytes = writer.try_pop().unwrap();
|
|
assert_eq!(resulting_bytes, expected_bytes);
|
|
assert_eq!(writer.try_pop(), None);
|
|
assert_eq!(writer.is_empty(), true);
|
|
}
|
|
|
|
#[test]
|
|
fn will_not_write_small_chunks_if_more_data_remains() {
|
|
let mut writer = MessageWriter::new();
|
|
// Because we also have to pass lengths (of chunk `CHUNK_LENGTH_FIELD` amd of message `4`),
|
|
// sending this will leave us with exactly 10 free bytes in the buffer
|
|
let long_message = "Q".repeat(UE_INPUT_BUFFER - CHUNK_LENGTH_FIELD - 4 - 10);
|
|
writer.push(&long_message);
|
|
writer.try_pop();
|
|
let short_message = "Hello, world!";
|
|
writer.push(&short_message);
|
|
// `MessageWriter` can neither send full message, nor a chunk of size 10
|
|
// (because it is too short)
|
|
assert_eq!(writer.try_pop(), None);
|
|
assert_eq!(writer.is_empty(), false);
|
|
}
|