stygian_proxy/storage/
mod.rs1use async_trait::async_trait;
4use uuid::Uuid;
5
6use crate::error::ProxyResult;
7use crate::types::{Proxy, ProxyRecord};
8
9#[async_trait]
35pub trait ProxyStoragePort: Send + Sync + 'static {
36 async fn add(&self, proxy: Proxy) -> ProxyResult<ProxyRecord>;
38
39 async fn remove(&self, id: Uuid) -> ProxyResult<()>;
41
42 async fn list(&self) -> ProxyResult<Vec<ProxyRecord>>;
44
45 async fn get(&self, id: Uuid) -> ProxyResult<ProxyRecord>;
47
48 async fn update_metrics(&self, id: Uuid, success: bool, latency_ms: u64) -> ProxyResult<()>;
53
54 async fn list_with_metrics(&self) -> ProxyResult<Vec<(ProxyRecord, Arc<ProxyMetrics>)>>;
60}
61
62pub type BoxedProxyStorage = Box<dyn ProxyStoragePort>;
64
65fn validate_proxy_url(url: &str) -> ProxyResult<()> {
72 use crate::error::ProxyError;
73
74 let (scheme, rest) = url
75 .split_once("://")
76 .ok_or_else(|| ProxyError::InvalidProxyUrl {
77 url: url.to_owned(),
78 reason: "missing scheme separator '://'".into(),
79 })?;
80
81 match scheme {
82 "http" | "https" => {}
83 #[cfg(feature = "socks")]
84 "socks4" | "socks5" => {}
85 other => {
86 return Err(ProxyError::InvalidProxyUrl {
87 url: url.to_owned(),
88 reason: format!("unsupported scheme '{other}'"),
89 });
90 }
91 }
92
93 let authority = rest.split('/').next().unwrap_or("");
95 let host_and_port = authority.split('@').next_back().unwrap_or("");
96
97 let (host, port_str) = if host_and_port.starts_with('[') {
99 let close = host_and_port.find(']').unwrap_or(host_and_port.len());
100 let after = &host_and_port[close + 1..];
101 let port = after.strip_prefix(':').unwrap_or("");
102 (&host_and_port[..=close], port)
103 } else {
104 match host_and_port.rsplit_once(':') {
105 Some((h, p)) => (h, p),
106 None => (host_and_port, ""),
107 }
108 };
109
110 if host.is_empty() || host == "[]" {
111 return Err(ProxyError::InvalidProxyUrl {
112 url: url.to_owned(),
113 reason: "empty host".into(),
114 });
115 }
116
117 if !port_str.is_empty() {
118 let port: u32 = port_str.parse().map_err(|_| ProxyError::InvalidProxyUrl {
119 url: url.to_owned(),
120 reason: format!("non-numeric port '{port_str}'"),
121 })?;
122 if port == 0 || port > 65535 {
123 return Err(ProxyError::InvalidProxyUrl {
124 url: url.to_owned(),
125 reason: format!("port {port} is out of range [1, 65535]"),
126 });
127 }
128 }
129
130 Ok(())
131}
132
133use std::collections::HashMap;
138use tokio::sync::RwLock;
139
140use crate::types::ProxyMetrics;
141use std::sync::Arc;
142
143type StoreMap = HashMap<Uuid, (ProxyRecord, Arc<ProxyMetrics>)>;
144
145#[derive(Debug, Default, Clone)]
169pub struct MemoryProxyStore {
170 inner: Arc<RwLock<StoreMap>>,
171}
172
173impl MemoryProxyStore {
174 pub async fn with_proxies(proxies: Vec<Proxy>) -> ProxyResult<Self> {
178 let store = Self::default();
179 for proxy in proxies {
180 store.add(proxy).await?;
181 }
182 Ok(store)
183 }
184}
185
186#[async_trait]
187impl ProxyStoragePort for MemoryProxyStore {
188 async fn add(&self, proxy: Proxy) -> ProxyResult<ProxyRecord> {
189 validate_proxy_url(&proxy.url)?;
190 let record = ProxyRecord::new(proxy);
191 let metrics = Arc::new(ProxyMetrics::default());
192 self.inner
193 .write()
194 .await
195 .insert(record.id, (record.clone(), metrics));
196 Ok(record)
197 }
198
199 async fn remove(&self, id: Uuid) -> ProxyResult<()> {
200 self.inner
201 .write()
202 .await
203 .remove(&id)
204 .map(|_| ())
205 .ok_or_else(|| crate::error::ProxyError::StorageError(format!("proxy {id} not found")))
206 }
207
208 async fn list(&self) -> ProxyResult<Vec<ProxyRecord>> {
209 Ok(self
210 .inner
211 .read()
212 .await
213 .values()
214 .map(|(r, _)| r.clone())
215 .collect())
216 }
217
218 async fn get(&self, id: Uuid) -> ProxyResult<ProxyRecord> {
219 self.inner
220 .read()
221 .await
222 .get(&id)
223 .map(|(r, _)| r.clone())
224 .ok_or_else(|| crate::error::ProxyError::StorageError(format!("proxy {id} not found")))
225 }
226
227 async fn list_with_metrics(&self) -> ProxyResult<Vec<(ProxyRecord, Arc<ProxyMetrics>)>> {
228 Ok(self
229 .inner
230 .read()
231 .await
232 .values()
233 .map(|(r, m)| (r.clone(), Arc::clone(m)))
234 .collect())
235 }
236
237 async fn update_metrics(&self, id: Uuid, success: bool, latency_ms: u64) -> ProxyResult<()> {
238 use std::sync::atomic::Ordering;
239
240 let metrics = self
241 .inner
242 .read()
243 .await
244 .get(&id)
245 .map(|(_, m)| Arc::clone(m))
246 .ok_or_else(|| {
247 crate::error::ProxyError::StorageError(format!("proxy {id} not found"))
248 })?;
249
250 metrics.requests_total.fetch_add(1, Ordering::Relaxed);
252 if success {
253 metrics.successes.fetch_add(1, Ordering::Relaxed);
254 } else {
255 metrics.failures.fetch_add(1, Ordering::Relaxed);
256 }
257 metrics
258 .total_latency_ms
259 .fetch_add(latency_ms, Ordering::Relaxed);
260 Ok(())
261 }
262}
263
264#[cfg(test)]
269mod tests {
270 use super::*;
271 use crate::types::ProxyType;
272 use std::sync::atomic::Ordering;
273
274 fn make_proxy(url: &str) -> Proxy {
275 Proxy {
276 url: url.into(),
277 proxy_type: ProxyType::Http,
278 username: None,
279 password: None,
280 weight: 1,
281 tags: vec![],
282 capabilities: crate::types::ProxyCapabilities::default(),
283 }
284 }
285
286 #[tokio::test]
287 async fn add_list_remove() -> crate::error::ProxyResult<()> {
288 let store = MemoryProxyStore::default();
289 let r1 = store.add(make_proxy("http://a.test:8080")).await?;
290 let r2 = store.add(make_proxy("http://b.test:8080")).await?;
291 let r3 = store.add(make_proxy("http://c.test:8080")).await?;
292 assert_eq!(store.list().await?.len(), 3);
293 store.remove(r2.id).await?;
294 let remaining = store.list().await?;
295 assert_eq!(remaining.len(), 2);
296 let ids: Vec<_> = remaining.iter().map(|r| r.id).collect();
297 assert!(ids.contains(&r1.id));
298 assert!(ids.contains(&r3.id));
299 Ok(())
300 }
301
302 #[tokio::test]
303 async fn invalid_url_rejected() -> std::result::Result<(), Box<dyn std::error::Error>> {
304 let store = MemoryProxyStore::default();
305 let err = store
306 .add(make_proxy("not-a-url"))
307 .await
308 .err()
309 .ok_or_else(|| std::io::Error::other("invalid URL should be rejected"))?;
310 assert!(matches!(
311 err,
312 crate::error::ProxyError::InvalidProxyUrl { .. }
313 ));
314 Ok(())
315 }
316
317 #[tokio::test]
318 async fn invalid_url_empty_host() -> std::result::Result<(), Box<dyn std::error::Error>> {
319 let store = MemoryProxyStore::default();
320 let err = store
321 .add(make_proxy("http://:8080"))
322 .await
323 .err()
324 .ok_or_else(|| std::io::Error::other("empty host URL should be rejected"))?;
325 assert!(matches!(
326 err,
327 crate::error::ProxyError::InvalidProxyUrl { .. }
328 ));
329 Ok(())
330 }
331
332 #[tokio::test]
333 async fn concurrent_metrics_updates() -> std::result::Result<(), Box<dyn std::error::Error>> {
334 use tokio::task::JoinSet;
335
336 let store = Arc::new(MemoryProxyStore::default());
337 let record = store
338 .add(make_proxy("http://proxy.test:3128"))
339 .await
340 .map_err(|e| std::io::Error::other(format!("failed to add proxy: {e}")))?;
341 let id = record.id;
342
343 let mut tasks = JoinSet::new();
344 for i in 0u64..50 {
345 let s = Arc::clone(&store);
346 tasks.spawn(async move { s.update_metrics(id, i % 2 == 0, i * 10).await });
347 }
348 while let Some(res) = tasks.join_next().await {
349 let inner = res.map_err(|e| std::io::Error::other(format!("join failed: {e}")))?;
350 inner.map_err(|e| std::io::Error::other(format!("update_metrics failed: {e}")))?;
351 }
352
353 let guard = store.inner.read().await;
355 let metrics = guard
356 .get(&id)
357 .map(|(_, m)| Arc::clone(m))
358 .ok_or_else(|| std::io::Error::other("missing metrics for inserted proxy"))?;
359 drop(guard);
360
361 let total = metrics.requests_total.load(Ordering::Relaxed);
362 let successes = metrics.successes.load(Ordering::Relaxed);
363 let failures = metrics.failures.load(Ordering::Relaxed);
364 assert_eq!(total, 50);
365 assert_eq!(successes + failures, 50);
366 Ok(())
367 }
368}