use anyhow::{Context, Result};
use std::time::{Duration, Instant};
use super::{SourceRejected, UpstreamFault};
fn endpoint() -> String {
std::env::var("OXIMG_GCS_ENDPOINT")
.ok()
.map(|v| v.trim_end_matches('/').to_string())
.filter(|v| !v.is_empty())
.unwrap_or_else(|| "https://storage.googleapis.com".to_string())
}
fn metadata_host() -> String {
std::env::var("GCE_METADATA_HOST")
.ok()
.filter(|v| !v.trim().is_empty())
.unwrap_or_else(|| "metadata.google.internal".to_string())
}
struct CachedToken {
bearer: String,
expires_at: Instant,
}
static TOKEN: tokio::sync::Mutex<Option<CachedToken>> = tokio::sync::Mutex::const_new(None);
async fn fetch_token() -> Result<CachedToken> {
let url = format!(
"http://{}/computeMetadata/v1/instance/service-accounts/default/token",
metadata_host()
);
let resp = super::fetch_client()
.get(&url)
.header("Metadata-Flavor", "Google")
.send()
.await
.context("GCP metadata server token request")?;
let body = resp.text().await.context("read metadata token response")?;
let v: serde_json::Value =
serde_json::from_str(&body).context("parse metadata token response")?;
let token = v["access_token"]
.as_str()
.context("metadata token response lacks access_token")?;
let expires_in = v["expires_in"].as_u64().unwrap_or(300).min(24 * 3600);
Ok(CachedToken {
bearer: format!("Bearer {token}"),
expires_at: Instant::now() + Duration::from_secs(expires_in.saturating_sub(60).max(10)),
})
}
async fn bearer(force: bool) -> Result<String> {
let mut guard = TOKEN.lock().await;
if !force
&& let Some(t) = guard.as_ref()
&& t.expires_at > Instant::now()
{
return Ok(t.bearer.clone());
}
let fresh = fetch_token().await?;
let bearer = fresh.bearer.clone();
*guard = Some(fresh);
Ok(bearer)
}
pub(crate) fn startup() -> Result<(), String> {
super::block_on_fetch(async { bearer(false).await.map(|_| ()) }).map_err(|e| {
format!(
"gs:// source needs GCP-attached credentials \
(metadata server at {:?} unreachable: {e:#}) — GKE Workload \
Identity, Cloud Run, and GCE provide them; service-account \
JSON keys are not supported (use the HTTP mode off-GCP)",
metadata_host()
)
})
}
const GCS_MAX_KEY_BYTES: usize = 1024;
fn decoded_key_len(key: &str) -> usize {
key.len() - 2 * key.bytes().filter(|b| *b == b'%').count()
}
fn retryable_status(code: u16) -> bool {
matches!(code, 401 | 429 | 500 | 502 | 503 | 504)
}
pub(crate) async fn fetch(bucket: &str, key: &str) -> Result<reqwest::Response> {
if decoded_key_len(key) > GCS_MAX_KEY_BYTES {
return Err(anyhow::Error::new(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!(
"object name is {} bytes, over the {GCS_MAX_KEY_BYTES}-byte GCS limit",
decoded_key_len(key)
),
)));
}
let url = format!("{}/{bucket}/{key}", endpoint());
let attempt = async |force_token: bool| -> Result<reqwest::Response> {
let bearer = bearer(force_token).await?;
super::fetch_client()
.get(&url)
.header("Authorization", &bearer)
.send()
.await
.map_err(anyhow::Error::new)
};
let resp = match attempt(false).await {
Ok(resp) if retryable_status(resp.status().as_u16()) => {
let force_token = resp.status().as_u16() == 401;
super::UPSTREAM_RETRIES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
tokio::time::sleep(Duration::from_millis(100)).await;
attempt(force_token).await.map_err(map_transport_err)?
}
Ok(resp) => resp,
Err(e) => {
let transient = e
.downcast_ref::<reqwest::Error>()
.is_some_and(|re| !re.is_timeout() && (re.is_connect() || re.is_request()));
if !transient {
return Err(map_transport_err(e));
}
super::UPSTREAM_RETRIES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
tokio::time::sleep(Duration::from_millis(100)).await;
attempt(false).await.map_err(map_transport_err)?
}
};
refuse_status(resp, bucket)
}
fn map_transport_err(e: anyhow::Error) -> anyhow::Error {
if e.downcast_ref::<reqwest::Error>()
.is_some_and(reqwest::Error::is_timeout)
{
return anyhow::Error::new(std::io::Error::new(std::io::ErrorKind::TimedOut, e));
}
e.context("fetch gcs object").context(UpstreamFault)
}
fn refuse_status(resp: reqwest::Response, bucket: &str) -> Result<reqwest::Response> {
let status = resp.status();
match status.as_u16() {
404 => Err(anyhow::Error::new(std::io::Error::new(
std::io::ErrorKind::NotFound,
"object not found in bucket",
))),
code @ (400 | 414) => Err(
anyhow::anyhow!("object store rejected the request ({code})").context(SourceRejected),
),
401 | 403 => Err(anyhow::Error::new(std::io::Error::new(
std::io::ErrorKind::PermissionDenied,
format!("access to bucket {bucket:?} denied (check the service account's roles)"),
))),
_ if status.is_redirection() => Err(anyhow::anyhow!(
"object store answered {status} (redirects are not followed)"
)
.context(UpstreamFault)),
_ if !status.is_success() => Err(anyhow::anyhow!("object store answered {status}")
.context("fetch gcs object")
.context(UpstreamFault)),
_ => Ok(resp),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn decoded_key_len_counts_escapes_as_one_byte() {
assert_eq!(decoded_key_len("photo.jpg"), 9);
assert_eq!(decoded_key_len("a%20b.jpg"), 7);
assert_eq!(decoded_key_len("%E4%B8%AD.jpg"), 7);
}
#[test]
fn key_length_boundary_is_the_documented_limit() {
let at = |n: usize| decoded_key_len(&"x".repeat(n)) > GCS_MAX_KEY_BYTES;
assert!(!at(1023));
assert!(!at(1024));
assert!(at(1025));
assert!(!(decoded_key_len(&"%20".repeat(341)) > GCS_MAX_KEY_BYTES));
}
}