Using async websockets in Rust with tokio, rmp-serde and varint encoding

Question

This is my first attempt at a real life app in Rust. I would mostly like to know if I am writing idiomatic Rust code, coming from Java and Python. I'll appreciate any pointers to improvements, though.

The code uses tokio_tungstenite to receive messages from a websocket server asynchronously. The connection as well as the handshake and a mandatory subscription call happen in a synchronous manner before.
Once subscribed, the websocket server will start sending varint encoded messagepack messages. I could not find a decent varint crate that suited my purposes, so I wrote I small module myself. Furthermore, I had to write another module that handles messagepack timestamp serialization and deserialization since the rmp-serde crate does not provide this functionality and the incoming messagepack messages contain messagepack style timestamps.

The websocket server occasionally sends a ping message amongst the actual messages. In this case, a custom pong message is being sent to the server containing the contents of the ping itself.
I am using channels to notify the pong sender task to send back the pong from within the loop that waits for new messages to come in.

Specific questions:

1 - The RealResponse struct represents one response message, however, only the data field has meaningful value. All other fields are either {}, null or []. Let me know if I gave those fields the correct type (like Option<()>).

2 - I have many constructs like

let a = match x {
   Ok(a) => a,
   Err(e) => handle error meaningfully
}

for which I would like to know if I could write them in a more concise way.

Code structure:

   |-- Cargo.toml
   |-- src
         |-- mp_ext
         |   |-- mod.rs
         |
         |-- varint
         |   |-- mod.rs
         |
         |-- lib.rs
         |-- main.rs

Cargo.toml looks like this:

[package]
name = "myproject"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
tokio = {version = "1.34", features = ["full"]}
tokio-tungstenite = {version = "0.20", features = ["native-tls"]}
futures-util = "0.3"
futures-sink = "0.3"
futures-channel = "0.3"
rmp-serde = "1.1.2"
serde = { version = "1", features = ["derive"] }
serde_bytes = "0.11"

Code:

main.rs:

#[tokio::main]
async fn main() {
    myproject::start().await;
}

lib.rs:

use std::collections::HashMap;
use std::vec;

use futures_channel::mpsc;
use futures_util::{
    stream::{SplitSink, SplitStream},
    SinkExt, StreamExt,
};
use serde::{Deserialize, Serialize};
use std::process;
use tokio::net::TcpStream;
use tokio_tungstenite::{
    connect_async,
    tungstenite::{self, handshake, http, Message},
    MaybeTlsStream, WebSocketStream,
};

mod mp_ext;
mod varint;

