From 7b82b6399f33cb84311cae4a6d07ade49ab0cca7 Mon Sep 17 00:00:00 2001
From: Jensenn <neil.jensen@pluto.tv>
Date: Thu, 7 Dec 2023 08:31:17 -0700
Subject: [PATCH] Add pssh box

---
 src/mp4box/mod.rs  |   6 ++-
 src/mp4box/moov.rs |  16 +++++++
 src/mp4box/pssh.rs | 109 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 130 insertions(+), 1 deletion(-)
 create mode 100644 src/mp4box/pssh.rs

diff --git a/src/mp4box/mod.rs b/src/mp4box/mod.rs
index 646993b..6a500c0 100644
--- a/src/mp4box/mod.rs
+++ b/src/mp4box/mod.rs
@@ -52,6 +52,7 @@
 //!     mvex
 //!         mehd
 //!         trex
+//!     pssh
 //! emsg
 //! moof
 //!     mfhd
@@ -98,6 +99,7 @@ pub(crate) mod moov;
 pub(crate) mod mp4a;
 pub(crate) mod mvex;
 pub(crate) mod mvhd;
+pub(crate) mod pssh;
 pub(crate) mod saio;
 pub(crate) mod saiz;
 pub(crate) mod schi;
@@ -152,6 +154,7 @@ pub use moov::MoovBox;
 pub use mp4a::Mp4aBox;
 pub use mvex::MvexBox;
 pub use mvhd::MvhdBox;
+pub use pssh::PsshBox;
 pub use saio::SaioBox;
 pub use saiz::SaizBox;
 pub use schi::SchiBox;
@@ -278,7 +281,8 @@ boxtype! {
     TencBox => 0x74656e63,
     SaizBox => 0x7361697a,
     SaioBox => 0x7361696f,
-    SencBox => 0x73656e63
+    SencBox => 0x73656e63,
+    PsshBox => 0x70737368
 }
 
 pub trait Mp4Box: Sized {
diff --git a/src/mp4box/moov.rs b/src/mp4box/moov.rs
index ac19381..a3c0741 100644
--- a/src/mp4box/moov.rs
+++ b/src/mp4box/moov.rs
@@ -20,6 +20,9 @@ pub struct MoovBox {
 
     #[serde(skip_serializing_if = "Option::is_none")]
     pub udta: Option<UdtaBox>,
+
+    #[serde(rename = "pssh")]
+    pub psshs: Vec<PsshBox>,
 }
 
 impl MoovBox {
@@ -38,6 +41,9 @@ impl MoovBox {
         if let Some(udta) = &self.udta {
             size += udta.box_size();
         }
+        for pssh in &self.psshs {
+            size += pssh.box_size();
+        }
         size
     }
 }
@@ -70,6 +76,7 @@ impl<R: Read + Seek> ReadBox<&mut R> for MoovBox {
         let mut udta = None;
         let mut mvex = None;
         let mut traks = Vec::new();
+        let mut psshs = Vec::new();
 
         let mut current = reader.stream_position()?;
         let end = start + size;
@@ -100,6 +107,10 @@ impl<R: Read + Seek> ReadBox<&mut R> for MoovBox {
                 BoxType::UdtaBox => {
                     udta = Some(UdtaBox::read_box(reader, s)?);
                 }
+                BoxType::PsshBox => {
+                    let pssh = PsshBox::read_box(reader, s)?;
+                    psshs.push(pssh);
+                }
                 _ => {
                     // XXX warn!()
                     skip_box(reader, s)?;
@@ -121,6 +132,7 @@ impl<R: Read + Seek> ReadBox<&mut R> for MoovBox {
             udta,
             mvex,
             traks,
+            psshs,
         })
     }
 }
@@ -140,6 +152,9 @@ impl<W: Write> WriteBox<&mut W> for MoovBox {
         if let Some(udta) = &self.udta {
             udta.write_box(writer)?;
         }
+        for pssh in &self.psshs {
+            pssh.write_box(writer)?;
+        }
         Ok(0)
     }
 }
@@ -158,6 +173,7 @@ mod tests {
             traks: vec![],
             meta: Some(MetaBox::default()),
             udta: Some(UdtaBox::default()),
+            psshs: vec![],
         };
 
         let mut buf = Vec::new();
diff --git a/src/mp4box/pssh.rs b/src/mp4box/pssh.rs
new file mode 100644
index 0000000..34f2019
--- /dev/null
+++ b/src/mp4box/pssh.rs
@@ -0,0 +1,109 @@
+use serde::Serialize;
+use std::io::{Read, Seek, Write};
+
+use crate::mp4box::*;
+
+#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize)]
+pub struct PsshBox {
+    pub version: u8,
+    pub system_id: [u8; 16],
+    pub kids: Vec<[u8; 16]>,
+    pub data: Vec<u8>,
+}
+
+impl PsshBox {
+    pub fn get_type(&self) -> BoxType {
+        BoxType::PsshBox
+    }
+
+    pub fn get_size(&self) -> u64 {
+        let mut s = HEADER_SIZE + HEADER_EXT_SIZE + 16;
+        if self.version > 0 {
+            s += 4 + (16 * self.kids.len() as u64);
+        }
+        s += 4 + self.data.len() as u64;
+        s
+    }
+}
+
+impl Mp4Box for PsshBox {
+    fn box_type(&self) -> BoxType {
+        self.get_type()
+    }
+
+    fn box_size(&self) -> u64 {
+        self.get_size()
+    }
+
+    fn to_json(&self) -> Result<String> {
+        Ok(serde_json::to_string(&self).unwrap())
+    }
+
+    fn summary(&self) -> Result<String> {
+        let s = format!(
+            "system_id={:02x?} data_size={}",
+            self.system_id,
+            self.data.len(),
+        );
+        Ok(s)
+    }
+}
+
+impl<R: Read + Seek> ReadBox<&mut R> for PsshBox {
+    fn read_box(reader: &mut R, size: u64) -> Result<Self> {
+        let start = box_start(reader)?;
+
+        let (version, _flags) = read_box_header_ext(reader)?;
+
+        let mut system_id = [0; 16];
+        reader.read_exact(&mut system_id)?;
+
+        let mut kids = Vec::new();
+        if version > 0 {
+            let kid_count = reader.read_u32::<BigEndian>()?;
+            kids.reserve(kid_count as usize);
+            for _ in 0..kid_count {
+                let mut kid = [0; 16];
+                reader.read_exact(&mut kid)?;
+                kids.push(kid);
+            }
+        }
+
+        let data_size = reader.read_u32::<BigEndian>()?;
+
+        let mut data = vec![0; data_size as usize];
+        reader.read_exact(&mut data)?;
+
+        skip_bytes_to(reader, start + size)?;
+
+        Ok(PsshBox {
+            version,
+            system_id,
+            kids,
+            data,
+        })
+    }
+}
+
+impl<W: Write> WriteBox<&mut W> for PsshBox {
+    fn write_box(&self, writer: &mut W) -> Result<u64> {
+        let size = self.box_size();
+        BoxHeader::new(self.box_type(), size).write(writer)?;
+
+        write_box_header_ext(writer, self.version, 0)?;
+
+        writer.write_all(&self.system_id)?;
+
+        if self.version > 0 {
+            writer.write_u32::<BigEndian>(self.kids.len() as u32)?;
+            for kid in &self.kids {
+                writer.write_all(kid)?;
+            }
+        }
+
+        writer.write_u32::<BigEndian>(self.data.len() as u32)?;
+        writer.write_all(&self.data)?;
+
+        Ok(size)
+    }
+}