AppFlowy-Cloud/libs/database/src/file/utils.rs

62 lines
1.3 KiB
Rust

use base64::alphabet::URL_SAFE;
use base64::engine::general_purpose::PAD;
use base64::engine::GeneralPurpose;
use base64::Engine;
use sha2::{Digest, Sha256};
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{self, AsyncRead, AsyncReadExt, ReadBuf};
pub const URL_SAFE_ENGINE: GeneralPurpose = GeneralPurpose::new(&URL_SAFE, PAD);
pub struct BlobStreamReader<R> {
reader: R,
hasher: Sha256,
}
impl<R> AsyncRead for BlobStreamReader<R>
where
R: AsyncRead + Unpin,
{
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let before = buf.filled().len();
let poll = Pin::new(&mut self.reader).poll_read(cx, buf);
let after = buf.filled().len();
if after > before {
self.hasher.update(&buf.filled()[before..after]);
}
poll
}
}
impl<R> BlobStreamReader<R>
where
R: AsyncRead + Unpin,
{
pub fn new(reader: R) -> Self {
Self {
reader,
hasher: Sha256::new(),
}
}
pub async fn finish(mut self) -> io::Result<(Vec<u8>, String)> {
let mut buffer = Vec::new();
let _ = self.read_to_end(&mut buffer).await?;
let hash = URL_SAFE_ENGINE.encode(self.hasher.finalize());
Ok((buffer, hash))
}
}
impl<R> AsRef<R> for BlobStreamReader<R>
where
R: AsyncRead + Unpin,
{
fn as_ref(&self) -> &R {
&self.reader
}
}