From dd0ce35bef92903e32968ba52d6f40234dc79024 Mon Sep 17 00:00:00 2001
From: Luca <Luca@hackerspace-bamberg.de>
Date: Fri, 28 Jun 2024 22:29:32 +0200
Subject: [PATCH] feat: implement serialize method on TPM2.net packet

---
 src/tpm2net.rs | 74 ++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 71 insertions(+), 3 deletions(-)

diff --git a/src/tpm2net.rs b/src/tpm2net.rs
index 3786b97..6bd735e 100644
--- a/src/tpm2net.rs
+++ b/src/tpm2net.rs
@@ -1,12 +1,15 @@
 pub const PACKET_START_BYTE: u8 = 0x9C;
 
-pub const PAYLOAD_SIZE_MAX: usize = 1490;
+pub const PAYLOAD_SIZE_MAX: usize = 1465;
+
+const OVERHEAD: usize = 7;
 
 pub use crate::common::{PacketType, PACKET_END_BYTE};
 
 #[derive(PartialEq)]
 #[non_exhaustive]
 pub enum Tpm2NetError {
+    BufferTooSmall,
     IndexOutOfBounds,
     PayloadTooLarge,
 }
@@ -63,9 +66,33 @@ impl<'a> Packet<'a> {
         self.packet_type
     }
 
-    pub fn size(&self) -> usize {
+    pub fn payload_size(&self) -> usize {
         self.payload.len()
     }
+
+    pub fn total_size(&self) -> usize {
+        self.payload.len() + OVERHEAD
+    }
+
+    pub fn serialize(&self, buf: &mut [u8]) -> Result<(), Tpm2NetError> {
+        if buf.len() < self.total_size() {
+            return Err(Tpm2NetError::BufferTooSmall);
+        }
+
+        buf[0] = PACKET_START_BYTE;
+        buf[1] = self.packet_type as u8;
+        buf[2] = ((self.payload.len() + 2) >> 8) as u8; // frame size includes packet number + packet count
+        buf[3] = ((self.payload.len() + 2) & 255) as u8;
+        buf[4] = self.index;
+        buf[5] = self.total;
+        buf[self.payload.len() + 6] = PACKET_END_BYTE;
+
+        for (i, b) in self.payload.iter().enumerate() {
+            buf[i + 6] = *b;
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -75,7 +102,7 @@ mod tests {
     #[test]
     fn empty_packet() {
         let packet = Packet::new(PacketType::Command);
-        assert_eq!(packet.size(), 0);
+        assert_eq!(packet.payload_size(), 0);
     }
 
     #[test]
@@ -116,4 +143,45 @@ mod tests {
         assert!(result.is_err());
         assert!(result.err().unwrap() == Tpm2NetError::IndexOutOfBounds);
     }
+
+    #[test]
+    fn serialize() {
+        let payload: [u8; 4] = [0xca, 0xfe, 0xba, 0xbe];
+        let packet = Packet::with_payload_and_index(PacketType::Data, &payload, 23, 42)
+            .ok()
+            .unwrap();
+
+        let mut buf = [0u8; 11];
+        let result = packet.serialize(&mut buf);
+        assert!(result.is_ok());
+        assert_eq!(
+            buf,
+            [
+                PACKET_START_BYTE,
+                PacketType::Data as u8,
+                0,
+                payload.len() as u8 + 2,
+                23,
+                42,
+                payload[0],
+                payload[1],
+                payload[2],
+                payload[3],
+                PACKET_END_BYTE
+            ]
+        );
+    }
+
+    #[test]
+    fn buffer_too_small() {
+        let payload: [u8; 4] = [0xca, 0xfe, 0xba, 0xbe];
+        let packet = Packet::with_payload(PacketType::Data, &payload)
+            .ok()
+            .unwrap();
+
+        let mut buf = [0u8; 10];
+        let result = packet.serialize(&mut buf);
+        assert!(result.is_err());
+        assert!(result.err().unwrap() == Tpm2NetError::BufferTooSmall);
+    }
 }