Skip to main content
KeyOS API Reference

crypto/
sha2.rs

1// SPDX-FileCopyrightText: 2024 Foundation Devices, Inc. <hello@foundation.xyz>
2// SPDX-License-Identifier: GPL-3.0-or-later
3
4#[cfg(keyos)]
5use server::xous::MemoryRange;
6use server::{permission_set, CheckedConn};
7
8use crate::error::CryptoError;
9use crate::messages::{ShaDrop, ShaGetContext, ShaSetContext, ShaUpdate};
10use crate::CryptoApi;
11
12permission_set!(pub trait ShaPermissions { ShaSetContext, ShaUpdate, ShaGetContext, ShaDrop });
13
14pub const SHA224_HASH_SIZE: usize = 28;
15pub const SHA256_HASH_SIZE: usize = 32;
16pub const SHA384_HASH_SIZE: usize = 48;
17pub const SHA512_HASH_SIZE: usize = 64;
18
19#[cfg(keyos)]
20const SW_THRESHOLD: usize = 0x1000;
21
22#[cfg(keyos)]
23const SCRATCH_CAP: usize = 128 * 1024;
24
25#[derive(Debug, Clone, Copy, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)]
26pub enum ShaAlgo {
27    Sha224,
28    Sha256,
29    Sha384,
30    Sha512,
31}
32
33impl From<usize> for ShaAlgo {
34    fn from(value: usize) -> Self {
35        match value {
36            0 => ShaAlgo::Sha224,
37            1 => ShaAlgo::Sha256,
38            2 => ShaAlgo::Sha384,
39            3 => ShaAlgo::Sha512,
40            _ => unreachable!(),
41        }
42    }
43}
44
45impl From<ShaAlgo> for usize {
46    fn from(value: ShaAlgo) -> Self {
47        match value {
48            ShaAlgo::Sha224 => 0,
49            ShaAlgo::Sha256 => 1,
50            ShaAlgo::Sha384 => 2,
51            ShaAlgo::Sha512 => 3,
52        }
53    }
54}
55
56impl ShaAlgo {
57    pub fn block_size(self) -> usize {
58        match self {
59            ShaAlgo::Sha224 | ShaAlgo::Sha256 => 64,
60            ShaAlgo::Sha384 | ShaAlgo::Sha512 => 128,
61        }
62    }
63
64    pub fn hash_size(self) -> usize {
65        match self {
66            ShaAlgo::Sha224 => SHA224_HASH_SIZE,
67            ShaAlgo::Sha256 => SHA256_HASH_SIZE,
68            ShaAlgo::Sha384 => SHA384_HASH_SIZE,
69            ShaAlgo::Sha512 => SHA512_HASH_SIZE,
70        }
71    }
72
73    pub fn initial_hash_state(self) -> [u8; 64] {
74        match self {
75            ShaAlgo::Sha224 => [
76                0xc1, 0x05, 0x9e, 0xd8, 0x36, 0x7c, 0xd5, 0x07, 0x30, 0x70, 0xdd, 0x17, 0xf7, 0x0e, 0x59,
77                0x39, 0xff, 0xc0, 0x0b, 0x31, 0x68, 0x58, 0x15, 0x11, 0x64, 0xf9, 0x8f, 0xa7, 0xbe, 0xfa,
78                0x4f, 0xa4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
79                0, 0, 0, 0, 0,
80            ],
81            ShaAlgo::Sha256 => [
82                0x6a, 0x09, 0xe6, 0x67, 0xbb, 0x67, 0xae, 0x85, 0x3c, 0x6e, 0xf3, 0x72, 0xa5, 0x4f, 0xf5,
83                0x3a, 0x51, 0x0e, 0x52, 0x7f, 0x9b, 0x05, 0x68, 0x8c, 0x1f, 0x83, 0xd9, 0xab, 0x5b, 0xe0,
84                0xcd, 0x19, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
85                0, 0, 0, 0, 0,
86            ],
87            ShaAlgo::Sha384 => [
88                0xcb, 0xbb, 0x9d, 0x5d, 0xc1, 0x05, 0x9e, 0xd8, 0x62, 0x9a, 0x29, 0x2a, 0x36, 0x7c, 0xd5,
89                0x07, 0x91, 0x59, 0x01, 0x5a, 0x30, 0x70, 0xdd, 0x17, 0x15, 0x2f, 0xec, 0xd8, 0xf7, 0x0e,
90                0x59, 0x39, 0x67, 0x33, 0x26, 0x67, 0xff, 0xc0, 0x0b, 0x31, 0x8e, 0xb4, 0x4a, 0x87, 0x68,
91                0x58, 0x15, 0x11, 0xdb, 0x0c, 0x2e, 0x0d, 0x64, 0xf9, 0x8f, 0xa7, 0x47, 0xb5, 0x48, 0x1d,
92                0xbe, 0xfa, 0x4f, 0xa4,
93            ],
94            ShaAlgo::Sha512 => [
95                0x6a, 0x09, 0xe6, 0x67, 0xf3, 0xbc, 0xc9, 0x08, 0xbb, 0x67, 0xae, 0x85, 0x84, 0xca, 0xa7,
96                0x3b, 0x3c, 0x6e, 0xf3, 0x72, 0xfe, 0x94, 0xf8, 0x2b, 0xa5, 0x4f, 0xf5, 0x3a, 0x5f, 0x1d,
97                0x36, 0xf1, 0x51, 0x0e, 0x52, 0x7f, 0xad, 0xe6, 0x82, 0xd1, 0x9b, 0x05, 0x68, 0x8c, 0x2b,
98                0x3e, 0x6c, 0x1f, 0x1f, 0x83, 0xd9, 0xab, 0xfb, 0x41, 0xbd, 0x6b, 0x5b, 0xe0, 0xcd, 0x19,
99                0x13, 0x7e, 0x21, 0x79,
100            ],
101        }
102    }
103}
104
105pub struct ShaStreamingContext<P: ShaPermissions> {
106    #[cfg_attr(not(keyos), allow(dead_code))]
107    conn: CheckedConn<P>,
108    algo: ShaAlgo,
109    accumulator: Vec<u8>,
110    bytes_compressed: u64,
111    hash_state: [u8; 64],
112    #[cfg(keyos)]
113    server_id: Option<usize>,
114    #[cfg(keyos)]
115    server_authoritative: bool,
116    #[cfg(keyos)]
117    scratch: Option<xous::DropDeallocate>,
118}
119
120pub type Sha256StreamingContext<P> = ShaStreamingContext<P>;
121
122impl<P: ShaPermissions> CryptoApi<P> {
123    pub fn sha2(&self, data: &[u8], algo: ShaAlgo) -> Result<Vec<u8>, CryptoError> {
124        let mut ctx = self.sha_init(algo);
125        ctx.update(data)?;
126        ctx.finalize()
127    }
128
129    pub fn sha224(&self, data: &[u8]) -> Result<[u8; SHA224_HASH_SIZE], CryptoError> {
130        Ok(self.sha2(data, ShaAlgo::Sha224)?.try_into().unwrap())
131    }
132
133    pub fn sha256(&self, data: &[u8]) -> Result<[u8; SHA256_HASH_SIZE], CryptoError> {
134        Ok(self.sha2(data, ShaAlgo::Sha256)?.try_into().unwrap())
135    }
136
137    pub fn sha384(&self, data: &[u8]) -> Result<[u8; SHA384_HASH_SIZE], CryptoError> {
138        Ok(self.sha2(data, ShaAlgo::Sha384)?.try_into().unwrap())
139    }
140
141    pub fn sha512(&self, data: &[u8]) -> Result<[u8; SHA512_HASH_SIZE], CryptoError> {
142        Ok(self.sha2(data, ShaAlgo::Sha512)?.try_into().unwrap())
143    }
144
145    pub fn sha_init(&self, algo: ShaAlgo) -> ShaStreamingContext<P> {
146        ShaStreamingContext {
147            conn: self.conn.clone(),
148            algo,
149            accumulator: Vec::new(),
150            bytes_compressed: 0,
151            hash_state: algo.initial_hash_state(),
152            #[cfg(keyos)]
153            server_id: None,
154            #[cfg(keyos)]
155            server_authoritative: false,
156            #[cfg(keyos)]
157            scratch: None,
158        }
159    }
160
161    pub fn sha256_init(&self) -> ShaStreamingContext<P> { self.sha_init(ShaAlgo::Sha256) }
162}
163
164impl<P: ShaPermissions> ShaStreamingContext<P> {
165    pub fn hash_size(&self) -> usize { self.algo.hash_size() }
166
167    pub fn update(&mut self, data: &[u8]) -> Result<(), CryptoError> {
168        if data.is_empty() {
169            return Ok(());
170        }
171        #[cfg(keyos)]
172        if self.accumulator.len() + data.len() >= SW_THRESHOLD {
173            self.push_state_to_server()?;
174            self.update_hw(data)?;
175            return Ok(());
176        }
177        self.update_sw(data)
178    }
179
180    fn update_sw(&mut self, data: &[u8]) -> Result<(), CryptoError> {
181        let bs = self.algo.block_size();
182        self.accumulator.extend_from_slice(data);
183        let num_blocks = self.accumulator.len() / bs;
184        if num_blocks > 0 {
185            #[cfg(keyos)]
186            self.fetch_server_state()?;
187            let blocks_end = num_blocks * bs;
188            sw_compress_blocks(&mut self.hash_state, self.algo, &self.accumulator[..blocks_end]);
189            self.accumulator.drain(..blocks_end);
190            self.bytes_compressed += blocks_end as u64;
191        }
192        Ok(())
193    }
194
195    #[cfg_attr(not(keyos), allow(unused_mut))]
196    pub fn finalize(mut self) -> Result<Vec<u8>, CryptoError> {
197        #[cfg(keyos)]
198        self.fetch_server_state()?;
199
200        let total_bits = (self.bytes_compressed + self.accumulator.len() as u64) * 8;
201        let pad = sha_padding(&self.accumulator, self.algo, total_bits);
202        sw_compress_blocks(&mut self.hash_state, self.algo, &pad);
203
204        Ok(self.hash_state[..self.algo.hash_size()].to_vec())
205    }
206
207    #[cfg(keyos)]
208    fn push_state_to_server(&mut self) -> Result<(), CryptoError> {
209        if !self.server_authoritative {
210            let id = self.conn.send_blocking_archive(ShaSetContext {
211                context_id: self.server_id,
212                algo: self.algo,
213                hash_state: self.hash_state,
214            })?;
215            self.server_id = Some(id);
216            self.server_authoritative = true;
217        }
218        Ok(())
219    }
220
221    #[cfg(keyos)]
222    fn fetch_server_state(&mut self) -> Result<(), CryptoError> {
223        if self.server_authoritative {
224            let snap =
225                self.conn.send_blocking_archive(ShaGetContext { context_id: self.server_id.unwrap() })?;
226            self.hash_state = snap.hash_state;
227            self.server_authoritative = false;
228        }
229        Ok(())
230    }
231
232    #[cfg(keyos)]
233    fn scratch_range(&mut self) -> Result<MemoryRange, CryptoError> {
234        if self.scratch.is_none() {
235            let mem = xous::map_memory(None, None, SCRATCH_CAP, xous::MemoryFlags::W)?;
236            self.scratch = Some(xous::DropDeallocate::new(mem));
237        }
238        Ok(**self.scratch.as_ref().unwrap())
239    }
240
241    #[cfg(keyos)]
242    fn update_hw(&mut self, data: &[u8]) -> Result<(), CryptoError> {
243        use xous::keyos::PAGE_SIZE;
244
245        let bs = self.algo.block_size();
246        // Fast path: directly lend the page-aligned prefix to HW, accumulate the tail in SW.
247        if self.accumulator.is_empty() && (data.as_ptr() as usize) % PAGE_SIZE == 0 && data.len() >= PAGE_SIZE
248        {
249            let hw_len = (data.len() / PAGE_SIZE) * PAGE_SIZE;
250            let mr = unsafe { MemoryRange::new(data.as_ptr() as usize, hw_len)? };
251            self.conn.lend_mut(ShaUpdate { context_id: self.server_id.unwrap(), buf: mr, length: hw_len })?;
252            self.bytes_compressed += hw_len as u64;
253            self.update_sw(&data[hw_len..])?;
254            return Ok(());
255        }
256
257        let mut scratch = self.scratch_range()?;
258
259        let mut remaining = data;
260        while self.accumulator.len() + remaining.len() >= bs {
261            let acc_len = self.accumulator.len();
262            let from_data = (SCRATCH_CAP - acc_len).min(remaining.len());
263            let send_len = (acc_len + from_data) / bs * bs;
264
265            scratch.as_slice_mut::<u8>()[..acc_len].copy_from_slice(&self.accumulator);
266            scratch.as_slice_mut::<u8>()[acc_len..send_len].copy_from_slice(&remaining[..send_len - acc_len]);
267
268            remaining = &remaining[send_len - acc_len..];
269            self.accumulator.clear();
270
271            self.conn.lend_mut(ShaUpdate {
272                context_id: self.server_id.unwrap(),
273                buf: scratch.subrange(0, send_len.next_multiple_of(PAGE_SIZE)).unwrap(),
274                length: send_len,
275            })?;
276            self.bytes_compressed += send_len as u64;
277        }
278        self.accumulator = remaining.to_vec();
279
280        Ok(())
281    }
282}
283
284impl<P: ShaPermissions> Drop for ShaStreamingContext<P> {
285    fn drop(&mut self) {
286        #[cfg(keyos)]
287        if let Some(id) = self.server_id {
288            self.conn.try_send_scalar(ShaDrop(id)).ok();
289        }
290    }
291}
292
293fn sw_compress_blocks(hash_state: &mut [u8; 64], algo: ShaAlgo, data: &[u8]) {
294    match algo {
295        ShaAlgo::Sha224 | ShaAlgo::Sha256 => {
296            let mut state = [0u32; 8];
297            for (i, chunk) in hash_state[..32].chunks_exact(4).enumerate() {
298                state[i] = u32::from_be_bytes(chunk.try_into().unwrap());
299            }
300            for block in data.chunks_exact(64) {
301                sha2::compress256(&mut state, core::slice::from_ref(block.try_into().unwrap()));
302            }
303            for (i, w) in state.iter().enumerate() {
304                hash_state[i * 4..i * 4 + 4].copy_from_slice(&w.to_be_bytes());
305            }
306        }
307        ShaAlgo::Sha384 | ShaAlgo::Sha512 => {
308            let mut state = [0u64; 8];
309            for (i, chunk) in hash_state[..64].chunks_exact(8).enumerate() {
310                state[i] = u64::from_be_bytes(chunk.try_into().unwrap());
311            }
312            for block in data.chunks_exact(128) {
313                sha2::compress512(&mut state, core::slice::from_ref(block.try_into().unwrap()));
314            }
315            for (i, w) in state.iter().enumerate() {
316                hash_state[i * 8..i * 8 + 8].copy_from_slice(&w.to_be_bytes());
317            }
318        }
319    }
320}
321
322fn sha_padding(acc: &[u8], algo: ShaAlgo, total_bits: u64) -> Vec<u8> {
323    let (bs, len_field): (usize, usize) = match algo {
324        ShaAlgo::Sha224 | ShaAlgo::Sha256 => (64, 8),
325        ShaAlgo::Sha384 | ShaAlgo::Sha512 => (128, 16),
326    };
327
328    let mut pad = Vec::with_capacity(bs * 2);
329    pad.extend_from_slice(acc);
330    pad.push(0x80);
331
332    let len_mod = (pad.len() + len_field) % bs;
333    let zeros = if len_mod == 0 { 0 } else { bs - len_mod };
334    pad.extend(core::iter::repeat(0u8).take(zeros));
335
336    if len_field == 8 {
337        pad.extend_from_slice(&total_bits.to_be_bytes());
338    } else {
339        pad.extend_from_slice(&(total_bits as u128).to_be_bytes());
340    }
341
342    debug_assert_eq!(pad.len() % bs, 0);
343    pad
344}