#[derive(Debug, PartialEq, Serialize)]
struct SubscribeRequest {
    r#type: u8,
    empty: HashMap<(), ()>,
    invocation_id: &'static str,
    endpoint: &'static str,
    identifiers: [Vec<&'static str>; 1],
}

#[derive(Debug, PartialEq, Deserialize)]
#[serde(rename_all = "PascalCase")]
struct RealData {
    push_identifier: String,
    data1: f64,
    data2: f64,
    data3: f64,
    data4: f64,
    #[serde(deserialize_with = "mp_ext::deserialize_mp_timestamp_sec")]
    data1_date_time: u64,
    #[serde(deserialize_with = "mp_ext::deserialize_mp_timestamp_sec")]
    data2_date_time: u64,
    #[serde(deserialize_with = "mp_ext::deserialize_mp_timestamp_sec")]
    data3_date_time: u64,
    #[serde(deserialize_with = "mp_ext::deserialize_mp_timestamp_sec")]
    data4_date_time: u64,
}

#[derive(Debug, PartialEq, Deserialize)]
struct RealResponse {
    r#type: u8,
    empty: HashMap<(), ()>,
    unknown: Option<()>,
    endpoint: String,
    data: [[RealData; 1]; 1],
    unknown2: [(); 0],
}

pub async fn start() {
    let request = http::Request::builder()
        .uri("wss://some_domain.com/real")
        .header("Host", "some_domain.com")
        .header("Origin", "https://www.some_domain.com")
        .header("Connection", "Upgrade")
        .header("Upgrade", "websocket")
        .header("pragma", "no-cache")
        .header("Cache-Control", "no-cache")
        .header("Sec-WebSocket-Version", 13)
        .header("Sec-WebSocket-Extensions", "permessage-deflate; client_max_window_bits")
        .header("Sec-WebSocket-Key", handshake::client::generate_key())
        .header("User-Agent", "some_agent")
        .header("Accept-Encoding", "gzip, deflate, br")
        .header("Accept-Language", "en-GB,en-US;q=0.9,en;q=0.8,de;q=0.7")
        .body(())
        .unwrap();

    // connect and split into sender and receiver
    let (ws_stream, _) = connect_async(request).await.unwrap();
    let (mut sender_ws, mut receiver_ws) = ws_stream.split();

    // send handshake and subscription request
    websocket_handshake(&mut sender_ws, &mut receiver_ws).await;
    subscribe(&mut sender_ws, &mut receiver_ws).await;

    // this channel is used to send pongs to the server
    let (sender_pong, receiver_pong) = mpsc::unbounded::<Vec<u8>>();

    // spawn the pong handler
    tokio::spawn(spawn_pong_handler(receiver_pong, sender_ws));

    // receive messages until connection has been closed
    loop {
        let response = receiver_ws.next().await;

        if response.is_none() {
            println!("Connection has been closed.");
            break;
        }

        // spawn a new task for every good response
        tokio::spawn(process_incoming_messages(
            response.unwrap(),
            sender_pong.clone(),
        ));
    }
}

async fn websocket_handshake(
    sender_ws: &mut SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
    receiver_ws: &mut SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
) {
    sender_ws
        .send(Message::Text(
            "{\"protocol\":\"messagepack\",\"version\":1}\x1e".to_string(),
        ))
        .await
        .expect("Unable to send handshake.");
    let response = receiver_ws
        .next()
        .await
        .expect("Unable to read handshake.")
        .expect("Unable to read handshake.");

    if let Ok("{}\x1e") = response.to_text() {
        println!("WebSocket handshake has been successfully completed.");
    } else {
        eprintln!("Bad handshake. Aborting!");
        process::exit(1);
    }
}

async fn subscribe(
    sender_ws: &mut SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
    receiver_ws: &mut SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
) {
    // send subscribe message for given IDs
    let subscribe_request = SubscribeRequest {
        r#type: 1,
        empty: HashMap::new(),
        invocation_id: "0",
        endpoint: "Subscribe",
        identifiers: [vec!["IDENTIFIER"]],
    };

    let subscribe_request = match rmp_serde::to_vec(&subscribe_request) {
        Ok(subscribe_request) => subscribe_request,
        Err(e) => {
            eprintln!("Error serializing subscribe request: {}.", e);
            process::exit(1);
        }
    };
    let subscribe_request = varint::encode(subscribe_request);
    sender_ws
        .send(Message::Binary(subscribe_request))
        .await
        .expect("Unable to send subscribe message.");

    let response = receiver_ws
        .next()
        .await
        .expect("Unable to read subscribe response.")
        .expect("Unable to read subscribe response.")
        .into_data();
    if response == vec![7, 149, 3, 128, 161, 48, 3, 192] {
        println!("Successfully sent subscription request.");
    } else {
        eprintln!("Bad subscription response. Aborting!");
        process::exit(1);
    }
}

async fn spawn_pong_handler(
    mut receiver_pong: mpsc::UnboundedReceiver<Vec<u8>>,
    mut sender_ws: SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
) {
    // The ping/pong handler
    loop {
        // For the pong we need to send the same data as the ping.
        let ping = match receiver_pong.next().await {
            Some(ping) => ping,
            // connection has been closed
            None => continue,
        };

        match sender_ws.send(Message::Binary(varint::encode(ping))).await {
            Ok(_) => println!("Pong "),
            Err(e) => eprintln!("Unable to send pong: {}.", e),
        };
    }
}

async fn process_incoming_messages(
    response: Result<Message, tungstenite::Error>,
    sender_pong: mpsc::UnboundedSender<Vec<u8>>,
) {
    // decode the responses (data or ping)
    let response = match response {
        Ok(res) => res,
        Err(e) => {
            eprintln!("Unable to read socket: {}.", e);
            process::exit(1);
        }
    };
    let response = match varint::decode(&response.into_data()) {
        Ok(response) => response,
        Err(e) => {
            eprintln!("Error decoding message: {}.", e);
            return;
        }
    };

    for message in response {
        match message.len() {
            2 => process_ping(&sender_pong, message),
            _ => process_message(message),
        }
    }
}

fn process_ping(sender_pong: &mpsc::UnboundedSender<Vec<u8>>, message: Vec<u8>) {
    println!("Ping!");
    if let Err(e) = sender_pong.unbounded_send(message) {
        eprintln!("Unable to send ping: {}.", e);
    }
}

fn process_message(message: Vec<u8>) {
    // the actual handler
    let response = match rmp_serde::from_slice::<RealResponse>(&message) {
        Ok(res) => res,
        Err(e) => {
            eprintln!("Error decoding message: {}.", e);
            return;
        }
    };
    println!("{:?}", response.data[0][0]);
    //todo!();
}

mod.rs in varint module (no need to check the logic):

pub fn encode(byte_data: Vec<u8>) -> Vec<u8> {
    let mut varint_length = byte_data.len();
    let mut varint_bytes = Vec::new();
    while varint_length > 0 {
        let mut byte = (varint_length & 127) as u8;
        varint_length >>= 7;
        if varint_length > 0 {
            byte |= 128;
        }
        varint_bytes.push(byte);
    }
    varint_bytes.extend(byte_data);
    varint_bytes
}

pub fn decode(byte_data: &[u8]) -> Result<Vec<Vec<u8>>, &'static str> {
    let mut decoded_messages: Vec<Vec<u8>> = Vec::new();
    let mut index: usize = 0;

