From f3dc6b118d19bb5e67e0e0608e06b217f25f8f5c Mon Sep 17 00:00:00 2001
From: Hubert <lesny.rumcajs@gmail.com>
Date: Fri, 24 Jun 2022 13:13:18 +0200
Subject: [PATCH] use arrayvec for packets (#4)

---
 Cargo.toml |  3 ++-
 src/lib.rs | 74 ++++++++++++++++++++++++++++++++++++++++++++----------
 2 files changed, 63 insertions(+), 14 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index c76306e..db8ed18 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "wakey"
-version = "0.2.0"
+version = "0.2.1"
 
 authors = ["Hubert Bugaj<lesny.rumcajs@gmail.com>"]
 edition = "2021"
@@ -14,4 +14,5 @@ categories = ["network-programming"]
 
 [dependencies]
 hex = "~0.3"
+arrayvec = "0.7.2"
 clap = { version = "3.1.18", features = ["derive"] }
diff --git a/src/lib.rs b/src/lib.rs
index 96c7e48..63da85a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -12,14 +12,20 @@
 use std::iter;
 use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};
 
+use arrayvec::ArrayVec;
+
 const MAC_SIZE: usize = 6;
 const MAC_PER_MAGIC: usize = 16;
-static HEADER: [u8; 6] = [0xFF; 6];
+const HEADER: [u8; 6] = [0xFF; 6];
+const PACKET_LEN: usize = HEADER.len() + MAC_SIZE * MAC_PER_MAGIC;
+
+type Packet = ArrayVec<u8, PACKET_LEN>;
 
 /// Wake-on-LAN packet
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct WolPacket {
     /// WOL packet bytes
-    packet: Vec<u8>,
+    packet: Packet,
 }
 
 impl WolPacket {
@@ -80,28 +86,39 @@ impl WolPacket {
     /// Converts string representation of MAC address (e.x. 00:01:02:03:04:05) to raw bytes.
     /// # Panic
     /// Panics when input MAC is invalid (i.e. contains non-byte characters)
-    fn mac_to_byte(data: &str, sep: char) -> Vec<u8> {
-        data.split(sep)
+    fn mac_to_byte(data: &str, sep: char) -> ArrayVec<u8, MAC_SIZE> {
+        let bytes = data
+            .split(sep)
             .flat_map(|x| hex::decode(x).expect("Invalid mac!"))
-            .collect()
+            .collect::<ArrayVec<u8, MAC_SIZE>>();
+
+        debug_assert_eq!(MAC_SIZE, bytes.len());
+
+        bytes
     }
 
     /// Extends the MAC address to fill the magic packet
-    fn extend_mac(mac: &[u8]) -> Vec<u8> {
-        iter::repeat(mac)
+    fn extend_mac(mac: &[u8]) -> ArrayVec<u8, { MAC_SIZE * MAC_PER_MAGIC }> {
+        let magic = iter::repeat(mac)
             .take(MAC_PER_MAGIC)
             .flatten()
-            .cloned()
-            .collect()
+            .copied()
+            .collect::<ArrayVec<u8, { MAC_SIZE * MAC_PER_MAGIC }>>();
+
+        debug_assert_eq!(MAC_SIZE * MAC_PER_MAGIC, magic.len());
+
+        magic
     }
 
     /// Creates bytes of the magic packet from MAC address
-    fn create_packet_bytes(mac: &[u8]) -> Vec<u8> {
-        let mut packet = Vec::with_capacity(HEADER.len() + MAC_SIZE * MAC_PER_MAGIC);
+    fn create_packet_bytes(mac: &[u8]) -> Packet {
+        let mut packet = Packet::new();
 
-        packet.extend(HEADER.iter());
+        packet.extend(HEADER);
         packet.extend(WolPacket::extend_mac(mac));
 
+        debug_assert_eq!(PACKET_LEN, packet.len());
+
         packet
     }
 }
@@ -121,12 +138,29 @@ mod tests {
         );
     }
 
+    #[test]
+    #[should_panic]
+    fn extend_mac_mac_too_long_test() {
+        let mac = vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07];
+        super::WolPacket::extend_mac(&mac);
+    }
+
+    #[test]
+    #[should_panic]
+    fn extend_mac_mac_too_short_test() {
+        let mac = vec![0x01, 0x02, 0x03, 0x04, 0x05];
+        super::WolPacket::extend_mac(&mac);
+    }
+
     #[test]
     fn mac_to_byte_test() {
         let mac = "01:02:03:04:05:06";
         let result = super::WolPacket::mac_to_byte(mac, ':');
 
-        assert_eq!(result, vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06]);
+        assert_eq!(
+            result.into_inner().unwrap(),
+            [0x01, 0x02, 0x03, 0x04, 0x05, 0x06]
+        );
     }
 
     #[test]
@@ -143,6 +177,20 @@ mod tests {
         super::WolPacket::mac_to_byte(mac, ':');
     }
 
+    #[test]
+    #[should_panic]
+    fn mac_to_byte_mac_too_long_test() {
+        let mac = "01:02:03:04:05:06:07";
+        super::WolPacket::mac_to_byte(mac, ':');
+    }
+
+    #[test]
+    #[should_panic]
+    fn mac_to_byte_mac_too_short_test() {
+        let mac = "01:02:03:04:05";
+        super::WolPacket::mac_to_byte(mac, ':');
+    }
+
     #[test]
     fn create_packet_bytes_test() {
         let bytes = super::WolPacket::create_packet_bytes(&[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF]);