systemprompt_files/services/upload/
service.rs

1//! [`FileUploadService`]: decode, validate, store, and record uploads.
2//!
3//! Decodes base64 payloads, enforces upload policy via [`FileValidator`],
4//! writes bytes to a traversal-checked storage path derived from the
5//! persistence mode, and records the file through [`FileRepository`], cleaning
6//! up the on-disk artefact if the database write fails.
7//!
8//! Copyright (c) systemprompt.io — Business Source License 1.1.
9//! See <https://systemprompt.io> for licensing details.
10
11use base64::Engine;
12use base64::engine::general_purpose::STANDARD;
13use sha2::{Digest, Sha256};
14use std::path::PathBuf;
15use systemprompt_identifiers::{ContextId, FileId, UserId};
16use tokio::fs;
17use tokio::io::AsyncWriteExt;
18use uuid::Uuid;
19
20use crate::config::{FilePersistenceMode, FilesConfig};
21use crate::models::{FileChecksums, FileMetadata};
22use crate::repository::{FileRepository, InsertFileRequest};
23
24use super::error::FileUploadError;
25use super::request::{FileUploadRequest, UploadedFile};
26use super::validator::{FileCategory, FileValidator};
27
28struct StoredArtifact<'a> {
29    file_id: &'a FileId,
30    storage_path: &'a std::path::Path,
31    public_url: &'a str,
32    size_bytes: u64,
33    sha256: String,
34}
35
36#[derive(Debug, Clone)]
37pub struct FileUploadService {
38    files_config: FilesConfig,
39    file_repository: FileRepository,
40    validator: FileValidator,
41}
42
43impl FileUploadService {
44    pub const fn new(file_repository: FileRepository, files_config: FilesConfig) -> Self {
45        let validator = FileValidator::new(*files_config.upload());
46
47        Self {
48            files_config,
49            file_repository,
50            validator,
51        }
52    }
53
54    pub const fn validator(&self) -> &FileValidator {
55        &self.validator
56    }
57
58    pub fn is_enabled(&self) -> bool {
59        let cfg = self.files_config.upload();
60        cfg.enabled && cfg.persistence_mode != FilePersistenceMode::Disabled
61    }
62
63    pub async fn upload_file(
64        &self,
65        request: FileUploadRequest,
66    ) -> Result<UploadedFile, FileUploadError> {
67        let upload_config = self.files_config.upload();
68
69        if upload_config.persistence_mode == FilePersistenceMode::Disabled {
70            return Err(FileUploadError::PersistenceDisabled);
71        }
72
73        let max_encoded_size = (upload_config.max_file_size_bytes as f64 * 1.34) as usize + 100;
74        if request.bytes_base64.len() > max_encoded_size {
75            return Err(FileUploadError::Base64TooLarge {
76                encoded_size: request.bytes_base64.len(),
77            });
78        }
79
80        let bytes = STANDARD.decode(&request.bytes_base64)?;
81        let size_bytes = bytes.len() as u64;
82
83        let category = self.validator.validate(&request.mime_type, size_bytes)?;
84
85        let file_id = FileId::new(Uuid::new_v4().to_string());
86        let extension = FileValidator::get_extension(&request.mime_type, request.name.as_deref());
87        let filename = format!("{}.{}", file_id.as_str(), extension);
88
89        let (storage_path, relative_path) = self.determine_storage_path(
90            &category,
91            &filename,
92            &request.context_id,
93            request.user_id.as_ref(),
94        )?;
95
96        if let Some(parent) = storage_path.parent() {
97            fs::create_dir_all(parent).await?;
98        }
99
100        let mut file = fs::File::create(&storage_path).await?;
101        file.write_all(&bytes).await?;
102        file.flush().await?;
103
104        let sha256 = hex::encode(Sha256::digest(&bytes));
105
106        let public_url = self.files_config.upload_url(&relative_path);
107
108        self.record_file(
109            StoredArtifact {
110                file_id: &file_id,
111                storage_path: &storage_path,
112                public_url: &public_url,
113                size_bytes,
114                sha256,
115            },
116            &request,
117        )
118        .await?;
119
120        Ok(UploadedFile {
121            file_id,
122            path: relative_path,
123            public_url,
124            size_bytes: size_bytes as i64,
125        })
126    }
127
128    async fn record_file(
129        &self,
130        artifact: StoredArtifact<'_>,
131        request: &FileUploadRequest,
132    ) -> Result<(), FileUploadError> {
133        let StoredArtifact {
134            file_id,
135            storage_path,
136            public_url,
137            size_bytes,
138            sha256,
139        } = artifact;
140
141        let metadata = FileMetadata::new().with_checksums(FileChecksums::new().with_sha256(sha256));
142
143        let mut insert_request = InsertFileRequest::new(
144            file_id.clone(),
145            storage_path.to_string_lossy().to_string(),
146            public_url.to_owned(),
147            request.mime_type.clone(),
148        )
149        .with_size(size_bytes as i64)
150        .with_metadata(metadata)
151        .with_context_id(request.context_id.clone());
152
153        if let Some(user_id) = request.user_id.clone() {
154            insert_request = insert_request.with_user_id(user_id);
155        }
156
157        if let Some(session_id) = request.session_id.clone() {
158            insert_request = insert_request.with_session_id(session_id);
159        }
160
161        if let Some(trace_id) = request.trace_id.clone() {
162            insert_request = insert_request.with_trace_id(trace_id);
163        }
164
165        if let Err(e) = self.file_repository.insert(insert_request).await {
166            if let Err(cleanup_err) = fs::remove_file(storage_path).await {
167                tracing::warn!(
168                    path = %storage_path.display(),
169                    error = %cleanup_err,
170                    "Failed to clean up uploaded file after database error"
171                );
172            }
173            return Err(FileUploadError::Database(e.to_string()));
174        }
175
176        Ok(())
177    }
178
179    fn determine_storage_path(
180        &self,
181        category: &FileCategory,
182        filename: &str,
183        context_id: &ContextId,
184        user_id: Option<&UserId>,
185    ) -> Result<(PathBuf, String), FileUploadError> {
186        let base = self.files_config.uploads();
187        let upload_config = self.files_config.upload();
188
189        let context_str = context_id.as_str();
190        Self::validate_path_inputs(context_str, filename, user_id)?;
191
192        let (full_path, relative) = match upload_config.persistence_mode {
193            FilePersistenceMode::ContextScoped => {
194                let rel = format!(
195                    "contexts/{}/{}/{}",
196                    context_str,
197                    category.storage_subdir(),
198                    filename
199                );
200                (base.join(&rel), rel)
201            },
202            FilePersistenceMode::UserLibrary => {
203                let user_dir =
204                    user_id.map_or_else(|| "anonymous".to_owned(), |u| u.as_str().to_owned());
205                let rel = format!(
206                    "users/{}/{}/{}",
207                    user_dir,
208                    category.storage_subdir(),
209                    filename
210                );
211                (base.join(&rel), rel)
212            },
213            FilePersistenceMode::Disabled => {
214                return Err(FileUploadError::PersistenceDisabled);
215            },
216        };
217
218        for component in std::path::Path::new(&relative).components() {
219            use std::path::Component;
220            match component {
221                Component::Normal(_) | Component::CurDir => {},
222                Component::ParentDir | Component::RootDir | Component::Prefix(_) => {
223                    return Err(FileUploadError::PathValidation(
224                        "Resolved path contains traversal or absolute component".to_owned(),
225                    ));
226                },
227            }
228        }
229
230        if !full_path.starts_with(&base) {
231            return Err(FileUploadError::PathValidation(
232                "Resolved path escapes upload directory".to_owned(),
233            ));
234        }
235
236        Ok((full_path, relative))
237    }
238
239    fn validate_path_inputs(
240        context_str: &str,
241        filename: &str,
242        user_id: Option<&UserId>,
243    ) -> Result<(), FileUploadError> {
244        if context_str.contains("..") || context_str.contains('\0') {
245            return Err(FileUploadError::PathValidation(
246                "Invalid context_id: contains path traversal sequence".to_owned(),
247            ));
248        }
249
250        if let Some(uid) = user_id {
251            let user_str = uid.as_str();
252            if user_str.contains("..") || user_str.contains('\0') {
253                return Err(FileUploadError::PathValidation(
254                    "Invalid user_id: contains path traversal sequence".to_owned(),
255                ));
256            }
257        }
258
259        if filename.contains('\0')
260            || filename.contains('/')
261            || filename.contains('\\')
262            || filename == ".."
263            || filename == "."
264            || filename.is_empty()
265        {
266            return Err(FileUploadError::PathValidation(
267                "Invalid filename: must be a single path component".to_owned(),
268            ));
269        }
270
271        Ok(())
272    }
273}
systemprompt_files/services/upload/service.rs

systemprompt_files/services/upload/
service.rs
