Added PacketWriter for easier packet creation.

grebneerg · Nov 17, 2018 · 8d61b2f · 8d61b2f
1 parent 18466c9
commit 8d61b2f
Show file tree

Hide file tree

Showing 4 changed files with 122 additions and 50 deletions.
diff --git a/src/bytes.rs b/src/bytes.rs
@@ -1,13 +1,19 @@
-use byteorder::{ByteOrder, NetworkEndian};
+use byteorder::{ByteOrder, NetworkEndian, ReadBytesExt, WriteBytesExt};
 use std::collections::VecDeque;
+use std::io::{Cursor, Write};
 
-/// [Packet] represents a data packet and allows for extracting components in order.
-pub struct Packet(VecDeque<u8>);
+/// [PacketReader] represents a recieved data packet and allows for extracting components in order.
+pub struct PacketReader(VecDeque<u8>);
 
-impl Packet {
+impl PacketReader {
     /// Creates a [Packet] from the given [Vec<u8>].
     pub fn from_vec(vec: Vec<u8>) -> Self {
-        Packet(VecDeque::from(vec))
+        PacketReader(VecDeque::from(vec))
+    }
+
+    /// Returns a [Vec] of the remaining bytes.
+    pub fn into_vec(self) -> Vec<u8> {
+        Vec::from(self.0)
     }
 
     /// Returns the number of bytes left in the [Packet].
@@ -83,3 +89,70 @@ pub fn extract_string_u16_size(bytes: &mut VecDeque<u8>) -> Option<String> {
         }
     }
 }
