From 08a6880f566c6bcaf40a26651a846aecfb4d67a1 Mon Sep 17 00:00:00 2001
From: Kenneth Knudsen <ksk@blackbird.online>
Date: Thu, 27 Jun 2024 15:09:08 +0200
Subject: [PATCH] Add feature flag to set max payload size

---
 mqttrust_core/Cargo.toml         | 13 +++++++++++--
 mqttrust_core/src/eventloop.rs   |  3 ++-
 mqttrust_core/src/lib.rs         |  4 +++-
 mqttrust_core/src/max_payload.rs | 15 +++++++++++++++
 4 files changed, 31 insertions(+), 4 deletions(-)
 create mode 100644 mqttrust_core/src/max_payload.rs

diff --git a/mqttrust_core/Cargo.toml b/mqttrust_core/Cargo.toml
index 16cade3..e1ea9c6 100644
--- a/mqttrust_core/Cargo.toml
+++ b/mqttrust_core/Cargo.toml
@@ -43,8 +43,17 @@ dns-lookup = "1.0.3"
 env_logger = "0.9.0"
 
 [features]
-default = []
+default = ["max_payload_size_4096"]
+max_payload_size_2048 = []
+max_payload_size_4096 = []
+max_payload_size_8192 = []
+
 
 std = []
 
-defmt-impl = ["defmt", "mqttrust/defmt-impl", "heapless/defmt-impl", "fugit/defmt"]
+defmt-impl = [
+    "defmt",
+    "mqttrust/defmt-impl",
+    "heapless/defmt-impl",
+    "fugit/defmt",
+]
diff --git a/mqttrust_core/src/eventloop.rs b/mqttrust_core/src/eventloop.rs
index b534663..403cb01 100644
--- a/mqttrust_core/src/eventloop.rs
+++ b/mqttrust_core/src/eventloop.rs
@@ -1,3 +1,4 @@
+use crate::max_payload::MAX_PAYLOAD_SIZE;
 use crate::options::Broker;
 use crate::packet::SerializedPacket;
 use crate::state::{MqttConnectionStatus, MqttState};
@@ -417,7 +418,7 @@ impl<S> NetworkHandle<S> {
 #[derive(Debug)]
 struct PacketBuffer {
     range: RangeTo<usize>,
-    buffer: Vec<u8, 4096>,
+    buffer: Vec<u8, { MAX_PAYLOAD_SIZE }>,
 }
 
 impl PacketBuffer {
diff --git a/mqttrust_core/src/lib.rs b/mqttrust_core/src/lib.rs
index df72d34..79e11d6 100644
--- a/mqttrust_core/src/lib.rs
+++ b/mqttrust_core/src/lib.rs
@@ -8,6 +8,7 @@ pub(crate) mod fmt;
 
 mod client;
 mod eventloop;
+mod max_payload;
 mod options;
 mod packet;
 mod state;
@@ -18,6 +19,7 @@ pub use client::Client;
 use core::convert::TryFrom;
 pub use eventloop::EventLoop;
 use heapless::{String, Vec};
+use max_payload::MAX_PAYLOAD_SIZE;
 pub use mqttrust::encoding::v4::{Pid, Publish, QoS, QosPid, Suback};
 pub use mqttrust::*;
 pub use options::{Broker, MqttOptions};
@@ -30,7 +32,7 @@ pub struct PublishNotification {
     pub qospid: QoS,
     pub retain: bool,
     pub topic_name: String<256>,
-    pub payload: Vec<u8, 4096>,
+    pub payload: Vec<u8, MAX_PAYLOAD_SIZE>,
 }
 
 /// Includes incoming packets from the network and other interesting events
diff --git a/mqttrust_core/src/max_payload.rs b/mqttrust_core/src/max_payload.rs
new file mode 100644
index 0000000..9854215
--- /dev/null
+++ b/mqttrust_core/src/max_payload.rs
@@ -0,0 +1,15 @@
+#[cfg(not(any(
+    feature = "max_payload_size_2048",
+    feature = "max_payload_size_4096",
+    feature = "max_payload_size_8192"
+)))]
+pub const MAX_PAYLOAD_SIZE: usize = 4096;
+
+#[cfg(feature = "max_payload_size_2048")]
+pub const MAX_PAYLOAD_SIZE: usize = 2048;
+
+#[cfg(feature = "max_payload_size_4096")]
+pub const MAX_PAYLOAD_SIZE: usize = 4096;
+
+#[cfg(feature = "max_payload_size_8192")]
+pub const MAX_PAYLOAD_SIZE: usize = 8192;