Rust 加密写入 与 解密读

use cipher::{
    block_padding::{Padding, Pkcs7},
    generic_array::GenericArray,
    BlockDecryptMut, BlockEncryptMut, Unsigned,
};

use std::io::{ErrorKind, Read, Write};
pub struct EncryptWriter<C: BlockEncryptMut, W: Write> {
    inner: W,
    cipher: C,
    last_block: GenericArray<u8, C::BlockSize>,
    pos: usize,
}
impl<C: BlockEncryptMut, W: Write> EncryptWriter<C, W> {
    pub fn new(inner: W, cipher: C) -> Self {
        Self {
            inner,
            cipher,
            last_block: GenericArray::default(),
            pos: 0,
        }
    }
}
impl<C: BlockEncryptMut, W: Write> Write for EncryptWriter<C, W> {
    fn write(&mut self, mut buf: &[u8]) -> std::io::Result<usize> {
        let block_size = C::block_size();
        let mut written = 0;
        if self.pos > 0 {
            let len = self.last_block.len() - self.pos;
            if len > buf.len() {
                self.last_block[self.pos..self.pos + buf.len()].copy_from_slice(buf);
                self.pos += buf.len();
                return Ok(buf.len());
            } else {
                self.last_block[self.pos..].copy_from_slice(&buf[..len]);
                let mut block = self.last_block.clone();
                self.cipher.encrypt_block_mut(&mut block);
                self.inner.write_all(&block)?;
                written += len;
                self.pos = 0;
                buf = &buf[len..];
            }
        }
        let full_blocks = buf.len() / block_size;
        let (buf, last) = buf.split_at(full_blocks * block_size);
        for chunk in buf.chunks(block_size) {
            let mut block = GenericArray::clone_from_slice(chunk);
            self.cipher.encrypt_block_mut(&mut block);
            self.inner.write_all(&block)?;
            written += block_size;
        }
        self.last_block[..last.len()].copy_from_slice(last);
        self.pos = last.len();
        written += last.len();
        Ok(written)
    }
    fn flush(&mut self) -> std::io::Result<()> {
        self.inner.flush()
    }
}

impl<C: BlockEncryptMut, W: Write> Drop for EncryptWriter<C, W> {
    fn drop(&mut self) {
        Pkcs7::pad(&mut self.last_block, self.pos);
        self.cipher.encrypt_block_mut(&mut self.last_block);
        self.inner.write_all(&self.last_block).unwrap();
        self.inner.flush().unwrap();
    }
}

pub struct DecryptoReader<C: BlockDecryptMut, R: Read> {
    inner: R,
    cipher: C,
    last_block: Option<GenericArray<u8, C::BlockSize>>,
    eof: bool,
}

impl<C: BlockDecryptMut, R: Read> DecryptoReader<C, R> {
    pub fn new(inner: R, cipher: C) -> Self {
        Self {
            inner,
            cipher,
            last_block: None,
            eof: false,
        }
    }
}

impl<C: BlockDecryptMut, R: Read> Read for DecryptoReader<C, R> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if self.eof {
            return Ok(0);
        }
        if buf.len() < C::BlockSize::USIZE {
            return Err(ErrorKind::InvalidInput.into());
        }
        let mut read_blocks = 0;
        if let Some(block) = self.last_block.take() {
            buf[..block.len()].copy_from_slice(&block);
            read_blocks += 1;
        }
        let block_size = C::BlockSize::USIZE;

        for i in buf.chunks_mut(block_size).skip(read_blocks) {
            let block = GenericArray::from_mut_slice(i);
            if let Err(e) = self.inner.read_exact(block) {
                if ErrorKind::UnexpectedEof == e.kind() {
                    self.eof = true;
                    break;
                }
                return Err(e);
            };
            self.cipher.decrypt_block_mut(block);
            read_blocks += 1;
        }

        if !self.eof {
            let mut block = GenericArray::default();
            if let Err(e) = self.inner.read_exact(&mut block) {
                if ErrorKind::UnexpectedEof == e.kind() {
                    self.eof = true;
                }
            } else {
                self.cipher.decrypt_block_mut(&mut block);
                self.last_block = Some(block);
            };
        }

        let mut read = read_blocks * block_size;
        if self.eof && read_blocks > 0 {
            let tmp = &mut buf[(read_blocks - 1) * block_size..];
            let block = GenericArray::<u8, C::BlockSize>::from_slice(&tmp[..block_size]);
            let result =
                Pkcs7::unpad(block).map_err(|e| std::io::Error::new(ErrorKind::InvalidData, e))?;
            read = (read_blocks - 1) * block_size + result.len();
        }

        Ok(read)
    }
}