Avoid use of Vec after DequeBuffer refactor

Lonami · Jun 29, 2024 · b466db5 · b466db5
1 parent 468a634
commit b466db5
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 4 deletions.
diff --git a/lib/grammers-mtproto/src/mtp/encrypted.rs b/lib/grammers-mtproto/src/mtp/encrypted.rs
@@ -8,6 +8,7 @@
 use super::{
     Deserialization, DeserializationFailure, DeserializeError, Mtp, RpcResult, RpcResultError,
 };
+use crate::utils::StackBuffer;
 use crate::{manual_tl, MsgId};
 use getrandom::getrandom;
 use grammers_crypto::{decrypt_data_v2, encrypt_data_v2, AuthKey, DequeBuffer};
@@ -261,7 +262,7 @@ impl Encrypted {
             // Prepend a container, setting its message ID and sequence number.
             // + 8 because it has to include the constructor ID and length (4 bytes each).
             let len = (buffer.len() + 8) as i32;
-            let mut header = Vec::with_capacity(MESSAGE_CONTAINER_HEADER_LEN);
+            let mut header = StackBuffer::<MESSAGE_CONTAINER_HEADER_LEN>::new();
 
             // Manually `serialize_msg` because the container body was already written.
             self.get_new_msg_id().serialize(&mut header);
@@ -271,16 +272,16 @@ impl Encrypted {
 
             manual_tl::MessageContainer::CONSTRUCTOR_ID.serialize(&mut header);
             (self.msg_count as i32).serialize(&mut header);
-            buffer.extend_front(&header);
+            buffer.extend_front(&header.into_inner());
         }
 
         {
             // Prepend the message header
-            let mut header = Vec::with_capacity(PLAIN_PACKET_HEADER_LEN);
+            let mut header = StackBuffer::<PLAIN_PACKET_HEADER_LEN>::new();
             self.get_current_salt().serialize(&mut header); // 8 bytes
 
             self.client_id.serialize(&mut header); // 8 bytes
-            buffer.extend_front(&header);
+            buffer.extend_front(&header.into_inner());
         }
 
         self.msg_count = 0;

diff --git a/lib/grammers-mtproto/src/utils.rs b/lib/grammers-mtproto/src/utils.rs
@@ -15,3 +15,31 @@ pub(crate) fn check_message_buffer(message: &[u8]) -> Result<(), DeserializeErro
         Ok(())
     }
 }
+
+/// Stack buffer to support extending from an iterator.
+pub(crate) struct StackBuffer<const N: usize> {
+    array: [u8; N],
+    pos: usize,
+}
+
+impl<const N: usize> StackBuffer<N> {
+    pub(crate) fn new() -> Self {
+        Self {
+            array: [0; N],
+            pos: 0,
+        }
+    }
+
+    pub(crate) fn into_inner(self) -> [u8; N] {
+        self.array
+    }
+}
+
+impl<const N: usize> Extend<u8> for StackBuffer<N> {
+    fn extend<T: IntoIterator<Item = u8>>(&mut self, iter: T) {
+        iter.into_iter().for_each(|x| {
+            self.array[self.pos] = x;
+            self.pos += 1;
+        });
+    }
+}