gitops_operator/registry/
registry.rs

1use anyhow::{Context, Result};
2use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
3use k8s_openapi::api::core::v1::Secret;
4use kube::{api::Api, Client as K8sClient};
5use reqwest::{
6    header::{HeaderMap, HeaderValue, ACCEPT, AUTHORIZATION, WWW_AUTHENTICATE},
7    Client,
8};
9use serde::Deserialize;
10use serde_json::Value;
11use tracing::{error, info};
12
13#[derive(Debug, Deserialize)]
14struct TokenResponse {
15    token: String,
16    access_token: Option<String>,
17}
18
19#[derive(Debug)]
20pub struct AuthChallenge {
21    pub realm: String,
22    pub service: String,
23    pub scope: String,
24}
25
26impl AuthChallenge {
27    pub fn from_header(header: &str) -> Option<Self> {
28        let mut realm = None;
29        let mut service = None;
30        let mut scope = None;
31
32        if let Some(bearer_str) = header.strip_prefix("Bearer ") {
33            for pair in bearer_str.split(',') {
34                let mut parts = pair.trim().splitn(2, '=');
35                if let (Some(key), Some(value)) = (parts.next(), parts.next()) {
36                    let value = value.trim_matches('"');
37                    match key {
38                        "realm" => realm = Some(value.to_string()),
39                        "service" => service = Some(value.to_string()),
40                        "scope" => scope = Some(value.to_string()),
41                        _ => {}
42                    }
43                }
44            }
45        }
46
47        match (realm, service, scope) {
48            (Some(realm), Some(service), Some(scope)) => Some(AuthChallenge {
49                realm,
50                service,
51                scope,
52            }),
53            _ => None,
54        }
55    }
56}
57
58#[derive(Debug)]
59pub struct RegistryChecker {
60    pub client: Client,
61    pub registry_url: String,
62    pub auth_token: Option<String>,
63    pub username: Option<String>,
64    pub password: Option<String>,
65}
66
67impl RegistryChecker {
68    pub async fn new(registry_url: String, auth_token: Option<String>) -> Result<Self> {
69        let mut headers = HeaderMap::new();
70        headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
71
72        info!("Creating HTTP client for registry checks");
73        let client = Client::builder()
74            .default_headers(headers)
75            .build()
76            .context("Failed to create HTTP client")?;
77
78        // If we have a Basic auth token, extract username and password
79        let (username, password) = if let Some(token) = &auth_token {
80            if let Some(credentials) = token.strip_prefix("Basic ") {
81                if let Ok(decoded) = BASE64.decode(credentials) {
82                    if let Ok(auth_str) = String::from_utf8(decoded) {
83                        let mut parts = auth_str.splitn(2, ':');
84                        (
85                            parts.next().map(String::from),
86                            parts.next().map(String::from),
87                        )
88                    } else {
89                        (None, None)
90                    }
91                } else {
92                    (None, None)
93                }
94            } else {
95                (None, None)
96            }
97        } else {
98            (None, None)
99        };
100
101        Ok(Self {
102            client,
103            registry_url,
104            auth_token,
105            username,
106            password,
107        })
108    }
109
110    pub async fn get_bearer_token(&self, challenge: &AuthChallenge) -> Result<String> {
111        let mut request = self
112            .client
113            .get(&challenge.realm)
114            .query(&[("service", &challenge.service), ("scope", &challenge.scope)]);
115
116        // Add basic auth if credentials are available
117        if let (Some(username), Some(password)) = (&self.username, &self.password) {
118            request = request.basic_auth(username, Some(password));
119        }
120
121        let response = request.send().await?;
122
123        if !response.status().is_success() {
124            error!("Failed to get bearer token: {}", response.status());
125            anyhow::bail!("Failed to get bearer token: {}", response.status());
126        }
127
128        let token_response: TokenResponse = response.json().await?;
129        Ok(token_response.access_token.unwrap_or(token_response.token))
130    }
131
132    #[tracing::instrument(name = "check_image", skip(self), fields())]
133    pub async fn check_image(&self, image: &str, tag: &str) -> Result<bool> {
134        let registry_url = match self.registry_url.as_str() {
135            url if url.ends_with("/v1/") => url.replace("/v1", "/v2"),
136            url if url.ends_with("/v2/") => url.to_string(),
137            url => format!("{}/v2", url.trim_end_matches('/')),
138        };
139
140        let url = format!("{}/{}/manifests/{}", registry_url, image, tag);
141        info!("Checking image: {}", url);
142
143        // First request - might result in 401 with auth challenge
144        let response = self
145            .client
146            .head(&url)
147            .header(
148                AUTHORIZATION,
149                self.auth_token.as_ref().unwrap_or(&String::new()),
150            )
151            .send()
152            .await?;
153
154        if response.status().as_u16() == 401 {
155            if let Some(auth_header) = response.headers().get(WWW_AUTHENTICATE) {
156                if let Some(challenge) =
157                    AuthChallenge::from_header(auth_header.to_str().unwrap_or_default())
158                {
159                    // Get bearer token and retry with it
160                    let token = self.get_bearer_token(&challenge).await?;
161                    let auth_value = format!("Bearer {}", token);
162
163                    let response = self
164                        .client
165                        .head(&url)
166                        .header(AUTHORIZATION, auth_value)
167                        .send()
168                        .await?;
169                    info!("registry checker status: {}", response.status());
170
171                    return Ok(response.status().is_success());
172                }
173            }
174        }
175
176        Ok(response.status().is_success())
177    }
178}
179
180#[tracing::instrument(name = "get_registry_auth_from_secret", skip(), fields())]
181pub async fn get_registry_auth_from_secret(
182    secret_name: &str,
183    namespace: &str,
184    registry_url: &str,
185) -> Result<String> {
186    let client = K8sClient::try_default().await?;
187    let secrets: Api<Secret> = Api::namespaced(client, namespace);
188
189    let secret = secrets
190        .get(secret_name)
191        .await
192        .context("Failed to get secret")?;
193
194    let data = secret
195        .data
196        .ok_or_else(|| anyhow::anyhow!("Secret data is empty"))?;
197
198    // Get the .dockerconfigjson data
199    let raw_data = data
200        .get(".dockerconfigjson")
201        .ok_or_else(|| anyhow::anyhow!(".dockerconfigjson not found in secret"))?;
202
203    let key_bytes = raw_data.0.clone();
204
205    let str_data = String::from_utf8(key_bytes).context("Failed to convert raw data to string")?;
206
207    // Parse the JSON
208    let config: Value = serde_json::from_str(&str_data)?;
209
210    // Extract auth token for the specified registry
211    let auth = config
212        .get("auths")
213        .and_then(|auths| auths.get(registry_url))
214        .and_then(|registry| registry.get("auth"))
215        .and_then(|auth| auth.as_str())
216        .ok_or_else(|| anyhow::anyhow!("Auth not found for registry {}", registry_url))?;
217
218    Ok(format!("Basic {}", auth))
219}