oma_fetch/
checksum.rs

1use faster_hex::{hex_decode, hex_string};
2use md5::Md5;
3use serde::{Deserialize, Serialize};
4use sha2::{Digest, Sha256, Sha512};
5use snafu::ResultExt;
6use std::{fmt::Display, fs::File, io, path::Path};
7
8#[derive(PartialEq, Eq, Clone, Debug, Serialize, Deserialize)]
9pub enum Checksum {
10    Sha256(Vec<u8>),
11    Sha512(Vec<u8>),
12    Md5(Vec<u8>),
13}
14
15#[derive(Clone, Debug)]
16pub enum ChecksumValidator {
17    Sha256((Vec<u8>, Sha256)),
18    Sha512((Vec<u8>, Sha512)),
19    Md5((Vec<u8>, Md5)),
20}
21
22impl ChecksumValidator {
23    pub fn update(&mut self, data: impl AsRef<[u8]>) {
24        match self {
25            ChecksumValidator::Sha256((_, v)) => v.update(data),
26            ChecksumValidator::Sha512((_, v)) => v.update(data),
27            ChecksumValidator::Md5((_, v)) => v.update(data),
28        }
29    }
30
31    pub fn finish(&self) -> bool {
32        match self {
33            ChecksumValidator::Sha256((c, v)) => c == &v.clone().finalize().to_vec(),
34            ChecksumValidator::Sha512((c, v)) => c == &v.clone().finalize().to_vec(),
35            ChecksumValidator::Md5((c, v)) => c == &v.clone().finalize().to_vec(),
36        }
37    }
38}
39
40#[derive(Debug, snafu::Snafu)]
41pub enum ChecksumError {
42    #[snafu(display("Failed to open file"))]
43    OpenFile { source: io::Error, path: Box<Path> },
44    #[snafu(display("Failed to checksum file"))]
45    Copy { source: io::Error },
46    #[snafu(display("Bad Length"))]
47    BadLength,
48    #[snafu(display("Failed to verify data"))]
49    Decode { source: faster_hex::Error },
50}
51
52pub type Result<T> = std::result::Result<T, ChecksumError>;
53
54impl Checksum {
55    pub fn from_file_sha256(path: &Path) -> Result<Self> {
56        let mut file = File::open(path).context(OpenFileSnafu {
57            path: Box::from(path),
58        })?;
59
60        let mut hasher = Sha256::new();
61        io::copy(&mut file, &mut hasher).context(CopySnafu)?;
62        let hash = hasher.finalize().to_vec();
63
64        Ok(Self::Sha256(hash))
65    }
66
67    /// This function does not do input sanitization, so do checks before!
68    pub fn from_sha256_str(s: &str) -> Result<Self> {
69        if s.len() != 64 {
70            return Err(ChecksumError::BadLength);
71        }
72
73        let from = s.as_bytes();
74        // dst 的长度必须是 src 的一半
75        let mut dst = vec![0; from.len() / 2];
76        hex_decode(from, &mut dst).context(DecodeSnafu)?;
77
78        Ok(Checksum::Sha256(dst))
79    }
80
81    /// This function does not do input sanitization, so do checks before!
82    pub fn from_sha512_str(s: &str) -> Result<Self> {
83        if s.len() != 128 {
84            return Err(ChecksumError::BadLength);
85        }
86
87        let from = s.as_bytes();
88        // dst 的长度必须是 src 的一半
89        let mut dst = vec![0; from.len() / 2];
90        hex_decode(from, &mut dst).context(DecodeSnafu)?;
91
92        Ok(Checksum::Sha512(dst))
93    }
94
95    /// This function does not do input sanitization, so do checks before!
96    pub fn from_md5_str(s: &str) -> Result<Self> {
97        if s.len() != 32 {
98            return Err(ChecksumError::BadLength);
99        }
100
101        let from = s.as_bytes();
102        // dst 的长度必须是 src 的一半
103        let mut dst = vec![0; from.len() / 2];
104        hex_decode(from, &mut dst).context(DecodeSnafu)?;
105
106        Ok(Checksum::Md5(dst))
107    }
108
109    pub fn get_validator(&self) -> ChecksumValidator {
110        match self {
111            Checksum::Sha256(c) => ChecksumValidator::Sha256((c.clone(), Sha256::new())),
112            Checksum::Sha512(c) => ChecksumValidator::Sha512((c.clone(), Sha512::new())),
113            Checksum::Md5(c) => ChecksumValidator::Md5((c.clone(), Md5::new())),
114        }
115    }
116
117    pub fn cmp_read(&self, mut r: Box<dyn std::io::Read>) -> Result<bool> {
118        match self {
119            Checksum::Sha256(hex) => {
120                let mut hasher = Sha256::new();
121                io::copy(&mut r, &mut hasher).context(CopySnafu)?;
122                let hash = hasher.finalize().to_vec();
123                Ok(hex == &hash)
124            }
125            Checksum::Sha512(hex) => {
126                let mut hasher = Sha512::new();
127                io::copy(&mut r, &mut hasher).context(CopySnafu)?;
128                let hash = hasher.finalize().to_vec();
129                Ok(hex == &hash)
130            }
131            Checksum::Md5(hex) => {
132                let mut hasher = Md5::new();
133                io::copy(&mut r, &mut hasher).context(CopySnafu)?;
134                let hash = hasher.finalize().to_vec();
135                Ok(hex == &hash)
136            }
137        }
138    }
139
140    pub fn cmp_file(&self, path: &Path) -> Result<bool> {
141        let file = File::open(path).context(OpenFileSnafu {
142            path: Box::from(path),
143        })?;
144
145        self.cmp_read(Box::new(file) as Box<dyn std::io::Read>)
146    }
147}
148
149impl Display for Checksum {
150    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
151        match self {
152            Checksum::Sha256(hex) => {
153                f.write_str("sha256::")?;
154                f.write_str(&hex_string(hex))
155            }
156            Checksum::Sha512(hex) => {
157                f.write_str("sha512::")?;
158                f.write_str(&hex_string(hex))
159            }
160            Checksum::Md5(hex) => {
161                f.write_str("md5::")?;
162                f.write_str(&hex_string(hex))
163            }
164        }
165    }
166}