Skip to main content

headless_lms_server/domain/
answer_files_archive.rs

1//! Streams every file-typed answer of one exercise to a teacher as a zip archive.
2//!
3//! Entries are named positionally rather than from `file_uploads.name`, so that a plugin which
4//! anonymizes filenames for its own views cannot leak the real ones through this export.
5
6use 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
23/// Extensions for the mime types answer files realistically arrive as. Anything else is written
24/// without an extension rather than guessed at.
25const 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
44/// The extension to give an archive entry, without the leading dot. `None` for a mime type we have
45/// no confident extension for, including `application/octet-stream`.
46fn 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
59/// `<user_id>/<submission_id>/<order_number><ext>`, so that the two things a teacher needs to
60/// identify an answer come from host-owned data only.
61fn 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
72/// Lowercased `[a-z0-9-]` rendering of user-authored text, for use inside a download filename.
73fn 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
90/// `attachment; filename="..."` for an exercise's archive.
91pub 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
99/// How many written chunks may sit between the archive writer and the client before the writer has
100/// to wait. The producer reads from the file store as fast as it is served, so without a bound a
101/// slow client makes the whole archive accumulate in memory.
102const ARCHIVE_CHANNEL_CAPACITY: usize = 8;
103
104/// Writes the archive bytes into the streaming response channel as they are produced, waiting when
105/// the client is not reading them fast enough.
106struct 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
136/// Starts the archive download for one exercise.
137///
138/// The status line is committed before the first byte is read from the file store, so an upload
139/// that has vanished from it is logged and skipped rather than failing the response.
140pub 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    // `FileStore`'s futures are `?Send`, so the writer task has to stay on this thread.
157    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}