Skip to main content

s3_wire/operation/
checksum.rs

1use base64::Engine as _;
2use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
3use sha2::{Digest as _, Sha256};
4
5use super::{Checksum, ChecksumAlgorithm};
6
7/// Failure to calculate a requested checksum locally.
8#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
9#[error("local calculation for {algorithm:?} is not implemented")]
10pub struct ChecksumCalculationError {
11    /// Algorithm that could not be calculated.
12    pub algorithm: ChecksumAlgorithm,
13}
14
15impl Checksum {
16    /// Calculates one checksum over `bytes` using standard S3 wire encoding.
17    ///
18    /// CRC32, CRC32C, CRC64NVME, and SHA-256 are supported. SHA-1 remains an
19    /// accepted wire value but is not calculated locally.
20    ///
21    /// # Errors
22    ///
23    /// Returns an error when local calculation is unavailable for `algorithm`.
24    pub fn calculate(
25        algorithm: ChecksumAlgorithm,
26        bytes: &[u8],
27    ) -> Result<Self, ChecksumCalculationError> {
28        let mut checksum = Self::default();
29        match algorithm {
30            ChecksumAlgorithm::Crc32 => {
31                checksum.crc32 = Some(BASE64_STANDARD.encode(crc32(bytes).to_be_bytes()));
32            }
33            ChecksumAlgorithm::Crc32c => {
34                checksum.crc32c = Some(BASE64_STANDARD.encode(crc32c(bytes).to_be_bytes()));
35            }
36            ChecksumAlgorithm::Crc64Nvme => {
37                checksum.crc64_nvme = Some(BASE64_STANDARD.encode(crc64_nvme(bytes).to_be_bytes()));
38            }
39            ChecksumAlgorithm::Sha256 => {
40                checksum.sha256 = Some(BASE64_STANDARD.encode(Sha256::digest(bytes)));
41            }
42            ChecksumAlgorithm::Sha1 => return Err(ChecksumCalculationError { algorithm }),
43        }
44        Ok(checksum)
45    }
46
47    pub(crate) async fn calculate_cooperatively(
48        algorithm: ChecksumAlgorithm,
49        bytes: &[u8],
50    ) -> Result<Self, ChecksumCalculationError> {
51        const CHUNK_SIZE: usize = 256 * 1024;
52        let mut checksum = Self::default();
53        match algorithm {
54            ChecksumAlgorithm::Crc32 | ChecksumAlgorithm::Crc32c => {
55                let table = match algorithm {
56                    ChecksumAlgorithm::Crc32 => &CRC32_TABLE,
57                    ChecksumAlgorithm::Crc32c => &CRC32C_TABLE,
58                    _ => unreachable!("matched CRC32 family"),
59                };
60                let mut value = u32::MAX;
61                for chunk in bytes.chunks(CHUNK_SIZE) {
62                    value = update_crc32_state(table, value, chunk);
63                    tokio::task::yield_now().await;
64                }
65                let encoded = BASE64_STANDARD.encode((!value).to_be_bytes());
66                match algorithm {
67                    ChecksumAlgorithm::Crc32 => checksum.crc32 = Some(encoded),
68                    ChecksumAlgorithm::Crc32c => checksum.crc32c = Some(encoded),
69                    _ => unreachable!("matched CRC32 family"),
70                }
71            }
72            ChecksumAlgorithm::Crc64Nvme => {
73                let mut value = u64::MAX;
74                for chunk in bytes.chunks(CHUNK_SIZE) {
75                    value = update_crc64_state(value, chunk);
76                    tokio::task::yield_now().await;
77                }
78                checksum.crc64_nvme = Some(BASE64_STANDARD.encode((!value).to_be_bytes()));
79            }
80            ChecksumAlgorithm::Sha256 => {
81                let mut digest = Sha256::new();
82                for chunk in bytes.chunks(CHUNK_SIZE) {
83                    digest.update(chunk);
84                    tokio::task::yield_now().await;
85                }
86                checksum.sha256 = Some(BASE64_STANDARD.encode(digest.finalize()));
87            }
88            ChecksumAlgorithm::Sha1 => return Err(ChecksumCalculationError { algorithm }),
89        }
90        Ok(checksum)
91    }
92}
93
94const CRC32_TABLE: [u32; 256] = crc32_table(0xedb8_8320);
95const CRC32C_TABLE: [u32; 256] = crc32_table(0x82f6_3b78);
96const CRC64_NVME_TABLE: [u64; 256] = crc64_table(0x9a6c_9329_ac4b_c9b5);
97
98fn crc32(bytes: &[u8]) -> u32 {
99    update_crc32(&CRC32_TABLE, bytes)
100}
101
102fn crc32c(bytes: &[u8]) -> u32 {
103    update_crc32(&CRC32C_TABLE, bytes)
104}
105
106fn update_crc32(table: &[u32; 256], bytes: &[u8]) -> u32 {
107    !update_crc32_state(table, u32::MAX, bytes)
108}
109
110fn update_crc32_state(table: &[u32; 256], mut value: u32, bytes: &[u8]) -> u32 {
111    for &byte in bytes {
112        let index = usize::from((value as u8) ^ byte);
113        value = table[index] ^ (value >> 8);
114    }
115    value
116}
117
118fn crc64_nvme(bytes: &[u8]) -> u64 {
119    !update_crc64_state(u64::MAX, bytes)
120}
121
122fn update_crc64_state(mut value: u64, bytes: &[u8]) -> u64 {
123    for &byte in bytes {
124        let index = usize::from((value as u8) ^ byte);
125        value = CRC64_NVME_TABLE[index] ^ (value >> 8);
126    }
127    value
128}
129
130const fn crc32_table(polynomial: u32) -> [u32; 256] {
131    let mut table = [0_u32; 256];
132    let mut index = 0;
133    while index < table.len() {
134        let mut value = index as u32;
135        let mut bit = 0;
136        while bit < 8 {
137            value = if value & 1 == 0 {
138                value >> 1
139            } else {
140                (value >> 1) ^ polynomial
141            };
142            bit += 1;
143        }
144        table[index] = value;
145        index += 1;
146    }
147    table
148}
149
150const fn crc64_table(polynomial: u64) -> [u64; 256] {
151    let mut table = [0_u64; 256];
152    let mut index = 0;
153    while index < table.len() {
154        let mut value = index as u64;
155        let mut bit = 0;
156        while bit < 8 {
157            value = if value & 1 == 0 {
158                value >> 1
159            } else {
160                (value >> 1) ^ polynomial
161            };
162            bit += 1;
163        }
164        table[index] = value;
165        index += 1;
166    }
167    table
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn standard_crc_check_values_match() {
176        assert_eq!(crc32(b"123456789"), 0xcbf4_3926);
177        assert_eq!(crc32c(b"123456789"), 0xe306_9283);
178        assert_eq!(crc64_nvme(b"123456789"), 0xae8b_1486_0a79_9888);
179    }
180
181    #[test]
182    fn calculates_s3_base64_values() {
183        assert_eq!(
184            Checksum::calculate(ChecksumAlgorithm::Crc64Nvme, b"123456789")
185                .unwrap()
186                .crc64_nvme
187                .as_deref(),
188            Some("rosUhgp5mIg=")
189        );
190        assert_eq!(
191            Checksum::calculate(ChecksumAlgorithm::Sha256, b"abc")
192                .unwrap()
193                .sha256
194                .as_deref(),
195            Some("ungWv48Bz+pBQUDeXa4iI7ADYaOWF3qctBD/YfIAFa0=")
196        );
197    }
198
199    #[tokio::test]
200    async fn cooperative_calculation_matches_synchronous_vectors() {
201        let bytes = vec![0x5a; 1024 * 1024 + 17];
202        for algorithm in [
203            ChecksumAlgorithm::Crc32,
204            ChecksumAlgorithm::Crc32c,
205            ChecksumAlgorithm::Crc64Nvme,
206            ChecksumAlgorithm::Sha256,
207        ] {
208            assert_eq!(
209                Checksum::calculate_cooperatively(algorithm, &bytes)
210                    .await
211                    .unwrap(),
212                Checksum::calculate(algorithm, &bytes).unwrap()
213            );
214        }
215    }
216}