random_image_server/
lib.rs

1use std::{convert::Infallible, fs, path::PathBuf, sync::Arc};
2
3use anyhow::{Result, anyhow};
4use http_body_util::Full;
5use hyper::{Request, Response, body::Bytes, service::service_fn};
6use hyper_util::{
7    rt::{TokioExecutor, TokioIo},
8    server::conn::auto,
9};
10use tokio::{
11    net::TcpListener,
12    sync::{RwLock, broadcast::Receiver},
13};
14use url::Url;
15
16use crate::config::{Config, ImageSource};
17use crate::state::ServerState;
18use crate::termination::Interrupted;
19
20pub mod cache;
21pub mod config;
22mod logging;
23pub mod state;
24pub use logging::init_logging;
25pub mod env;
26pub mod termination;
27
28pub const ALLOWED_IMAGE_EXTENSIONS: &[&str] = &["jpg", "jpeg", "png", "webp", "gif"];
29
30/// The main server structure
31pub struct ImageServer {
32    pub config: Config,
33    pub state: Arc<RwLock<ServerState>>,
34}
35
36impl ImageServer {
37    /// Create a new `ImageServer` instance with default configuration
38    #[must_use]
39    pub fn new() -> Self {
40        Self {
41            config: Config::default(),
42            state: Arc::new(RwLock::new(ServerState::default())),
43        }
44    }
45
46    /// Create a new `ImageServer` instance with custom configuration
47    #[must_use]
48    pub fn with_config(config: Config) -> Self {
49        Self {
50            state: Arc::new(RwLock::new(ServerState::with_config(&config))),
51            config,
52        }
53    }
54
55    /// Populate the cache with the configured images
56    ///
57    /// # Errors
58    ///
59    /// Returns an error if the image file does not exist, is not a file, or has an unsupported extension.
60    pub async fn populate_cache(&self) {
61        // This method can be implemented to load images from configured sources
62        // and populate the cache. For now, it is a placeholder.
63        tracing::info!("Populating cache with configured images...");
64
65        for source in &self.config.server.sources {
66            match source {
67                ImageSource::Url(url) => {
68                    tracing::info!("Loading image from URL: {url}");
69                    let key = cache::CacheKey::ImageUrl(url.clone());
70                    // fetch the image from the URL and store it in the cache
71                    match read_image_from_url(url).await {
72                        Ok(image) => {
73                            let set_result = self.state.write().await.cache.set(key, image);
74                            if let Err(err) = set_result {
75                                tracing::error!("Failed to store image in cache: {err}");
76                            }
77                        }
78                        Err(e) => {
79                            tracing::error!("Failed to read image from URL {url}: {e}");
80                        }
81                    }
82                }
83                ImageSource::Path(path) if path.is_file() => {
84                    let path = path.canonicalize().unwrap_or_else(|_| {
85                        tracing::warn!("Failed to canonicalize path: {}", path.display());
86                        path.clone()
87                    });
88                    if path.extension().is_some_and(|ext| {
89                        ALLOWED_IMAGE_EXTENSIONS.contains(&ext.to_string_lossy().as_ref())
90                    }) {
91                        tracing::info!("Loading image from file path: {}", path.display());
92                        // read the image file from the path and store it in the cache
93                        let Ok(image) = read_image_from_path(&path) else {
94                            tracing::error!("Failed to read image file: {}", path.display());
95                            continue;
96                        };
97                        let key = cache::CacheKey::ImagePath(path.clone());
98                        let set_result = self.state.write().await.cache.set(key, image);
99                        if let Err(err) = set_result {
100                            tracing::error!("Failed to store image in cache: {err}");
101                        }
102                    } else {
103                        tracing::warn!("Unsupported image file extension: {}", path.display());
104                    }
105                }
106                ImageSource::Path(path) if path.is_dir() => {
107                    let path = path.canonicalize().unwrap_or_else(|_| {
108                        tracing::warn!("Failed to canonicalize path: {}", path.display());
109                        path.clone()
110                    });
111
112                    tracing::info!("Loading images from directory: {}", path.display());
113                    // Read all image files in the directory and store them in the cache
114                    let mut state = self.state.write().await;
115                    walkdir::WalkDir::new(&path)
116                        .into_iter()
117                        .filter_map(Result::ok)
118                        .filter(|e| e.file_type().is_file())
119                        .filter(|e| {
120                            e.path()
121                                .extension()
122                                .and_then(|ext| ext.to_str())
123                                .is_some_and(|ext| ALLOWED_IMAGE_EXTENSIONS.contains(&ext))
124                        })
125                        .for_each(|entry| {
126                            let path = entry.path().to_path_buf();
127                            tracing::info!("Loading image from file: {}", path.display());
128                            // read the image file and store it in the cache
129                            match read_image_from_path(&path) {
130                                Ok(image) => {
131                                    let key = cache::CacheKey::ImagePath(path.clone());
132                                    let set_result = state.cache.set(key, image);
133                                    if let Err(err) = set_result {
134                                        tracing::error!("Failed to store image in cache: {err}");
135                                    }
136                                }
137                                Err(e) => {
138                                    tracing::error!(
139                                        "Failed to read image from path {}: {e}",
140                                        path.display(),
141                                    );
142                                }
143                            }
144                        });
145                }
146                ImageSource::Path(path) => {
147                    tracing::warn!("Unsupported image path: {}", path.display());
148                }
149            }
150        }
151    }
152
153    /// Start the server
154    ///
155    /// # Errors
156    ///
157    /// Returns an error if the server fails to start or encounters an unexpected error.
158    pub async fn start(&self, mut interrupt_rx: Receiver<Interrupted>) -> Result<()> {
159        let addr = self.config.socket_addr()?;
160        let listener = TcpListener::bind(addr).await?;
161        tracing::info!("Server running on http://{addr}");
162        tracing::debug!("Configuration: {:?}", self.config);
163
164        // Populate the cache with images from configured sources
165        self.populate_cache().await;
166        if self.state.read().await.cache.size() == 0 {
167            tracing::error!("No images found in cache, please check your configuration");
168            return Err(anyhow!(
169                "No images found in cache, please check your configuration"
170            ));
171        }
172
173        let executor = auto::Builder::new(TokioExecutor::new());
174        let graceful = hyper_util::server::graceful::GracefulShutdown::new();
175
176        loop {
177            tokio::select! {
178                Ok((stream, _addr)) = listener.accept() => {
179                    let io = TokioIo::new(stream);
180
181                    // Clone state for the handler
182                    let state = self.state.clone();
183                    let service = service_fn(move |req| {
184                        handle_request(req, state.clone())
185                    });
186
187                    // watch this connection
188                    let conn = executor.serve_connection(io, service);
189                    let fut = graceful.watch(conn.into_owned());
190
191                    // Spawn a new task to handle the connection
192                    tokio::spawn(async move {
193                        if let Err(e) = fut.await {
194                            tracing::error!("Failed to serve connection: {e}");
195                        }
196                    });
197                },
198
199                _ = interrupt_rx.recv() => {
200                    drop(listener);
201                    tracing::info!("Received termination signal, shutting down server");
202                    break;
203                }
204            };
205        }
206
207        // Start the shutdown and wait for any existing connections to close
208        tokio::select! {
209            () = graceful.shutdown() => {
210                tracing::info!("All connections gracefully closed");
211            }
212            () = tokio::time::sleep(std::time::Duration::from_secs(5)) => {
213                tracing::warn!("Timed out waiting for all connections to close");
214            }
215        }
216
217        Ok(())
218    }
219}
220
221impl Default for ImageServer {
222    fn default() -> Self {
223        Self::new()
224    }
225}
226
227/// Read an image file from the given path and return it as a `CacheValue`
228///
229/// # Errors
230///
231/// Returns an error if the file does not exist, is not a file, or has an unsupported extension.
232pub fn read_image_from_path(path: &PathBuf) -> Result<cache::CacheValue> {
233    let path_display = path.display();
234    if !path.exists() || !path.is_file() {
235        return Err(anyhow!("Image file does not exist: {path_display}"));
236    }
237    let Some(ext) = path.extension().and_then(|ext| ext.to_str()) else {
238        return Err(anyhow!("Image file has no extension: {path_display}"));
239    };
240    if !ALLOWED_IMAGE_EXTENSIONS.contains(&ext) {
241        return Err(anyhow!(
242            "Unsupported image file extension: {}",
243            path.display()
244        ));
245    }
246
247    let image_data = fs::read(path).map_err(|e| anyhow!("Failed to read image file: {e}"))?;
248    let content_type = mime_guess::from_path(path)
249        .first()
250        .ok_or_else(|| anyhow!("Failed to determine content type for image file: {path_display}"))?
251        .to_string();
252    Ok(cache::CacheValue {
253        data: image_data,
254        content_type,
255    })
256}
257
258/// Fetch an image from a URL and return it as a `CacheValue`
259///
260/// # Errors
261///
262/// Returns an error if the image cannot be fetched or if the content type is unsupported.
263pub async fn read_image_from_url(url: &Url) -> Result<cache::CacheValue> {
264    let response = reqwest::get(url.as_str())
265        .await
266        .map_err(|e| anyhow!("Failed to fetch image from URL: {e}"))?;
267
268    if !response.status().is_success() {
269        return Err(anyhow!(
270            "Failed to fetch image, status: {}",
271            response.status()
272        ));
273    }
274
275    let content_type = response
276        .headers()
277        .get("Content-Type")
278        .and_then(|v| v.to_str().ok())
279        .ok_or_else(|| anyhow!("Failed to get Content-Type header from response"))?
280        .to_string();
281
282    if !ALLOWED_IMAGE_EXTENSIONS.contains(&content_type.split('/').next_back().unwrap_or("")) {
283        return Err(anyhow!("Unsupported image content type: {content_type}"));
284    }
285
286    let data = response
287        .bytes()
288        .await
289        .map_err(|e| anyhow!("Failed to read image bytes from response: {e}"))?;
290
291    Ok(cache::CacheValue {
292        data: data.to_vec(),
293        content_type,
294    })
295}
296
297/// Handle incoming HTTP requests
298///
299/// # Errors
300///
301/// should be Infallible
302pub async fn handle_request(
303    req: Request<hyper::body::Incoming>,
304    state: Arc<RwLock<ServerState>>,
305) -> Result<Response<Full<Bytes>>, Infallible> {
306    match req.uri().path() {
307        "/" => Ok(Response::new(Full::new(Bytes::from(
308            "Welcome to the Random Image Server!",
309        )))),
310        "/health" => Ok(Response::new(Full::new(Bytes::from("OK")))),
311        "/random" => match handle_random_image(state).await {
312            Ok(response) => Ok(response),
313            Err(err) => {
314                tracing::error!("Failed to get random image: {err}");
315                let mut not_found = Response::new(Full::new(Bytes::from("Not Found")));
316                *not_found.status_mut() = hyper::StatusCode::NOT_FOUND;
317                Ok(not_found)
318            }
319        },
320        "/sequential" => match handle_sequential_image(state).await {
321            Ok(response) => Ok(response),
322            Err(err) => {
323                tracing::error!("Failed to get sequential image: {err}");
324                let mut not_found = Response::new(Full::new(Bytes::from("Not Found")));
325                *not_found.status_mut() = hyper::StatusCode::NOT_FOUND;
326                Ok(not_found)
327            }
328        },
329        _ => {
330            let mut not_found = Response::new(Full::new(Bytes::from("Not Found")));
331            *not_found.status_mut() = hyper::StatusCode::NOT_FOUND;
332            Ok(not_found)
333        }
334    }
335}
336
337/// Handle random image serving
338///
339/// # Errors
340///
341/// Returns an error if no images are configured or if the image cannot be found in the cache.
342pub async fn handle_random_image(state: Arc<RwLock<ServerState>>) -> Result<Response<Full<Bytes>>> {
343    let state = state.read().await;
344
345    // get a random image from the cache
346    state.cache.get_random().map_or_else(
347        || {
348            Err(anyhow!(
349                "Failed to retrieve a random image, perhaps no images are configured"
350            ))
351        },
352        |image| {
353            let body = Full::new(Bytes::from(image.data));
354            let mut response = Response::new(body);
355            *response.status_mut() = hyper::StatusCode::OK;
356            response
357                .headers_mut()
358                .insert(hyper::header::CONTENT_TYPE, image.content_type.parse()?);
359            Ok(response)
360        },
361    )
362}
363
364/// Handle sequential image serving
365///
366/// # Errors
367///
368/// Returns an error if no images are configured or if the image cannot be found in the cache.
369pub async fn handle_sequential_image(
370    state: Arc<RwLock<ServerState>>,
371) -> Result<Response<Full<Bytes>>> {
372    let mut state = state.write().await;
373
374    if state.cache.is_empty() {
375        return Err(anyhow!("No image sources configured"));
376    }
377
378    let current_index = state.current_index % state.cache.size();
379    let source = state.cache.keys()[current_index].clone();
380    state.current_index = (current_index + 1) % state.cache.size();
381
382    // Fetch the image from the cache or source
383    if let Some(image) = state.cache.get(source.clone()) {
384        let body = Full::new(Bytes::from(image.data));
385        let mut response = Response::new(body);
386        *response.status_mut() = hyper::StatusCode::OK;
387        response
388            .headers_mut()
389            .insert(hyper::header::CONTENT_TYPE, image.content_type.parse()?);
390        Ok(response)
391    } else {
392        state.cache.remove(&source);
393        drop(state);
394        Err(anyhow!("Image not found in cache"))
395    }
396}
397
398#[cfg(test)]
399mod tests {
400    use super::*;
401    use crate::termination::create_termination;
402    use pretty_assertions::assert_eq;
403    use rstest::rstest;
404
405    #[test]
406    fn test_allowed_image_extensions() {
407        assert!(ALLOWED_IMAGE_EXTENSIONS.contains(&"jpg"));
408        assert!(ALLOWED_IMAGE_EXTENSIONS.contains(&"jpeg"));
409        assert!(ALLOWED_IMAGE_EXTENSIONS.contains(&"png"));
410        assert!(ALLOWED_IMAGE_EXTENSIONS.contains(&"webp"));
411        assert!(ALLOWED_IMAGE_EXTENSIONS.contains(&"gif"));
412        assert_eq!(ALLOWED_IMAGE_EXTENSIONS.len(), 5);
413    }
414
415    #[rstest]
416    #[tokio::test]
417    #[timeout(std::time::Duration::from_secs(2))]
418    async fn test_start_stop_server() {
419        let mut server = ImageServer::default();
420        let port = 0;
421        server.config.server.port = port;
422        server.config.server.sources = vec![ImageSource::Path(PathBuf::from("assets"))];
423
424        let (mut terminator, interrupt_rx) = create_termination();
425        terminator.terminate(Interrupted::UserInt).unwrap();
426        server.start(interrupt_rx).await.unwrap();
427    }
428}
random_image_server/lib.rs

random_image_server/
lib.rs
