Skip to main content

pb_mapper_core/
codec.rs

1use std::mem::size_of;
2
3use ring::aead::{
4    AES_256_GCM, Aad, BoundKey, NONCE_LEN, Nonce, NonceSequence, OpeningKey, SealingKey, Tag,
5    UnboundKey,
6};
7
8#[derive(Clone, Copy, Default)]
9struct Counter(u32);
10
11impl Counter {
12    fn advance(&mut self) {
13        self.0 += 1;
14    }
15
16    const fn size() -> usize {
17        size_of::<u32>()
18    }
19
20    fn to_bytes(self) -> [u8; 4] {
21        self.0.to_be_bytes()
22    }
23}
24
25#[derive(Clone, Copy, Default)]
26struct CounterNonceSequence(Counter, [u8; NONCE_LEN]);
27
28type RingResult<T> = Result<T, ring::error::Unspecified>;
29
30impl NonceSequence for CounterNonceSequence {
31    // called once for each seal operation
32    fn advance(&mut self) -> RingResult<Nonce> {
33        let nonce_bytes = &mut self.1;
34
35        let bytes = self.0.to_bytes();
36        nonce_bytes[NONCE_LEN - Counter::size()..].copy_from_slice(&bytes);
37
38        self.0.advance(); // advance the counter
39        Ok(Nonce::assume_unique_for_key(*nonce_bytes))
40    }
41}
42
43pub struct Aes256GcmCodec {
44    seal: Aes256GcmEnCodec,
45    open: Aes256GcmDeCodec,
46}
47
48impl Aes256GcmCodec {
49    pub fn try_new(key: &[u8]) -> RingResult<Self> {
50        Ok(Self {
51            seal: Aes256GcmEnCodec::try_new(key)?,
52            open: Aes256GcmDeCodec::try_new(key)?,
53        })
54    }
55
56    pub fn encrypt(&mut self, data: &mut [u8]) -> RingResult<Tag> {
57        self.seal.encrypt(data)
58    }
59
60    pub fn decrypt_with_tag(
61        &mut self,
62        decrypeted_data: &[u8],
63        tag: Tag,
64    ) -> RingResult<(Vec<u8>, usize)> {
65        self.open.decrypt_with_tag(decrypeted_data, tag)
66    }
67
68    pub fn decrypt<'a>(&mut self, data: &'a mut [u8]) -> RingResult<&'a mut [u8]> {
69        self.open.decrypt(data)
70    }
71}
72
73#[derive(Debug)]
74pub struct Aes256GcmEnCodec {
75    seal: SealingKey<CounterNonceSequence>,
76}
77
78impl Aes256GcmEnCodec {
79    pub fn try_new(key: &[u8]) -> RingResult<Self> {
80        let counter = CounterNonceSequence::default();
81        Ok(Self {
82            seal: SealingKey::new(UnboundKey::new(&AES_256_GCM, key)?, counter),
83        })
84    }
85
86    pub fn new(key: UnboundKey) -> Self {
87        let counter = CounterNonceSequence::default();
88        Self {
89            seal: SealingKey::new(key, counter),
90        }
91    }
92}
93
94impl Encryptor for Aes256GcmEnCodec {
95    fn encrypt(&mut self, data: &mut [u8]) -> RingResult<Tag> {
96        self.seal.seal_in_place_separate_tag(Aad::empty(), data)
97    }
98}
99
100#[derive(Debug)]
101pub struct Aes256GcmDeCodec {
102    open: OpeningKey<CounterNonceSequence>,
103}
104
105impl Aes256GcmDeCodec {
106    pub fn try_new(key: &[u8]) -> RingResult<Self> {
107        let counter = CounterNonceSequence::default();
108        Ok(Self {
109            open: OpeningKey::new(UnboundKey::new(&AES_256_GCM, key)?, counter),
110        })
111    }
112
113    pub fn new(key: UnboundKey) -> Self {
114        let counter = CounterNonceSequence::default();
115        Self {
116            open: OpeningKey::new(key, counter),
117        }
118    }
119}
120
121impl Decryptor for Aes256GcmDeCodec {
122    fn decrypt_with_tag(
123        &mut self,
124        decrypeted_data: &[u8],
125        tag: Tag,
126    ) -> RingResult<(Vec<u8>, usize)> {
127        let mut new_data = [decrypeted_data, tag.as_ref()].concat();
128        let new_data_len = self.decrypt(&mut new_data)?.len();
129        Ok((new_data, new_data_len))
130    }
131
132    fn decrypt<'a>(&mut self, data: &'a mut [u8]) -> RingResult<&'a mut [u8]> {
133        self.open.open_in_place(Aad::empty(), data)
134    }
135}
136
137pub trait Encryptor: 'static {
138    fn encrypt(&mut self, data: &mut [u8]) -> RingResult<Tag>;
139}
140
141pub trait Decryptor {
142    fn decrypt_with_tag(
143        &mut self,
144        decrypeted_data: &[u8],
145        tag: Tag,
146    ) -> RingResult<(Vec<u8>, usize)>;
147    fn decrypt<'a>(&mut self, data: &'a mut [u8]) -> RingResult<&'a mut [u8]>;
148}
149
150#[cfg(test)]
151mod tests {
152    use std::slice::from_raw_parts_mut;
153    use std::time::Instant;
154
155    use crate::codec::Aes256GcmCodec;
156
157    struct Timer {
158        ins: Instant,
159        hint: String,
160    }
161
162    impl Timer {
163        fn new_with_hint(hint: String) -> Self {
164            Self {
165                ins: Instant::now(),
166                hint,
167            }
168        }
169    }
170
171    impl Drop for Timer {
172        fn drop(&mut self) {
173            println!("{} consume time:{:?}", self.hint, self.ins.elapsed());
174        }
175    }
176
177    fn write_slice(slice: &[u8]) {
178        let ptr = slice.as_ptr() as *mut u8;
179        let mut_slice = unsafe { from_raw_parts_mut(ptr, slice.len()) };
180        mut_slice[1] = b'a';
181        println!("{mut_slice:?}")
182    }
183
184    #[test]
185    fn test_write_mut_slice() {
186        let vec: Vec<u8> = "bbc".into();
187        write_slice(&vec);
188        assert_eq!(vec[1], b'a');
189    }
190
191    #[test]
192    fn test_codec() {
193        const TEST_KEY: [u8; 32] = [0x42; 32];
194        let data = String::from(
195            "fdafas反对fdasfasfasfsdafdasfsdfasd范德萨发顺🤣❤️😁😍👍👍丰十大大师傅士大夫大撒发射点发士大夫大师傅大师傅士大夫士大夫阿斯蒂芬大师傅阿斯顿法大师傅看叫阿三的发就可是大家发开始打客服开始大幅喀什的开发点卡收费就开始打客服就是的咖啡肯定撒法开始打客服就是的咖啡就开始大幅扣税的急啊看发叫阿三的发生的开发就是大家可是大家发看大数据开发大数据开发大家ask发就是的咖啡的萨芬就卡死的房价开始打家开发商的JFK上的飞机卡上的纠纷开始打飞机宽带技术开发就开始大家开发建设的卡JFK大数据风控静安寺的看法角度看萨芬卡上的纠纷看静安寺的看法角度思考积分可是大家发卡是大家看法就大肆砍伐尽快打算减肥肯定是积分开始大幅技术大咖积分开始打飞机扣税的急啊看发的技术开发就是JFK十大福克斯大家开发大撒发射点幅度萨芬撒旦发发收范德萨发顺丰士大夫十大阿斯蒂芬大师傅阿斯顿附件是的客服对接撒巨大石块积分的课时费阿斯蒂芬法大师傅大师傅十大法大师傅阿斯蒂芬阿斯顿法大师傅阿斯蒂芬大师傅阿斯顿法大师傅大师傅阿斯蒂芬阿斯蒂芬士大夫阿斯蒂芬大师傅的萨芬打算减肥上岛咖啡加快速度大数据开发就是打客服看大数据开发就开始减肥卡萨丁JFK是大家看法加快速度JFK技术大咖积分喀什的开发独守空房技术大咖积分空手道解放扣税的开发商的开发接口是大家看法角度看是否扣税的急啊看发生的开发的快速减肥开始大幅就是打客服卡上的纠纷啊撒旦解放扣税的急啊看发加快速度点卡JFK啥的但是法大师傅技术大咖积分卡萨丁就反馈是大家看法啊是大家看法卡上的纠纷可是大家发喀什的开发大卡司喀什的开发就是打客服法大师傅士大夫的式咖啡机上岛咖啡就是的咖啡艰苦大师傅看上雕刻技法喀什的开发上岛咖啡就喀什的开发就是打客服卡上的纠纷技术的咖啡机肯定撒开发啊十大科技开发速度加啊反馈就是的咖啡开始大幅大师傅似的十大放假啊上岛咖啡就可是大家发空间的是否撒旦士大夫的撒娇开发是大家看法大肆砍伐就喀什的开发氨基酸的考虑非军事对抗疗法金克拉撒旦发艰苦拉萨的飞机喀什打开发就可是大家发可是大家看附件卡上的纠纷卡刷点卡技术的咖啡机可是大家发卡是大家看法静安寺的看法就可是大家发卡萨丁就开发商的急啊看飞机迪斯科发技术的咖啡机可是大家发看电视剧开发商大开始打到发大水发大水",
196        );
197        let mut cryption = Aes256GcmCodec::try_new(&TEST_KEY).unwrap();
198        let mut out_buf = data.as_bytes().to_vec();
199        let tag = {
200            let _timer = Timer::new_with_hint("Encrypt".into());
201            cryption.encrypt(&mut out_buf).unwrap()
202        };
203        println!("tag:{:?}", tag.as_ref());
204        let (decrypeted_data, len) = {
205            let _timer = Timer::new_with_hint("Decrypt".into());
206            cryption.decrypt_with_tag(&out_buf, tag).unwrap()
207        };
208        assert_eq!(
209            data,
210            String::from_utf8(decrypeted_data[..len].to_vec()).unwrap()
211        );
212    }
213}