feat: config works

Cargo.toml

@@ -21,3 +21,7 @@ p384 = "0.13.0"
 p521 = "0.13.3"
 once_cell = "1.19.0"
 bytes = "1.6.0"
+serde = { version = "1.0.197", features = ["derive"] }
+serde_json = "1.0.115"
+tokio-dns-unofficial = "0.4.0"
+trust-dns-resolver = { version = "0.23.2", features = ["tokio"] }

proxy_config.json

@@ -0,0 +1,18 @@
+{
+  "groups": [
+    {
+      "port": 4430,
+      "lookup_dns": true,
+      "tls": {
+        "intermediate_cert": "__ignore_intermediate_cert.pem",
+        "intermediate_key": "__ignore_intermediate_pri_key.pem"
+      },
+      "proxy_map": {
+        "hatter.ink": {
+          "address": "101.132.122.240:443",
+          "tls": true
+        }
+      }
+    }
+  ]
+}

src/app.rs

@@ -6,16 +6,54 @@ use http::HeaderName;
 use pingora::{Error, ErrorType};
 use pingora::http::ResponseHeader;
 use pingora::prelude::{HttpPeer, ProxyHttp, Result, Session};
+use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
+use trust_dns_resolver::proto::rr::RData;
+use trust_dns_resolver::TokioAsyncResolver;
 
 use super::service::HostConfig;
 
 pub struct ProxyApp {
+    tls: bool,
+    lookup_dns: bool,
     host_configs: Vec<HostConfig>,
+    tokio_async_resolver: TokioAsyncResolver,
 }
 
 impl ProxyApp {
-    pub fn new(host_configs: Vec<HostConfig>) -> Self {
+    pub fn new(tls: bool, lookup_dns: bool, host_configs: Vec<HostConfig>) -> Self {
-        ProxyApp { host_configs }
+        let tokio_async_resolver = TokioAsyncResolver::tokio(
+            ResolverConfig::default(),
+            ResolverOpts::default(),
+        );
+        ProxyApp {
+            tls,
+            lookup_dns,
+            host_configs,
+            tokio_async_resolver,
+        }
+    }
+
+    async fn lookup_ipv4(&self, hostname: &str) -> Option<String> {
+        let ips = self.tokio_async_resolver.ipv4_lookup(hostname).await;
+        log::debug!("lookup {} --> {:#?}", hostname, ips);
+        match ips {
+            Ok(ips) => {
+                let records = ips.as_lookup().records();
+                if records.len() > 0 {
+                    let record = &records[0];
+                    if let Some(rdata) = record.data() {
+                        match rdata {
+                            RData::A(a) => {
+                                return Some(a.0.to_string());
+                            }
+                            _ => {}
+                        }
+                    }
+                }
+            }
+            Err(_) => {}
+        }
+        None
     }
 }
 
@@ -45,15 +83,30 @@ impl ProxyHttp for ProxyApp {
         let host_config = self
             .host_configs
             .iter()
-            .find(|x| x.proxy_hostname == hostname)
+            .find(|x| x.proxy_hostname == hostname);
-            .unwrap_or_else(|| panic!("find config for: {} failed", hostname));
+        if let Some(host_config) = host_config {
-        let peer = HttpPeer::new(
+            let peer = HttpPeer::new(
-            host_config.proxy_addr.as_str(),
+                host_config.proxy_addr.as_str(),
-            host_config.proxy_tls,
+                host_config.proxy_tls,
-            host_config.proxy_hostname.clone(),
+                host_config.proxy_hostname.clone(),
-        );
+            );
-        log::info!("Find peer: {:?}", peer._address);
+            log::info!("Find peer: {} --> {:?}", hostname, host_config.proxy_addr);
-        Ok(Box::new(peer))
+            return Ok(Box::new(peer));
+        }
+
+        if self.lookup_dns {
+            if let Some(address) = self.lookup_ipv4(&hostname).await {
+                let peer = HttpPeer::new(
+                    format!("{}:{}", address, if self.tls { 443 } else { 80 }),
+                    self.tls,
+                    hostname.to_string(),
+                );
+                log::info!("Find peer: {} --> {:?}", hostname, address);
+                return Ok(Box::new(peer));
+            }
+        }
+
+        panic!("Cannot find peer: {}", hostname);
     }
 
     async fn request_filter(&self, session: &mut Session, _ctx: &mut Self::CTX) -> Result<bool>

src/cert.rs

@@ -1,61 +1,35 @@
-use std::collections::HashMap;
 use std::fs;
 
-use once_cell::sync::Lazy;
 use rcgen::{Certificate, CertificateParams, DnType, ExtendedKeyUsagePurpose, IsCa, KeyPair, KeyUsagePurpose};
 use time::{Duration, OffsetDateTime};
-use tokio::sync::RwLock;
-
-const INTERMEDIATE_CERT_ENV_VAR: &str = "INTERMEDIATE_CERT";
-const INTERMEDIATE_KEY_ENV_VAR: &str = "INTERMEDIATE_KEY";
-
-static INTERMEDIATE_CA: Lazy<(Certificate, String)> = Lazy::new(|| {
-    let cert_fn = std::env::var(INTERMEDIATE_CERT_ENV_VAR)
-        .unwrap_or("__ignore_intermediate_cert.pem".to_string());
-    let key_fn = std::env::var(INTERMEDIATE_KEY_ENV_VAR)
-        .unwrap_or("__ignore_intermediate_pri_key.pem".to_string());
-    let cert_pem = fs::read_to_string(cert_fn).expect("Read cert file failed");
-    let key_pem = fs::read_to_string(key_fn).expect("Read key file failed");
-    let key_pem = parse_pkcs8(&key_pem);
-    let key_pair = KeyPair::from_pem(&key_pem).expect("Parse keypair failed");
-    // 底层逻辑限制，P256 与 SHA256 搭配，P384 与 SHA384 搭配
-    let certificate_params = CertificateParams::from_ca_cert_pem(&cert_pem, key_pair)
-        .expect("Cert and keypair mismatch");
-    let cert = Certificate::from_params(certificate_params)
-        .expect("Parse cert params failed");
-    (cert, cert_pem)
-});
 
 #[derive(Debug, Clone)]
 pub struct Cert {
     pub cert_pem: String,
-    pub intermediate_pem: String,
     pub key_pem: String,
 }
 
-static CERTIFICATE_CACHE_MAP: Lazy<RwLock<HashMap<String, Cert>>> = Lazy::new(|| {
+pub fn load_certificate(cert_fn: &str, key_fn: &str) -> Result<(Certificate, String), String> {
-    RwLock::new(HashMap::new())
+    let cert_pem = fs::read_to_string(cert_fn).map_err(|e| format!("Read file: {} failed: {}", cert_fn, e))?;
-});
+    let key_pem = fs::read_to_string(key_fn).map_err(|e| format!("Read file: {} failed: {}", key_fn, e))?;
+    let key_pem = parse_pkcs8(&key_pem);
+    let key_pair = KeyPair::from_pem(&key_pem).map_err(|e| format!("Parse key: {} failed: {}", key_fn, e))?;
+    // 底层逻辑限制，P256 与 SHA256 搭配，P384 与 SHA384 搭配
+    let certificate_params = CertificateParams::from_ca_cert_pem(&cert_pem, key_pair)
+        .map_err(|e| format!("Cert and keypair match failed: {}", e))?;
+    let cert = Certificate::from_params(certificate_params)
+        .map_err(|e| format!("Parse cert params failed: {}", e))?;
+    Ok((cert, cert_pem))
+}
 
-pub async fn issue_certificate(domain: &str) -> Cert {
+pub fn issue_certificate(intermediate_certificate: &Certificate, domain: &str) -> Result<Cert, String> {
-    {
+    let cert = new_end_entity(domain)?;
-        if let Some(cert) = CERTIFICATE_CACHE_MAP.read().await.get(domain) {
+    let cert_pem = cert.serialize_pem_with_signer(intermediate_certificate).map_err(|e| format!("Sign cert failed: {}", e))?;
-            return cert.clone();
-        }
-    }
-    let cert = new_end_entity(domain);
-    let cert_pem = cert.serialize_pem_with_signer(&INTERMEDIATE_CA.0).expect("Sign cert failed");
     let key_pem = cert.serialize_private_key_pem();
-    let intermediate_pem = INTERMEDIATE_CA.1.clone();
+    Ok(Cert {
-    let cert = Cert {
         cert_pem,
-        intermediate_pem,
         key_pem,
-    };
+    })
-    {
-        CERTIFICATE_CACHE_MAP.write().await.insert(domain.to_string(), cert.clone());
-    }
-    cert
 }
 
 fn parse_pkcs8(pem: &str) -> String {
@@ -89,9 +63,9 @@ fn parse_pkcs8(pem: &str) -> String {
     pem.to_string()
 }
 
-fn new_end_entity(domain: &str) -> Certificate {
+fn new_end_entity(domain: &str) -> Result<Certificate, String> {
     let mut params = CertificateParams::new(vec![domain.into()]);
-    let (start, end) = validity_period();
+    let (start, end) = validity_period()?;
     params.distinguished_name.push(DnType::CommonName, domain);
     params.use_authority_key_identifier_extension = true;
     params.key_usages.push(KeyUsagePurpose::DigitalSignature);
@@ -100,11 +74,11 @@ fn new_end_entity(domain: &str) -> Certificate {
     params.extended_key_usages.push(ExtendedKeyUsagePurpose::ClientAuth);
     params.not_before = start;
     params.not_after = end;
-    Certificate::from_params(params).expect("Generate domain certificate failed")
+    Ok(Certificate::from_params(params).map_err(|e| format!("New cert failed: {}", e))?)
 }
 
-fn validity_period() -> (OffsetDateTime, OffsetDateTime) {
+fn validity_period() -> Result<(OffsetDateTime, OffsetDateTime), String> {
-    let start = OffsetDateTime::now_utc().checked_sub(Duration::hours(1)).unwrap();
+    let start = OffsetDateTime::now_utc().checked_sub(Duration::hours(1)).expect("SHOULD NOT HAPPEN!");
-    let end = OffsetDateTime::now_utc().checked_add(Duration::days(90)).unwrap();
+    let end = OffsetDateTime::now_utc().checked_add(Duration::days(90)).expect("SHOULD NOT HAPPEN!");
-    (start, end)
+    Ok((start, end))
 }

src/config.rs

@@ -0,0 +1,46 @@
+use std::collections::HashMap;
+use std::fs;
+
+use serde::Deserialize;
+use serde::Serialize;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ProxyConfig {
+    pub prometheus: Option<bool>,
+    pub groups: Vec<ProxyGroup>,
+}
+
+impl ProxyConfig {
+    pub fn load(file_name: &str) -> Result<Self, String> {
+        let file_content = match fs::read_to_string(file_name) {
+            Ok(file_content) => file_content,
+            Err(e) => return Err(format!("read file: {} failed: {}", file_name, e)),
+        };
+        let proxy_config: ProxyConfig = match serde_json::from_str(&file_content) {
+            Ok(proxy_config) => proxy_config,
+            Err(e) => return Err(format!("parse file: {} failed: {}", file_name, e)),
+        };
+        Ok(proxy_config)
+    }
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ProxyGroup {
+    pub port: u16,
+    pub lookup_dns: Option<bool>,
+    pub tls: Option<ProxyTls>,
+    pub proxy_map: Option<HashMap<String, ProxyItem>>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ProxyItem {
+    pub address: String,
+    pub tls: Option<bool>,
+    pub sni: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ProxyTls {
+    pub intermediate_cert: String,
+    pub intermediate_key: String,
+}

src/main.rs

@@ -6,11 +6,13 @@ use pingora::{
 use pretty_env_logger::env_logger::Builder;
 use structopt::StructOpt;
 
+use crate::config::ProxyConfig;
 use crate::service::HostConfig;
 
 mod app;
 mod service;
 mod cert;
+mod config;
 
 pub fn main() {
     init_logger();
@@ -19,38 +21,54 @@ pub fn main() {
     let mut my_server = Server::new(opt).unwrap();
     my_server.bootstrap();
 
-    let proxy_service_tcp = service::proxy_service_tcp(
+    let proxy_config = ProxyConfig::load("proxy_config.json").unwrap_or_else(|e| {
-        &my_server.configuration,
+        panic!("Load proxy_config.json failed: {}", e);
-        "0.0.0.0:8800",
+    });
-        vec![],
-    );
 
-    let proxy_service_ssl2 = service::proxy_service_tls(
+    let mut services: Vec<Box<dyn Service>> = vec![];
-        &my_server.configuration,
+    for group in &proxy_config.groups {
-        "0.0.0.0:4430",
+        let listen_address = format!("0.0.0.0:{}", group.port);
-        vec![
+        let mut host_configs = vec![];
-            HostConfig {
+        if let Some(proxy_map) = &group.proxy_map {
-                proxy_addr: "101.132.122.240:443".to_owned(),
+            for (hostname, proxy_item) in proxy_map {
-                proxy_tls: true,
+                host_configs.push(HostConfig {
-                proxy_hostname: "hatter.ink".to_owned(),
+                    proxy_addr: proxy_item.address.clone(),
-            },
+                    proxy_tls: proxy_item.tls.unwrap_or(false),
-            HostConfig {
+                    proxy_hostname: proxy_item.sni.clone().unwrap_or_else(|| hostname.clone()),
-                proxy_addr: "1.1.1.1:443".to_owned(),
+                });
-                proxy_tls: true,
+            }
-                proxy_hostname: "one.one.one.one".to_owned(),
+        }
-            },
+        let lookup_tls = group.lookup_dns.unwrap_or(false);
-        ],
+        match &group.tls {
-    );
+            None => {
+                let proxy_service_tcp = service::proxy_service_tcp(
+                    &my_server.configuration,
+                    &listen_address,
+                    lookup_tls,
+                    host_configs,
+                );
+                services.push(Box::new(proxy_service_tcp));
+            }
+            Some(proxy_tls) => {
+                let proxy_service_ssl = service::proxy_service_tls(
+                    &my_server.configuration,
+                    &listen_address,
+                    lookup_tls,
+                    &proxy_tls,
+                    host_configs,
+                );
+                services.push(Box::new(proxy_service_ssl));
+            }
+        }
+    }
 
-    let mut prometheus_service_http = ListeningService::prometheus_http_service();
+    if let Some(true) = proxy_config.prometheus {
-    prometheus_service_http.add_tcp("127.0.0.1:6150");
+        let mut prometheus_service_http = ListeningService::prometheus_http_service();
+        prometheus_service_http.add_tcp("127.0.0.1:6150");
+        services.push(Box::new(prometheus_service_http));
+    }
 
     info!("start listen...");
-    let services: Vec<Box<dyn Service>> = vec![
-        Box::new(proxy_service_tcp),
-        Box::new(proxy_service_ssl2),
-        Box::new(prometheus_service_http),
-    ];
     my_server.add_services(services);
     my_server.run_forever();
 }

src/service.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -10,15 +11,41 @@ use pingora::tls::ext;
 use pingora::tls::pkey::PKey;
 use pingora::tls::ssl::{NameType, SslRef};
 use pingora::tls::x509::X509;
+use rcgen::Certificate;
+use tokio::sync::RwLock;
 
 use crate::app::ProxyApp;
 use crate::cert;
+use crate::cert::Cert;
+use crate::config::ProxyTls;
 
-struct Callback();
+struct Callback {
+    intermediate_certificate: Certificate,
+    intermediate_certificate_pem: String,
+    certificate_cache_map: RwLock<HashMap<String, Cert>>,
+}
 
 impl Callback {
-    fn new() -> Self {
+    fn new(proxy_tls: &ProxyTls) -> Result<Self, String> {
-        Self()
+        let (cert, cert_pem) = cert::load_certificate(&proxy_tls.intermediate_cert, &proxy_tls.intermediate_key)?;
+        Ok(Self {
+            intermediate_certificate: cert,
+            intermediate_certificate_pem: cert_pem,
+            certificate_cache_map: Default::default(),
+        })
+    }
+
+    async fn issue_certificate(&self, hostname: &str) -> Result<Cert, String> {
+        {
+            if let Some(cert) = self.certificate_cache_map.read().await.get(hostname) {
+                return Ok(cert.clone());
+            }
+        }
+        let cert = cert::issue_certificate(&self.intermediate_certificate, &hostname)?;
+        {
+            self.certificate_cache_map.write().await.insert(hostname.to_string(), cert.clone());
+        }
+        Ok(cert)
     }
 }
 
@@ -28,22 +55,26 @@ impl TlsAccept for Callback {
         let sni_provided = ssl.servername(NameType::HOST_NAME).expect("get sni failed").to_string();
         log::info!("SNI provided: {}", sni_provided);
 
-        let cert = cert::issue_certificate(&sni_provided).await;
+        let cert = self.issue_certificate(&sni_provided).await
+            .unwrap_or_else(|e| panic!("Issue certificate failed: {}", e));
 
         let x509_cert = X509::from_pem(cert.cert_pem.as_bytes())
-            .unwrap_or_else(|_| panic!("parse cert: {} failed", cert.cert_pem));
+            .unwrap_or_else(|e| panic!("parse cert: {} failed: {}", cert.cert_pem, e));
-        let x509_intermediate_cert = X509::from_pem(cert.intermediate_pem.as_bytes())
+        let x509_intermediate_cert = X509::from_pem(self.intermediate_certificate_pem.as_bytes())
-            .unwrap_or_else(|_| panic!("parse intermediate cert: {} failed", cert.intermediate_pem));
+            .unwrap_or_else(|e| panic!("parse intermediate cert: {} failed: {}", self.intermediate_certificate_pem, e));
         let private_key = PKey::private_key_from_pem(cert.key_pem.as_bytes())
-            .unwrap_or_else(|_| panic!("parse key: {} failed", cert.key_pem));
+            .unwrap_or_else(|e| panic!("parse key: {} failed: {}", cert.key_pem, e));
 
-        ext::ssl_use_certificate(ssl, &x509_cert).unwrap_or_else(|_| panic!("apply certificate for: {} failed", sni_provided));
+        ext::ssl_use_certificate(ssl, &x509_cert)
-        ext::ssl_add_chain_cert(ssl, &x509_intermediate_cert).unwrap_or_else(|_| panic!("apply intermediate certificate for: {} failed", sni_provided));
+            .unwrap_or_else(|e| panic!("apply certificate for: {} failed: {}", sni_provided, e));
-        ext::ssl_use_private_key(ssl, &private_key).unwrap_or_else(|_| panic!("apply key for: {} failed", sni_provided));
+        ext::ssl_add_chain_cert(ssl, &x509_intermediate_cert)
+            .unwrap_or_else(|e| panic!("apply intermediate 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));
     }
 }
 
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 pub struct HostConfig {
     pub proxy_addr: String,
     pub proxy_tls: bool,
@@ -53,9 +84,10 @@ pub struct HostConfig {
 pub fn proxy_service_tcp(
     server_conf: &Arc<ServerConf>,
     listen_addr: &str,
+    lookup_dns: bool,
     host_configs: Vec<HostConfig>,
 ) -> impl pingora::services::Service {
-    let proxy_app = ProxyApp::new(host_configs.clone());
+    let proxy_app = ProxyApp::new(false, lookup_dns, host_configs);
     let mut service = http_proxy_service(server_conf, proxy_app);
 
     service.add_tcp(listen_addr);
@@ -66,12 +98,14 @@ pub fn proxy_service_tcp(
 pub fn proxy_service_tls(
     server_conf: &Arc<ServerConf>,
     listen_addr: &str,
+    lookup_dns: bool,
+    proxy_tls: &ProxyTls,
     host_configs: Vec<HostConfig>,
 ) -> impl pingora::services::Service {
-    let proxy_app = ProxyApp::new(host_configs.clone());
+    let proxy_app = ProxyApp::new(true, lookup_dns, host_configs);
     let mut service = http_proxy_service(server_conf, proxy_app);
 
-    let cb = Box::new(Callback::new());
+    let cb = Box::new(Callback::new(proxy_tls).unwrap());
     let tls_settings = TlsSettings::with_callbacks(cb).unwrap();
     service.add_tls_with_settings(listen_addr, None, tls_settings);