use std::time::Duration;
use bytes::Bytes;
use prost::Message;
use crate::framing::MessageType;
use crate::proto;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TapDirection {
Sent,
Received,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TapPacing {
pub delay: Duration,
pub tokens_left: u32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TapPluginData {
pub sender: Option<u32>,
pub receivers: Vec<u32>,
pub data_id: String,
pub data: Bytes,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TappedFrame {
pub direction: TapDirection,
pub kind: u16,
pub size: usize,
pub plugin_data: Option<TapPluginData>,
pub pacing: Option<TapPacing>,
}
impl TappedFrame {
pub(crate) fn new(
direction: TapDirection,
kind: u16,
payload: &Bytes,
pacing: Option<TapPacing>,
) -> Self {
let plugin_data = (kind == MessageType::PluginDataTransmission as u16)
.then(|| proto::PluginDataTransmission::decode(payload.clone()).ok())
.flatten()
.map(|m| TapPluginData {
sender: m.sender_session,
receivers: m.receiver_sessions,
data_id: m.data_id.unwrap_or_default(),
data: m.data.map(Bytes::from).unwrap_or_default(),
});
Self {
direction,
kind,
size: payload.len(),
plugin_data,
pacing,
}
}
}
pub trait PacketTap: Send + Sync {
fn enabled(&self) -> bool;
fn on_frame(&self, frame: TappedFrame);
}
#[cfg(test)]
mod tests {
use super::*;
use crate::framing::Frame;
#[test]
fn plugin_data_is_decoded_and_other_messages_are_not() {
let m = proto::PluginDataTransmission {
sender_session: Some(3),
receiver_sessions: vec![4, 5],
data: Some(b"{}".to_vec()),
data_id: Some("mumcord.test".into()),
};
let frame = Frame::from_message(MessageType::PluginDataTransmission, &m);
let tapped = TappedFrame::new(TapDirection::Sent, frame.kind, &frame.payload, None);
assert_eq!(tapped.size, frame.payload.len());
let data = tapped.plugin_data.expect("decoded");
assert_eq!(data.sender, Some(3));
assert_eq!(data.receivers, vec![4, 5]);
assert_eq!(data.data_id, "mumcord.test");
assert_eq!(&data.data[..], b"{}");
let text = proto::TextMessage {
message: "secret".into(),
..Default::default()
};
let frame = Frame::from_message(MessageType::TextMessage, &text);
let tapped = TappedFrame::new(TapDirection::Received, frame.kind, &frame.payload, None);
assert_eq!(tapped.plugin_data, None);
assert_eq!(tapped.kind, MessageType::TextMessage as u16);
}
}