async rustls rewrite

Cargo.toml

@@ -4,11 +4,13 @@ version = "0.1.2"
 edition = "2021"
 
 [dependencies]
-openssl = "0.10.72"
+tokio = { version = "1.44.2", features = ["full"] }
+tokio-io-timeout = "1.2.0"
+tokio-rustls = "0.26.2"
+rustls = "0.23.25"
+wildmatch = "2.4.0"
 serde_yml = "0.0.12"
+serde_json = "1.0.140"
 log = "0.4.27"
 colog = "1.3.0"
-threadpool = "1.8.1"
+ignore-result = "0.2.0"
-wildcard_ex = "0.1.2"
-websocket = "0.27.1"
-serde_json = "1.0.140"

shell.nix

@@ -0,0 +1,8 @@
+with import <nixpkgs> { };
+
+mkShell {
+  nativeBuildInputs = [
+    openssl
+    pkg-config
+  ];
+}

src/flowgate.rs

@@ -1,3 +1,3 @@
 pub mod config;
 pub mod server;
-pub mod ssl_cert;
+pub mod tls;

src/flowgate/config.rs

@@ -1,15 +1,17 @@
-use std::{fs, net::TcpStream, time::Duration};
+use std::{fs, time::Duration};
+
+use tokio::net::TcpStream;
 
 use serde_yml::{Number, Value};
-use wildcard_ex::is_match_simple;
+use wildmatch::WildMatch;
 
