rustfs_rio/
encrypt_reader.rs

1// Copyright 2024 RustFS Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::HashReaderDetector;
16use crate::HashReaderMut;
17use crate::compress_index::{Index, TryGetIndex};
18use crate::{EtagResolvable, Reader};
19use aes_gcm::aead::Aead;
20use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
21use pin_project_lite::pin_project;
22use rustfs_utils::{put_uvarint, put_uvarint_len};
23use std::pin::Pin;
24use std::task::{Context, Poll};
25use tokio::io::{AsyncRead, ReadBuf};
26
27pin_project! {
28    /// A reader wrapper that encrypts data on the fly using AES-256-GCM.
29    /// This is a demonstration. For production, use a secure and audited crypto library.
30    #[derive(Debug)]
31    pub struct EncryptReader<R> {
32        #[pin]
33        pub inner: R,
34        key: [u8; 32],   // AES-256-GCM key
35        nonce: [u8; 12], // 96-bit nonce for GCM
36        buffer: Vec<u8>,
37        buffer_pos: usize,
38        finished: bool,
39    }
40}
41
42impl<R> EncryptReader<R>
43where
44    R: Reader,
45{
46    pub fn new(inner: R, key: [u8; 32], nonce: [u8; 12]) -> Self {
47        Self {
48            inner,
49            key,
50            nonce,
51            buffer: Vec::new(),
52            buffer_pos: 0,
53            finished: false,
54        }
55    }
56}
57
58impl<R> AsyncRead for EncryptReader<R>
59where
60    R: AsyncRead + Unpin + Send + Sync,
61{
62    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
63        let mut this = self.project();
64        // Serve from buffer if any
65        if *this.buffer_pos < this.buffer.len() {
66            let to_copy = std::cmp::min(buf.remaining(), this.buffer.len() - *this.buffer_pos);
67            buf.put_slice(&this.buffer[*this.buffer_pos..*this.buffer_pos + to_copy]);
68            *this.buffer_pos += to_copy;
69            if *this.buffer_pos == this.buffer.len() {
70                this.buffer.clear();
71                *this.buffer_pos = 0;
72            }
73            return Poll::Ready(Ok(()));
74        }
75        if *this.finished {
76            return Poll::Ready(Ok(()));
77        }
78        // Read a fixed block size from inner
79        let block_size = 8 * 1024;
80        let mut temp = vec![0u8; block_size];
81        let mut temp_buf = ReadBuf::new(&mut temp);
82        match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
83            Poll::Pending => Poll::Pending,
84            Poll::Ready(Ok(())) => {
85                let n = temp_buf.filled().len();
86                if n == 0 {
87                    // EOF, write end header
88                    let mut header = [0u8; 8];
89                    header[0] = 0xFF; // type: end
90                    *this.buffer = header.to_vec();
91                    *this.buffer_pos = 0;
92                    *this.finished = true;
93                    let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
94                    buf.put_slice(&this.buffer[..to_copy]);
95                    *this.buffer_pos += to_copy;
96                    Poll::Ready(Ok(()))
97                } else {
98                    // Encrypt the chunk
99                    let cipher = Aes256Gcm::new_from_slice(this.key).expect("key");
100                    let nonce = Nonce::from_slice(this.nonce);
101                    let plaintext = &temp_buf.filled()[..n];
102                    let plaintext_len = plaintext.len();
103                    let crc = crc32fast::hash(plaintext);
104                    let ciphertext = cipher
105                        .encrypt(nonce, plaintext)
106                        .map_err(|e| std::io::Error::other(format!("encrypt error: {e}")))?;
107                    let int_len = put_uvarint_len(plaintext_len as u64);
108                    let clen = int_len + ciphertext.len() + 4;
109                    // Header: 8 bytes
110                    // 0: type (0 = encrypted, 0xFF = end)
111                    // 1-3: length (little endian u24, ciphertext length)
112                    // 4-7: CRC32 of ciphertext (little endian u32)
113                    let mut header = [0u8; 8];
114                    header[0] = 0x00; // 0 = encrypted
115                    header[1] = (clen & 0xFF) as u8;
116                    header[2] = ((clen >> 8) & 0xFF) as u8;
117                    header[3] = ((clen >> 16) & 0xFF) as u8;
118                    header[4] = (crc & 0xFF) as u8;
119                    header[5] = ((crc >> 8) & 0xFF) as u8;
120                    header[6] = ((crc >> 16) & 0xFF) as u8;
121                    header[7] = ((crc >> 24) & 0xFF) as u8;
122                    let mut out = Vec::with_capacity(8 + int_len + ciphertext.len());
123                    out.extend_from_slice(&header);
124                    let mut plaintext_len_buf = vec![0u8; int_len];
125                    put_uvarint(&mut plaintext_len_buf, plaintext_len as u64);
126                    out.extend_from_slice(&plaintext_len_buf);
127                    out.extend_from_slice(&ciphertext);
128                    *this.buffer = out;
129                    *this.buffer_pos = 0;
130                    let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
131                    buf.put_slice(&this.buffer[..to_copy]);
132                    *this.buffer_pos += to_copy;
133                    Poll::Ready(Ok(()))
134                }
135            }
136            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
137        }
138    }
139}
140
141impl<R> EtagResolvable for EncryptReader<R>
142where
143    R: EtagResolvable,
144{
145    fn try_resolve_etag(&mut self) -> Option<String> {
146        self.inner.try_resolve_etag()
147    }
148}
149
150impl<R> HashReaderDetector for EncryptReader<R>
151where
152    R: EtagResolvable + HashReaderDetector,
153{
154    fn is_hash_reader(&self) -> bool {
155        self.inner.is_hash_reader()
156    }
157
158    fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
159        self.inner.as_hash_reader_mut()
160    }
161}
162
163impl<R> TryGetIndex for EncryptReader<R>
164where
165    R: TryGetIndex,
166{
167    fn try_get_index(&self) -> Option<&Index> {
168        self.inner.try_get_index()
169    }
170}
171
172pin_project! {
173    /// A reader wrapper that decrypts data on the fly using AES-256-GCM.
174    /// This is a demonstration. For production, use a secure and audited crypto library.
175#[derive(Debug)]
176    pub struct DecryptReader<R> {
177        #[pin]
178        pub inner: R,
179        key: [u8; 32],   // AES-256-GCM key
180        nonce: [u8; 12], // 96-bit nonce for GCM
181        buffer: Vec<u8>,
182        buffer_pos: usize,
183        finished: bool,
184        // For block framing
185        header_buf: [u8; 8],
186        header_read: usize,
187        header_done: bool,
188        ciphertext_buf: Option<Vec<u8>>,
189        ciphertext_read: usize,
190        ciphertext_len: usize,
191    }
192}
193
194impl<R> DecryptReader<R>
195where
196    R: Reader,
197{
198    pub fn new(inner: R, key: [u8; 32], nonce: [u8; 12]) -> Self {
199        Self {
200            inner,
201            key,
202            nonce,
203            buffer: Vec::new(),
204            buffer_pos: 0,
205            finished: false,
206            header_buf: [0u8; 8],
207            header_read: 0,
208            header_done: false,
209            ciphertext_buf: None,
210            ciphertext_read: 0,
211            ciphertext_len: 0,
212        }
213    }
214}
215
216impl<R> AsyncRead for DecryptReader<R>
217where
218    R: AsyncRead + Unpin + Send + Sync,
219{
220    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
221        let mut this = self.project();
222        // Serve from buffer if any
223        if *this.buffer_pos < this.buffer.len() {
224            let to_copy = std::cmp::min(buf.remaining(), this.buffer.len() - *this.buffer_pos);
225            buf.put_slice(&this.buffer[*this.buffer_pos..*this.buffer_pos + to_copy]);
226            *this.buffer_pos += to_copy;
227            if *this.buffer_pos == this.buffer.len() {
228                this.buffer.clear();
229                *this.buffer_pos = 0;
230            }
231            return Poll::Ready(Ok(()));
232        }
233        if *this.finished {
234            return Poll::Ready(Ok(()));
235        }
236        // Read header (8 bytes), support partial header read
237        while !*this.header_done && *this.header_read < 8 {
238            let mut temp = [0u8; 8];
239            let mut temp_buf = ReadBuf::new(&mut temp[0..8 - *this.header_read]);
240            match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
241                Poll::Pending => return Poll::Pending,
242                Poll::Ready(Ok(())) => {
243                    let n = temp_buf.filled().len();
244                    if n == 0 {
245                        break;
246                    }
247                    this.header_buf[*this.header_read..*this.header_read + n].copy_from_slice(&temp_buf.filled()[..n]);
248                    *this.header_read += n;
249                }
250                Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
251            }
252            if *this.header_read < 8 {
253                return Poll::Pending;
254            }
255        }
256        if !*this.header_done && *this.header_read == 8 {
257            *this.header_done = true;
258        }
259        if !*this.header_done {
260            return Poll::Pending;
261        }
262        let typ = this.header_buf[0];
263        let len = (this.header_buf[1] as usize) | ((this.header_buf[2] as usize) << 8) | ((this.header_buf[3] as usize) << 16);
264        let crc = (this.header_buf[4] as u32)
265            | ((this.header_buf[5] as u32) << 8)
266            | ((this.header_buf[6] as u32) << 16)
267            | ((this.header_buf[7] as u32) << 24);
268        *this.header_read = 0;
269        *this.header_done = false;
270        if typ == 0xFF {
271            *this.finished = true;
272            return Poll::Ready(Ok(()));
273        }
274        // Read ciphertext block (len bytes), support partial read
275        if this.ciphertext_buf.is_none() {
276            *this.ciphertext_len = len - 4; // 4 bytes for CRC32
277            *this.ciphertext_buf = Some(vec![0u8; *this.ciphertext_len]);
278            *this.ciphertext_read = 0;
279        }
280        let ciphertext_buf = this.ciphertext_buf.as_mut().unwrap();
281        while *this.ciphertext_read < *this.ciphertext_len {
282            let mut temp_buf = ReadBuf::new(&mut ciphertext_buf[*this.ciphertext_read..]);
283            match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
284                Poll::Pending => return Poll::Pending,
285                Poll::Ready(Ok(())) => {
286                    let n = temp_buf.filled().len();
287                    if n == 0 {
288                        break;
289                    }
290                    *this.ciphertext_read += n;
291                }
292                Poll::Ready(Err(e)) => {
293                    this.ciphertext_buf.take();
294                    *this.ciphertext_read = 0;
295                    *this.ciphertext_len = 0;
296                    return Poll::Ready(Err(e));
297                }
298            }
299        }
300        if *this.ciphertext_read < *this.ciphertext_len {
301            return Poll::Pending;
302        }
303        // Parse uvarint for plaintext length
304        let (plaintext_len, uvarint_len) = rustfs_utils::uvarint(&ciphertext_buf[0..16]);
305        let ciphertext = &ciphertext_buf[uvarint_len as usize..];
306
307        // Decrypt
308        let cipher = Aes256Gcm::new_from_slice(this.key).expect("key");
309        let nonce = Nonce::from_slice(this.nonce);
310        let plaintext = cipher
311            .decrypt(nonce, ciphertext)
312            .map_err(|e| std::io::Error::other(format!("decrypt error: {e}")))?;
313        if plaintext.len() != plaintext_len as usize {
314            this.ciphertext_buf.take();
315            *this.ciphertext_read = 0;
316            *this.ciphertext_len = 0;
317            return Poll::Ready(Err(std::io::Error::other("Plaintext length mismatch")));
318        }
319        // CRC32 check
320        let actual_crc = crc32fast::hash(&plaintext);
321        if actual_crc != crc {
322            this.ciphertext_buf.take();
323            *this.ciphertext_read = 0;
324            *this.ciphertext_len = 0;
325            return Poll::Ready(Err(std::io::Error::other("CRC32 mismatch")));
326        }
327        *this.buffer = plaintext;
328        *this.buffer_pos = 0;
329        // Clear block state for next block
330        this.ciphertext_buf.take();
331        *this.ciphertext_read = 0;
332        *this.ciphertext_len = 0;
333        let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
334        buf.put_slice(&this.buffer[..to_copy]);
335        *this.buffer_pos += to_copy;
336        Poll::Ready(Ok(()))
337    }
338}
339
340impl<R> EtagResolvable for DecryptReader<R>
341where
342    R: EtagResolvable,
343{
344    fn try_resolve_etag(&mut self) -> Option<String> {
345        self.inner.try_resolve_etag()
346    }
347}
348
349impl<R> HashReaderDetector for DecryptReader<R>
350where
351    R: EtagResolvable + HashReaderDetector,
352{
353    fn is_hash_reader(&self) -> bool {
354        self.inner.is_hash_reader()
355    }
356
357    fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
358        self.inner.as_hash_reader_mut()
359    }
360}
361
362#[cfg(test)]
363mod tests {
364    use std::io::Cursor;
365
366    use crate::WarpReader;
367
368    use super::*;
369    use rand::RngCore;
370    use tokio::io::{AsyncReadExt, BufReader};
371
372    #[tokio::test]
373    async fn test_encrypt_decrypt_reader_aes256gcm() {
374        let data = b"hello sse encrypt";
375        let mut key = [0u8; 32];
376        let mut nonce = [0u8; 12];
377        rand::rng().fill_bytes(&mut key);
378        rand::rng().fill_bytes(&mut nonce);
379
380        let reader = BufReader::new(&data[..]);
381        let encrypt_reader = EncryptReader::new(WarpReader::new(reader), key, nonce);
382
383        // Encrypt
384        let mut encrypt_reader = encrypt_reader;
385        let mut encrypted = Vec::new();
386        encrypt_reader.read_to_end(&mut encrypted).await.unwrap();
387
388        // Decrypt using DecryptReader
389        let reader = Cursor::new(encrypted.clone());
390        let decrypt_reader = DecryptReader::new(WarpReader::new(reader), key, nonce);
391        let mut decrypt_reader = decrypt_reader;
392        let mut decrypted = Vec::new();
393        decrypt_reader.read_to_end(&mut decrypted).await.unwrap();
394
395        assert_eq!(&decrypted, data);
396    }
397
398    #[tokio::test]
399    async fn test_decrypt_reader_only() {
400        // Encrypt some data first
401        let data = b"test decrypt only";
402        let mut key = [0u8; 32];
403        let mut nonce = [0u8; 12];
404        rand::rng().fill_bytes(&mut key);
405        rand::rng().fill_bytes(&mut nonce);
406
407        // Encrypt
408        let reader = BufReader::new(&data[..]);
409        let encrypt_reader = EncryptReader::new(WarpReader::new(reader), key, nonce);
410        let mut encrypt_reader = encrypt_reader;
411        let mut encrypted = Vec::new();
412        encrypt_reader.read_to_end(&mut encrypted).await.unwrap();
413
414        // Now test DecryptReader
415
416        let reader = Cursor::new(encrypted.clone());
417        let decrypt_reader = DecryptReader::new(WarpReader::new(reader), key, nonce);
418        let mut decrypt_reader = decrypt_reader;
419        let mut decrypted = Vec::new();
420        decrypt_reader.read_to_end(&mut decrypted).await.unwrap();
421
422        assert_eq!(&decrypted, data);
423    }
424
425    #[tokio::test]
426    async fn test_encrypt_decrypt_reader_large() {
427        use rand::Rng;
428        let size = 1024 * 1024;
429        let mut data = vec![0u8; size];
430        rand::rng().fill(&mut data[..]);
431        let mut key = [0u8; 32];
432        let mut nonce = [0u8; 12];
433        rand::rng().fill_bytes(&mut key);
434        rand::rng().fill_bytes(&mut nonce);
435
436        let reader = std::io::Cursor::new(data.clone());
437        let encrypt_reader = EncryptReader::new(WarpReader::new(reader), key, nonce);
438        let mut encrypt_reader = encrypt_reader;
439        let mut encrypted = Vec::new();
440        encrypt_reader.read_to_end(&mut encrypted).await.unwrap();
441
442        let reader = std::io::Cursor::new(encrypted.clone());
443        let decrypt_reader = DecryptReader::new(WarpReader::new(reader), key, nonce);
444        let mut decrypt_reader = decrypt_reader;
445        let mut decrypted = Vec::new();
446        decrypt_reader.read_to_end(&mut decrypted).await.unwrap();
447
448        assert_eq!(&decrypted, &data);
449    }
450}
rustfs_rio/encrypt_reader.rs

rustfs_rio/
encrypt_reader.rs
