Skip to main content

headless_lms_utils/file_store/
local_file_store.rs

1use std::path::{Path, PathBuf};
2
3use async_trait::async_trait;
4use bytes::Bytes;
5use futures::{Stream, StreamExt};
6
7use tokio::{
8    fs::{self, OpenOptions},
9    io::{self, AsyncWriteExt, BufWriter},
10};
11use tokio_util::io::ReaderStream;
12
13use super::{FileStore, GenericPayload, generate_cache_folder_dir, path_to_str};
14use crate::prelude::*;
15
16#[derive(Debug, Clone)]
17pub struct LocalFileStore {
18    pub base_path: PathBuf,
19    pub base_url: String,
20    pub cache_files_path: PathBuf,
21}
22
23impl LocalFileStore {
24    pub fn new(base_path: PathBuf, base_url: String) -> UtilResult<Self> {
25        if base_path.exists() {
26            if !base_path.is_dir() {
27                return Err(UtilError::new(
28                    UtilErrorType::Other,
29                    "Base path should be a folder".to_string(),
30                    None,
31                ));
32            }
33        } else {
34            std::fs::create_dir_all(&base_path)?;
35        }
36        let cache_files_path = generate_cache_folder_dir()?;
37        Ok(Self {
38            base_path,
39            base_url,
40            cache_files_path,
41        })
42    }
43}
44#[async_trait(?Send)]
45impl FileStore for LocalFileStore {
46    async fn upload(&self, path: &Path, contents: Vec<u8>, _mime_type: &str) -> UtilResult<()> {
47        let full_path = self.base_path.join(path);
48        if let Some(parent) = full_path.parent() {
49            fs::create_dir_all(parent).await?;
50        }
51        fs::write(full_path, contents).await?;
52        Ok(())
53    }
54
55    async fn download(&self, path: &Path) -> UtilResult<Vec<u8>> {
56        let full_path = self.base_path.join(path);
57        Ok(fs::read(full_path).await?)
58    }
59
60    async fn delete(&self, path: &Path) -> UtilResult<()> {
61        let full_path = self.base_path.join(path);
62        fs::remove_file(full_path).await?;
63        Ok(())
64    }
65
66    async fn get_direct_download_url(&self, path: &Path) -> UtilResult<String> {
67        let full_path = self.base_path.join(path);
68        if !full_path.exists() {
69            return Err(UtilError::new(
70                UtilErrorType::Other,
71                "File does not exist.".to_string(),
72                None,
73            ));
74        }
75        let path_str = path_to_str(path)?;
76        if self.base_url.ends_with('/') {
77            return Ok(format!("{}{}", self.base_url, path_str));
78        }
79        Ok(format!("{}/{}", self.base_url, path_str))
80    }
81
82    /// The dev upload handler turns these parameters into the attachment header.
83    async fn get_direct_attachment_download_url(
84        &self,
85        path: &Path,
86        file_name: &str,
87    ) -> UtilResult<String> {
88        let encoded_file_name =
89            percent_encoding::utf8_percent_encode(file_name, percent_encoding::NON_ALPHANUMERIC);
90        Ok(format!(
91            "{}?download=1&name={encoded_file_name}",
92            self.get_direct_download_url(path).await?
93        ))
94    }
95
96    async fn upload_stream(
97        &self,
98        path: &Path,
99        mut contents: GenericPayload,
100        _mime_type: &str,
101    ) -> UtilResult<()> {
102        let full_path = self.base_path.join(path);
103        let parent = full_path.parent().ok_or_else(|| {
104            UtilError::new(
105                UtilErrorType::Other,
106                "Media path did not have a parent folder".to_string(),
107                None,
108            )
109        })?;
110        if parent.exists() {
111            if !parent.is_dir() {
112                return Err(UtilError::new(
113                    UtilErrorType::Other,
114                    "Base path should be a folder".to_string(),
115                    None,
116                ));
117            }
118        } else {
119            fs::create_dir_all(&parent).await?;
120        }
121        let file = OpenOptions::new()
122            .truncate(true)
123            .create(true)
124            .write(true)
125            .open(full_path)
126            .await?;
127
128        let mut buf_writer = BufWriter::new(file);
129
130        while let Some(bytes_res) = contents.next().await {
131            let bytes =
132                bytes_res.map_err(|e| UtilError::new(UtilErrorType::Other, e.to_string(), None))?;
133            buf_writer.write_all(&bytes).await?;
134        }
135
136        buf_writer.flush().await?;
137
138        Ok(())
139    }
140
141    async fn download_stream(
142        &self,
143        path: &Path,
144    ) -> UtilResult<Box<dyn Stream<Item = std::io::Result<Bytes>>>> {
145        let full_path = self.base_path.join(path);
146        let file = fs::File::open(full_path).await?;
147        let reader = io::BufReader::new(file);
148        let stream = ReaderStream::new(reader);
149        Ok(Box::new(stream))
150    }
151
152    fn get_cache_files_folder_path(&self) -> UtilResult<&Path> {
153        Ok(&self.cache_files_path)
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use std::{env, path::Path};
160
161    use tempfile::TempDir;
162
163    use super::LocalFileStore;
164    use crate::file_store::FileStore;
165
166    #[tokio::test]
167    async fn upload_download_delete_works() {
168        // TODO: Audit that the environment access only happens in single-threaded code.
169        unsafe { env::set_var("HEADLESS_LMS_CACHE_FILES_PATH", "/tmp") };
170        let dir = TempDir::new().expect("Failed to create a temp dir");
171        let base_path = dir.path().to_path_buf();
172        let local_file_store =
173            LocalFileStore::new(base_path.clone(), "http://localhost:3000".to_string())
174                .expect("Could not create local file storage");
175
176        let path1 = Path::new("file1");
177        let test_file_contents = "Test file contents".as_bytes().to_vec();
178        // Put content to storage and read it back
179        local_file_store
180            .upload(path1, test_file_contents.clone(), "text/plain")
181            .await
182            .expect("Failed to put a file into local file storage.");
183        let retrivied_file = local_file_store
184            .download(path1)
185            .await
186            .expect("Failed to retrieve a file from local file storage");
187        assert_eq!(test_file_contents, retrivied_file);
188
189        local_file_store
190            .delete(path1)
191            .await
192            .expect("Failed to delete a file");
193
194        // After deletion getting the file should fail
195        let retrivied_file2 = local_file_store.download(path1).await;
196        assert!(retrivied_file2.is_err());
197    }
198
199    #[tokio::test]
200    async fn get_download_url_works() {
201        // TODO: Audit that the environment access only happens in single-threaded code.
202        unsafe { env::set_var("HEADLESS_LMS_CACHE_FILES_PATH", "/tmp") };
203        let dir = TempDir::new().expect("Failed to create a temp dir");
204        let base_path = dir.path().to_path_buf();
205        let local_file_store =
206            LocalFileStore::new(base_path.clone(), "http://localhost:3000".to_string())
207                .expect("Could not create local file storage");
208        let test_file_contents = "Test file contents 2".as_bytes().to_vec();
209        let path1 = Path::new("file1");
210        local_file_store
211            .upload(path1, test_file_contents.clone(), "text/plain")
212            .await
213            .expect("Failed to put a file into local file storage.");
214        let url = local_file_store
215            .get_direct_download_url(path1)
216            .await
217            .expect("Failed to get a download url");
218        let expected_url = format!("http://localhost:3000/{}", path1.to_string_lossy());
219        assert_eq!(url, expected_url);
220
221        let nonexistant_file = Path::new("does-not-exist");
222        let res = local_file_store
223            .get_direct_download_url(nonexistant_file)
224            .await;
225        assert!(res.is_err());
226    }
227}