gway 0.1.1

use std::{collections::BTreeMap, fmt::Debug, sync::Arc};

use dashmap::DashMap;
use faststr::FastStr;
use parking_lot::RwLock;
use rustls::{crypto::ring::sign::any_supported_type, sign::CertifiedKey};
use sub_host::sub_host;
use tokio::time;

use crate::{
  cert::Cert,
  error::{Error, Result},
};

#[cfg(feature = "cert_dir")]
mod dir;

#[cfg(feature = "cert_dir")]
pub use dir::CertDir;

pub trait LoadCert: Send + Sync + 'static + Debug {
  fn load(&self, host: impl Into<FastStr>) -> Result<Option<Cert>>;
}

#[derive(Debug)]
pub struct CertLoader<L: LoadCert> {
  // 证书缓存
  pub host_cert: DashMap<FastStr, Arc<CertifiedKey>>,
  // 证书过期时间
  pub expire: RwLock<BTreeMap<i64, Vec<FastStr>>>,
  // 证书加载器
  pub loader: L,
}

impl<L: LoadCert> CertLoader<L> {
  pub fn new(loader: L) -> Arc<Self> {
    let s = Arc::new(Self {
      host_cert: DashMap::new(),
      expire: RwLock::new(BTreeMap::new()),
      loader,
    });

    let s2 = s.clone();
    tokio::spawn(async move {
      let mut interval = time::interval(time::Duration::from_secs(86400));
      loop {
        interval.tick().await;
        s2.rm_expired(2);
      }
    });

    s
  }

  pub fn rm_expired(&self, days: i64) {
    let mut to_rm_expire = Vec::new();
    for (expire, hosts) in self.expire.read().iter() {
      if *expire < chrono::Utc::now().timestamp() + days * 24 * 60 * 60 {
        for host in hosts {
          self.host_cert.remove(host);
        }
        to_rm_expire.push(*expire);
      } else {
        break;
      }
    }

    if !to_rm_expire.is_empty() {
      let mut expire_write = self.expire.write();
      for expire in to_rm_expire {
        expire_write.remove(&expire);
      }
    }
  }

  pub fn get(&self, host: impl Into<String>) -> Result<Arc<CertifiedKey>> {
    let host = host.into();
    let mut h = host.clone();
    let mut wait_check_sub = true;
    loop {
      if let Some(c) = self.host_cert.get::<FastStr>(&h.clone().into()) {
        return Ok(c.clone());
      }

      if let Some(cert) = self.loader.load(h.clone())? {
        let cert = Arc::new(cert);
        let key =
          any_supported_type(&cert.key).map_err(|e| Error::PrivateKeyUnsupported(e.to_string()))?;
        let certified_key = Arc::new(CertifiedKey::new(cert.fullchain.clone(), key));
        if let Some(domain) = cert.host_li.first() {
          self.host_cert.insert(domain.clone(), certified_key.clone());
          self
            .expire
            .write()
            .entry(cert.expire)
            .or_default()
            .push(domain.clone());
        }
        return Ok(certified_key);
      }

      if wait_check_sub {
        wait_check_sub = false;
      } else {
        break;
      }
      if let Some(sub) = sub_host(h) {
        h = sub;
      } else {
        break;
      }
    }

    Err(Error::CertNotFound(host))
  }
}

impl<L: LoadCert> rustls::server::ResolvesServerCert for CertLoader<L> {
  fn resolve(
    &self,
    client_hello: rustls::server::ClientHello,
  ) -> Option<Arc<rustls::sign::CertifiedKey>> {
    let host = FastStr::from(client_hello.server_name()?.to_string());
    self.get(host).ok()
  }
}