encryptable_tokio_fs/tokio/
encryptable_fs.rs

1use std::io::ErrorKind;
2use crate::crypto::cryptor::{CryptorAsyncReader, CryptorAsyncWriter};
3use std::path::Path;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6use tokio::io;
7pub use tokio::fs::*;
8use tokio::io::{AsyncReadExt, AsyncWriteExt};
9
10pub type Key = [u8; crate::crypto::cryptor::KEY_LEN];
11
12static mut ENCRYPTION_KEY: Option<Key> = None;
13
14pub fn set_key(key: Key) -> Option<Key> {
15    let old_key = unsafe {
16        #[allow(static_mut_refs)]
17        ENCRYPTION_KEY.replace(key)
18    };
19    std::sync::atomic::fence(std::sync::atomic::Ordering::SeqCst);
20    old_key
21}
22
23fn get_key() -> &'static Option<Key> {
24    unsafe {
25        #[allow(static_mut_refs)]
26        &ENCRYPTION_KEY
27    }
28}
29
30// Tokio FS API replacements
31////////////////////////////
32
33/// Encryptable replacement for [tokio::fs::read()]
34pub async fn read(path: impl AsRef<Path>) -> io::Result<Vec<u8>> {
35    match get_key() {
36        None => tokio::fs::read(path).await,
37        Some(_key) => {
38            let file = File::open(path).await?;
39            let mut reader = io::BufReader::new(file);
40            let mut contents = Vec::new();
41            reader.read_to_end(&mut contents).await?;
42            Ok(contents)
43        },
44    }
45}
46
47/// Encryptable replacement for [tokio::fs::write()]
48pub async fn write(path: impl AsRef<Path>, contents: impl AsRef<[u8]>) -> io::Result<()> {
49    match get_key() {
50        None => tokio::fs::write(path, contents).await,
51        Some(_key) => {
52            let file = File::create(path).await?;
53            let mut writer = io::BufWriter::new(file);
54            writer.write_all(contents.as_ref()).await?;
55            writer.shutdown().await?;
56            Ok(())
57        },
58    }
59}
60
61/// Encryptable replacement for [tokio::fs::File]
62pub enum File {
63    Plain { tokio_file: tokio::fs::File },
64    EncryptedReader { cryptor_async_reader: CryptorAsyncReader<'static, tokio::fs::File> },
65    EncryptedWriter { cryptor_async_writer: CryptorAsyncWriter<tokio::fs::File> },
66}
67
68impl File {
69
70    /// Encryptable replacement for [tokio::fs::File::open()]
71    pub async fn open(path: impl AsRef<Path>) -> io::Result<File> {
72        Ok(
73            match get_key() {
74                None => File::Plain { tokio_file: tokio::fs::File::open(path).await? },
75                Some(key) => File::EncryptedReader { cryptor_async_reader: CryptorAsyncReader::new(tokio::fs::File::open(path).await?, key) }
76            }
77        )
78    }
79
80    /// Encryptable replacement for [tokio::fs::File::create()]
81    pub async fn create(path: impl AsRef<Path>) -> io::Result<File> {
82        Ok(
83            match get_key() {
84                None => File::Plain { tokio_file: tokio::fs::File::create(path).await? },
85                Some(key) => File::EncryptedWriter { cryptor_async_writer: CryptorAsyncWriter::new(tokio::fs::File::create(path).await?, key) }
86            }
87        )
88    }
89
90}
91
92impl io::AsyncRead for File {
93    fn poll_read(
94        self: Pin<&mut Self>,
95        cx: &mut Context<'_>,
96        dst: &mut io::ReadBuf<'_>,
97    ) -> Poll<io::Result<()>> {
98        match self.get_mut() {
99            Self::Plain { tokio_file } => std::pin::pin!(tokio_file).poll_read(cx, dst),
100            Self::EncryptedReader { cryptor_async_reader } => std::pin::pin!(cryptor_async_reader).poll_read(cx, dst),
101            File::EncryptedWriter { .. } => Poll::Ready(Err(io::Error::new(ErrorKind::Unsupported, "Attempted to read from a file opened exclusively for writing"))),
102        }
103    }
104}
105
106impl io::AsyncWrite for File {
107    fn poll_write(
108        self: Pin<&mut Self>,
109        cx: &mut Context<'_>,
110        src: &[u8],
111    ) -> Poll<std::io::Result<usize>> {
112        match self.get_mut() {
113            Self::Plain { tokio_file } => std::pin::pin!(tokio_file).poll_write(cx, src),
114            File::EncryptedWriter { cryptor_async_writer } => std::pin::pin!(cryptor_async_writer).poll_write(cx, src),
115            Self::EncryptedReader { .. } => Poll::Ready(Err(io::Error::new(ErrorKind::Unsupported, "Attempted to write on a file opened exclusively for reading"))),
116        }
117    }
118
119    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
120        match self.get_mut() {
121            Self::Plain { tokio_file } => std::pin::pin!(tokio_file).poll_flush(cx),
122            File::EncryptedWriter { cryptor_async_writer } => std::pin::pin!(cryptor_async_writer).poll_flush(cx),
123            Self::EncryptedReader { .. } => Poll::Ready(Err(io::Error::new(ErrorKind::Unsupported, "Attempted to write (flush) a file opened exclusively for reading"))),
124        }
125    }
126
127    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
128        match self.get_mut() {
129            Self::Plain { tokio_file } => std::pin::pin!(tokio_file).poll_shutdown(cx),
130            File::EncryptedWriter { cryptor_async_writer } => std::pin::pin!(cryptor_async_writer).poll_shutdown(cx),
131            Self::EncryptedReader { .. } => Poll::Ready(Err(io::Error::new(ErrorKind::Unsupported, "Attempted to write (shutdown) on a file opened exclusively for reading"))),
132        }
133    }
134
135}