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:
houseme
2026-04-07 08:33:46 +08:00
committed by GitHub
parent 8d27170ce4
commit 32bf8f5bf3
84 changed files with 15932 additions and 3592 deletions
+139 -43
View File
@@ -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;