feat: include timestamp_received when opted into (#53)

* feat: include `timestamp_received` when opted into

Signed-off-by: Skyler Ross <skyler@launchbadge.com>

* fix: stop replacing the new user_properties with the old one.

* chore: apply suggestions

---------

Signed-off-by: Skyler Ross <skyler@launchbadge.com>
Co-authored-by: Felipe Pessoa <felipe.pessoa@launchbadge.com>

Author: izik1

Cargo.lock

@@ -5,6 +5,10 @@ members = [".", "rumqttd-protocol", "fuzz"]
 # Rename our in-tree copy of `rumqttd` to prevent confusion with the published crate.
 rumqttd-protocol = { path = "rumqttd-protocol", package = "rumqttd" }
 
+[workspace.dependencies.tokio]
+version = "1.35.1"
+features = ["full"]
+
 [package]
 name = "foxmq"
 version = "0.1.0"
@@ -25,6 +29,7 @@ argon2 = { version = "0.5.3", features = ["std"] }
 bytes = "1.5.0"
 color-eyre = "0.6.2"
 dialoguer = "0.11.0"
+tokio = { workspace = true }
 dotenvy = "0.15.7"
 hex = "0.4.3"
 num_enum = "0.7.2"
@@ -43,14 +48,12 @@ arbitrary = { version = "1", optional = true, features = ["derive"] }
 # TODO: upgrade RusTLS in TCE and message-queue [TG-461]
 tokio-rustls = "0.24.1"
 rustls-pemfile = "1.0.4"
+time = { version = "0.3.36", features = ["formatting"] }
 
 [dependencies.clap]
 version = "4.5.0"
 features = ["derive", "env"]
 
-[dependencies.tokio]
-version = "1.35.1"
-features = ["full"]
 
 [dependencies.tracing-subscriber]
 version = "0.3.18"

Cargo.toml

@@ -29,7 +29,9 @@ use crate::mqtt::connect::ConnectPacket;
 use crate::mqtt::mailbox::OpenMailbox;
 use crate::mqtt::packets::{IncomingPacketSet, IncomingSub, IncomingUnsub, PacketId};
 use crate::mqtt::publish::ValidateError;
-use crate::mqtt::router::{FilterProperties, RouterConnection, RouterMessage, SubscriptionId};
+use crate::mqtt::router::{
+    FilterProperties, RouterConnection, RouterMessage, SubscribeRequest, SubscriptionId,
+};
 use crate::mqtt::session::{Session, SessionStore};
 use crate::mqtt::trie::Filter;
 use crate::mqtt::KeepAlive;
@@ -414,6 +416,7 @@ impl<S: MqttSocket> Connection<S> {
                 mail.delivery_meta,
                 Some(mail.packet_id),
                 &mail.subscription_ids,
+                mail.include_broker_timestamps,
             ))
             .await?;
 
@@ -426,6 +429,7 @@ impl<S: MqttSocket> Connection<S> {
                 PublishMeta::new(QoS::AtMostOnce, mail.retain(), false),
                 None,
                 mail.subscription_ids(),
+                mail.include_broker_timestamps,
             ))
             .await?;
         }
@@ -833,12 +837,25 @@ impl<S: MqttSocket> Connection<S> {
             has_valid_filters = true;
         }
 