+
+/// [PacketWriter] represents a new data packet and allows for adding components in order.
+pub struct PacketWriter(Cursor<Vec<u8>>);
+
+impl PacketWriter {
+    /// Creates a new [PacketWriter]
+    pub fn new() -> Self {
+        PacketWriter(Cursor::new(Vec::new()))
+    }
+
+    /// Returns a [Vec<u8>] representing the packet.
+    pub fn into_vec(self) -> Vec<u8> {
+        self.0.into_inner()
+    }
+
+    /// Returns the length of the packet in bytes.
+    pub fn len(&self) -> usize {
+        self.0.get_ref().len()
+    }
+
+    /// Writes one [u8] to the end of the packet.
+    pub fn write_u8(&mut self, val: u8) {
+        self.0.write(&[val.to_be()]).unwrap();
+    }
+
+    /// Writes one [i8] to the end of the packet.
+    pub fn write_i8(&mut self, val: i8) {
+        self.0.write(&[val.to_be() as u8]).unwrap();
+    }
+
+    /// Writes one [u16] to the end of the packet.
+    pub fn write_u16(&mut self, val: u16) {
+        self.0.write_u16::<NetworkEndian>(val).unwrap();
+    }
+
+    /// Writes one [i16] to the end of the packet.
+    pub fn write_i16(&mut self, val: i16) {
+        self.0.write_i16::<NetworkEndian>(val).unwrap();
+    }
+
+    /// Writes one [u32] to the end of the packet.
+    pub fn write_u32(&mut self, val: u32) {
+        self.0.write_u32::<NetworkEndian>(val).unwrap();
+    }
+
+    /// Writes one [f32] to the end of the packet.
+    pub fn write_f32(&mut self, val: f32) {
+        self.0.write_f32::<NetworkEndian>(val).unwrap();
+    }
+
+    /// Writes a [String] to the end of the packet.
+    pub fn write_string(&mut self, val: String) {
+        self.0.write(val.as_bytes()).unwrap();
+    }
+
+    pub fn append_packet(&mut self, other: PacketWriter) {
+        self.0.write(other.into_vec().as_ref()).unwrap();
+    }
+
+    pub fn write_slice(&mut self, slice: &[u8]) {
+        self.0.write(slice).unwrap();
+    }
+
+    pub fn write_vec(&mut self, vec: Vec<u8>) {
+        self.write_slice(vec.as_ref());
+    }
+}
diff --git a/src/joystick.rs b/src/joystick.rs
@@ -1,3 +1,5 @@
+use bytes::PacketWriter;
+
 #[derive(Clone)]
 enum JoystickType {
     Unknown = -1,
@@ -61,14 +63,14 @@ impl Joystick {
     }
 
     pub fn udp_tag(&self) -> Vec<u8> {
-        let mut tag: Vec<u8> = Vec::new();
+        let mut tag = PacketWriter::new();
 
-        tag.push(self.axes.len() as u8);
+        tag.write_u8(self.axes.len() as u8);
         for axis in &self.axes {
-            tag.push(*axis as u8); // this might work
+            tag.write_u8(*axis as u8); // this might work
         }
 
-        tag.push(self.buttons.len() as u8);
+        tag.write_u8(self.buttons.len() as u8);
         let mut index = 7;
         let mut byte: u8 = 0;
         for button in &self.buttons {
@@ -77,21 +79,21 @@ impl Joystick {
             }
             index -= 1;
             if index < 0 {
-                tag.push(byte);
+                tag.write_u8(byte);
                 byte = 0;
                 index = 7;
             }
         }
         if index != 7 {
-            tag.push(byte);
+            tag.write_u8(byte);
         }
 
-        tag.push(self.povs.len() as u8);
+        tag.write_u8(self.povs.len() as u8);
         for pov in &self.povs {
-            tag.push(((pov >> 8) & 0xff as i16) as u8);
-            tag.push((pov & 0xff) as u8);
+            tag.write_u8(((pov >> 8) & 0xff as i16) as u8);
+            tag.write_u8((pov & 0xff) as u8);
         }
 
-        tag
+        tag.into_vec()
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -10,6 +10,7 @@ mod joystick;
 mod messages;
 mod states;
 
+use bytes::PacketWriter;
 use joystick::Joystick;
 use states::{Alliance, RobotMode};
 
@@ -33,44 +34,43 @@ impl DriverStation {
     }
 
     fn udp_packet(&mut self) -> Vec<u8> {
-        let mut packet: Vec<u8> = Vec::new();
+        let mut packet = PacketWriter::new();
 
         // Packet number in case they arrive out of order
-        packet.push(((self.sequence_num >> 8) & 0xff) as u8);
-        packet.push((self.sequence_num & 0xff) as u8);
+        packet.write_u16(self.sequence_num);
         self.sequence_num += 1;
 
-        packet.push(0x01); // comm version
-        packet.push(self.control_byte()); // control byte
-        packet.push(0); // TODO: actually restart code or rio with this byte.
-        packet.push(self.alliance.to_position_u8()); // alliance
+        packet.write_u8(0x01); // comm version
+        packet.write_u8(self.control_byte()); // control byte
+        packet.write_u8(0); // TODO: actually restart code or rio with this byte.
+        packet.write_u8(self.alliance.to_position_u8()); // alliance
 
         // joystick tags
         for stick in &self.joysticks {
             if let Some(stick) = stick {
                 let mut tag = stick.udp_tag();
-                packet.push(tag.len() as u8 + 1); // size
-                packet.push(0x0c); // id
-                packet.extend(tag); // joystick tag info
+                packet.write_u8(tag.len() as u8 + 1); // size
+                packet.write_u8(0x0c); // id
+                packet.write_vec(tag); // joystick tag info
             } else {
                 // Empty joystick tag
-                packet.push(0x01); // size
-                packet.push(0x0c); // id
+                packet.write_u8(0x01); // size
+                packet.write_u8(0x0c); // id
             }
         }
 
         // datetime and timezone
         if self.request_time {
             // timezone
-            packet.push(TIMEZONE.len() as u8 + 1); // size
-            packet.push(0x10); // id
-            packet.extend(TIMEZONE.as_bytes());
+            packet.write_u8(TIMEZONE.len() as u8 + 1); // size
+            packet.write_u8(0x10); // id
+            packet.write_slice(TIMEZONE.as_bytes());
 
             // date and time
-            packet.extend(date_packet());
+            packet.write_vec(date_packet());
         }
 
-        packet
+        packet.into_vec()
     }
 
     fn control_byte(&self) -> u8 {
@@ -90,23 +90,20 @@ impl DriverStation {
 }
 
 fn date_packet() -> Vec<u8> {
-    let mut packet = Vec::new();
+    let mut packet = PacketWriter::new();
     let now = Utc::now();
-    packet.push(11); // size
-    packet.push(0x0f); // id
+    packet.write_u8(11); // size
+    packet.write_u8(0x0f); // id
     let nanos = now.nanosecond();
     let micros = nanos / 1000;
-    packet.push(((micros >> 24) & 0xff) as u8);
-    packet.push(((micros >> 16) & 0xff) as u8);
-    packet.push(((micros >> 8) & 0xff) as u8);
-    packet.push(((micros >> 0) & 0xff) as u8);
-    packet.push(now.second() as u8);
-    packet.push(now.minute() as u8);
-    packet.push(now.hour() as u8);
-    packet.push(now.day() as u8); // should this be day0?
-    packet.push(now.month0() as u8);
-    packet.push((now.year() - 1900) as u8);
-    packet
+    packet.write_u32(micros);
+    packet.write_u8(now.second() as u8);
+    packet.write_u8(now.minute() as u8);
+    packet.write_u8(now.hour() as u8);
+    packet.write_u8(now.day() as u8); // should this be day0?
+    packet.write_u8(now.month0() as u8);
+    packet.write_u8((now.year() - 1900) as u8);
+    packet.into_vec()
 }
 
 impl Default for DriverStation {

diff --git a/src/messages.rs b/src/messages.rs
@@ -1,6 +1,6 @@
 use std::convert::From;
 
-use bytes::Packet;
+use bytes::PacketReader;
 use states::RobotMode;
 
 pub(crate) struct Trace {
@@ -53,15 +53,15 @@ struct RioUdpPacket {
     trace: Trace,
     battery_voltage: f32,
     request_date: bool,
-    tags: Packet, // TODO: parse tags from this Packet
+    tags: PacketReader, // TODO: parse tags from this Packet
 }
 
 impl RioUdpPacket {
     fn from_bytes(bytes: Vec<u8>) -> Option<Self> {
         if bytes.len() > 8 {
             None
         } else {
-            let mut packet = Packet::from_vec(bytes);
+            let mut packet = PacketReader::from_vec(bytes);
             Some(RioUdpPacket {
                 sequence_num: packet.next_u16().unwrap(),
                 comm_version: packet.next_u8().unwrap(),
@@ -130,7 +130,7 @@ impl RioTcpPacket {
             None
         } else {
             use self::RioTcpPacket::*;
-            let mut packet = Packet::from_vec(bytes);
+            let mut packet = PacketReader::from_vec(bytes);
             match packet.next_u8().unwrap() {
                 0x00 => Some(RadioEvent({
                     let size = packet.len();