bsdiff_android/
bsdf2_writer.rs

1use std::io::{self, Write};
2use bzip2::write::BzEncoder;
3use bzip2::Compression as BzCompression;
4
5const BSDIFF_MAGIC: &[u8; 8] = b"BSDIFF40";
6const BSDF2_MAGIC: &[u8; 5] = b"BSDF2";
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum CompressionAlgorithm {
10    None = 0,
11    Bz2 = 1,
12    Brotli = 2,
13}
14
15fn compress(alg: CompressionAlgorithm, data: &[u8]) -> io::Result<Vec<u8>> {
16    match alg {
17        CompressionAlgorithm::None => Ok(data.to_vec()),
18        CompressionAlgorithm::Bz2 => {
19            let mut encoder = BzEncoder::new(Vec::new(), BzCompression::best());
20            encoder.write_all(data)?;
21            encoder.finish()
22        }
23        CompressionAlgorithm::Brotli => {
24            let mut compressed = Vec::new();
25            {
26                let mut encoder = brotli::CompressorWriter::new(
27                    &mut compressed,
28                    4096,  // buffer size
29                    11,    // quality (11 = max)
30                    20,    // lg_window_size (matches Android kBrotliDefaultLgwin)
31                );
32                encoder.write_all(data)?;
33                encoder.flush()?;
34            } // encoder drops here, finalizing compression
35            Ok(compressed)
36        }
37    }
38}
39
40/// encode signed integer in bspatch sign-magnitude format
41#[inline]
42fn encode_int64(x: i64, buf: &mut [u8]) {
43    if x >= 0 {
44        buf.copy_from_slice(&x.to_le_bytes());
45    } else {
46        let tmp = ((-x) as u64) | (1u64 << 63);
47        buf.copy_from_slice(&tmp.to_le_bytes());
48    }
49}
50
51/// Control entry matching
52#[derive(Debug, Clone, Copy)]
53pub struct ControlEntry {
54    pub diff_size: i64,
55    pub extra_size: i64,
56    pub offset_increment: i64,
57}
58
59/// BSDF2 patch writer
60pub struct Bsdf2Writer {
61    ctrl_data: Vec<u8>,
62    diff_data: Vec<u8>,
63    extra_data: Vec<u8>,
64    ctrl_alg: CompressionAlgorithm,
65    diff_alg: CompressionAlgorithm,
66    extra_alg: CompressionAlgorithm,
67    written_output: u64,
68}
69
70impl Bsdf2Writer {
71    /// Create a new BSDF2 writer with specified compression for each stream
72    pub fn new(
73        ctrl_alg: CompressionAlgorithm,
74        diff_alg: CompressionAlgorithm,
75        extra_alg: CompressionAlgorithm,
76    ) -> Self {
77        Self {
78            ctrl_data: Vec::new(),
79            diff_data: Vec::new(),
80            extra_data: Vec::new(),
81            ctrl_alg,
82            diff_alg,
83            extra_alg,
84            written_output: 0,
85        }
86    }
87    pub fn new_legacy() -> Self {
88        Self::new(
89            CompressionAlgorithm::Bz2,
90            CompressionAlgorithm::Bz2,
91            CompressionAlgorithm::Bz2,
92        )
93    }
94    pub fn add_control_entry(&mut self, entry: ControlEntry) -> io::Result<()> {
95        let mut buf = [0u8; 24];
96        encode_int64(entry.diff_size, &mut buf[0..8]);
97        encode_int64(entry.extra_size, &mut buf[8..16]);
98        encode_int64(entry.offset_increment, &mut buf[16..24]);
99
100        self.ctrl_data.extend_from_slice(&buf);
101        self.written_output += (entry.diff_size + entry.extra_size) as u64;
102        Ok(())
103    }
104
105    /// Write diff stream data
106    pub fn write_diff_stream(&mut self, data: &[u8]) -> io::Result<()> {
107        self.diff_data.extend_from_slice(data);
108        Ok(())
109    }
110    pub fn write_extra_stream(&mut self, data: &[u8]) -> io::Result<()> {
111        self.extra_data.extend_from_slice(data);
112        Ok(())
113    }
114    pub fn close<W: Write>(&mut self, writer: &mut W) -> io::Result<()> {
115        // Compress all streams
116        let ctrl_compressed = compress(self.ctrl_alg, &self.ctrl_data)?;
117        let diff_compressed = compress(self.diff_alg, &self.diff_data)?;
118        let extra_compressed = compress(self.extra_alg, &self.extra_data)?;
119
120        // Write header
121        let is_legacy = self.ctrl_alg == CompressionAlgorithm::Bz2
122            && self.diff_alg == CompressionAlgorithm::Bz2
123            && self.extra_alg == CompressionAlgorithm::Bz2;
124
125        self.write_header(
126            writer,
127            is_legacy,
128            ctrl_compressed.len() as u64,
129            diff_compressed.len() as u64,
130        )?;
131
132        // Write compressed streams
133        writer.write_all(&ctrl_compressed)?;
134        writer.write_all(&diff_compressed)?;
135        writer.write_all(&extra_compressed)?;
136
137        Ok(())
138    }
139
140    fn write_header<W: Write>(
141        &self,
142        writer: &mut W,
143        is_legacy: bool,
144        ctrl_size: u64,
145        diff_size: u64,
146    ) -> io::Result<()> {
147        let mut header = [0u8; 32];
148
149        if is_legacy {
150            header[0..8].copy_from_slice(BSDIFF_MAGIC);
151        } else {
152            header[0..5].copy_from_slice(BSDF2_MAGIC);
153            header[5] = self.ctrl_alg as u8;
154            header[6] = self.diff_alg as u8;
155            header[7] = self.extra_alg as u8;
156        }
157
158        encode_int64(ctrl_size as i64, &mut header[8..16]);
159        encode_int64(diff_size as i64, &mut header[16..24]);
160        encode_int64(self.written_output as i64, &mut header[24..32]);
161
162        writer.write_all(&header)?;
163        Ok(())
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170
171    #[test]
172    fn test_encode_int64_positive() {
173        let mut buf = [0u8; 8];
174        encode_int64(42, &mut buf);
175        assert_eq!(buf, [42, 0, 0, 0, 0, 0, 0, 0]);
176    }
177
178    #[test]
179    fn test_encode_int64_negative() {
180        let mut buf = [0u8; 8];
181        encode_int64(-42, &mut buf);
182        assert_eq!(buf, [42, 0, 0, 0, 0, 0, 0, 0x80]);
183    }
184
185    #[test]
186    fn test_encode_int64_zero() {
187        let mut buf = [0u8; 8];
188        encode_int64(0, &mut buf);
189        assert_eq!(buf, [0; 8]);
190    }
191
192    #[test]
193    fn test_writer_creation() {
194        let writer = Bsdf2Writer::new(
195            CompressionAlgorithm::Brotli,
196            CompressionAlgorithm::Brotli,
197            CompressionAlgorithm::Bz2,
198        );
199        assert_eq!(writer.ctrl_alg, CompressionAlgorithm::Brotli);
200        assert_eq!(writer.written_output, 0);
201    }
202
203    #[test]
204    fn test_legacy_writer() {
205        let writer = Bsdf2Writer::new_legacy();
206        assert_eq!(writer.ctrl_alg, CompressionAlgorithm::Bz2);
207        assert_eq!(writer.diff_alg, CompressionAlgorithm::Bz2);
208        assert_eq!(writer.extra_alg, CompressionAlgorithm::Bz2);
209    }
210}