Skip to main content

orion_accessor/addr/access_ctrl/
unit.rs

1use crate::prelude::*;
2use crate::{
3    addr::{
4        GitRepository, HttpResource,
5        access_ctrl::{auth::AuthConfig, rule::Rule},
6        proxy::ProxyConfig,
7    },
8    opt::OptionFrom,
9    timeout::TimeoutConfig,
10};
11use derive_more::From;
12use getset::Getters;
13use log::info;
14use serde_derive::{Deserialize, Serialize};
15
16#[derive(Debug, Clone, Serialize, Deserialize, Getters, PartialEq)]
17#[getset(get = "pub")]
18pub struct Unit {
19    rules: Vec<Rule>,
20    #[serde(skip_serializing_if = "Option::is_none")]
21    auth: Option<AuthConfig>,
22    #[serde(skip_serializing_if = "Option::is_none")]
23    timeout: Option<TimeoutConfig>,
24    #[serde(skip_serializing_if = "Option::is_none")]
25    proxy: Option<ProxyConfig>,
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize, From)]
29pub enum RedirectResult {
30    Origin(String),
31    Direct(String, Option<AuthConfig>),
32}
33
34impl RedirectResult {
35    pub fn path(&self) -> &str {
36        match self {
37            RedirectResult::Origin(path) => path,
38            RedirectResult::Direct(path, _) => path,
39        }
40    }
41
42    pub fn is_proxy(&self) -> bool {
43        match self {
44            RedirectResult::Origin(_) => false,
45            RedirectResult::Direct(_, _) => true,
46        }
47    }
48}
49
50impl Unit {
51    pub fn new(rules: Vec<Rule>, auth: Option<AuthConfig>, proxy: Option<ProxyConfig>) -> Self {
52        Self {
53            rules,
54            auth,
55            timeout: Some(TimeoutConfig::http_simple()),
56            proxy,
57        }
58    }
59
60    pub fn add_rule(&mut self, rule: Rule) {
61        self.rules.push(rule);
62    }
63
64    pub fn set_auth(&mut self, auth: AuthConfig) {
65        self.auth = Some(auth);
66    }
67
68    pub fn set_timeout_config(&mut self, config: TimeoutConfig) {
69        self.timeout = Some(config);
70    }
71
72    pub fn timeout_config_mut(&mut self) -> &mut Option<TimeoutConfig> {
73        &mut self.timeout
74    }
75
76    pub fn redirect(&self, input: &str) -> RedirectResult {
77        for rule in &self.rules {
78            let result = rule.replace(input);
79            if let Some(result) = result {
80                return RedirectResult::Direct(result, self.auth.clone());
81            }
82        }
83        RedirectResult::Origin(input.to_string())
84    }
85
86    pub fn direct_http_addr(&self, input: &HttpResource) -> Option<HttpResource> {
87        for rule in &self.rules {
88            let result = rule.replace(input.url());
89            if let Some(result) = result {
90                let mut direct = input.clone();
91                direct.set_url(result);
92                if let Some(auth) = self.auth() {
93                    direct.set_username(auth.username().clone().to_opt());
94                    direct.set_password(auth.password().clone().to_opt());
95                }
96                return Some(direct);
97            }
98        }
99        None
100    }
101
102    pub fn direct_git_addr(&self, input: &GitRepository) -> Option<GitRepository> {
103        for rule in &self.rules {
104            let result = rule.replace(input.repo());
105            if let Some(result) = result {
106                info!(
107                    target: "git",
108                    "redirect to {result}, origin: {}",
109                    input.repo()
110                );
111                let mut direct = input.clone();
112                direct.set_repo(result);
113                if let Some(auth) = self.auth() {
114                    direct.set_username(auth.username().clone().to_opt());
115                    direct.set_token(auth.password().clone().to_opt());
116                }
117                return Some(direct);
118            }
119        }
120        None
121    }
122
123    pub fn make_example() -> Self {
124        Self {
125            rules: vec![Rule::new(
126                "https://github.com/example/*",
127                "https://mirror.example.com/*",
128            )],
129            auth: Some(AuthConfig::new(
130                "username".to_string(),
131                "password".to_string(),
132            )),
133            timeout: Some(TimeoutConfig::http_large_file()),
134            proxy: None,
135        }
136    }
137}
138
139impl EnvEvalable<Unit> for Unit {
140    fn env_eval(self, dict: &EnvDict) -> Unit {
141        Unit {
142            rules: self
143                .rules
144                .into_iter()
145                .map(|rule| rule.env_eval(dict))
146                .collect(),
147            auth: self.auth.map(|auth| auth.env_eval(dict)),
148            timeout: self.timeout,
149            proxy: self.proxy.map(|x| x.env_eval(dict)),
150        }
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use crate::addr::access_ctrl::auth::AuthConfig;
158    use crate::addr::access_ctrl::rule::Rule;
159    use crate::addr::proxy::ProxyConfig;
160
161    #[test]
162    fn test_unit_new() {
163        let rules = vec![Rule::new("https://example.com/*", "https://mirror.com/*")];
164        let auth = AuthConfig::new("user", "pass");
165        let proxy = ProxyConfig::new("http://proxy.example.com:8080");
166
167        let unit = Unit::new(rules.clone(), Some(auth.clone()), Some(proxy.clone()));
168
169        assert_eq!(unit.rules(), &rules);
170        assert_eq!(unit.auth(), &Some(auth));
171        assert!(unit.timeout().is_some());
172        assert_eq!(unit.proxy(), &Some(proxy));
173    }
174
175    #[test]
176    fn test_unit_add_rule() {
177        let mut unit = Unit::new(vec![], None, None);
178        let rule = Rule::new("https://example.com/*", "https://mirror.com/*");
179
180        unit.add_rule(rule.clone());
181
182        assert_eq!(unit.rules().len(), 1);
183        assert_eq!(unit.rules()[0], rule);
184    }
185
186    #[test]
187    fn test_unit_set_auth() {
188        let mut unit = Unit::new(vec![], None, None);
189        let auth = AuthConfig::new("user", "pass");
190
191        unit.set_auth(auth.clone());
192
193        assert_eq!(unit.auth(), &Some(auth));
194    }
195
196    #[test]
197    fn test_unit_set_timeout_config() {
198        let mut unit = Unit::new(vec![], None, None);
199        let timeout_config = TimeoutConfig::http_large_file();
200
201        unit.set_timeout_config(timeout_config.clone());
202
203        assert_eq!(unit.timeout(), &Some(timeout_config));
204    }
205
206    #[test]
207    fn test_unit_redirect_with_match() {
208        let rules = vec![Rule::new("https://example.com/*", "https://mirror.com/")];
209        let unit = Unit::new(rules, None, None);
210
211        let result = unit.redirect("https://example.com/file.txt");
212
213        match result {
214            RedirectResult::Direct(path, auth) => {
215                assert_eq!(path, "https://mirror.com/file.txt");
216                assert_eq!(auth, None);
217            }
218            RedirectResult::Origin(_) => panic!("Expected RedirectResult::Direct"),
219        }
220    }
221
222    #[test]
223    fn test_unit_redirect_without_match() {
224        let rules = vec![Rule::new("https://example.com/*", "https://mirror.com/")];
225        let unit = Unit::new(rules, None, None);
226
227        let result = unit.redirect("https://other.com/file.txt");
228
229        match result {
230            RedirectResult::Origin(path) => {
231                assert_eq!(path, "https://other.com/file.txt");
232            }
233            _ => panic!("Expected RedirectResult::Origin"),
234        }
235    }
236
237    #[test]
238    fn test_unit_redirect_with_auth() {
239        let rules = vec![Rule::new("https://example.com/*", "https://mirror.com/")];
240        let auth = AuthConfig::new("user", "pass");
241        let unit = Unit::new(rules, Some(auth.clone()), None);
242
243        let result = unit.redirect("https://example.com/file.txt");
244
245        match result {
246            RedirectResult::Direct(path, result_auth) => {
247                assert_eq!(path, "https://mirror.com/file.txt");
248                assert_eq!(result_auth, Some(auth));
249            }
250            RedirectResult::Origin(_) => panic!("Expected RedirectResult::Direct"),
251        }
252    }
253
254    #[test]
255    fn test_unit_direct_address_redirection() {
256        use crate::addr::{GitRepository, HttpResource};
257
258        // Test HTTP address redirection
259        let http_rules = vec![Rule::new(
260            "https://github.com/*",
261            "https://github-mirror.com/",
262        )];
263        let auth = AuthConfig::new("test-user", "test-pass");
264        let unit = Unit::new(http_rules, Some(auth.clone()), None);
265
266        let http_addr =
267            HttpResource::from("https://github.com/user/repo").with_credentials("user", "pass");
268        let redirected_http = unit.direct_http_addr(&http_addr);
269
270        assert!(redirected_http.is_some());
271        if let Some(redirected) = redirected_http {
272            assert_eq!(redirected.url(), "https://github-mirror.com/user/repo");
273            assert_eq!(*redirected.username(), Some("test-user".to_string()));
274            assert_eq!(*redirected.password(), Some("test-pass".to_string()));
275        }
276
277        // Test HTTP address without match (should return None)
278        let no_match_http = HttpResource::from("https://gitlab.com/user/repo");
279        let no_redirect_result = unit.direct_http_addr(&no_match_http);
280        assert!(no_redirect_result.is_none());
281
282        // Test Git repository redirection
283        let git_rules = vec![Rule::new("git@gitlab.com:*", "git@gitlab-mirror.com:")];
284        let git_auth = AuthConfig::new("git-user", "git-token");
285        let git_unit = Unit::new(git_rules, Some(git_auth.clone()), None);
286
287        let git_repo = GitRepository::from("git@gitlab.com:user/repo.git");
288        let redirected_git = git_unit.direct_git_addr(&git_repo);
289
290        assert!(redirected_git.is_some());
291        if let Some(redirected) = redirected_git {
292            assert_eq!(redirected.repo(), "git@gitlab-mirror.com:user/repo.git");
293            assert_eq!(*redirected.username(), Some("git-user".to_string()));
294            assert_eq!(*redirected.token(), Some("git-token".to_string()));
295        }
296
297        // Test Git repository without match (should return None)
298        let no_match_git = GitRepository::from("git@github.com:user/repo.git");
299        let no_git_redirect = git_unit.direct_git_addr(&no_match_git);
300        assert!(no_git_redirect.is_none());
301    }
302
303    #[test]
304    fn test_unit_edge_cases_and_environment_evaluation() {
305        // Test timeout_config_mut method
306        let mut unit = Unit::new(vec![], None, None);
307        let original_timeout = unit.timeout().clone();
308
309        let mut_timeout_config = unit.timeout_config_mut();
310        *mut_timeout_config = Some(TimeoutConfig::new());
311
312        // Since TimeoutConfig::new() creates default values, we'll test that
313        // the configuration can be modified rather than comparing specific values
314        assert!(unit.timeout().is_some());
315        assert!(original_timeout.is_some());
316
317        // Test RedirectResult methods
318        let origin_result = RedirectResult::Origin("https://example.com".to_string());
319        let direct_result = RedirectResult::Direct("https://mirror.com".to_string(), None);
320
321        assert!(!origin_result.is_proxy());
322        assert!(direct_result.is_proxy());
323        assert_eq!(origin_result.path(), "https://example.com");
324        assert_eq!(direct_result.path(), "https://mirror.com");
325
326        // Test make_example method
327        let example_unit = Unit::make_example();
328        assert_eq!(example_unit.rules().len(), 1);
329        assert!(example_unit.auth().is_some());
330        assert!(example_unit.timeout().is_some());
331        assert!(example_unit.proxy().is_none());
332
333        let example_rule = example_unit.rules()[0].clone();
334        assert_eq!(example_rule.pattern(), "https://github.com/example/*");
335        assert_eq!(example_rule.target(), "https://mirror.example.com/*");
336
337        // Test environment evaluation
338        let mut env_dict = std::collections::HashMap::new();
339        env_dict.insert("TEST_USER".to_string(), "env-user".to_string());
340        env_dict.insert("TEST_PASS".to_string(), "env-pass".to_string());
341        let env_dict: EnvDict = env_dict.into();
342
343        // Create a unit with environment variables in auth
344        let rules = vec![];
345        let auth = Some(AuthConfig::new(
346            "${TEST_USER}".to_string(),
347            "${TEST_PASS}".to_string(),
348        ));
349        let proxy = None;
350        let test_unit = Unit::new(rules, auth, proxy);
351
352        // Apply environment evaluation
353        let evaluated_unit = test_unit.env_eval(&env_dict);
354
355        // Note: The actual evaluation behavior depends on the AuthConfig implementation
356        // This test verifies the environment evaluation mechanism works
357        assert!(evaluated_unit.auth().is_some());
358
359        // Test multiple rules with priority (first match wins)
360        let mut multi_rule_unit = Unit::new(vec![], None, None);
361        multi_rule_unit.add_rule(Rule::new("https://test.com/*", "https://first-mirror.com/"));
362        multi_rule_unit.add_rule(Rule::new(
363            "https://test.com/api/*",
364            "https://second-mirror.com/",
365        ));
366
367        // First rule should match
368        let result1 = multi_rule_unit.redirect("https://test.com/file.txt");
369        assert_eq!(result1.path(), "https://first-mirror.com/file.txt");
370
371        // More specific rule should match when it comes first
372        let mut specific_rule_unit = Unit::new(vec![], None, None);
373        specific_rule_unit.add_rule(Rule::new(
374            "https://test.com/api/*",
375            "https://second-mirror.com/",
376        ));
377        specific_rule_unit.add_rule(Rule::new("https://test.com/*", "https://first-mirror.com/"));
378
379        let result2 = specific_rule_unit.redirect("https://test.com/api/endpoint");
380        assert_eq!(result2.path(), "https://second-mirror.com/endpoint");
381
382        // Test with empty rules
383        let empty_rule_unit = Unit::new(vec![], None, None);
384        let result3 = empty_rule_unit.redirect("https://any-url.com/file.txt");
385        match result3 {
386            RedirectResult::Origin(path) => {
387                assert_eq!(path, "https://any-url.com/file.txt");
388            }
389            _ => panic!("Expected RedirectResult::Origin for empty rules"),
390        }
391    }
392}