garbg Public

Watch 0 Fork 0 Star 0
Rust · 6877 bytes Raw Blame History
1 //! S3 source provider
2 //!
3 //! Supports fetching wallpapers from S3 buckets and S3-compatible storage.
4 //!
5 //! URI format: `s3://bucket/prefix`
6 //!
7 //! Authentication is handled via:
8 //! - AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables
9 //! - AWS credentials file (~/.aws/credentials)
10 //! - IAM instance roles (when running on EC2)
11 //!
12 //! For S3-compatible endpoints (MinIO, etc.):
13 //! - Set AWS_ENDPOINT_URL environment variable
14
15 #![cfg(feature = "s3")]
16
17 use anyhow::{Context, Result};
18 use async_trait::async_trait;
19 use aws_sdk_s3::Client;
20 use std::path::Path;
21
22 use super::{FetchedImage, MediaType, SourceProvider, WallpaperEntry};
23 use crate::media::ImageLoader;
24
25 /// Provider for S3 and S3-compatible object storage
26 pub struct S3Provider {
27 client: Client,
28 }
29
30 impl S3Provider {
31 /// Create a new S3 provider
32 ///
33 /// Loads credentials from environment or AWS config files.
34 pub async fn new() -> Result<Self> {
35 let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
36 let client = Client::new(&config);
37
38 Ok(Self { client })
39 }
40
41 /// Create with a custom endpoint URL (for S3-compatible services like MinIO)
42 pub async fn with_endpoint(endpoint_url: &str) -> Result<Self> {
43 let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
44
45 let s3_config = aws_sdk_s3::config::Builder::from(&config)
46 .endpoint_url(endpoint_url)
47 .force_path_style(true) // Required for most S3-compatible services
48 .build();
49
50 let client = Client::from_conf(s3_config);
51
52 Ok(Self { client })
53 }
54
55 /// Parse an s3:// URI into (bucket, prefix)
56 fn parse_uri(uri: &str) -> Result<(String, String)> {
57 let path = uri
58 .strip_prefix("s3://")
59 .context("Invalid S3 URI")?;
60
61 let parts: Vec<&str> = path.splitn(2, '/').collect();
62 let bucket = parts[0].to_string();
63 let prefix = parts.get(1).map(|s| s.to_string()).unwrap_or_default();
64
65 if bucket.is_empty() {
66 anyhow::bail!("S3 URI must include a bucket: s3://bucket/prefix");
67 }
68
69 Ok((bucket, prefix))
70 }
71
72 /// Determine media type from key (path)
73 fn media_type_from_key(key: &str) -> MediaType {
74 let ext = Path::new(key)
75 .extension()
76 .and_then(|e| e.to_str())
77 .map(|e| e.to_lowercase());
78
79 match ext.as_deref() {
80 Some("gif") => MediaType::AnimatedImage,
81 Some("mp4" | "webm") => MediaType::Video,
82 _ => MediaType::StaticImage,
83 }
84 }
85
86 /// Check if a key is a supported image format
87 fn is_image_key(key: &str) -> bool {
88 ImageLoader::is_supported_format(Path::new(key))
89 }
90 }
91
92 #[async_trait]
93 impl SourceProvider for S3Provider {
94 fn id(&self) -> &str {
95 "s3"
96 }
97
98 fn can_handle(&self, uri: &str) -> bool {
99 uri.starts_with("s3://")
100 }
101
102 async fn list(&self, uri: &str) -> Result<Vec<WallpaperEntry>> {
103 let (bucket, prefix) = Self::parse_uri(uri)?;
104
105 tracing::debug!("Listing S3 objects: bucket={}, prefix={}", bucket, prefix);
106
107 let mut entries = Vec::new();
108 let mut continuation_token: Option<String> = None;
109
110 // Paginate through results
111 loop {
112 let mut request = self.client
113 .list_objects_v2()
114 .bucket(&bucket)
115 .prefix(&prefix);
116
117 if let Some(token) = &continuation_token {
118 request = request.continuation_token(token);
119 }
120
121 let response = request
122 .send()
123 .await
124 .with_context(|| format!("Failed to list S3 bucket: {}", bucket))?;
125
126 if let Some(contents) = response.contents {
127 for object in contents {
128 if let Some(key) = object.key {
129 // Skip directory markers
130 if key.ends_with('/') {
131 continue;
132 }
133
134 // Only include supported image formats
135 if !Self::is_image_key(&key) {
136 continue;
137 }
138
139 let name = Path::new(&key)
140 .file_name()
141 .and_then(|n| n.to_str())
142 .unwrap_or(&key)
143 .to_string();
144
145 entries.push(WallpaperEntry {
146 uri: format!("s3://{}/{}", bucket, key),
147 name,
148 media_type: Self::media_type_from_key(&key),
149 size: object.size.map(|s| s as u64),
150 metadata: Default::default(),
151 });
152 }
153 }
154 }
155
156 // Check for more results
157 if response.is_truncated == Some(true) {
158 continuation_token = response.next_continuation_token;
159 } else {
160 break;
161 }
162 }
163
164 tracing::debug!("Found {} images in S3", entries.len());
165 Ok(entries)
166 }
167
168 async fn fetch(&self, entry: &WallpaperEntry) -> Result<FetchedImage> {
169 let (bucket, key) = Self::parse_uri(&entry.uri)?;
170
171 tracing::debug!("Fetching S3 object: bucket={}, key={}", bucket, key);
172
173 let response = self.client
174 .get_object()
175 .bucket(&bucket)
176 .key(&key)
177 .send()
178 .await
179 .with_context(|| format!("Failed to fetch S3 object: {}", entry.uri))?;
180
181 let bytes = response
182 .body
183 .collect()
184 .await
185 .context("Failed to read S3 object body")?
186 .into_bytes();
187
188 let image = ImageLoader::load_bytes(&bytes, None)?;
189
190 Ok(FetchedImage {
191 image,
192 uri: entry.uri.clone(),
193 media_type: entry.media_type,
194 })
195 }
196 }
197
198 #[cfg(test)]
199 mod tests {
200 use super::*;
201
202 #[test]
203 fn test_parse_uri() {
204 let (bucket, prefix) = S3Provider::parse_uri("s3://my-bucket/wallpapers/").unwrap();
205 assert_eq!(bucket, "my-bucket");
206 assert_eq!(prefix, "wallpapers/");
207
208 let (bucket, prefix) = S3Provider::parse_uri("s3://bucket").unwrap();
209 assert_eq!(bucket, "bucket");
210 assert_eq!(prefix, "");
211
212 let (bucket, prefix) = S3Provider::parse_uri("s3://bucket/path/to/file.png").unwrap();
213 assert_eq!(bucket, "bucket");
214 assert_eq!(prefix, "path/to/file.png");
215 }
216
217 #[test]
218 fn test_invalid_uri() {
219 assert!(S3Provider::parse_uri("http://example.com").is_err());
220 assert!(S3Provider::parse_uri("s3://").is_err());
221 }
222 }
223