use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use tokio_rustls::rustls::crypto::aws_lc_rs::sign::any_supported_type;
use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer};
use tokio_rustls::rustls::server::{ClientHello, ResolvesServerCert};
use tokio_rustls::rustls::sign::CertifiedKey;
use crate::error::{ProxyError, ProxyResult};
use crate::tls::cert_store::CertificateStore;
pub struct SniCertResolver {
keys: RwLock<HashMap<String, Arc<CertifiedKey>>>,
fallback: Option<Arc<CertifiedKey>>,
}
impl SniCertResolver {
pub fn new(fallback: Option<Arc<CertifiedKey>>) -> Self {
Self {
keys: RwLock::new(HashMap::new()),
fallback,
}
}
pub fn certified_key(cert_pem: &[u8], key_pem: &[u8]) -> ProxyResult<Arc<CertifiedKey>> {
crate::tls::server::install_crypto_provider();
let certs: Vec<CertificateDer<'static>> =
rustls_pemfile::certs(&mut std::io::BufReader::new(cert_pem))
.filter_map(|r| r.ok())
.collect();
if certs.is_empty() {
return Err(ProxyError::Tls("no certificates found in PEM".into()));
}
let key: PrivateKeyDer<'static> =
rustls_pemfile::private_key(&mut std::io::BufReader::new(key_pem))
.map_err(|e| ProxyError::Tls(format!("could not parse private key: {e}")))?
.ok_or_else(|| ProxyError::Tls("no private key found in PEM".into()))?;
let signing_key = any_supported_type(&key)
.map_err(|e| ProxyError::Tls(format!("unsupported private key: {e}")))?;
Ok(Arc::new(CertifiedKey::new(certs, signing_key)))
}
pub async fn refresh(&self, store: &CertificateStore) -> ProxyResult<usize> {
let certificates = store.load_all().await?;
let mut fresh = HashMap::with_capacity(certificates.len());
for cert in certificates {
match Self::certified_key(&cert.cert_pem, &cert.key_pem) {
Ok(key) => {
fresh.insert(normalise(&cert.domain), key);
}
Err(e) => tracing::warn!("Certificate for {} will not load: {}", cert.domain, e),
}
}
let count = fresh.len();
*self.keys.write().unwrap_or_else(|e| e.into_inner()) = fresh;
tracing::info!("Serving {} certificate(s) by SNI", count);
Ok(count)
}
fn lookup(&self, name: &str) -> Option<Arc<CertifiedKey>> {
let keys = self.keys.read().unwrap_or_else(|e| e.into_inner());
let name = normalise(name);
if let Some(key) = keys.get(&name) {
return Some(key.clone());
}
let (_, parent) = name.split_once('.')?;
if parent.contains('.') {
if let Some(key) = keys.get(&format!("*.{parent}")) {
return Some(key.clone());
}
}
None
}
}
fn normalise(name: &str) -> String {
name.trim_end_matches('.').to_ascii_lowercase()
}
impl ResolvesServerCert for SniCertResolver {
fn resolve(&self, client_hello: ClientHello<'_>) -> Option<Arc<CertifiedKey>> {
match client_hello.server_name() {
Some(name) => self.lookup(name).or_else(|| self.fallback.clone()),
None => self.fallback.clone(),
}
}
}
impl std::fmt::Debug for SniCertResolver {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let keys = self.keys.read().unwrap_or_else(|e| e.into_inner());
f.debug_struct("SniCertResolver")
.field("domains", &keys.keys().collect::<Vec<_>>())
.field("has_fallback", &self.fallback.is_some())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
const CERT: &[u8] = include_bytes!("testdata/valid.pem");
const KEY: &[u8] = include_bytes!("testdata/valid.key.pem");
fn resolver_with(domains: &[&str]) -> SniCertResolver {
let resolver = SniCertResolver::new(None);
let key = SniCertResolver::certified_key(CERT, KEY).unwrap();
let keys = domains
.iter()
.map(|d| (normalise(d), key.clone()))
.collect();
*resolver.keys.write().unwrap() = keys;
resolver
}
#[test]
fn an_exact_name_is_found() {
let resolver = resolver_with(&["a.example.com"]);
assert!(resolver.lookup("a.example.com").is_some());
assert!(resolver.lookup("b.example.com").is_none());
}
#[test]
fn sni_is_matched_case_insensitively() {
let resolver = resolver_with(&["a.example.com"]);
assert!(resolver.lookup("A.Example.COM").is_some());
}
#[test]
fn a_certificate_stored_with_uppercase_or_a_trailing_dot_is_served() {
let resolver = resolver_with(&["A.Example.COM", "b.example.com.", "*.Wild.Example.com"]);
assert!(resolver.lookup("a.example.com").is_some());
assert!(resolver.lookup("A.EXAMPLE.COM").is_some());
assert!(resolver.lookup("b.example.com").is_some());
assert!(resolver.lookup("b.example.com.").is_some());
assert!(resolver.lookup("x.wild.example.com").is_some());
}
#[test]
fn a_wildcard_covers_exactly_one_label() {
let resolver = resolver_with(&["*.example.com"]);
assert!(resolver.lookup("a.example.com").is_some());
assert!(resolver.lookup("a.b.example.com").is_none());
assert!(resolver.lookup("example.com").is_none());
}
#[test]
fn a_wildcard_does_not_swallow_a_public_suffix() {
let resolver = resolver_with(&["*.com"]);
assert!(resolver.lookup("example.com").is_none());
}
#[test]
fn the_fallback_answers_when_nothing_matches() {
let fallback = SniCertResolver::certified_key(CERT, KEY).unwrap();
let resolver = SniCertResolver::new(Some(fallback));
assert!(resolver.lookup("anything.example.com").is_none());
assert!(resolver.fallback.is_some());
}
#[test]
fn a_pem_that_will_not_parse_is_an_error_not_a_panic() {
assert!(SniCertResolver::certified_key(b"nonsense", b"nonsense").is_err());
}
}