linkle/format/
utils.rs

1use serde::de::{Unexpected, Visitor};
2use serde::{Deserialize, Deserializer, Serialize, Serializer};
3use sha2::{Digest, Sha256};
4use std::fmt;
5use std::fs::File;
6use std::io::{Read, Seek, SeekFrom};
7
8pub fn align(size: usize, padding: usize) -> usize {
9    ((size as usize) + padding) & !padding
10}
11
12pub fn add_padding(vec: &mut Vec<u8>, padding: usize) {
13    let real_size = vec.len();
14    vec.resize(align(real_size, padding), 0);
15}
16
17pub fn check_string_or_truncate(string: &mut String, name: &str, size: usize) {
18    if string.len() >= size {
19        println!("Warning: Truncating {} to 0x{:x}", name, size - 1);
20        string.truncate(size);
21    }
22}
23
24pub fn get_segment_data(
25    file: &mut File,
26    header: &elf::types::ProgramHeader,
27) -> std::io::Result<Vec<u8>> {
28    let mut data = vec![0; header.filesz as usize];
29    file.seek(SeekFrom::Start(header.offset))?;
30    file.read_exact(&mut data)?;
31    Ok(data)
32}
33
34pub fn compress_lz4(uncompressed_data: &mut Vec<u8>) -> std::io::Result<Vec<u8>> {
35    lz4::block::compress(&uncompressed_data[..], None, false)
36}
37
38pub fn compress_blz(uncompressed_data: &mut Vec<u8>) -> blz_nx::BlzResult<Vec<u8>> {
39    let mut compressed_data =
40        vec![0; blz_nx::get_worst_compression_buffer_size(uncompressed_data.len())];
41    let res = blz_nx::compress_raw(&mut uncompressed_data[..], &mut compressed_data[..])?;
42    compressed_data.resize(res, 0);
43    Ok(compressed_data)
44}
45
46pub fn calculate_sha256(data: &[u8]) -> std::io::Result<Vec<u8>> {
47    let mut hasher = Sha256::default();
48    hasher.update(data);
49    Ok(Vec::from(hasher.finalize().as_slice()))
50}
51
52#[derive(Default)]
53pub struct HexOrNum(pub u64);
54
55impl fmt::Debug for HexOrNum {
56    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
57        f.write_fmt(format_args!("{:#010x}", self.0))
58    }
59}
60
61impl<'de> Deserialize<'de> for HexOrNum {
62    fn deserialize<D>(deserializer: D) -> Result<HexOrNum, D::Error>
63    where
64        D: Deserializer<'de>,
65    {
66        struct HexOrNumVisitor;
67
68        impl<'a> Visitor<'a> for HexOrNumVisitor {
69            type Value = u64;
70
71            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
72                write!(formatter, "an integer or a hex-formatted string")
73            }
74
75            fn visit_u64<E>(self, v: u64) -> Result<u64, E>
76            where
77                E: serde::de::Error,
78            {
79                Ok(v)
80            }
81
82            fn visit_str<E>(self, v: &str) -> Result<u64, E>
83            where
84                E: serde::de::Error,
85            {
86                if v.starts_with("0x") {
87                    u64::from_str_radix(&v[2..], 16)
88                        .map_err(|_| E::invalid_value(Unexpected::Str(v), &"a hex-encoded string"))
89                } else {
90                    Err(E::invalid_value(
91                        Unexpected::Str(v),
92                        &"a hex-encoded string",
93                    ))
94                }
95            }
96        }
97
98        let num = Deserializer::deserialize_any(deserializer, HexOrNumVisitor)?;
99        Ok(HexOrNum(num))
100    }
101}
102
103impl Serialize for HexOrNum {
104    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
105    where
106        S: Serializer,
107    {
108        serializer.collect_str(&format_args!("{:#010x}", self.0))
109    }
110}