1use crate::*;
2
3use ::aes::{Aes256, BlockCipher};
4use block_modes::block_padding::Pkcs7;
5use block_modes::{BlockMode, Cbc};
6use generic_array::typenum::{marker_traits::Unsigned, U48, U8};
7use generic_array::GenericArray;
8use sha2::Digest;
9use std::fmt;
10use std::fmt::{Display, Formatter};
11use std::hash::Hash;
12use std::str::FromStr;
13use base58::{FromBase58, ToBase58};
14
15#[derive(Clone, Eq, PartialEq)]
17pub struct AesKey(GenericArray<u8, U48>);
18
19impl std::fmt::Debug for AesKey {
20 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21 write!(f, "{}", self.to_base58())
22 }
23}
24impl std::fmt::Display for AesKey {
25 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26 write!(f, "{}", self.to_base58())
27 }
28}
29
30impl From<GenericArray<u8, U48>> for AesKey {
31 fn from(aes_key: GenericArray<u8, U48>) -> Self {
32 Self(aes_key)
33 }
34}
35
36impl From<AesKey> for GenericArray<u8, U48> {
37 fn from(aes_key: AesKey) -> Self {
38 aes_key.0
39 }
40}
41
42impl AsRef<GenericArray<u8, U48>> for AesKey {
43 fn as_ref(&self) -> &GenericArray<u8, U48> {
44 &self.0
45 }
46}
47
48impl RawFixedBytes for AesKey {
49 fn raw_bytes() -> Option<usize> {
50 Some(U48::to_usize())
51 }
52}
53
54impl RawEncode for AesKey {
55 fn raw_measure(&self, _purpose: &Option<RawEncodePurpose>) -> BuckyResult<usize> {
56 Ok(U48::to_usize())
57 }
58 fn raw_encode<'a>(
59 &self,
60 buf: &'a mut [u8],
61 _purpose: &Option<RawEncodePurpose>,
62 ) -> BuckyResult<&'a mut [u8]> {
63 let bytes = Self::raw_bytes().unwrap();
64 if buf.len() < bytes {
65 let msg = format!(
66 "not enough buffer for encode AesKey, except={}, got={}",
67 bytes,
68 buf.len()
69 );
70 error!("{}", msg);
71
72 return Err(BuckyError::new(BuckyErrorCode::OutOfLimit, msg));
73 }
74 unsafe {
75 std::ptr::copy(self.0.as_slice().as_ptr(), buf.as_mut_ptr(), bytes);
76 }
77
78 Ok(&mut buf[bytes..])
79 }
80}
81
82impl<'de> RawDecode<'de> for AesKey {
83 fn raw_decode(buf: &'de [u8]) -> BuckyResult<(Self, &'de [u8])> {
84 let bytes = Self::raw_bytes().unwrap();
85 if buf.len() < bytes {
86 let msg = format!(
87 "not enough buffer for decode AesKey, except={}, got={}",
88 bytes,
89 buf.len()
90 );
91 error!("{}", msg);
92
93 return Err(BuckyError::new(BuckyErrorCode::OutOfLimit, msg));
94 }
95 let mut hash = GenericArray::default();
96 unsafe {
97 std::ptr::copy(buf.as_ptr(), hash.as_mut_slice().as_mut_ptr(), bytes);
98 }
99 Ok((Self(hash), &buf[bytes..]))
100 }
101}
102
103impl From<&[u8; 48]> for AesKey {
104 fn from(v: &[u8; 48]) -> Self {
105 Self(GenericArray::clone_from_slice(v))
106 }
107}
108
109impl From<Vec<u8>> for AesKey {
110 fn from(v: Vec<u8>) -> Self {
111 let ar: [u8; 48] = v.try_into().unwrap_or_else(|v: Vec<u8>| {
112 panic!(
113 "AesKey expected a Vec of length {} but it was {}",
114 48,
115 v.len()
116 )
117 });
118
119 Self(GenericArray::clone_from_slice(&ar))
120 }
121}
122
123impl Default for AesKey {
124 fn default() -> Self {
125 Self(GenericArray::default())
126 }
127}
128
129impl AesKey {
130 pub fn len(&self) -> usize {
131 self.0.len()
132 }
133
134 pub fn to_base58(&self) -> String {
135 self.0.as_slice().to_base58()
136 }
137
138 pub fn from_base58(s: &str) -> BuckyResult<Self> {
139 let buf = s.from_base58().map_err(|e| {
140 let msg = format!(
141 "convert base58 str to AesKey buf failed, str={}, {:?}",
142 s, e
143 );
144 error!("{}", msg);
145 BuckyError::new(BuckyErrorCode::InvalidFormat, msg)
146 })?;
147
148 if buf.len() != 48 {
149 let msg = format!(
150 "convert base58 str to AesKey failed, len unmatch: str={}",
151 s
152 );
153 return Err(BuckyError::new(BuckyErrorCode::InvalidFormat, msg));
154 }
155
156 Ok(Self::from(buf))
157 }
158
159 pub fn proxy(n: u64) -> AesKey {
160 let mut key = [0u8; 48];
161 for i in 0..3 {
162 key[i * 8..(i + 1) * 8].copy_from_slice(&n.to_be_bytes());
163 }
164 for i in 3..6 {
165 let r = rand::random::<u64>();
166 key[i * 8..(i + 1) * 8].copy_from_slice(&r.to_be_bytes());
167 }
168 AesKey::from(&key)
169 }
170
171 pub fn random() -> AesKey {
172 let mut key = [0u8; 48];
173 for i in 0..6 {
174 let r = rand::random::<u64>();
175 key[i * 8..(i + 1) * 8].copy_from_slice(&r.to_be_bytes());
176 }
177 AesKey::from(&key)
178 }
179
180 pub fn mix_hash(&self, salt: Option<u64>) -> KeyMixHash {
181 let mut sha = sha2::Sha256::new();
182 sha.input(self.0.as_slice());
183 if let Some(salt) = salt {
184 sha.input(&salt.to_le_bytes());
185 }
186
187 let hash = sha.result();
188 let mut mix_hash =
189 GenericArray::from_slice(&hash.as_slice()[..KeyMixHash::raw_bytes().unwrap()]).clone();
190 mix_hash[0] = mix_hash[0] & 0x7f;
191 KeyMixHash(mix_hash)
192 }
193
194 pub fn padded_len(in_len: usize) -> usize {
195 let block_size = <Aes256 as BlockCipher>::BlockSize::to_usize();
196 block_size * ((in_len / block_size) + 1)
197 }
198
199 pub fn encrypt(
200 &self,
201 in_buf: &[u8],
202 out: &mut [u8],
203 in_len: usize,
204 ) -> Result<usize, BuckyError> {
205 out[..in_len].copy_from_slice(&in_buf[..in_len]);
206 self.inplace_encrypt(out, in_len)
207 }
208
209 pub fn decrypt(
210 &self,
211 in_buf: &[u8],
212 out: &mut [u8],
213 in_len: usize,
214 ) -> Result<usize, BuckyError> {
215 out[..in_len].copy_from_slice(&in_buf[..in_len]);
216 self.inplace_decrypt(out, in_len)
217 }
218
219 pub fn inplace_encrypt(&self, inout: &mut [u8], in_len: usize) -> Result<usize, BuckyError> {
220 let key = self.0.as_slice();
232 let cipher = Cbc::<Aes256, Pkcs7>::new_from_slices(&key[0..32], &key[32..]).unwrap();
233
234 match cipher.encrypt(inout, in_len) {
235 Ok(buf) => Ok(buf.len()),
236 Err(e) => {
237 let msg = format!(
238 "AesKey::inplace_encrypt error, inout={}, in_len={}, {}",
239 inout.len(),
240 in_len,
241 e
242 );
243 error!("{}", msg);
244
245 Err(BuckyError::new(BuckyErrorCode::OutOfLimit, msg))
246 }
247 }
248 }
249
250 pub fn inplace_decrypt(&self, inout: &mut [u8], in_len: usize) -> Result<usize, BuckyError> {
251 let key = self.0.as_slice();
254 let cipher = Cbc::<Aes256, Pkcs7>::new_from_slices(&key[0..32], &key[32..]).unwrap();
255 match cipher.decrypt(&mut inout[..in_len]) {
256 Ok(buf) => Ok(buf.len()),
257 Err(e) => {
258 let msg = format!(
259 "AesKey::inplace_decrypt error, inout={}, in_len={}, {}",
260 inout.len(),
261 in_len,
262 e
263 );
264 error!("{}", msg);
265
266 Err(BuckyError::new(BuckyErrorCode::OutOfLimit, msg))
267 }
268 }
269 }
270
271 pub fn as_slice(&self) -> &[u8] {
272 self.0.as_slice()
273 }
274
275 pub fn as_mut_slice(&mut self) -> &mut [u8] {
276 self.0.as_mut_slice()
277 }
278}
279
280impl FromStr for AesKey {
281 type Err = BuckyError;
282 fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
283 Self::from_base58(s)
284 }
285}
286#[derive(Eq, PartialEq, Hash, Clone, Ord, PartialOrd, Debug)]
288pub struct KeyMixHash(GenericArray<u8, U8>);
289
290impl AsRef<GenericArray<u8, U8>> for KeyMixHash {
291 fn as_ref(&self) -> &GenericArray<u8, U8> {
292 &self.0
293 }
294}
295
296impl AsMut<GenericArray<u8, U8>> for KeyMixHash {
297 fn as_mut(&mut self) -> &mut GenericArray<u8, U8> {
298 &mut self.0
299 }
300}
301
302impl Display for KeyMixHash {
303 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
304 write!(f, "{}", hex::encode(self.0.as_slice()))
305 }
306}
307
308impl RawFixedBytes for KeyMixHash {
309 fn raw_bytes() -> Option<usize> {
310 Some(U8::to_usize())
311 }
312}
313
314impl RawEncode for KeyMixHash {
315 fn raw_measure(&self, _purpose: &Option<RawEncodePurpose>) -> BuckyResult<usize> {
316 Ok(U8::to_usize())
317 }
318 fn raw_encode<'a>(
319 &self,
320 buf: &'a mut [u8],
321 _purpose: &Option<RawEncodePurpose>,
322 ) -> BuckyResult<&'a mut [u8]> {
323 let bytes = Self::raw_bytes().unwrap();
324 if buf.len() < bytes {
325 let msg = format!(
326 "not enough buffer for encode KeyMixHash, except={}, got={}",
327 bytes,
328 buf.len()
329 );
330 error!("{}", msg);
331
332 return Err(BuckyError::new(BuckyErrorCode::OutOfLimit, msg));
333 }
334 unsafe {
335 std::ptr::copy(self.0.as_slice().as_ptr(), buf.as_mut_ptr(), bytes);
336 }
337
338 Ok(&mut buf[bytes..])
339 }
340}
341
342impl<'de> RawDecode<'de> for KeyMixHash {
343 fn raw_decode(buf: &'de [u8]) -> BuckyResult<(Self, &'de [u8])> {
344 let bytes = Self::raw_bytes().unwrap();
345 if buf.len() < bytes {
346 let msg = format!(
347 "not enough buffer for decode KeyMixHash, except={}, got={}",
348 bytes,
349 buf.len()
350 );
351 error!("{}", msg);
352
353 return Err(BuckyError::new(BuckyErrorCode::OutOfLimit, msg));
354 }
355 let mut hash = GenericArray::default();
356 unsafe {
357 std::ptr::copy(buf.as_ptr(), hash.as_mut_slice().as_mut_ptr(), bytes);
358 }
359 Ok((Self(hash), &buf[bytes..]))
360 }
361}
362
363#[cfg(test)]
364mod test_aes {
365 use super::AesKey;
366 use generic_array::typenum::U48;
367 use generic_array::GenericArray;
368
369 #[test]
370 fn test() {
371 let key = b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0A\x0B\x0C\x0D\x0E\x0F\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0A\x0B\x0C\x0D\x0E\x0F\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0A\x0B\x0C\x0D\x0E\x0E";
372 let array = GenericArray::<u8, U48>::clone_from_slice(key);
373 let aes_key = AesKey(array);
374 let mut data: [u8; 128] = [30; 128];
375 let d = b"Some Crypto Text11111dsfasdfsdsdSome Crypto Text1111";
376 data[..d.len()].copy_from_slice(d);
377
378 assert!(aes_key.inplace_encrypt(data.as_mut(), d.len()).is_ok());
379 assert!(aes_key
380 .inplace_decrypt(data.as_mut(), (d.len() / 16 + 1) * 16)
381 .is_ok());
382 assert_eq!(
383 String::from_utf8(data[..d.len()].to_vec()),
384 String::from_utf8(d.to_vec())
385 );
386 }
387}