aria_engine/
auth.rs

1//! Instance-level Engine auth (in-memory; does not write config.yml).
2
3use std::sync::Mutex;
4use std::time::Duration;
5use thiserror::Error;
6
7pub const INTL_CLOUD: &str = "https://gateway.ariacompute.com";
8pub const INTL_SITE: &str = "https://ariacompute.com";
9pub const INTL_UPGRADE: &str = "https://github.com/ariacompute";
10pub const CN_CLOUD: &str = "https://gateway.ariacompute.cn";
11pub const CN_SITE: &str = "https://ariacompute.cn";
12pub const CN_UPGRADE: &str = "https://gitee.com/ariacompute";
13
14#[derive(Debug, Error, Clone, PartialEq, Eq)]
15pub enum AuthError {
16    #[error("invalid hybrid_mode: {0}")]
17    InvalidHybridMode(String),
18    #[error("invalid hybrid_execution: {0}")]
19    InvalidHybridExecution(String),
20    #[error("invalid compute: {0}")]
21    InvalidCompute(String),
22    #[error("hybrid_semantic_timeout_ms / cache_size must be positive integers")]
23    InvalidTimeoutOrCache,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct AuthConfig {
28    pub cloud_api_key: String,
29    pub cloud_url: String,
30    pub site_url: String,
31    pub upgrade_url: String,
32    pub hybrid_mode: String,
33    pub hybrid_execution: String,
34    pub hybrid_semantic: bool,
35    pub hybrid_semantic_timeout_ms: i32,
36    pub hybrid_semantic_cache_size: i32,
37    pub compute: String,
38    pub hf_token: String,
39    pub modelscope_api_token: String,
40}
41
42impl Default for AuthConfig {
43    fn default() -> Self {
44        Self {
45            cloud_api_key: String::new(),
46            cloud_url: String::new(),
47            site_url: String::new(),
48            upgrade_url: String::new(),
49            hybrid_mode: "balance".into(),
50            hybrid_execution: "hybrid".into(),
51            hybrid_semantic: true,
52            hybrid_semantic_timeout_ms: 800,
53            hybrid_semantic_cache_size: 512,
54            compute: "auto".into(),
55            hf_token: String::new(),
56            modelscope_api_token: String::new(),
57        }
58    }
59}
60
61/// Partial merge. `None` fields are omitted.
62#[derive(Debug, Clone, Default)]
63pub struct AuthUpdates {
64    pub cloud_api_key: Option<String>,
65    pub cloud_url: Option<String>,
66    pub site_url: Option<String>,
67    pub upgrade_url: Option<String>,
68    pub hybrid_mode: Option<String>,
69    pub hybrid_execution: Option<String>,
70    pub hybrid_semantic: Option<bool>,
71    pub hybrid_semantic_timeout_ms: Option<i32>,
72    pub hybrid_semantic_cache_size: Option<i32>,
73    pub compute: Option<String>,
74    pub hf_token: Option<String>,
75    pub modelscope_api_token: Option<String>,
76}
77
78fn gateway_region(url: &str) -> Option<&'static str> {
79    let lower = url.to_ascii_lowercase();
80    if lower.contains("ariacompute.cn") || lower.contains("gitee.com/ariacompute") {
81        Some("cn")
82    } else if lower.contains("ariacompute.com") || lower.contains("github.com/ariacompute") {
83        Some("intl")
84    } else {
85        None
86    }
87}
88
89fn pair_urls(region: &str) -> (&'static str, &'static str, &'static str) {
90    if region == "cn" {
91        (CN_CLOUD, CN_SITE, CN_UPGRADE)
92    } else {
93        (INTL_CLOUD, INTL_SITE, INTL_UPGRADE)
94    }
95}
96
97/// Fill missing cloud/site/upgrade URLs from a provided TLD.
98pub fn fill_auth_urls(mut cfg: AuthConfig) -> AuthConfig {
99    let region = gateway_region(&cfg.site_url)
100        .or_else(|| gateway_region(&cfg.cloud_url))
101        .or_else(|| gateway_region(&cfg.upgrade_url));
102    let Some(region) = region else {
103        return cfg;
104    };
105    let (cloud, site, upgrade) = pair_urls(region);
106    if cfg.cloud_url.is_empty() {
107        cfg.cloud_url = cloud.into();
108    }
109    if cfg.site_url.is_empty() {
110        cfg.site_url = site.into();
111    }
112    if cfg.upgrade_url.is_empty() {
113        cfg.upgrade_url = upgrade.into();
114    }
115    cfg
116}
117
118fn locale_prefers_cn() -> bool {
119    let lang = format!(
120        "{}{}",
121        std::env::var("LANG").unwrap_or_default(),
122        std::env::var("LC_ALL").unwrap_or_default()
123    )
124    .to_ascii_lowercase();
125    lang.contains("zh") || lang.contains(".cn") || lang.starts_with("cn")
126}
127
128type ProbeFn = fn(&str, &str) -> bool;
129
130fn default_probe_dashboard(site_url: &str, api_key: &str) -> bool {
131    let url = format!(
132        "{}/api/dashboard/models",
133        site_url.trim_end_matches('/')
134    );
135    let resp = ureq::get(&url)
136        .set("User-Agent", "aria-engine-sdk/0.1.0")
137        .set("Authorization", &format!("Bearer {api_key}"))
138        .timeout(Duration::from_secs(10))
139        .call();
140    match resp {
141        Ok(r) => (200..300).contains(&r.status()),
142        Err(ureq::Error::Status(code, _)) => (200..300).contains(&code),
143        Err(_) => false,
144    }
145}
146
147static PROBE_DASHBOARD: Mutex<ProbeFn> = Mutex::new(default_probe_dashboard);
148
149#[cfg(test)]
150pub fn set_probe_dashboard(f: ProbeFn) {
151    *PROBE_DASHBOARD.lock().unwrap() = f;
152}
153
154#[cfg(test)]
155pub fn reset_probe_dashboard() {
156    *PROBE_DASHBOARD.lock().unwrap() = default_probe_dashboard;
157}
158
159fn probe_dashboard(site_url: &str, api_key: &str) -> bool {
160    let f = *PROBE_DASHBOARD.lock().unwrap();
161    f(site_url, api_key)
162}
163
164/// Match CLI detect: probe Dashboard with the key, else locale fallback.
165pub fn detect_gateway_pair(api_key: &str) -> (&'static str, &'static str, &'static str) {
166    let key = api_key.trim();
167    let (first, second) = if locale_prefers_cn() {
168        ("cn", "intl")
169    } else {
170        ("intl", "cn")
171    };
172    for region in [first, second] {
173        let (cloud, site, upgrade) = pair_urls(region);
174        if !key.is_empty() && probe_dashboard(site, key) {
175            return (cloud, site, upgrade);
176        }
177    }
178    pair_urls(first)
179}
180
181/// Merge `updates` into `existing`. Validates; does not mutate `existing`.
182pub fn apply_auth(existing: &AuthConfig, updates: &AuthUpdates) -> Result<AuthConfig, AuthError> {
183    let mut out = existing.clone();
184    if let Some(v) = &updates.cloud_api_key {
185        out.cloud_api_key = v.clone();
186    }
187    if let Some(v) = &updates.cloud_url {
188        out.cloud_url = v.clone();
189    }
190    if let Some(v) = &updates.site_url {
191        out.site_url = v.clone();
192    }
193    if let Some(v) = &updates.upgrade_url {
194        out.upgrade_url = v.clone();
195    }
196    if let Some(v) = &updates.hybrid_mode {
197        out.hybrid_mode = v.clone();
198    }
199    if let Some(v) = &updates.hybrid_execution {
200        out.hybrid_execution = v.clone();
201    }
202    if let Some(v) = updates.hybrid_semantic {
203        out.hybrid_semantic = v;
204    }
205    if let Some(v) = updates.hybrid_semantic_timeout_ms {
206        out.hybrid_semantic_timeout_ms = v;
207    }
208    if let Some(v) = updates.hybrid_semantic_cache_size {
209        out.hybrid_semantic_cache_size = v;
210    }
211    if let Some(v) = &updates.compute {
212        out.compute = v.clone();
213    }
214    if let Some(v) = &updates.hf_token {
215        out.hf_token = v.clone();
216    }
217    if let Some(v) = &updates.modelscope_api_token {
218        out.modelscope_api_token = v.clone();
219    }
220    match out.hybrid_mode.as_str() {
221        "cost" | "balance" | "intelligence" => {}
222        other => return Err(AuthError::InvalidHybridMode(other.into())),
223    }
224    match out.hybrid_execution.as_str() {
225        "hybrid" | "device" | "cloud" => {}
226        other => return Err(AuthError::InvalidHybridExecution(other.into())),
227    }
228    match out.compute.as_str() {
229        "auto" | "cpu" | "cuda" => {}
230        other => return Err(AuthError::InvalidCompute(other.into())),
231    }
232    if out.hybrid_semantic_timeout_ms <= 0 || out.hybrid_semantic_cache_size <= 0 {
233        return Err(AuthError::InvalidTimeoutOrCache);
234    }
235    out = fill_auth_urls(out);
236    if !out.cloud_api_key.is_empty()
237        && (out.cloud_url.is_empty() || out.site_url.is_empty() || out.upgrade_url.is_empty())
238    {
239        let (cloud, site, upgrade) = detect_gateway_pair(&out.cloud_api_key);
240        if out.cloud_url.is_empty() {
241            out.cloud_url = cloud.into();
242        }
243        if out.site_url.is_empty() {
244            out.site_url = site.into();
245        }
246        if out.upgrade_url.is_empty() {
247            out.upgrade_url = upgrade.into();
248        }
249    }
250    Ok(out)
251}
aria_engine/auth.rs

aria_engine/
auth.rs