    while index < byte_data.len() {
        let mut varint_bytes: usize = 0;
        let mut varint_length: usize = 0;

        loop {
            let current_byte = byte_data[index + varint_bytes];
            varint_length |= ((current_byte & 127) as usize) << (varint_bytes * 7);
            varint_bytes += 1;
            if current_byte & 128 == 0 {
                break;
            }
        }

        if varint_length > 0x7FFFFFFF {
            return Err("Messages bigger than 2GB are not supported.");
        }

        index += varint_bytes;

        if byte_data.len() >= index + varint_length {
            let message = byte_data[index..index + varint_length].to_vec();
            decoded_messages.push(message);
            index += varint_length;
        } else {
            return Err("Incomplete message.");
        }
    }

    Ok(decoded_messages)
}

mod.rs in mp_ext (no need to check the logic of the timestamp conversions):

use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_bytes::ByteBuf;

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename = "_ExtStruct")]
struct ExtStruct((i8, ByteBuf));

#[allow(dead_code)]
pub fn serialize_mp_timestamp<S>(timestamp: &Timestamp, serializer: S) -> Result<S::Ok, S::Error>
where
    S: Serializer,
{
    ExtStruct((-1, timestamp.as_mp_bytes())).serialize(serializer)
}

pub fn deserialize_mp_timestamp<'de, D>(deserializer: D) -> Result<Timestamp, D::Error>
where
    D: Deserializer<'de>,
{
    let ExtStruct((ext_type, buf)) = Deserialize::deserialize(deserializer)?;
    if ext_type != -1 {
        return Err(serde::de::Error::custom("Invalid extension type (!=-1)"));
    }

    Timestamp::from_mp_bytes(buf).map_err(serde::de::Error::custom)
}

pub fn deserialize_mp_timestamp_sec<'de, D>(deserializer: D) -> Result<u64, D::Error>
where
    D: Deserializer<'de>,
{
    deserialize_mp_timestamp(deserializer).map(|ts| ts.seconds)
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Timestamp {
    pub seconds: u64,
    pub nanoseconds: u32,
}

impl Timestamp {
    fn from_mp_bytes(bytes: ByteBuf) -> Result<Self, &'static str> {
        let bytes = bytes.into_vec();
        let len = bytes.len();
        let (seconds, nanoseconds) = match len {
            4 => (u32::from_be_bytes(bytes.try_into().unwrap()) as u64, 0),
            8 => {
                let data64 = u64::from_be_bytes(bytes.try_into().unwrap());
                (data64 & 0x00000003ffffffff, (data64 >> 34) as u32)
            }
            12 => {
                let (nanoseconds_bytes, seconds_bytes) = bytes.split_at(4);
                (
                    u64::from_be_bytes(seconds_bytes.try_into().unwrap()),
                    u32::from_be_bytes(nanoseconds_bytes.try_into().unwrap()),
                )
            }
            _ => return Err("Timestamp can only be created from 32, 64, or 96-bit byte objects"),
        };

        Ok(Self {
            seconds,
            nanoseconds,
        })
    }

    fn as_mp_bytes(&self) -> ByteBuf {
        if (self.seconds >> 34) == 0 {
            let mut data = [0u8; 8];

            let data64 = (self.nanoseconds as u64) << 34 | self.seconds;

            if data64 & 0xFFFFFFFF00000000 == 0 {
                data[..4].copy_from_slice(&data64.to_be_bytes()[..4]);
                data[4..].fill(0);
            } else {
                data.copy_from_slice(&data64.to_be_bytes());
            }

            return ByteBuf::from(data);
        }

        let mut data = [0u8; 12];
        data[..4].copy_from_slice(&self.nanoseconds.to_be_bytes());
        data[4..].copy_from_slice(&self.seconds.to_be_bytes());

        ByteBuf::from(data)
    }
}

[ I did not manage to find a decent example online for neither, tokio_tungstenite nor messagepack style timestamps. If anything, this code might help others struggling with those crates. ]

Answer 1

First of all, kudos to using cargo fmt. It makes for an easy read.

Use clippy

If you're not doing it yet, I suggest you use cargo clippy with the following options:

• -W clippy::pedantic
• -W clippy::nursery
• -W clippy::unwrap_used
• -W clippy::cargo

It will provide you with a stricter set of lints, though you need to be aware that it also may produce false-positives in rare cases.

RealResponse

Regarding RealResponse, why does it contain the other fields if they are unused? Serde will skip unknown fields by default, so you can just leave unused fields out of the struct.

Repeating match statements

If the repeated match statements bother you, you may want to have a look at unwrap_or() and unwrap_or_else() respectively. With the options above, clippy may already suggest those to you in appropriate cases. That being said, there is generally nothing wrong with the match syntax as you use it - I'd rather consider this a matter of taste.

Long types

The types SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message> and SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>> are rather long. You may want to consider to type-alias those for abbreviation.