use crate::config::basic_auth_header;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::time::{Duration, Instant};
struct CachedToken {
token: String,
expires_at: Instant,
}
pub struct DockerAuth {
tokens: RwLock<HashMap<String, CachedToken>>,
client: reqwest::Client,
}
impl DockerAuth {
pub fn new(timeout: u64) -> Self {
Self {
tokens: RwLock::new(HashMap::new()),
client: reqwest::Client::builder()
.timeout(Duration::from_secs(timeout))
.build()
.unwrap_or_default(),
}
}
pub async fn get_token(
&self,
registry_url: &str,
name: &str,
www_authenticate: Option<&str>,
basic_auth: Option<&str>,
) -> Option<String> {
let cache_key = format!("{}:{}", registry_url, name);
{
let tokens = self.tokens.read();
if let Some(cached) = tokens.get(&cache_key) {
if cached.expires_at > Instant::now() {
return Some(cached.token.clone());
}
}
}
let www_auth = www_authenticate?;
let token = self.fetch_token(www_auth, name, basic_auth).await?;
{
let mut tokens = self.tokens.write();
tokens.insert(
cache_key,
CachedToken {
token: token.clone(),
expires_at: Instant::now() + Duration::from_secs(300),
},
);
}
Some(token)
}
async fn fetch_token(
&self,
www_authenticate: &str,
name: &str,
basic_auth: Option<&str>,
) -> Option<String> {
let params = parse_www_authenticate(www_authenticate)?;
let realm = params.get("realm")?;
let service = params.get("service").map(|s| s.as_str()).unwrap_or("");
let scope = format!("repository:{}:pull", name);
let url = format!("{}?service={}&scope={}", realm, service, scope);
tracing::debug!(url = %url, "Fetching auth token");
let mut request = self.client.get(&url);
if let Some(credentials) = basic_auth {
request = request.header("Authorization", basic_auth_header(credentials));
tracing::debug!("Using basic auth for token request");
}
let response = request.send().await.ok()?;
if !response.status().is_success() {
tracing::warn!(status = %response.status(), "Token request failed");
return None;
}
let json: serde_json::Value = response.json().await.ok()?;
json.get("token")
.or_else(|| json.get("access_token"))
.and_then(|v| v.as_str())
.map(String::from)
}
}
impl Default for DockerAuth {
fn default() -> Self {
Self::new(60)
}
}
fn parse_www_authenticate(header: &str) -> Option<HashMap<String, String>> {
let header = header
.strip_prefix("Bearer ")
.or_else(|| header.strip_prefix("bearer "))?;
let mut params = HashMap::new();
for part in header.split(',') {
let part = part.trim();
if let Some((key, value)) = part.split_once('=') {
let value = value.trim_matches('"');
params.insert(key.to_string(), value.to_string());
}
}
Some(params)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_www_authenticate() {
let header = r#"Bearer realm="https://auth.docker.io/token",service="registry.docker.io",scope="repository:library/alpine:pull""#;
let params = parse_www_authenticate(header).unwrap();
assert_eq!(
params.get("realm"),
Some(&"https://auth.docker.io/token".to_string())
);
assert_eq!(
params.get("service"),
Some(&"registry.docker.io".to_string())
);
}
#[test]
fn test_parse_www_authenticate_lowercase() {
let header = r#"bearer realm="https://ghcr.io/token",service="ghcr.io""#;
let params = parse_www_authenticate(header).unwrap();
assert_eq!(
params.get("realm"),
Some(&"https://ghcr.io/token".to_string())
);
}
}