mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-11 15:46:53 +00:00
feat(storage): add direct chunk GET fast path (#2351)
Signed-off-by: houseme <housemecn@gmail.com> Co-authored-by: heihutu <heihutu@gmail.com> Co-authored-by: cxymds <Cxymds@qq.com>
This commit is contained in:
@@ -13,6 +13,7 @@
|
||||
// limitations under the License.
|
||||
|
||||
use crate::compress_index::{Index, TryGetIndex};
|
||||
use crate::{BlockReadable, BoxReadBlockFuture, Reader};
|
||||
use aes_gcm::aead::Aead;
|
||||
use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
|
||||
use pin_project_lite::pin_project;
|
||||
@@ -59,6 +60,69 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
fn encrypt_segment_bytes(cipher: &Aes256Gcm, nonce_bytes: &[u8; 12], plaintext: &[u8]) -> std::io::Result<Vec<u8>> {
|
||||
let nonce = Nonce::try_from(nonce_bytes.as_slice()).map_err(|_| Error::other("invalid nonce length"))?;
|
||||
let plaintext_len = plaintext.len();
|
||||
let crc = {
|
||||
let mut hasher = crc_fast::Digest::new(crc_fast::CrcAlgorithm::Crc32IsoHdlc);
|
||||
hasher.update(plaintext);
|
||||
hasher.finalize() as u32
|
||||
};
|
||||
let ciphertext = cipher
|
||||
.encrypt(&nonce, plaintext)
|
||||
.map_err(|e| Error::other(format!("encrypt error: {e}")))?;
|
||||
let int_len = put_uvarint_len(plaintext_len as u64);
|
||||
let clen = int_len + ciphertext.len() + 4;
|
||||
let mut header = [0u8; 8];
|
||||
header[0] = 0x00;
|
||||
header[1] = (clen & 0xFF) as u8;
|
||||
header[2] = ((clen >> 8) & 0xFF) as u8;
|
||||
header[3] = ((clen >> 16) & 0xFF) as u8;
|
||||
header[4] = (crc & 0xFF) as u8;
|
||||
header[5] = ((crc >> 8) & 0xFF) as u8;
|
||||
header[6] = ((crc >> 16) & 0xFF) as u8;
|
||||
header[7] = ((crc >> 24) & 0xFF) as u8;
|
||||
debug!(
|
||||
"encrypt block header typ=0 len={} header={:?} plaintext_len={} ciphertext_len={}",
|
||||
clen,
|
||||
header,
|
||||
plaintext_len,
|
||||
ciphertext.len()
|
||||
);
|
||||
let mut out = Vec::with_capacity(8 + int_len + ciphertext.len());
|
||||
out.extend_from_slice(&header);
|
||||
let mut plaintext_len_buf = vec![0u8; int_len];
|
||||
put_uvarint(&mut plaintext_len_buf, plaintext_len as u64);
|
||||
out.extend_from_slice(&plaintext_len_buf);
|
||||
out.extend_from_slice(&ciphertext);
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
impl<R> EncryptReader<R>
|
||||
where
|
||||
R: Reader,
|
||||
{
|
||||
fn copy_buffered(&mut self, buf: &mut [u8]) -> usize {
|
||||
if self.buffer_pos >= self.buffer.len() || buf.is_empty() {
|
||||
return 0;
|
||||
}
|
||||
|
||||
let to_copy = buf.len().min(self.buffer.len() - self.buffer_pos);
|
||||
buf[..to_copy].copy_from_slice(&self.buffer[self.buffer_pos..self.buffer_pos + to_copy]);
|
||||
self.buffer_pos += to_copy;
|
||||
if self.buffer_pos == self.buffer.len() {
|
||||
self.buffer.clear();
|
||||
self.buffer_pos = 0;
|
||||
}
|
||||
to_copy
|
||||
}
|
||||
|
||||
fn encrypt_segment(&self, plaintext: &[u8]) -> std::io::Result<Vec<u8>> {
|
||||
let nonce = derive_block_nonce(&self.base_nonce, self.block_index);
|
||||
encrypt_segment_bytes(&self.cipher, &nonce, plaintext)
|
||||
}
|
||||
}
|
||||
|
||||
impl<R> AsyncRead for EncryptReader<R>
|
||||
where
|
||||
R: AsyncRead + Unpin + Send + Sync,
|
||||
@@ -97,49 +161,8 @@ where
|
||||
*this.buffer_pos += to_copy;
|
||||
Poll::Ready(Ok(()))
|
||||
} else {
|
||||
// Encrypt the chunk
|
||||
let block_nonce = derive_block_nonce(this.base_nonce, *this.block_index);
|
||||
let nonce = Nonce::try_from(block_nonce.as_slice()).map_err(|_| Error::other("invalid nonce length"))?;
|
||||
let plaintext = &this.read_buffer[..n];
|
||||
let plaintext_len = plaintext.len();
|
||||
let crc = {
|
||||
let mut hasher = crc_fast::Digest::new(crc_fast::CrcAlgorithm::Crc32IsoHdlc);
|
||||
hasher.update(plaintext);
|
||||
hasher.finalize() as u32
|
||||
};
|
||||
let ciphertext = this
|
||||
.cipher
|
||||
.encrypt(&nonce, plaintext)
|
||||
.map_err(|e| Error::other(format!("encrypt error: {e}")))?;
|
||||
let int_len = put_uvarint_len(plaintext_len as u64);
|
||||
let clen = int_len + ciphertext.len() + 4;
|
||||
// Header: 8 bytes
|
||||
// 0: type (0 = encrypted, 0xFF = end)
|
||||
// 1-3: length (little endian u24, ciphertext length)
|
||||
// 4-7: CRC32 of ciphertext (little endian u32)
|
||||
let mut header = [0u8; 8];
|
||||
header[0] = 0x00; // 0 = encrypted
|
||||
header[1] = (clen & 0xFF) as u8;
|
||||
header[2] = ((clen >> 8) & 0xFF) as u8;
|
||||
header[3] = ((clen >> 16) & 0xFF) as u8;
|
||||
header[4] = (crc & 0xFF) as u8;
|
||||
header[5] = ((crc >> 8) & 0xFF) as u8;
|
||||
header[6] = ((crc >> 16) & 0xFF) as u8;
|
||||
header[7] = ((crc >> 24) & 0xFF) as u8;
|
||||
debug!(
|
||||
"encrypt block header typ=0 len={} header={:?} plaintext_len={} ciphertext_len={}",
|
||||
clen,
|
||||
header,
|
||||
plaintext_len,
|
||||
ciphertext.len()
|
||||
);
|
||||
let mut out = Vec::with_capacity(8 + int_len + ciphertext.len());
|
||||
out.extend_from_slice(&header);
|
||||
let mut plaintext_len_buf = [0u8; 10];
|
||||
let encoded_len = put_uvarint(&mut plaintext_len_buf, plaintext_len as u64);
|
||||
out.extend_from_slice(&plaintext_len_buf[..encoded_len]);
|
||||
out.extend_from_slice(&ciphertext);
|
||||
*this.buffer = out;
|
||||
*this.buffer = encrypt_segment_bytes(this.cipher, &block_nonce, &this.read_buffer[..n])?;
|
||||
*this.buffer_pos = 0;
|
||||
*this.block_index += 1;
|
||||
let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
|
||||
@@ -164,6 +187,50 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
impl<R> BlockReadable for EncryptReader<R>
|
||||
where
|
||||
R: Reader,
|
||||
{
|
||||
fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> {
|
||||
Box::pin(async move {
|
||||
if buf.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let mut written = self.copy_buffered(buf);
|
||||
while written < buf.len() {
|
||||
if self.finished {
|
||||
break;
|
||||
}
|
||||
|
||||
let mut plaintext = vec![0u8; 8 * 1024];
|
||||
let n = match self.inner.read_block(&mut plaintext).await {
|
||||
Ok(n) => n,
|
||||
Err(err) if err.kind() == std::io::ErrorKind::UnexpectedEof => 0,
|
||||
Err(err) => return Err(err),
|
||||
};
|
||||
if n == 0 {
|
||||
self.buffer = [0xFF, 0, 0, 0, 0, 0, 0, 0].to_vec();
|
||||
self.buffer_pos = 0;
|
||||
self.finished = true;
|
||||
} else {
|
||||
self.buffer = self.encrypt_segment(&plaintext[..n])?;
|
||||
self.buffer_pos = 0;
|
||||
}
|
||||
|
||||
let copied = self.copy_buffered(&mut buf[written..]);
|
||||
written += copied;
|
||||
|
||||
if copied == 0 && self.finished {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(written)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pin_project! {
|
||||
/// A reader wrapper that decrypts data on the fly using AES-256-GCM.
|
||||
/// This is a demonstration. For production, use a secure and audited crypto library.
|
||||
@@ -486,7 +553,7 @@ mod tests {
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use crate::HardLimitReader;
|
||||
use crate::{BlockReadable, HardLimitReader, WarpReader};
|
||||
|
||||
use super::*;
|
||||
use futures::StreamExt;
|
||||
@@ -674,6 +741,35 @@ mod tests {
|
||||
assert_eq!(&decrypted, &data);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_encrypt_reader_read_block_round_trips() {
|
||||
let data = b"hello sse encrypt via blocks";
|
||||
let mut key = [0u8; 32];
|
||||
let mut nonce = [0u8; 12];
|
||||
rand::rng().fill_bytes(&mut key);
|
||||
rand::rng().fill_bytes(&mut nonce);
|
||||
|
||||
let reader = Cursor::new(data.to_vec());
|
||||
let mut encrypt_reader = EncryptReader::new(WarpReader::new(reader), key, nonce);
|
||||
let mut encrypted = Vec::new();
|
||||
let mut buf = [0u8; 17];
|
||||
|
||||
loop {
|
||||
let n = encrypt_reader.read_block(&mut buf).await.unwrap();
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
encrypted.extend_from_slice(&buf[..n]);
|
||||
}
|
||||
|
||||
let reader = Cursor::new(encrypted);
|
||||
let mut decrypt_reader = DecryptReader::new(WarpReader::new(reader), key, nonce);
|
||||
let mut decrypted = Vec::new();
|
||||
decrypt_reader.read_to_end(&mut decrypted).await.unwrap();
|
||||
|
||||
assert_eq!(&decrypted, data);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_decrypt_reader_large_with_small_chunks() {
|
||||
let size = 1024 * 1024;
|
||||
|
||||
Reference in New Issue
Block a user