Skip to main content

stygian_proxy/storage/
mod.rs

1//! Storage port and in-memory adapter for proxy records.
2
3use async_trait::async_trait;
4use uuid::Uuid;
5
6use crate::error::ProxyResult;
7use crate::types::{Proxy, ProxyRecord};
8
9/// Abstract storage interface for persisting and querying proxy records.
10///
11/// Implementors must be `Send + Sync + 'static` to support concurrent access
12/// across async tasks. The trait is object-safe via [`macro@async_trait`].
13///
14/// # Example
15/// ```rust,no_run
16/// use stygian_proxy::storage::ProxyStoragePort;
17/// use stygian_proxy::types::{Proxy, ProxyCapabilities, ProxyType};
18/// use uuid::Uuid;
19///
20/// async fn demo(store: &dyn ProxyStoragePort) {
21///     let proxy = Proxy {
22///         url: "http://proxy.example.com:8080".into(),
23///         proxy_type: ProxyType::Http,
24///         username: None,
25///         password: None,
26///         weight: 1,
27///         tags: vec![],
28///         capabilities: ProxyCapabilities::default(),
29///     };
30///     let record = store.add(proxy).await.unwrap();
31///     let _ = store.get(record.id).await.unwrap();
32/// }
33/// ```
34#[async_trait]
35pub trait ProxyStoragePort: Send + Sync + 'static {
36    /// Add a new proxy to the store and return its [`ProxyRecord`].
37    async fn add(&self, proxy: Proxy) -> ProxyResult<ProxyRecord>;
38
39    /// Remove a proxy by its UUID. Returns an error if the ID is not found.
40    async fn remove(&self, id: Uuid) -> ProxyResult<()>;
41
42    /// Return all stored proxy records.
43    async fn list(&self) -> ProxyResult<Vec<ProxyRecord>>;
44
45    /// Fetch a single proxy record by UUID.
46    async fn get(&self, id: Uuid) -> ProxyResult<ProxyRecord>;
47
48    /// Record the outcome of a request through a proxy.
49    ///
50    /// - `success`: whether the request succeeded.
51    /// - `latency_ms`: elapsed time in milliseconds.
52    async fn update_metrics(&self, id: Uuid, success: bool, latency_ms: u64) -> ProxyResult<()>;
53
54    /// Return all stored proxy records paired with their live metrics reference.
55    ///
56    /// Used by [`ProxyManager`](crate::manager::ProxyManager) when building
57    /// [`ProxyCandidate`](crate::strategy::ProxyCandidate) slices so that
58    /// latency-aware strategies (e.g. least-used) see up-to-date counters.
59    async fn list_with_metrics(&self) -> ProxyResult<Vec<(ProxyRecord, Arc<ProxyMetrics>)>>;
60}
61
62/// Convenience alias for a heap-allocated, type-erased [`ProxyStoragePort`].
63pub type BoxedProxyStorage = Box<dyn ProxyStoragePort>;
64
65// ─────────────────────────────────────────────────────────────────────────────
66// URL validation helper
67// ─────────────────────────────────────────────────────────────────────────────
68
69/// Validate a proxy URL: scheme must be recognised, host must be non-empty,
70/// and the explicit port (if present) must be in [1, 65535].
71fn 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    // Strip any path/query, then strip user:pass@ if present.
94    let authority = rest.split('/').next().unwrap_or("");
95    let host_and_port = authority.split('@').next_back().unwrap_or("");
96
97    // Split host from port, handling IPv6 brackets.
98    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
133// ─────────────────────────────────────────────────────────────────────────────
134// MemoryProxyStore
135// ─────────────────────────────────────────────────────────────────────────────
136
137use 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/// In-memory implementation of [`ProxyStoragePort`].
146///
147/// Uses a `tokio::sync::RwLock`-guarded `HashMap` for thread-safe access.
148/// Metrics are updated via atomic operations, so only a **read** lock is
149/// needed for [`update_metrics`](MemoryProxyStore::update_metrics) calls —
150/// write contention stays low even under heavy concurrent load.
151///
152/// # Example
153/// ```
154/// # tokio_test::block_on(async {
155/// use stygian_proxy::storage::{MemoryProxyStore, ProxyStoragePort};
156/// use stygian_proxy::types::{Proxy, ProxyCapabilities, ProxyType};
157///
158/// let store = MemoryProxyStore::default();
159/// let proxy = Proxy { url: "http://proxy.example.com:8080".into(), proxy_type: ProxyType::Http,
160///                     username: None, password: None, weight: 1, tags: vec![],
161///                     capabilities: ProxyCapabilities::default() };
162/// let record = store.add(proxy).await.unwrap();
163/// assert_eq!(store.list().await.unwrap().len(), 1);
164/// store.remove(record.id).await.unwrap();
165/// assert!(store.list().await.unwrap().is_empty());
166/// # })
167/// ```
168#[derive(Debug, Default, Clone)]
169pub struct MemoryProxyStore {
170    inner: Arc<RwLock<StoreMap>>,
171}
172
173impl MemoryProxyStore {
174    /// Build a store pre-populated with `proxies`, validating each URL.
175    ///
176    /// Returns an error on the first invalid URL encountered.
177    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        // Lock released before the atomic updates — no long critical section.
251        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// ─────────────────────────────────────────────────────────────────────────────
265// Tests
266// ─────────────────────────────────────────────────────────────────────────────
267
268#[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        // Verify totals are internally consistent.
354        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}