feat: v0.2.0-rc, optimize code

src/app.rs

@@ -19,13 +19,13 @@ use super::service::HostConfig;
 pub struct ProxyApp {
     tls: bool,
     lookup_dns: bool,
-    host_configs: Vec<HostConfig>,
+    host_config_map: HashMap<String, HostConfig>,
     dns_resolver: TokioAsyncResolver,
-    dns_resolver_cache_map: RwLock<HashMap<String, String>>,
+    dns_resolver_cache: RwLock<HashMap<String, String>>,
 }
 
 impl ProxyApp {
-    pub fn new(tls: bool, lookup_dns: bool, host_configs: Vec<HostConfig>) -> Self {
+    pub fn new(tls: bool, lookup_dns: bool, host_config_map: HashMap<String, HostConfig>) -> Self {
         let dns_resolver = TokioAsyncResolver::tokio(
             ResolverConfig::default(),
             ResolverOpts::default(),
@@ -33,16 +33,16 @@ impl ProxyApp {
         ProxyApp {
             tls,
             lookup_dns,
-            host_configs,
+            host_config_map,
             dns_resolver,
-            dns_resolver_cache_map: Default::default(),
+            dns_resolver_cache: Default::default(),
         }
     }
 
     // just only support IPv4
     async fn lookup_ipv4(&self, hostname: &str) -> Option<String> {
         {
-            if let Some(ipv4_address) = self.dns_resolver_cache_map.read().await.get(hostname) {
+            if let Some(ipv4_address) = self.dns_resolver_cache.read().await.get(hostname) {
                 log::info!("DNS cached {} --> {}", hostname, ipv4_address);
                 return Some(ipv4_address.to_string());
             }
@@ -55,7 +55,7 @@ impl ProxyApp {
                     if let Some(RData::A(a)) = record.data() {
                         let ipv4_address = a.0.to_string();
                         {
-                            self.dns_resolver_cache_map.write().await
+                            self.dns_resolver_cache.write().await
                                 .insert(hostname.to_string(), ipv4_address.clone());
                         }
                         log::info!("DNS found {} --> {}", hostname, ipv4_address);
@@ -90,51 +90,48 @@ impl ProxyHttp for ProxyApp {
     fn new_ctx(&self) {}
 
     async fn upstream_peer(&self, session: &mut Session, _ctx: &mut ()) -> Result<Box<HttpPeer>> {
-        let host_header = session
+        let host_header = match session.get_header(HeaderName::from_static("host")) {
-            .get_header(HeaderName::from_static("host"))
+            None => return Err(Error::new(ErrorType::HTTPStatus(400))),
-            .unwrap()
+            Some(host_header) => host_header,
-            .to_str()
+        };
-            .expect("get host from http header failed");
+        let host = match host_header.to_str() {
-        log::info!("host header: {host_header}");
+            Ok(host) => host,
+            Err(_) => return Err(Error::new(ErrorType::HTTPStatus(400))),
+        };
+        log::info!("Find host header: {}", host);
 
-        let hostname = if host_header.contains(':') {
+        let hostname = if host.contains(':') {
-            host_header.chars().take_while(|c| c != &':').collect()
+            host.chars().take_while(|c| c != &':').collect()
         } else {
-            host_header.to_string()
+            host.to_string()
         };
 
-        if hostname == "localhost" {
+        if hostname == "127.0.0.1" || hostname == "localhost" {
-            return Err(Error::new(ErrorType::CustomCode("bad host", 400)));
+            return Err(Error::new(ErrorType::HTTPStatus(404)));
         }
 
-        let host_config = self
+        if let Some(host_config) = self.host_config_map.get(&hostname) {
-            .host_configs
+            log::info!("Find peer: {} --> {}", hostname, host_config.proxy_addr);
-            .iter()
+            return Ok(Box::new(HttpPeer::new(
-            .find(|x| x.proxy_hostname == hostname);
-        if let Some(host_config) = host_config {
-            let peer = HttpPeer::new(
                 host_config.proxy_addr.as_str(),
                 host_config.proxy_tls,
-                host_config.proxy_hostname.clone(),
+                host_config.proxy_servername.clone(),
-            );
+            )));
-            log::info!("Find peer: {} --> {}", hostname, host_config.proxy_addr);
-            return Ok(Box::new(peer));
         }
 
         if self.lookup_dns {
             if let Some(address) = self.lookup_ipv4(&hostname).await {
                 let peer_addr = format!("{}:{}", address, if self.tls { 443 } else { 80 });
-                let peer = HttpPeer::new(
+                log::info!("DNS peer: {} --> {}", hostname, peer_addr);
+                return Ok(Box::new(HttpPeer::new(
                     &peer_addr,
                     self.tls,
                     hostname.to_string(),
-                );
+                )));
-                log::info!("Generate peer: {} --> {}", hostname, peer_addr);
-                return Ok(Box::new(peer));
             }
         }
 
-        panic!("Cannot find peer: {}", hostname);
+        Err(Error::new(ErrorType::CustomCode("bad host", 400)))
     }
 
     async fn request_filter(&self, session: &mut Session, _ctx: &mut Self::CTX) -> Result<bool>

src/cert.rs

@@ -25,11 +25,14 @@ pub fn load_certificate(cert_fn: &str, key_fn: &str) -> Result<(Certificate, Str
     Ok((cert, cert_pem))
 }
 
-pub fn issue_certificate(intermediate_certificate: &Certificate, domain: &str) -> Result<Cert, String> {
+pub fn issue_certificate(issuer_certificate: &Certificate, domain: &str) -> Result<Cert, String> {
     let cert = new_end_entity(domain)?;
     log::info!("New certificate for: {} -> {}", domain, hex::encode(cert.get_key_identifier()));
-    let cert_pem = cert.serialize_pem_with_signer(intermediate_certificate).map_err(|e| format!("Sign cert failed: {}", e))?;
+
+    let cert_pem = cert.serialize_pem_with_signer(issuer_certificate)
+        .map_err(|e| format!("Sign cert failed: {}", e))?;
     let key_pem = cert.serialize_private_key_pem();
+
     Ok(Cert {
         cert_pem,
         key_pem,
@@ -71,7 +74,9 @@ fn new_end_entity(domain: &str) -> Result<Certificate, String> {
 }
 
 fn validity_period() -> Result<(OffsetDateTime, OffsetDateTime), String> {
-    let start = OffsetDateTime::now_utc().checked_sub(Duration::hours(1)).expect("SHOULD NOT HAPPEN!");
+    let start = OffsetDateTime::now_utc().checked_sub(Duration::hours(1))
-    let end = OffsetDateTime::now_utc().checked_add(Duration::days(90)).expect("SHOULD NOT HAPPEN!");
+        .expect("SHOULD NOT HAPPEN!");
+    let end = OffsetDateTime::now_utc().checked_add(Duration::days(90))
+        .expect("SHOULD NOT HAPPEN!");
     Ok((start, end))
 }

src/main.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use log::LevelFilter;
 use pingora::{
     server::{configuration::Opt, Server},
@@ -32,6 +34,10 @@ pub fn main() {
         services.push(Box::new(prometheus_service_http));
     }
 
+    if services.is_empty() {
+        panic!("No services is configured!");
+    }
+
     log::info!("start listen...");
     my_server.add_services(services);
     my_server.run_forever();
@@ -47,8 +53,9 @@ fn build_services(server: &Server, proxy_config: &ProxyConfig) -> Vec<Box<dyn Se
     let mut services: Vec<Box<dyn Service>> = vec![];
     for group in &proxy_config.groups {
         let listen_address = format!("0.0.0.0:{}", group.port);
-        let host_configs = build_host_configs(group);
+        let host_config_map = build_host_config_map(group);
         let lookup_dns = group.lookup_dns.unwrap_or(false);
+        
         log::info!("Listen at: {}, tls: {}, lookup_dns: {}", listen_address, group.tls.is_some(), lookup_dns);
 
         match &group.tls {
@@ -57,7 +64,7 @@ fn build_services(server: &Server, proxy_config: &ProxyConfig) -> Vec<Box<dyn Se
                     &server.configuration,
                     &listen_address,
                     lookup_dns,
-                    host_configs,
+                    host_config_map,
                 )));
             }
             Some(proxy_tls) => {
@@ -66,7 +73,7 @@ fn build_services(server: &Server, proxy_config: &ProxyConfig) -> Vec<Box<dyn Se
                     &listen_address,
                     lookup_dns,
                     proxy_tls,
-                    host_configs,
+                    host_config_map,
                 )));
             }
         }
@@ -74,16 +81,17 @@ fn build_services(server: &Server, proxy_config: &ProxyConfig) -> Vec<Box<dyn Se
     services
 }
 
-fn build_host_configs(group: &ProxyGroup) -> Vec<HostConfig> {
+fn build_host_config_map(group: &ProxyGroup) -> HashMap<String, HostConfig> {
-    let mut host_configs = vec![];
+    let mut host_config_map = HashMap::new();
     if let Some(proxy_map) = &group.proxy_map {
         for (hostname, proxy_item) in proxy_map {
-            host_configs.push(HostConfig {
+            let host_config = HostConfig {
                 proxy_addr: proxy_item.address.clone(),
                 proxy_tls: proxy_item.tls.unwrap_or(false),
-                proxy_hostname: proxy_item.sni.clone().unwrap_or_else(|| hostname.clone()),
+                proxy_servername: proxy_item.sni.clone().unwrap_or_else(|| hostname.clone()),
-            });
+            };
+            host_config_map.insert(hostname.to_string(), host_config);
         }
     }
-    host_configs
+    host_config_map
 }

src/service.rs

@@ -17,30 +17,30 @@ use crate::cert::Cert;
 use crate::config::ProxyTls;
 
 struct Callback {
-    intermediate_certificate: Certificate,
+    issuer_cert: Certificate,
-    intermediate_certificate_pem: String,
+    issuer_cert_pem: String,
-    certificate_cache_map: RwLock<HashMap<String, Cert>>,
+    issued_cert_cache: RwLock<HashMap<String, Cert>>,
 }
 
 impl Callback {
     fn new(proxy_tls: &ProxyTls) -> Result<Self, String> {
         let (cert, cert_pem) = cert::load_certificate(&proxy_tls.issuer_cert, &proxy_tls.issuer_key)?;
         Ok(Self {
-            intermediate_certificate: cert,
+            issuer_cert: cert,
-            intermediate_certificate_pem: cert_pem,
+            issuer_cert_pem: cert_pem,
-            certificate_cache_map: Default::default(),
+            issued_cert_cache: Default::default(),
         })
     }
 
     async fn issue_certificate(&self, hostname: &str) -> Result<Cert, String> {
         {
-            if let Some(cert) = self.certificate_cache_map.read().await.get(hostname) {
+            if let Some(cert) = self.issued_cert_cache.read().await.get(hostname) {
                 return Ok(cert.clone());
             }
         }
-        let cert = cert::issue_certificate(&self.intermediate_certificate, hostname)?;
+        let cert = cert::issue_certificate(&self.issuer_cert, hostname)?;
         {
-            self.certificate_cache_map.write().await.insert(hostname.to_string(), cert.clone());
+            self.issued_cert_cache.write().await.insert(hostname.to_string(), cert.clone());
         }
         Ok(cert)
     }
@@ -49,7 +49,8 @@ impl Callback {
 #[async_trait]
 impl TlsAccept for Callback {
     async fn certificate_callback(&self, ssl: &mut SslRef) -> () {
-        let sni_provided = ssl.servername(NameType::HOST_NAME).expect("get sni failed").to_string();
+        let sni_provided = ssl.servername(NameType::HOST_NAME)
+            .unwrap_or("127.0.0.1").to_string();
         log::info!("SNI provided: {}", sni_provided);
 
         let cert = self.issue_certificate(&sni_provided).await
@@ -57,15 +58,15 @@ impl TlsAccept for Callback {
 
         let x509_cert = X509::from_pem(cert.cert_pem.as_bytes())
             .unwrap_or_else(|e| panic!("parse cert: {} failed: {}", cert.cert_pem, e));
-        let x509_intermediate_cert = X509::from_pem(self.intermediate_certificate_pem.as_bytes())
+        let x509_intermediate_cert = X509::from_pem(self.issuer_cert_pem.as_bytes())
-            .unwrap_or_else(|e| panic!("parse intermediate cert: {} failed: {}", self.intermediate_certificate_pem, e));
+            .unwrap_or_else(|e| panic!("parse issuer cert: {} failed: {}", self.issuer_cert_pem, e));
         let private_key = PKey::private_key_from_pem(cert.key_pem.as_bytes())
             .unwrap_or_else(|e| panic!("parse key: {} failed: {}", cert.key_pem, e));
 
         ext::ssl_use_certificate(ssl, &x509_cert)
             .unwrap_or_else(|e| panic!("apply certificate for: {} failed: {}", sni_provided, e));
         ext::ssl_add_chain_cert(ssl, &x509_intermediate_cert)
-            .unwrap_or_else(|e| panic!("apply intermediate certificate for: {} failed: {}", sni_provided, e));
+            .unwrap_or_else(|e| panic!("apply issuer certificate for: {} failed: {}", sni_provided, e));
         ext::ssl_use_private_key(ssl, &private_key)
             .unwrap_or_else(|e| panic!("apply key for: {} failed: {}", sni_provided, e));
     }
@@ -75,16 +76,16 @@ impl TlsAccept for Callback {
 pub struct HostConfig {
     pub proxy_addr: String,
     pub proxy_tls: bool,
-    pub proxy_hostname: String,
+    pub proxy_servername: String,
 }
 
 pub fn proxy_service_tcp(
     server_conf: &Arc<ServerConf>,
     listen_addr: &str,
     lookup_dns: bool,
-    host_configs: Vec<HostConfig>,
+    host_config_map: HashMap<String, HostConfig>,
 ) -> impl pingora::services::Service {
-    let proxy_app = ProxyApp::new(false, lookup_dns, host_configs);
+    let proxy_app = ProxyApp::new(false, lookup_dns, host_config_map);
     let mut service = http_proxy_service(server_conf, proxy_app);
 
     service.add_tcp(listen_addr);
@@ -97,13 +98,17 @@ pub fn proxy_service_tls(
     listen_addr: &str,
     lookup_dns: bool,
     proxy_tls: &ProxyTls,
-    host_configs: Vec<HostConfig>,
+    host_config_map: HashMap<String, HostConfig>,
 ) -> impl pingora::services::Service {
-    let proxy_app = ProxyApp::new(true, lookup_dns, host_configs);
+    let proxy_app = ProxyApp::new(true, lookup_dns, host_config_map);
     let mut service = http_proxy_service(server_conf, proxy_app);
 
-    let cb = Box::new(Callback::new(proxy_tls).unwrap());
+    let cb = Box::new(Callback::new(proxy_tls).unwrap_or_else(|e| {
-    let tls_settings = TlsSettings::with_callbacks(cb).unwrap();
+        panic!("Init SSL callback failed: {}", e);
+    }));
+    let tls_settings = TlsSettings::with_callbacks(cb).unwrap_or_else(|e| {
+        panic!("Init SSL settings failed: {}", e);
+    });
     service.add_tls_with_settings(listen_addr, None, tls_settings);
 
     service