loco_protocol/secure/
client.rs

1/*
2 * Created on Sun Sep 10 2023
3 *
4 * Copyright (c) storycraft. Licensed under the MIT Licence.
5 */
6
7pub use rsa;
8
9use alloc::{boxed::Box, collections::VecDeque, vec::Vec};
10use core::mem;
11
12use aes::cipher::{AsyncStreamCipher, Key, KeyIvInit};
13use arrayvec::ArrayVec;
14use rand::thread_rng;
15use rsa::{Oaep, RsaPublicKey};
16use serde::{Deserialize, Serialize};
17use sha1::Sha1;
18
19use super::SecurePacket;
20
21type Aes128CfbEnc = cfb_mode::Encryptor<aes::Aes128>;
22type Aes128CfbDec = cfb_mode::Decryptor<aes::Aes128>;
23
24#[derive(Debug)]
25/// IO-free client secure layer
26pub struct LocoClientSecureLayer {
27    key: Key<aes::Aes128>,
28
29    read_state: ReadState,
30
31    encrypt_buffer: Vec<u8>,
32
33    /// Read buffer for layer
34    pub read_buffer: VecDeque<u8>,
35
36    /// Write buffer for layer
37    pub write_buffer: VecDeque<u8>,
38}
39
40impl LocoClientSecureLayer {
41    /// Create new [`LocoClientSecureLayer`] with given encrypt key
42    pub fn new(encrypt_key: [u8; 16]) -> Self {
43        Self {
44            key: encrypt_key.into(),
45
46            read_state: ReadState::Pending,
47
48            encrypt_buffer: Vec::new(),
49
50            read_buffer: VecDeque::new(),
51            write_buffer: VecDeque::new(),
52        }
53    }
54
55    pub const fn read_state(&self) -> &ReadState {
56        &self.read_state
57    }
58
59    /// Write handshake packet to [`LocoClientSecureLayer::write_buffer`] using given public key
60    pub fn handshake(&mut self, key: &RsaPublicKey) {
61        #[derive(Serialize)]
62        struct RawHandshakeHeader {
63            encrypted_key_size: u32,
64            key_type: u32,
65            encrypt_type: u32,
66        }
67
68        let encrypted_key = key
69            .encrypt(
70                &mut thread_rng(),
71                Oaep::new_with_mgf_hash::<Sha1, Sha1>(),
72                self.key.as_slice(),
73            )
74            .unwrap();
75
76        bincode::serialize_into(
77            &mut self.write_buffer,
78            &RawHandshakeHeader {
79                encrypted_key_size: encrypted_key.len() as u32,
80                key_type: 15,    // RSA OAEP SHA1 MGF1 SHA1
81                encrypt_type: 2, // AES_CFB128 NOPADDING
82            },
83        )
84        .unwrap();
85
86        self.write_buffer.extend(encrypted_key);
87    }
88
89    /// Try to read single [`SecurePacket`] from [`LocoClientSecureLayer::read_buffer`]
90    pub fn read(&mut self) -> Option<SecurePacket<Box<[u8]>>> {
91        loop {
92            match mem::replace(&mut self.read_state, ReadState::Corrupted) {
93                ReadState::Pending => {
94                    if self.read_buffer.len() < 20 {
95                        self.read_state = ReadState::Pending;
96                        return None;
97                    }
98
99                    let raw_header = {
100                        let buf = self.read_buffer.drain(..20).collect::<ArrayVec<u8, 20>>();
101
102                        bincode::deserialize::<RawHeader>(&buf).unwrap()
103                    };
104
105                    self.read_state = ReadState::Header(raw_header);
106                }
107
108                ReadState::Header(raw_header) => {
109                    let size = raw_header.size as usize - 16;
110
111                    if self.read_buffer.len() < size {
112                        self.read_state = ReadState::Header(raw_header);
113                        return None;
114                    }
115
116                    let mut data = self.read_buffer.drain(..size).collect::<Box<[u8]>>();
117                    Aes128CfbDec::new(&self.key, &raw_header.iv.into()).decrypt(&mut data);
118
119                    self.read_state = ReadState::Pending;
120                    return Some(SecurePacket {
121                        iv: raw_header.iv,
122                        data,
123                    });
124                }
125
126                ReadState::Corrupted => unreachable!(),
127            }
128        }
129    }
130
131    /// Write single [`SecurePacket`] to [`LocoClientSecureLayer::write_buffer`]
132    pub fn send(&mut self, packet: SecurePacket<impl AsRef<[u8]>>) {
133        let encrypted_data = {
134            let data = packet.data.as_ref();
135
136            self.encrypt_buffer.extend(data);
137            Aes128CfbEnc::new(&self.key, &packet.iv.into()).encrypt(&mut self.encrypt_buffer);
138
139            &mut self.encrypt_buffer
140        };
141
142        bincode::serialize_into(
143            &mut self.write_buffer,
144            &RawHeader {
145                size: 16 + encrypted_data.len() as u32,
146                iv: packet.iv,
147            },
148        )
149        .unwrap();
150
151        self.write_buffer.extend(encrypted_data.drain(..));
152    }
153}
154
155#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
156pub struct RawHeader {
157    /// Data size including iv
158    pub size: u32,
159
160    /// Encrypted data IV
161    pub iv: [u8; 16],
162}
163
164#[derive(Debug, Clone, PartialEq)]
165pub enum ReadState {
166    /// Client is waiting for packet
167    Pending,
168
169    /// Client read packet header and waiting for more data
170    Header(RawHeader),
171
172    /// Client corrupted and cannot continue
173    Corrupted,
174}
loco_protocol/secure/client.rs

loco_protocol/secure/
client.rs
