headless_lms_server/domain/
answer_files_archive.rs1use std::{
7 io,
8 path::Path,
9 pin::Pin,
10 task::{Context, Poll},
11};
12
13use async_zip::{Compression, ZipEntryBuilder, base::write::ZipFileWriter};
14use bytes::Bytes;
15use futures::{StreamExt, io::AsyncWrite, io::AsyncWriteExt, ready};
16use models::exercise_task_submission_files::ExerciseAnswerFile;
17use tokio_stream::wrappers::ReceiverStream;
18use tokio_util::sync::PollSender;
19
20use super::{authorization::AuthorizationToken, csv_export::make_authorized_streamable};
21use crate::prelude::*;
22
23const MIME_EXTENSIONS: &[(&str, &str)] = &[
26 ("application/gzip", "gz"),
27 ("application/json", "json"),
28 ("application/pdf", "pdf"),
29 ("application/x-tar", "tar"),
30 ("application/x-zstd-compressed-tar", "tar.zst"),
31 ("application/zip", "zip"),
32 ("application/zstd", "zst"),
33 ("image/gif", "gif"),
34 ("image/jpeg", "jpg"),
35 ("image/png", "png"),
36 ("image/svg+xml", "svg"),
37 ("image/webp", "webp"),
38 ("text/csv", "csv"),
39 ("text/html", "html"),
40 ("text/markdown", "md"),
41 ("text/plain", "txt"),
42];
43
44fn extension_for_mime(mime: &str) -> Option<&'static str> {
47 let essence = mime
48 .split(';')
49 .next()
50 .unwrap_or(mime)
51 .trim()
52 .to_ascii_lowercase();
53 MIME_EXTENSIONS
54 .iter()
55 .find(|(candidate, _)| *candidate == essence)
56 .map(|(_, extension)| *extension)
57}
58
59fn entry_name(file: &ExerciseAnswerFile) -> String {
62 let stem = format!(
63 "{}/{}/{}",
64 file.user_id, file.exercise_task_submission_id, file.order_number
65 );
66 match extension_for_mime(&file.mime) {
67 Some(extension) => format!("{stem}.{extension}"),
68 None => stem,
69 }
70}
71
72fn filename_component(text: &str) -> String {
74 let mut component = String::new();
75 for character in text.chars() {
76 if character.is_ascii_alphanumeric() {
77 component.extend(character.to_lowercase());
78 } else if !component.ends_with('-') {
79 component.push('-');
80 }
81 }
82 let trimmed = component.trim_matches('-');
83 if trimmed.is_empty() {
84 "answers".to_string()
85 } else {
86 trimmed.to_string()
87 }
88}
89
90pub fn content_disposition(course_or_exam_name: &str, exercise_name: &str) -> String {
92 format!(
93 "attachment; filename=\"{}-{}-answers.zip\"",
94 filename_component(course_or_exam_name),
95 filename_component(exercise_name)
96 )
97}
98
99const ARCHIVE_CHANNEL_CAPACITY: usize = 8;
103
104struct ArchiveSink {
107 sender: PollSender<ControllerResult<Bytes>>,
108 authorization_token: AuthorizationToken,
109}
110
111impl AsyncWrite for ArchiveSink {
112 fn poll_write(
113 self: Pin<&mut Self>,
114 cx: &mut Context<'_>,
115 buf: &[u8],
116 ) -> Poll<io::Result<usize>> {
117 let sink = self.get_mut();
118 ready!(sink.sender.poll_reserve(cx))
119 .map_err(|error| io::Error::other(error.to_string()))?;
120 let token = sink.authorization_token;
121 sink.sender
122 .send_item(token.authorized_ok(Bytes::copy_from_slice(buf)))
123 .map_err(|error| io::Error::other(error.to_string()))?;
124 Poll::Ready(Ok(buf.len()))
125 }
126
127 fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
128 Poll::Ready(Ok(()))
129 }
130
131 fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
132 Poll::Ready(Ok(()))
133 }
134}
135
136pub async fn stream_exercise_answer_files(
141 pool: web::Data<PgPool>,
142 file_store: web::Data<dyn FileStore>,
143 exercise_id: Uuid,
144 content_disposition: String,
145 token: AuthorizationToken,
146) -> ControllerResult<HttpResponse> {
147 let (sender, receiver) =
148 tokio::sync::mpsc::channel::<ControllerResult<Bytes>>(ARCHIVE_CHANNEL_CAPACITY);
149 let mut conn = pool.acquire().await?;
150 let files = models::exercise_task_submission_files::get_answer_files_by_exercise_id(
151 &mut conn,
152 exercise_id,
153 )
154 .await?;
155 drop(conn);
156 actix_web::rt::spawn(async move {
158 let sink = ArchiveSink {
159 sender: PollSender::new(sender),
160 authorization_token: token,
161 };
162 if let Err(error) = write_archive(file_store.as_ref(), files, sink).await {
163 tracing::error!("Failed to write answer file archive: {}", error);
164 }
165 });
166
167 token.authorized_ok(
168 HttpResponse::Ok()
169 .append_header(("Content-Disposition", content_disposition))
170 .append_header(("Content-Type", "application/zip"))
171 .streaming(make_authorized_streamable(ReceiverStream::new(receiver))),
172 )
173}
174
175async fn write_archive(
176 file_store: &dyn FileStore,
177 files: Vec<ExerciseAnswerFile>,
178 sink: ArchiveSink,
179) -> anyhow::Result<()> {
180 let mut archive = ZipFileWriter::new(sink);
181 for file in files {
182 let contents = match file_store.download_stream(Path::new(&file.path)).await {
183 Ok(contents) => contents,
184 Err(error) => {
185 tracing::warn!(
186 "Skipping answer file {} missing from the file store: {}",
187 file.path,
188 error
189 );
190 continue;
191 }
192 };
193 let mut contents = Box::into_pin(contents);
194 let mut entry = archive
195 .write_entry_stream(ZipEntryBuilder::new(
196 entry_name(&file).into(),
197 Compression::Deflate,
198 ))
199 .await?;
200 while let Some(chunk) = contents.next().await {
201 entry.write_all(&chunk?).await?;
202 }
203 entry.close().await?;
204 }
205 archive.close().await?;
206 Ok(())
207}
208
209#[cfg(test)]
210mod test {
211 use super::*;
212
213 #[test]
214 fn maps_known_mime_types_to_extensions() {
215 assert_eq!(extension_for_mime("application/zip"), Some("zip"));
216 assert_eq!(extension_for_mime("text/plain"), Some("txt"));
217 assert_eq!(extension_for_mime("TEXT/PLAIN; charset=utf-8"), Some("txt"));
218 assert_eq!(extension_for_mime("image/jpeg"), Some("jpg"));
219 assert_eq!(
220 extension_for_mime("application/x-zstd-compressed-tar"),
221 Some("tar.zst")
222 );
223 }
224
225 #[test]
226 fn leaves_an_unknown_mime_type_without_an_extension() {
227 assert_eq!(extension_for_mime("application/octet-stream"), None);
228 assert_eq!(extension_for_mime("application/x-made-up"), None);
229 assert_eq!(extension_for_mime(""), None);
230 }
231
232 #[test]
233 fn names_entries_positionally() {
234 let user_id = Uuid::parse_str("00000000-0000-0000-0000-000000000001").unwrap();
235 let submission_id = Uuid::parse_str("00000000-0000-0000-0000-000000000002").unwrap();
236 let file = ExerciseAnswerFile {
237 user_id,
238 exercise_task_submission_id: submission_id,
239 created_at: Utc::now(),
240 path: "uploads/secret-real-name.png".to_string(),
241 mime: "image/png".to_string(),
242 order_number: 3,
243 };
244 assert_eq!(
245 entry_name(&file),
246 format!("{user_id}/{submission_id}/3.png")
247 );
248
249 let unknown = ExerciseAnswerFile {
250 mime: "application/octet-stream".to_string(),
251 ..file
252 };
253 assert_eq!(entry_name(&unknown), format!("{user_id}/{submission_id}/3"));
254
255 let two_part_extension = ExerciseAnswerFile {
256 mime: "application/x-zstd-compressed-tar".to_string(),
257 ..unknown
258 };
259 assert_eq!(
260 entry_name(&two_part_extension),
261 format!("{user_id}/{submission_id}/3.tar.zst")
262 );
263 }
264
265 #[test]
266 fn builds_a_filename_from_course_and_exercise_names() {
267 assert_eq!(
268 content_disposition("Introduction to Everything", "Exercise 1: Loops!"),
269 "attachment; filename=\"introduction-to-everything-exercise-1-loops-answers.zip\""
270 );
271 assert_eq!(
272 content_disposition("", " "),
273 "attachment; filename=\"answers-answers-answers.zip\""
274 );
275 }
276}