Implement GlobMap for cert resolution

proxy/src/certs.rs

@@ -2,7 +2,7 @@ use rustls::{
     server::{ClientHello, ResolvesServerCert},
     sign::CertifiedKey,
 };
-use std::{io, sync::Arc};
+use std::{collections::BTreeMap, io, ops::Bound, sync::Arc};
 
 /// App-level configuration structs for TLS certificates.
 pub mod config {
@@ -72,6 +72,17 @@ pub mod config {
         /// Proxy server's private key.
         pub private_key: TlsKey,
     }
+
+    impl TlsServer {
+        pub fn into_certified_key(
+            self,
+        ) -> Result<rustls::sign::CertifiedKey, rustls::sign::SignError> {
+            Ok(rustls::sign::CertifiedKey::new(
+                self.certificate.0,
+                rustls::sign::any_supported_type(&self.private_key.0)?,
+            ))
+        }
+    }
 }
 
 /// Parse TLS certificate chain from a byte buffer.
@@ -100,7 +111,7 @@ fn parse_key(buf: &mut impl io::BufRead) -> io::Result<rustls::PrivateKey> {
 }
 
 /// Extract domain names from a certificate: first CN, then SANs.
-/// Further reading: https://www.rfc-editor.org/rfc/rfc4985
+/// Further reading: <https://www.rfc-editor.org/rfc/rfc4985>.
 fn certificate_names(cert: &rustls::Certificate) -> anyhow::Result<Vec<String>> {
     use x509_parser::{extensions::GeneralName, x509::AttributeTypeAndValue};
 
@@ -133,18 +144,80 @@ fn certificate_names(cert: &rustls::Certificate) -> anyhow::Result<Vec<String>>
     Ok(names)
 }
 
+#[derive(Debug)]
+struct GlobMapBuilder<V> {
+    builder: globset::GlobSetBuilder,
+    values: BTreeMap<usize, V>,
+}
+
+impl<V> GlobMapBuilder<V> {
+    fn new() -> Self {
+        Self {
+            builder: globset::GlobSetBuilder::new(),
+            values: Default::default(),
+        }
+    }
+
+    fn add(&mut self, globs: impl IntoIterator<Item = globset::Glob>, value: V) -> &mut Self {
+        let mut cnt = 0;
+        for glob in globs {
+            self.builder.add(glob);
+            cnt += 1;
+        }
+
+        if cnt > 0 {
+            let offset = self.values.last_key_value().map(|(k, _)| *k).unwrap_or(0);
+            self.values.insert(offset + cnt, value);
+        }
+
+        self
+    }
+
+    fn build(self) -> Result<GlobMap<V>, globset::Error> {
+        Ok(GlobMap {
+            set: self.builder.build()?,
+            values: self.values,
+        })
+    }
+}
+
+/// Maps a set of matching globs to an arbitrary value.
+/// See the tests below in case this description doesn't help.
+#[derive(Debug)]
+struct GlobMap<V> {
+    /// An ordered set of all loaded globs.
+    set: globset::GlobSet,
+    /// Store single value per range of globs.
+    values: BTreeMap<usize, V>,
+}
+
+impl<V> GlobMap<V> {
+    fn query(&self, text: &str) -> Vec<&V> {
+        let indices = self.set.matches(text);
+        let mut res = Vec::with_capacity(indices.len());
+
+        for i in indices {
+            let mut range = self.values.range((Bound::Excluded(i), Bound::Unbounded));
+            let (_, value) = range.next().expect("invariant: entry must exist");
+            res.push(value);
+        }
+
+        res
+    }
+}
+
 struct CertResolverEntry {
     raw: Arc<rustls::sign::CertifiedKey>,
     names: Vec<String>,
 }
 
 pub struct CertResolver {
-    certs: Vec<CertResolverEntry>,
+    storage: GlobMap<CertResolverEntry>,
 }
 
 impl CertResolver {
     pub fn new(config: config::TlsServers) -> anyhow::Result<Self> {
-        let mut builder = globset::GlobSetBuilder::new();
+        let mut builder = GlobMapBuilder::new();
         for server in config.0 {
             let Some(cert) = server.certificate.0.first() else {
                 tracing::warn!("found empty certificate, skipping");
@@ -152,15 +225,81 @@ impl CertResolver {
             };
 
             let names = certificate_names(cert)?;
-            dbg!(names);
+            let globs = names
+                .iter()
+                .map(|s| globset::Glob::new(s))
+                .collect::<Result<Vec<_>, _>>()?;
+
+            let entry = CertResolverEntry {
+                raw: Arc::new(server.into_certified_key()?),
+                names,
+            };
+
+            builder.add(globs, entry);
         }
 
-        todo!()
+        Ok(Self {
+            storage: builder.build()?,
+        })
     }
 }
 
 impl ResolvesServerCert for CertResolver {
     fn resolve(&self, message: ClientHello) -> Option<Arc<CertifiedKey>> {
-        todo!("implement")
+        let name = message.server_name()?;
+        let certs = self.storage.query(name);
+        let first = certs.first()?;
+
+        // TODO: warn if there's more than one match!
+        Some(first.raw.clone())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use globset::Glob;
+
+    #[test]
+    fn check_glob_map_basic() -> anyhow::Result<()> {
+        let mut builder = GlobMapBuilder::new();
+        builder
+            .add([Glob::new("*.localhost")?], 0)
+            .add([Glob::new("bar.localhost")?], 1)
+            .add([Glob::new("*.foo.localhost")?], 2);
+
+        let map = builder.build()?;
+
+        assert!(map.query("random").is_empty());
+        assert!(map.query("localhost").is_empty());
+        assert_eq!(map.query("foo.localhost"), [&0]);
+        assert_eq!(map.query("bar.localhost"), [&0, &1]);
+        assert_eq!(map.query("project.foo.localhost"), [&0, &2]);
+
+        Ok(())
+    }
+
+    #[test]
+    fn check_glob_map() -> anyhow::Result<()> {
+        let mut builder = GlobMapBuilder::new();
+        builder
+            .add(
+                [
+                    Glob::new("*.neon.tech")?,
+                    Glob::new("*.neon.internal.tech")?,
+                ],
+                "neon",
+            )
+            .add([Glob::new("*.localhost")?], "mock");
+
+        let map = builder.build()?;
+
+        assert!(map.query("random").is_empty());
+        assert!(map.query("localhost").is_empty());
+        assert_eq!(map.query("ep-1.neon.tech"), [&"neon"]);
+        assert_eq!(map.query("ep-1.neon.internal.tech"), [&"neon"]);
+        assert_eq!(map.query("ep-1.foo.localhost"), [&"mock"]);
+
+        Ok(())
     }
 }