72 lines
2.3 KiB
Rust
72 lines
2.3 KiB
Rust
use std::sync::Arc;
|
|
|
|
use axum::extract::State;
|
|
use axum::http::header::HeaderValue;
|
|
use axum::http::{header, HeaderMap, StatusCode, Uri};
|
|
use axum::response::IntoResponse;
|
|
use metrics::histogram;
|
|
use tracing::{info, warn};
|
|
|
|
use crate::state::{AppState, SharedState};
|
|
|
|
/// Fetch a JPEG 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(shared): State<Arc<SharedState>>,
|
|
headers: HeaderMap,
|
|
uri: Uri,
|
|
) -> impl IntoResponse {
|
|
let state = shared.load_state();
|
|
let qs = uri.query().unwrap_or_default();
|
|
let auth = headers.get(header::AUTHORIZATION);
|
|
let is_og = qs.contains("og=1");
|
|
|
|
let t0 = std::time::Instant::now();
|
|
let result = fetch_screenshot_bytes(&state, qs, auth).await;
|
|
let kind = if is_og { "og" } else { "export" };
|
|
histogram!("screenshot_duration_seconds", "kind" => kind).record(t0.elapsed().as_secs_f64());
|
|
|
|
match result {
|
|
Ok(bytes) => (
|
|
StatusCode::OK,
|
|
[
|
|
(header::CONTENT_TYPE, "image/jpeg"),
|
|
(header::CACHE_CONTROL, "public, max-age=86400"),
|
|
],
|
|
bytes,
|
|
)
|
|
.into_response(),
|
|
Err(err) => {
|
|
warn!("{err}");
|
|
(StatusCode::BAD_GATEWAY, "Screenshot service error").into_response()
|
|
}
|
|
}
|
|
}
|