Skip to main content

object_storage_proxy/parsers/
cos_map.rs

1use std::collections::HashMap;
2
3use pyo3::exceptions::PyKeyError;
4use pyo3::prelude::*;
5use pyo3::{PyResult, Python};
6
7use crate::credentials::models::BucketCredential;
8
9/// Represents a COS map item with its properties.
10///
11/// This struct is used to store the configuration for a COS bucket.
12/// Each bucket is identified by its name, and the properties include:
13/// - `host`: The host address of the COS service.
14/// - `port`: The port number for the COS service.
15/// - `region`: The region where the COS service is located.
16/// - `api_key`: The API key for accessing the COS service (optional).
17/// - `access_key`: The access key for the COS service (optional).
18/// - `secret_key`: The secret key for the COS service (optional).
19/// - `ttl`: The time-to-live for the COS bucket (optional).
20///
21#[pyclass]
22#[derive(Debug, Clone)]
23pub struct CosMapItem {
24    pub host: String,
25    pub port: u16,
26    pub region: Option<String>,
27    pub api_key: Option<String>,
28    pub access_key: Option<String>,
29    pub secret_key: Option<String>,
30    pub ttl: Option<u64>,
31    pub tls: Option<bool>,
32    pub addressing_style: Option<String>,
33}
34
35impl CosMapItem {
36    /// Returns `true` if all three HMAC fields are present.
37    pub fn has_hmac(&self) -> bool {
38        self.access_key.is_some() && self.secret_key.is_some()
39    }
40
41    /// Returns `true` if an API‑key is present.
42    pub fn has_api_key(&self) -> bool {
43        self.api_key.is_some()
44    }
45
46    /// Ensure that **some** credential (HMAC or API key) is populated.
47    ///
48    /// * If HMAC pair exists → OK
49    /// * Else if api_key exists → OK
50    /// * Else it calls the supplied async `fetcher(bucket)` which
51    ///   should return one of the accepted formats (see [`BucketCredential`]).
52    ///   The struct is then updated in‑place.
53    pub async fn ensure_credentials<F, Fut>(
54        &mut self,
55        bucket: &str,
56        fetcher: Option<F>,
57    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>
58    where
59        F: FnOnce(String) -> Fut + Send,
60        Fut: std::future::Future<Output = Result<String, Box<dyn std::error::Error + Send + Sync>>>
61            + Send,
62    {
63        if self.has_hmac() || self.has_api_key() {
64            return Ok(());
65        }
66
67        let Some(fetch) = fetcher else {
68            return Err("missing credentials and no fetcher provided".into());
69        };
70
71        let raw_creds = fetch(bucket.to_owned()).await?;
72        match BucketCredential::parse(&raw_creds) {
73            BucketCredential::Hmac {
74                access_key,
75                secret_key,
76            } => {
77                self.access_key = Some(access_key);
78                self.secret_key = Some(secret_key);
79            }
80            BucketCredential::ApiKey(k) => {
81                self.api_key = Some(k);
82            }
83        }
84        Ok(())
85    }
86}
87
88pub(crate) fn parse_cos_map(
89    py: Python,
90    cos_dict: &PyObject,
91) -> PyResult<HashMap<String, CosMapItem>> {
92    let raw_map: HashMap<String, HashMap<String, PyObject>> = cos_dict.extract(py)?;
93    let mut map = HashMap::new();
94
95    for (bucket, inner_map) in raw_map {
96        let host_obj = inner_map
97            .get("host")
98            .ok_or_else(|| PyKeyError::new_err("Missing 'host' in COS map entry"))?;
99        let host: String = host_obj.extract(py)?;
100
101        let port_obj = inner_map
102            .get("port")
103            .ok_or_else(|| PyKeyError::new_err("Missing 'port' in COS map entry"))?;
104        let port: u16 = port_obj.extract(py)?;
105
106        let region = inner_map.get("region").map(|v| v.extract(py)).transpose()?;
107
108        // Optional: api_key (allow 'api_key' or 'apikey')
109        let api_key =
110            if let Some(val) = inner_map.get("api_key").or_else(|| inner_map.get("apikey")) {
111                Some(val.extract(py)?)
112            } else {
113                None
114            };
115        let ttl = if let Some(val) = inner_map
116            .get("ttl")
117            .or_else(|| inner_map.get("time-to-live"))
118        {
119            Some(val.extract(py)?)
120        } else {
121            None
122        };
123
124        let tls = inner_map
125            .get("tls")
126            .or_else(|| inner_map.get("is_tls_enabled"))
127            .map(|v| v.extract(py))
128            .transpose()?;
129
130        let access_key = inner_map
131            .get("access_key")
132            .or_else(|| inner_map.get("accessKey"))
133            .map(|v| v.extract(py))
134            .transpose()?;
135
136        let secret_key = inner_map
137            .get("secret_key")
138            .or_else(|| inner_map.get("secretKey"))
139            .map(|v| v.extract(py))
140            .transpose()?;
141
142        let addressing_style = inner_map
143            .get("addressing_style")
144            .or_else(|| inner_map.get("addressingStyle"))
145            .map(|v| v.extract(py))
146            .transpose()?;
147
148        map.insert(
149            bucket.clone(),
150            CosMapItem {
151                host,
152                port,
153                region,
154                api_key,
155                access_key,
156                secret_key,
157                ttl,
158                tls,
159                addressing_style,
160            },
161        );
162    }
163
164    Ok(map)
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170
171    fn base_item() -> CosMapItem {
172        CosMapItem {
173            host: "s3.example.com".into(),
174            port: 443,
175            region: Some("us-east-1".into()),
176            api_key: None,
177            access_key: None,
178            secret_key: None,
179            ttl: None,
180            tls: Some(true),
181            addressing_style: None,
182        }
183    }
184
185    #[test]
186    fn has_hmac_requires_both_keys() {
187        let mut item = base_item();
188        assert!(!item.has_hmac());
189
190        item.access_key = Some("AK".into());
191        assert!(
192            !item.has_hmac(),
193            "only access_key should not satisfy has_hmac"
194        );
195
196        item.secret_key = Some("SK".into());
197        assert!(item.has_hmac());
198    }
199
200    #[test]
201    fn has_api_key_reflects_presence() {
202        let mut item = base_item();
203        assert!(!item.has_api_key());
204        item.api_key = Some("my-api-key".into());
205        assert!(item.has_api_key());
206    }
207
208    #[tokio::test]
209    async fn ensure_credentials_noop_when_hmac_present() {
210        let mut item = base_item();
211        item.access_key = Some("AK".into());
212        item.secret_key = Some("SK".into());
213
214        // fetcher should NOT be called
215        let called = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
216        let called2 = called.clone();
217        let fetcher = move |_bucket: String| {
218            let c = called2.clone();
219            async move {
220                c.store(true, std::sync::atomic::Ordering::SeqCst);
221                Ok::<String, Box<dyn std::error::Error + Send + Sync>>("unused".into())
222            }
223        };
224
225        item.ensure_credentials("my-bucket", Some(fetcher))
226            .await
227            .unwrap();
228        assert!(!called.load(std::sync::atomic::Ordering::SeqCst));
229    }
230
231    #[tokio::test]
232    async fn ensure_credentials_noop_when_api_key_present() {
233        let mut item = base_item();
234        item.api_key = Some("apikey123".into());
235
236        item.ensure_credentials(
237            "my-bucket",
238            None::<
239                fn(
240                    String,
241                )
242                    -> std::future::Ready<Result<String, Box<dyn std::error::Error + Send + Sync>>>,
243            >,
244        )
245        .await
246        .unwrap();
247    }
248
249    #[tokio::test]
250    async fn ensure_credentials_fetches_hmac_when_missing() {
251        let mut item = base_item();
252        let creds = r#"{"access_key": "FETCHED_AK", "secret_key": "FETCHED_SK"}"#;
253        let creds_str = creds.to_string();
254        let fetcher = move |_bucket: String| {
255            let c = creds_str.clone();
256            async move { Ok::<String, Box<dyn std::error::Error + Send + Sync>>(c) }
257        };
258
259        item.ensure_credentials("my-bucket", Some(fetcher))
260            .await
261            .unwrap();
262
263        assert_eq!(item.access_key.as_deref(), Some("FETCHED_AK"));
264        assert_eq!(item.secret_key.as_deref(), Some("FETCHED_SK"));
265    }
266
267    #[tokio::test]
268    async fn ensure_credentials_fetches_api_key_when_missing() {
269        let mut item = base_item();
270        let fetcher = |_bucket: String| async move {
271            Ok::<String, Box<dyn std::error::Error + Send + Sync>>("my-fetched-api-key".into())
272        };
273
274        item.ensure_credentials("my-bucket", Some(fetcher))
275            .await
276            .unwrap();
277
278        assert_eq!(item.api_key.as_deref(), Some("my-fetched-api-key"));
279    }
280
281    #[tokio::test]
282    async fn ensure_credentials_errors_without_fetcher() {
283        let mut item = base_item();
284        let result = item
285            .ensure_credentials(
286                "my-bucket",
287                None::<
288                    fn(
289                        String,
290                    ) -> std::future::Ready<
291                        Result<String, Box<dyn std::error::Error + Send + Sync>>,
292                    >,
293                >,
294            )
295            .await;
296        assert!(result.is_err());
297    }
298}