Deploy again

server-rs/src/routes/screenshot.rs

@@ -1,55 +1,65 @@
 use std::sync::Arc;
 
+use axum::http::header::HeaderValue;
 use axum::http::{header, HeaderMap, StatusCode, Uri};
 use axum::response::IntoResponse;
 use tracing::{info, warn};
 
 use crate::state::AppState;
 
+/// Fetch a PNG screenshot from the screenshot service.
+/// Used by both the `/api/screenshot` proxy and the xlsx export.
+pub async fn fetch_screenshot_bytes(
+    state: &AppState,
+    query_string: &str,
+    auth_header: Option<&HeaderValue>,
+) -> Result<Vec<u8>, String> {
+    let url = format!("{}/screenshot?{}", state.screenshot_url, query_string);
+    info!("Fetching screenshot from: {}", url);
+
+    let mut req = state.http_client.get(&url);
+    if let Some(auth) = auth_header {
+        req = req.header(header::AUTHORIZATION, auth);
+    }
+
+    match req.send().await {
+        Ok(resp) if resp.status().is_success() => resp
+            .bytes()
+            .await
+            .map(|b| b.to_vec())
+            .map_err(|err| format!("Failed to read screenshot response: {err}")),
+        Ok(resp) => {
+            let status = resp.status();
+            let body = resp.text().await.unwrap_or_default();
+            Err(format!(
+                "Screenshot service returned {status}: {body}"
+            ))
+        }
+        Err(err) => Err(format!("Failed to reach screenshot service: {err}")),
+    }
+}
+
 pub async fn get_screenshot(
     state: Arc<AppState>,
     headers: HeaderMap,
     uri: Uri,
 ) -> impl IntoResponse {
-    let screenshot_base = &state.screenshot_url;
+    let qs = uri.query().unwrap_or_default();
+    let auth = headers.get(header::AUTHORIZATION);
 
-    let qs = uri
-        .query()
-        .map(|q| format!("?{q}"))
-        .unwrap_or_default();
-    let url = format!("{screenshot_base}/screenshot{qs}");
-    info!("Proxying screenshot request to: {}", url);
-
-    let mut req = state.http_client.get(&url);
-    if let Some(auth) = headers.get(header::AUTHORIZATION) {
-        req = req.header(header::AUTHORIZATION, auth);
-    }
-
-    match req.send().await {
-        Ok(resp) if resp.status().is_success() => match resp.bytes().await {
-            Ok(bytes) => (
-                StatusCode::OK,
-                [
-                    (header::CONTENT_TYPE, "image/png"),
-                    (header::CACHE_CONTROL, "public, max-age=86400"),
-                ],
-                bytes,
-            )
-                .into_response(),
-            Err(err) => {
-                warn!("Failed to read screenshot response: {}", err);
-                (StatusCode::BAD_GATEWAY, "Failed to read screenshot").into_response()
-            }
-        },
-        Ok(resp) => {
-            let status = resp.status();
-            let body = resp.text().await.unwrap_or_default();
-            warn!("Screenshot service returned status {}: {}", status, body);
-            (StatusCode::BAD_GATEWAY, "Screenshot service error").into_response()
-        }
+    match fetch_screenshot_bytes(&state, qs, auth).await {
+        Ok(bytes) => (
+            StatusCode::OK,
+            [
+                (header::CONTENT_TYPE, "image/png"),
+                (header::CACHE_CONTROL, "public, max-age=86400"),
+            ],
+            bytes,
+        )
+            .into_response(),
         Err(err) => {
-            warn!("Failed to reach screenshot service: {}", err);
-            (StatusCode::BAD_GATEWAY, "Screenshot service unavailable").into_response()
+            warn!("{err}");
+            (StatusCode::BAD_GATEWAY, "Screenshot service error").into_response()
         }
     }
 }