use super::conn::Conn;
use super::ssrf::Scheme;
use std::collections::{HashMap, VecDeque};
use std::sync::{LazyLock, Mutex};
use std::time::{Duration, Instant};
const MAX_PER_HOST: usize = 8;
const MAX_GLOBAL: usize = 256;
const IDLE_TIMEOUT: Duration = Duration::from_secs(30);
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct PoolKey {
pub(crate) scheme: Scheme,
pub(crate) host: String,
pub(crate) port: u16,
}
struct PooledConn {
stream: Conn,
inserted: Instant,
}
struct Pool {
idle: HashMap<PoolKey, VecDeque<PooledConn>>,
total: usize,
}
impl Pool {
fn new() -> Self {
Self {
idle: HashMap::new(),
total: 0,
}
}
}
static POOL: LazyLock<Mutex<Pool>> = LazyLock::new(|| Mutex::new(Pool::new()));
pub(crate) fn checkout(key: &PoolKey) -> Option<Conn> {
let mut pool = POOL.lock().unwrap();
let now = Instant::now();
loop {
let conn = pool.idle.get_mut(key).and_then(|q| q.pop_front())?;
pool.total = pool.total.saturating_sub(1);
if now.duration_since(conn.inserted) > IDLE_TIMEOUT {
continue;
}
if is_alive(conn.stream.raw_fd()) {
return Some(conn.stream);
}
}
}
pub(crate) fn release(key: PoolKey, stream: Conn, keep_alive: bool) {
if !keep_alive {
return;
}
let mut pool = POOL.lock().unwrap();
if pool.total >= MAX_GLOBAL {
return;
}
let entries = pool.idle.entry(key).or_default();
if entries.len() >= MAX_PER_HOST {
return;
}
entries.push_front(PooledConn {
stream,
inserted: Instant::now(),
});
pool.total += 1;
}
fn is_alive(fd: std::os::fd::RawFd) -> bool {
let mut pfd = libc::pollfd {
fd,
events: libc::POLLIN,
revents: 0,
};
let rc = unsafe { libc::poll(&mut pfd, 1, 0) };
if rc < 0 {
return false;
}
if rc == 0 {
return true;
}
if pfd.revents & (libc::POLLHUP | libc::POLLERR | libc::POLLNVAL) != 0 {
return false;
}
false
}
#[cfg(test)]
pub(crate) fn clear_for_test() {
let mut pool = POOL.lock().unwrap();
pool.idle.clear();
pool.total = 0;
}