-use super::ssl_cert::SslCert;
+use super::tls::TlsCertificate;
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct SiteConfig {
-    pub domain: String,
+    pub domain: WildMatch,
     pub host: String,
-    pub ssl: Option<SslCert>,
+    pub ssl: Option<TlsCertificate>,
     pub enable_keep_alive: bool,
     pub support_keep_alive: bool,
     pub ip_forwarding: IpForwarding,
@@ -17,12 +19,12 @@ pub struct SiteConfig {
 }
 
 impl SiteConfig {
-    pub fn connect(&self) -> Option<TcpStream> {
+    pub async fn connect(&self) -> Option<TcpStream> {
-        TcpStream::connect(self.host.clone()).ok()
+        TcpStream::connect(self.host.clone()).await.ok()
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub enum IpForwarding {
     Simple,
     Header(String),
@@ -46,7 +48,7 @@ impl IpForwarding {
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct Config {
     pub sites: Vec<SiteConfig>,
     pub http_host: String,
@@ -80,12 +82,12 @@ impl Config {
         let sites_yaml = doc["sites"].as_sequence()?;
 
         for s in sites_yaml {
-            let mut cert: Option<SslCert> = None;
+            let mut cert: Option<TlsCertificate> = None;
             let s = s.as_mapping()?;
 
             if s.contains_key("ssl_cert") {
                 cert = Some(
-                    SslCert::new(
+                    TlsCertificate::new(
                         s.get("ssl_cert")?.as_str()?,
                         s.get("ssl_key")?.as_str()?,
                     )?,
@@ -93,7 +95,7 @@ impl Config {
             }
             
             let site = SiteConfig {
-                domain: s.get("domain")?.as_str()?.to_string(),
+                domain: WildMatch::new(&s.get("domain")?.as_str()?.to_string()),
                 host: s.get("host")?.as_str()?.to_string(),
                 ssl: cert,
                 enable_keep_alive: s.get("enable_keep_alive")
@@ -126,7 +128,7 @@ impl Config {
 
     pub fn get_site(&self, domain: &str) -> Option<&SiteConfig> {
         for i in &self.sites {
-            if is_match_simple(&i.domain, domain) {
+            if i.domain.matches(domain) {
                 return Some(i);
             }
         }

src/flowgate/server.rs

@@ -1,30 +1,23 @@
 use std::{
-    io::{
+    error::Error, 
-        Read, 
+    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}, 
-        Write
-    }, 
-    net::{
-        IpAddr, 
-        Ipv4Addr, 
-        Ipv6Addr, 
-        SocketAddr, 
-        SocketAddrV4, 
-        SocketAddrV6, 
-        TcpListener, 
-        TcpStream
-    }, 
     str::FromStr, 
-    sync::{
+    sync::Arc
-        Arc, 
+};
-        RwLock
+
-    }, 
+use tokio::sync::RwLock;
-    thread, 
+
-    time::Duration
+use ignore_result::Ignore;
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt}, 
+    net::{TcpListener, TcpStream}
 };
 
 use log::info;
-use openssl::ssl::SslStream;
+use tokio_io_timeout::TimeoutStream;
-use threadpool::ThreadPool;
+use tokio_rustls::TlsAcceptor;
+
+use crate::tls::create_server_config;
 
 use super::config::{
     Config,
@@ -40,15 +33,9 @@ pub trait Closeable {
     fn close(&mut self);
 }
 
-impl Closeable for SslStream<TcpStream> {
-    fn close(&mut self) {
-        let _ = self.shutdown();
-    }
-}
-
 impl Closeable for TcpStream {
     fn close(&mut self) {
-        let _ = self.shutdown(std::net::Shutdown::Both);
+        let _ = self.shutdown();
     }
 }
 
@@ -64,157 +51,132 @@ impl FlowgateServer {
         FlowgateServer { config }
     }
 
-    pub fn start(&self) {
+    pub async fn start(&self) {
-        let pool = ThreadPool::new(self.config.read().unwrap().threadpool_size);
+        tokio::spawn({ 
-        let pool = Arc::new(pool);
+            let config = self.config.clone(); 
 
-        thread::spawn({
+            async move { 
-            let config = Arc::clone(&self.config);
+                Self::run_http(config).await.ignore(); 
-            let pool = Arc::clone(&pool);
-            
-            move || {
-                Self::run_http(config, pool)
             } 
         });
 
-        thread::spawn({
+        tokio::spawn({ 
-            let config = Arc::clone(&self.config);
+            let config = self.config.clone(); 
-            let pool = Arc::clone(&pool);
 
-            move || {
+            async move { 
-                Self::run_https(config, pool)
+                Self::run_https(config).await.ignore(); 
             } 
         });
     }
 
-    pub fn run_http(
+    pub async fn run_http(
-        config: Arc<RwLock<Config>>,
+        config: Arc<RwLock<Config>>
-        pool: Arc<ThreadPool>
+    ) -> Result<(), Box<dyn Error>> {
-    ) -> Option<()> {
+        let listener = TcpListener::bind(&config.read().await.http_host).await?;
-        let listener = TcpListener::bind(&config.read().ok()?.http_host).ok()?;
 
-        info!("HTTP server runned on {}", &config.read().ok()?.http_host);
+        info!("HTTP server runned on {}", &config.read().await.http_host);
+
+        loop {
+            let Ok((stream, addr)) = listener.accept().await else { break }; 
 
-        for stream in listener.incoming() {
-            pool.execute({
             let config = config.clone();
 
-                move || {
+            tokio::spawn(async move {
-                    let Ok(mut stream) = stream else { return };
+                let mut stream = TimeoutStream::new(stream);
 
-                    let Ok(_) = stream.set_write_timeout(Some(Duration::from_secs(10))) else { return };
+                stream.set_write_timeout(Some(config.read().await.connection_timeout));
-                    let Ok(_) = stream.set_read_timeout(Some(Duration::from_secs(10))) else { return };
+                stream.set_read_timeout(Some(config.read().await.connection_timeout));
 
-                    let Ok(addr) = stream.peer_addr() else { return };
+                let mut stream = Box::pin(stream);
 
                 Self::accept_stream(
                     config,
                     &mut stream,
                     addr,
                     false
-                    );
+                ).await;
-                }
             });
         }
 
-        Some(())
+        Ok(())
     }
 
-    pub fn run_https(
+    pub async fn run_https(
-        config: Arc<RwLock<Config>>,
+        config: Arc<RwLock<Config>>
-        pool: Arc<ThreadPool>
+    ) -> Result<(), Box<dyn Error>> {
-    ) -> Option<()> {
+        let listener = TcpListener::bind(&config.read().await.https_host).await?;
-        use openssl::ssl::{NameType, SniError, SslAcceptor, SslAlert, SslMethod, SslRef};
+        let acceptor = TlsAcceptor::from(Arc::new(create_server_config(config.clone()).await));
 
-        let listener = TcpListener::bind(&config.read().ok()?.https_host).ok()?;
+        info!("HTTPS server runned on {}", &config.read().await.http_host);
 
-        let mut cert = SslAcceptor::mozilla_intermediate(SslMethod::tls()).ok()?;
+        loop {
+            let Ok((stream, addr)) = listener.accept().await else { break }; 
 
-        cert.set_servername_callback(Box::new({
             let config = config.clone();
+            let acceptor = acceptor.clone();
 
-                move |ssl: &mut SslRef, _: &mut SslAlert| -> Result<(), SniError> {
+            tokio::spawn(async move {
-                    let servname = ssl.servername(NameType::HOST_NAME).ok_or(SniError::NOACK)?;
+                let mut stream = TimeoutStream::new(stream);
-                    let c = config.read().unwrap();
-                    let cert = c.get_site(servname).ok_or(SniError::NOACK)?;
-                    ssl.set_ssl_context(&cert.ssl.as_ref().ok_or(SniError::NOACK)?.get_context()).ok().ok_or(SniError::NOACK)
-                }
-            }
-        ));
 
-        let cert = cert.build();
+                stream.set_write_timeout(Some(config.read().await.connection_timeout));
+                stream.set_read_timeout(Some(config.read().await.connection_timeout));
 
-        info!("HTTPS server runned on {}", &config.read().ok()?.https_host);
+                let Ok(mut stream) = acceptor.accept(Box::pin(stream)).await else { return };
-
-        for stream in listener.incoming() {
-            pool.execute({
-                let config = config.clone();
-                let cert = cert.clone();
-
-                move || {
-                    let Ok(stream) = stream else { return };
-                    
-                    let Ok(_) = stream.set_write_timeout(Some(config.read().unwrap().connection_timeout)) else { return };
-                    let Ok(_) = stream.set_read_timeout(Some(config.read().unwrap().connection_timeout)) else { return };
-
-                    let Ok(addr) = stream.peer_addr() else { return };
-
-                    let Ok(mut stream) = cert.accept(stream) else { return };
 
                 Self::accept_stream(
                     config,
                     &mut stream,
                     addr,
-                        true
+                    false
-                    );
+                ).await;
-                }
             });
         }
 
-        Some(())
+        Ok(())
     }
 
-    pub fn accept_stream(
+    pub async fn accept_stream(
         config: Arc<RwLock<Config>>, 
-        stream: &mut (impl Read + Write + Closeable), 
+        stream: &mut (impl AsyncReadExt + AsyncWriteExt + Unpin), 
         addr: SocketAddr,
         https: bool
     ) -> Option<()> {
-        let mut conn = Self::read_request(config.clone(), stream, addr, https, None)?;
+        let mut conn = read_request(config.clone(), stream, addr, https, None).await?;
 
         if conn.keep_alive && conn.config.enable_keep_alive {
             loop {
                 if !conn.config.support_keep_alive {
                     conn.stream.close();
-                    conn.stream = conn.config.connect()?;
+                    conn.stream = conn.config.connect().await?;
                 }
-                conn = Self::read_request(config.clone(), stream, addr, https, Some(conn))?;
+                conn = read_request(config.clone(), stream, addr, https, Some(conn)).await?;
             }
         }
 
         conn.stream.close();
-        stream.close();
+        stream.shutdown().await.ok()?;
 
         Some(())
     }
+}
 
-    fn read_request(
+async fn read_request(
     config: Arc<RwLock<Config>>, 
-        stream: &mut (impl Read + Write + Closeable), 
+    stream: &mut (impl AsyncReadExt + AsyncWriteExt + Unpin), 
     addr: SocketAddr,
     https: bool,
     conn: Option<Connection>
 ) -> Option<Connection> {
     let mut addr = addr;
 
-        match &config.read().ok()?.incoming_ip_forwarding {
+    match &config.read().await.incoming_ip_forwarding {
         IpForwarding::Simple => {
             let mut header = Vec::new();
 
             {
                 let mut buf = [0; 1];
 
-                    while let Ok(1) = stream.read(&mut buf) {
+                while let Ok(1) = stream.read(&mut buf).await {
                     let byte = buf[0];
                     if byte == b'\n' { break }
                     header.push(byte);
@@ -225,20 +187,20 @@ impl FlowgateServer {
         },
         IpForwarding::Modern => {
             let mut ipver = [0; 1];
-                stream.read(&mut ipver).ok()?;
+            stream.read(&mut ipver).await.ok()?;
             addr = match ipver[0] {
                 0x01 => {
                     let mut octets = [0; 4];
-                        stream.read(&mut octets).ok()?;
+                    stream.read(&mut octets).await.ok()?;
                     let mut port = [0; 2];
-                        stream.read(&mut port).ok()?;
+                    stream.read(&mut port).await.ok()?;
                     let port = u16::from_be_bytes(port);
                     SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::from(octets), port))
                 }, 0x02 => {
                     let mut octets = [0; 16];
-                        stream.read(&mut octets).ok()?;
+                    stream.read(&mut octets).await.ok()?;
                     let mut port = [0; 2];
-                        stream.read(&mut port).ok()?;
+                    stream.read(&mut port).await.ok()?;
                     let port = u16::from_be_bytes(port);
                     SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::from(octets), port, 0, 0))
                 }, _ => { return None },
@@ -253,7 +215,7 @@ impl FlowgateServer {
         let mut buf = [0; 1];
         let mut counter = 0;
 
-            while let Ok(1) = stream.read(&mut buf) {
+        while let Ok(1) = stream.read(&mut buf).await {
             let byte = buf[0];
             head.push(byte);
 
@@ -290,7 +252,7 @@ impl FlowgateServer {
         .map(|o| o.contains(&"chunked".to_string()))
         .unwrap_or(false);
     
-        if let IpForwarding::Header(header) = &config.read().ok()?.incoming_ip_forwarding {
+    if let IpForwarding::Header(header) = &config.read().await.incoming_ip_forwarding {
         if let Some(ip) = headers.iter().find(|o| o.0 == header).map(|o| o.1) {
             addr = SocketAddr::from_str(ip).ok()?;
         }
@@ -308,10 +270,10 @@ impl FlowgateServer {
             }
         }
 
-            let site = config.read().ok()?.get_site(&host)?.clone();
+        let site = config.read().await.get_site(&host)?.clone();
 
         Connection {
-                stream: site.connect()?,
+            stream: site.connect().await?,
             config: site,
             keep_alive,
             host
@@ -398,28 +360,27 @@ impl FlowgateServer {
         }
     }
 
-        conn.stream.write_all(&reqbuf).ok()?;
+    conn.stream.write_all(&reqbuf).await.ok()?;
 
     if content_length > 0 {
         let mut read = 0usize;
         let mut buf = vec![0; 4096];
-            while let Ok(size) = stream.read(&mut buf) {
+        while let Ok(size) = stream.read(&mut buf).await {
             if size == 0 { break }
             read += size;
             buf.truncate(size);
-                conn.stream.write_all(&buf).ok()?;
+            conn.stream.write_all(&buf).await.ok()?;
             buf = vec![0; 4096];
             if read >= content_length { break }
         }
     } else if is_chunked {
-            
         loop {
             let mut length = Vec::new();
             {
                 let mut buf = [0; 1];
                 let mut counter = 0;
 
-                    while let Ok(1) = stream.read(&mut buf) {
+                while let Ok(1) = stream.read(&mut buf).await {
                     let byte = buf[0];
                     length.push(byte);
 
@@ -429,16 +390,16 @@ impl FlowgateServer {
                         _ => 0,
                     };
                 }
-                    conn.stream.write_all(&length).ok()?;
+                conn.stream.write_all(&length).await.ok()?;
 
                 length.truncate(length.len() - 2);
             }
             let length = String::from_utf8(length).ok()?;
             let length = usize::from_str_radix(length.as_str(), 16).ok()?;
             let mut data = vec![0u8; length+2];
-                stream.read_exact(&mut data).ok()?;
+            stream.read_exact(&mut data).await.ok()?;
 
-                conn.stream.write_all(&data).ok()?;
+            conn.stream.write_all(&data).await.ok()?;
             if length == 0 {
                 break;
             }
@@ -452,11 +413,11 @@ impl FlowgateServer {
             let mut buf = [0; 1];
             let mut counter = 0;
 
-                while let Ok(1) = conn.stream.read(&mut buf) {
+            while let Ok(1) = conn.stream.read(&mut buf).await {
                 let byte = buf[0];
                 head.push(byte);
 
-                    stream.write_all(&buf).ok()?;
+                stream.write_all(&buf).await.ok()?;
 
                 counter = match (counter, byte) {
                     (0, b'\r') => 1,
@@ -497,11 +458,11 @@ impl FlowgateServer {
         if content_length > 0 {
             let mut read = 0usize;
             let mut buf = vec![0; 4096];
-                while let Ok(size) = conn.stream.read(&mut buf) {
+            while let Ok(size) = conn.stream.read(&mut buf).await {
                 if size == 0 { break }
                 read += size;
                 buf.truncate(size);
-                    stream.write_all(&buf).ok()?;
+                stream.write_all(&buf).await.ok()?;
                 buf = vec![0; 4096];
                 if read == content_length { break }
             }
@@ -512,7 +473,7 @@ impl FlowgateServer {
                     let mut buf = [0; 1];
                     let mut counter = 0;
 
-                        while let Ok(1) = conn.stream.read(&mut buf) {
+                    while let Ok(1) = conn.stream.read(&mut buf).await {
                         let byte = buf[0];
                         length.push(byte);
 
@@ -522,16 +483,16 @@ impl FlowgateServer {
                             _ => 0,
                         };
                     }
-                        stream.write_all(&length).ok()?;
+                    stream.write_all(&length).await.ok()?;
 
                     length.truncate(length.len() - 2);
                 }
                 let length = String::from_utf8(length).ok()?;
                 let length = usize::from_str_radix(length.as_str(), 16).ok()?;
                 let mut data = vec![0u8; length+2];
-                    conn.stream.read_exact(&mut data).ok()?;
+                conn.stream.read_exact(&mut data).await.ok()?;
 
-                    stream.write_all(&data).ok()?;
+                stream.write_all(&data).await.ok()?;
                 if length == 0 {
                     break;
                 }
@@ -539,10 +500,10 @@ impl FlowgateServer {
         }
     } else {
         let mut buf = vec![0;1024];
-            while let Ok(n) = conn.stream.read(&mut buf) {
+        while let Ok(n) = conn.stream.read(&mut buf).await {
             if n == 0 { break }
             buf.truncate(n);
-                stream.write_all(&buf).ok()?;
+            stream.write_all(&buf).await.ok()?;
             buf = vec![0;1024];
         }
     }
@@ -551,4 +512,3 @@ impl FlowgateServer {
 
     Some(conn)
 }
-}

src/flowgate/ssl_cert.rs

@@ -1,28 +0,0 @@
-use openssl::ssl::SslContext;
-
-#[derive(Clone)]
-pub struct SslCert {
-    context: SslContext,
-}
-
-fn generate_ctx(cert_file: &str, key_file: &str) -> Option<SslContext> {
-    use openssl::ssl::{SslFiletype, SslMethod};
-
-    let mut ctx = SslContext::builder(SslMethod::tls()).ok().unwrap();
-    ctx.set_private_key_file(&key_file, SslFiletype::PEM).ok().unwrap();
-    ctx.set_certificate_file(&cert_file, SslFiletype::PEM).ok().unwrap();
-    ctx.check_private_key().ok()?;
-    Some(ctx.build())
-}
-
-impl SslCert {
-    pub fn new(cert_file: &str, key_file: &str) -> Option<SslCert> {
-        Some(SslCert {
-            context: generate_ctx(cert_file, key_file)?
-        })
-    }
-
-    pub fn get_context(&self) -> SslContext {
-        self.context.clone()
-    }
-}

src/flowgate/tls.rs

@@ -0,0 +1,66 @@
+use std::sync::Arc;
+
+use rustls::{
+    crypto::aws_lc_rs::sign::any_supported_type, 
+    pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer}, 
+    server::{ClientHello, ResolvesServerCert}, 
+    sign::CertifiedKey, 
+    ServerConfig
+};
+
+use tokio::{runtime::Handle, sync::RwLock};
+
+use super::config::Config;
+
+
+#[derive(Clone, Debug)]
+pub struct TlsCertificate {
+    key: CertifiedKey
+}
+
+impl TlsCertificate {
+    pub fn new(cert_file: &str, key_file: &str) -> Option<TlsCertificate> {
+        let certs = CertificateDer::pem_file_iter(cert_file)
+            .unwrap()
+            .map(|cert| cert.unwrap())
+            .collect();
+        let private_key = PrivateKeyDer::from_pem_file(key_file).unwrap();
+        let key = CertifiedKey::new(certs, any_supported_type(&private_key).ok()?);
+
+        Some(Self { key })
+    }
+
+    pub fn get_key(&self) -> CertifiedKey {
+        self.key.clone()
+    }
+}
+
+#[derive(Debug)]
+struct ResolvesServerCertWildcard {
+    config: Arc<RwLock<Config>>,
+    handle: Handle
+}
+
+impl ResolvesServerCertWildcard {
+    pub async fn new(config: Arc<RwLock<Config>>) -> Self {
+        Self { config, handle: Handle::current() }
+    }
+}
+
+impl ResolvesServerCert for ResolvesServerCertWildcard {
+    fn resolve(&self, client_hello: ClientHello<'_>) -> Option<Arc<CertifiedKey>> {
+        if let Some(cert) = client_hello.server_name()
+                .and_then(|name| self.handle.block_on(self.config.read()).get_site(name).cloned())
+                .and_then(|site| site.ssl) {
+            Some(Arc::new(cert.get_key()))
+        } else {
+            None
+        }
+    }
+}
+
+pub async fn create_server_config(config: Arc<RwLock<Config>>) -> ServerConfig {
+    ServerConfig::builder()
+        .with_no_client_auth()
+        .with_cert_resolver(Arc::new(ResolvesServerCertWildcard::new(config).await))
+}

src/main.rs

@@ -1,16 +1,19 @@
-use std::{fs, path::Path, sync::{Arc, RwLock}};
+use std::{fs, path::Path, sync::Arc};
 
 use flowgate::{config::Config, server::FlowgateServer};
+use ignore_result::Ignore;
+use tokio::sync::RwLock;
 
-fn main() {
+#[tokio::main]
+async fn main() {
     colog::init();
 
     if !Path::new("conf.yml").exists() {
-        let _ = fs::write("conf.yml", include_bytes!("../conf.yml"));
+        fs::write("conf.yml", include_bytes!("../conf.yml")).ignore();
     }
 
     let config = Arc::new(RwLock::new(Config::parse("conf.yml").unwrap()));
     let server = FlowgateServer::new(config.clone());
 
-    server.start();
+    server.start().await;
 }