+        let include_broker_timestamps = sub_props.as_ref().is_some_and(|props| {
+            props.user_properties.iter().any(|(k, v)| {
+                k == "include_broker_timestamps" && v.parse::<bool>().unwrap_or(false)
+            })
+        });
+
         if has_valid_filters {
             self.incoming_packets
                 .insert_sub(packet_id, IncomingSub { return_codes })
                 .expect("BUG: we should have checked `.contains()` above");
 
-            router.subscribe(packet_id, sub_id, filters).await;
+            router
+                .subscribe(SubscribeRequest {
+                    packet_id,
+                    sub_id,
+                    filters,
+                    include_broker_timestamps,
+                })
+                .await;
         } else {
             self.send(Packet::SubAck(
                 SubAck {

src/mqtt/broker/connection.rs

@@ -57,6 +57,7 @@ pub struct OrderedMail {
     pub delivery_meta: PublishMeta,
     pub subscription_ids: Vec<SubscriptionId>,
     pub publish: Arc<PublishTrasaction>,
+    pub include_broker_timestamps: bool,
 }
 
 /// A QoS 0 PUBLISH.
@@ -65,6 +66,7 @@ pub struct OrderedMail {
 pub struct UnorderedMail {
     pub kind: UnorderedMailKind,
     pub publish: Arc<PublishTrasaction>,
+    pub include_broker_timestamps: bool,
 }
 
 /// A compact way to represent the RETAIN flag for a QoS 0 publish.
@@ -98,6 +100,7 @@ pub struct Release(#[allow(dead_code)] pub PacketId);
 struct Delivery {
     delivery_meta: PublishMeta,
     sub_id: Option<SubscriptionId>,
+    include_broker_timestamps: bool,
     publish: Arc<PublishTrasaction>,
 }
 
@@ -154,6 +157,7 @@ impl MailSender {
         subscription_qos: QoS,
         retain: bool,
         sub_id: Option<SubscriptionId>,
+        include_broker_timestamps: bool,
         publish: Arc<PublishTrasaction>,
     ) -> bool {
         let effective_qos = cmp::min(subscription_qos, publish.meta.qos());
@@ -173,6 +177,7 @@ impl MailSender {
             .send(Delivery {
                 delivery_meta: PublishMeta::new(effective_qos, retain, false),
                 sub_id,
+                include_broker_timestamps,
                 publish,
             })
             .is_ok()
@@ -238,6 +243,7 @@ impl OpenMailbox<'_> {
             // See doc comment on `UnorderedMailKind` for details.
             kind: UnorderedMailKind::NotRetained { subscription_ids },
             publish,
+            include_broker_timestamps: _,
         }) = self.unordered_mail.back_mut()
         {
             // If this is a duplicate delivery, coalesce into one PUBLISH.
@@ -265,6 +271,7 @@ impl OpenMailbox<'_> {
                 }
             },
             publish: delivery.publish,
+            include_broker_timestamps: delivery.include_broker_timestamps,
         });
     }
 
@@ -305,6 +312,7 @@ impl OpenMailbox<'_> {
             packet_id: self.mailbox.next_packet_id.wrapping_increment(),
             subscription_ids: Vec::from_iter(delivery.sub_id),
             publish: delivery.publish,
+            include_broker_timestamps: delivery.include_broker_timestamps,
         });
     }
 

src/mqtt/mailbox.rs

@@ -1,6 +1,7 @@
 use bytes::Bytes;
 use std::num::NonZeroU32;
 use std::ops::Not;
+use time::format_description::well_known::Rfc3339;
 
 use tashi_collections::FnvHashMap;
 
@@ -350,6 +351,7 @@ pub fn txn_to_packet(
     delivery_meta: PublishMeta,
     packet_id: Option<PacketId>,
     sub_ids: &[SubscriptionId],
+    include_broker_timestamps: bool,
 ) -> Packet {
     Packet::Publish(
         Publish::with_all(
@@ -369,13 +371,27 @@ pub fn txn_to_packet(
                 };
             }
 
+            let mut user_properties =
+                clone_prop!(user_properties).map_or(vec![], |props: UserProperties| props.0);
+
+            if include_broker_timestamps {
+                let timestamp_received =
+                    time::OffsetDateTime::from_unix_timestamp(txn.timestamp_received.0 as i64)
+                        .expect("Time overflow");
+                let timestamp_received = timestamp_received
+                    .format(&Rfc3339)
+                    .expect("formatting error");
+
+                user_properties.push(("timestamp_received".to_owned(), timestamp_received))
+            }
+
             PublishProperties {
                 payload_format_indicator: clone_prop!(payload_format_indicator),
                 message_expiry_interval: clone_prop!(message_expiry_interval),
                 topic_alias: None,
                 response_topic: clone_prop!(response_topic),
                 correlation_data: clone_prop!(correlation_data).map(|bytes| bytes.0),
-                user_properties: clone_prop!(user_properties).map_or(vec![], |props| props.0),
+                user_properties,
 
                 // TODO: now that `rumqttd-protocol` is in-tree, just change the type there
                 subscription_identifiers: sub_ids.iter().copied().map(Into::into).collect(),