rpfixer/
lib.rs

1use adler::Adler32;
2use crc32fast::Hasher;
3use miniz_oxide::inflate::core::{
4    decompress,
5    inflate_flags::{
6        TINFL_FLAG_IGNORE_ADLER32, TINFL_FLAG_PARSE_ZLIB_HEADER,
7        TINFL_FLAG_USING_NON_WRAPPING_OUTPUT_BUF,
8    },
9    DecompressorOxide,
10};
11use std::io::Cursor;
12use std::{
13    fmt::Debug,
14    io::{Read, Write},
15    usize,
16};
17
18struct Chunk {
19    kind: [u8; 4],
20    data: Vec<u8>,
21    crc: [u8; 4],
22}
23
24impl Debug for Chunk {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        f.debug_struct("Chunk")
27            .field("kind", &self.kind)
28            .field("data len", &self.data.len())
29            .field("crc", &self.crc)
30            .finish()
31    }
32}
33
34impl Chunk {
35    fn kind_to_string(&self) -> String {
36        String::from_utf8_lossy(&self.kind).to_string()
37    }
38}
39
40pub fn fix(mut bytes: Vec<u8>) -> Vec<u8> {
41    let mut bufread = Cursor::new(bytes);
42
43    // read the png header
44    let mut header = [0; 8];
45    bufread.read_exact(&mut header).unwrap();
46
47    let mut chunks = Vec::new();
48
49    let mut lenbuf = [0; 4];
50    while let Ok(_) = bufread.read_exact(&mut lenbuf) {
51        let len = u32::from_be_bytes(lenbuf);
52
53        let mut kind = [0; 4];
54        bufread.read_exact(&mut kind).unwrap();
55
56        let data = {
57            let mut data = vec![0; len as usize];
58            bufread.read_exact(&mut data).unwrap();
59            data
60        };
61
62        let mut crc = [0; 4];
63        bufread.read_exact(&mut crc).unwrap();
64
65        let mut chunk = Chunk { kind, data, crc };
66        // println!("{:?}", chunk);
67
68        // recode the compressed image data
69        if chunk.kind == *b"IDAT" {
70            // println!("Decompressing IDAT chunk");
71            let mut decompressor = DecompressorOxide::new();
72            decompressor.init();
73            let mut buf = vec![0; 1024 * 1024 * 1024]; // this could probably be smaller
74            let data = decompress(
75                &mut decompressor,
76                &chunk.data,
77                &mut buf,
78                0,
79                TINFL_FLAG_IGNORE_ADLER32
80                    | TINFL_FLAG_PARSE_ZLIB_HEADER
81                    | TINFL_FLAG_USING_NON_WRAPPING_OUTPUT_BUF,
82            );
83
84            // println!(
85            //     "Decompressed IDAT chunk status {:?}, bytes read {}, bytes outputted {}",
86            //     data.0, data.1, data.2
87            // );
88
89            let _ = buf.split_off(data.2);
90
91            let mut adler = Adler32::new();
92            adler.write_slice(&buf);
93
94            let csum = adler.checksum().to_be_bytes();
95
96            // replace the last 4 bytes of the data with the new checksum
97            let data_len = chunk.data.len();
98            chunk.data[data_len - 4..].copy_from_slice(&csum);
99            // println!("Corrected Adler32 checksum");
100        }
101
102        let mut hasher = Hasher::new();
103        hasher.update(&chunk.kind);
104        hasher.update(&chunk.data);
105        let checksum = hasher.finalize();
106
107        if checksum != u32::from_be_bytes(chunk.crc) {
108            // println!("CRC error in chunk {:?}", chunk.kind_to_string());
109            chunk.crc = checksum.to_be_bytes();
110            // println!("Corrected CRC");
111        }
112
113        chunks.push(chunk);
114    }
115
116    // new cursor
117    let mut newcursor = Cursor::new(Vec::new());
118    // write a new header
119    newcursor
120        .write_all(&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A])
121        .unwrap();
122
123    for chunk in chunks {
124        newcursor
125            .write_all(&u32::to_be_bytes(chunk.data.len() as u32))
126            .unwrap();
127        newcursor.write_all(&chunk.kind).unwrap();
128        newcursor.write_all(&chunk.data).unwrap();
129        newcursor.write_all(&chunk.crc).unwrap();
130    }
131    let finalres: Vec<u8> = newcursor.into_inner();
132    finalres
133}
rpfixer/lib.rs

rpfixer/
lib.rs
