MeexReay/flowgate
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

async rustls rewrite

Browse Source
This commit is contained in:
MeexReay 2025-04-07 00:15:10 +03:00
parent 4a0c00d421
commit 37c6122f87
9 changed files with 870 additions and 1274 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1247
Cargo.lock generated
View File

File diff suppressed because it is too large Load Diff

12
Cargo.toml
View File

@ -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"
wildcard_ex = "0.1.2"
websocket = "0.27.1"
serde_json = "1.0.140"
ignore-result = "0.2.0"

8
shell.nix Executable file
View File

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

2
src/flowgate.rs
View File

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

30
src/flowgate/config.rs
View File

@ -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> {
TcpStream::connect(self.host.clone()).ok()
pub async fn connect(&self) -> Option<TcpStream> {
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);
}
}

740
src/flowgate/server.rs
View File

@ -1,30 +1,23 @@
use std::{
io::{
Read,
Write
},
net::{
IpAddr,
Ipv4Addr,
Ipv6Addr,
SocketAddr,
SocketAddrV4,
SocketAddrV6,
TcpListener,
TcpStream
},
error::Error,
net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
str::FromStr,
sync::{
Arc,
RwLock
},
thread,
time::Duration
sync::Arc
};
use tokio::sync::RwLock;
use ignore_result::Ignore;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream}
};
use log::info;
use openssl::ssl::SslStream;
use threadpool::ThreadPool;
use tokio_io_timeout::TimeoutStream;
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,199 +51,374 @@ impl FlowgateServer {
FlowgateServer { config }
}
pub fn start(&self) {
let pool = ThreadPool::new(self.config.read().unwrap().threadpool_size);
let pool = Arc::new(pool);
pub async fn start(&self) {
tokio::spawn({
let config = self.config.clone();
thread::spawn({
let config = Arc::clone(&self.config);
let pool = Arc::clone(&pool);
move || {
Self::run_http(config, pool)
}
async move {
Self::run_http(config).await.ignore();
}
});
thread::spawn({
let config = Arc::clone(&self.config);
let pool = Arc::clone(&pool);
move || {
Self::run_https(config, pool)
}
tokio::spawn({
let config = self.config.clone();
async move {
Self::run_https(config).await.ignore();
}
});
}
pub fn run_http(
config: Arc<RwLock<Config>>,
pool: Arc<ThreadPool>
) -> Option<()> {
let listener = TcpListener::bind(&config.read().ok()?.http_host).ok()?;
pub async fn run_http(
config: Arc<RwLock<Config>>
) -> Result<(), Box<dyn Error>> {
let listener = TcpListener::bind(&config.read().await.http_host).await?;
info!("HTTP server runned on {}", &config.read().ok()?.http_host);
info!("HTTP server runned on {}", &config.read().await.http_host);
for stream in listener.incoming() {
pool.execute({
let config = config.clone();
loop {
let Ok((stream, addr)) = listener.accept().await else { break };
move || {
let Ok(mut stream) = stream else { return };
let config = config.clone();
let Ok(_) = stream.set_write_timeout(Some(Duration::from_secs(10))) else { return };
let Ok(_) = stream.set_read_timeout(Some(Duration::from_secs(10))) else { return };
tokio::spawn(async move {
let mut stream = TimeoutStream::new(stream);
let Ok(addr) = stream.peer_addr() else { return };
stream.set_write_timeout(Some(config.read().await.connection_timeout));
stream.set_read_timeout(Some(config.read().await.connection_timeout));
Self::accept_stream(
config,
&mut stream,
addr,
false
);
}
let mut stream = Box::pin(stream);
Self::accept_stream(
config,
&mut stream,
addr,
false
).await;
});
}
Some(())
Ok(())
}
pub fn run_https(
config: Arc<RwLock<Config>>,
pool: Arc<ThreadPool>
) -> Option<()> {
use openssl::ssl::{NameType, SniError, SslAcceptor, SslAlert, SslMethod, SslRef};
pub async fn run_https(
config: Arc<RwLock<Config>>
) -> Result<(), Box<dyn Error>> {
let listener = TcpListener::bind(&config.read().await.https_host).await?;
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 config = config.clone();
let acceptor = acceptor.clone();
move |ssl: &mut SslRef, _: &mut SslAlert| -> Result<(), SniError> {
let servname = ssl.servername(NameType::HOST_NAME).ok_or(SniError::NOACK)?;
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)
}
}
));
tokio::spawn(async move {
let mut stream = TimeoutStream::new(stream);
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
);
}
Self::accept_stream(
config,
&mut stream,
addr,
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(
config: Arc<RwLock<Config>>,
stream: &mut (impl Read + Write + Closeable),
addr: SocketAddr,
https: bool,
conn: Option<Connection>
) -> Option<Connection> {
let mut addr = addr;
async fn read_request(
config: Arc<RwLock<Config>>,
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 {
IpForwarding::Simple => {
let mut header = Vec::new();
match &config.read().await.incoming_ip_forwarding {
IpForwarding::Simple => {
let mut header = Vec::new();
{
let mut buf = [0; 1];
{
let mut buf = [0; 1];
while let Ok(1) = stream.read(&mut buf) {
let byte = buf[0];
if byte == b'\n' { break }
header.push(byte);
}
while let Ok(1) = stream.read(&mut buf).await {
let byte = buf[0];
if byte == b'\n' { break }
header.push(byte);
}
}
addr = SocketAddr::from_str(&String::from_utf8(header).ok()?).ok()?;
},
IpForwarding::Modern => {
let mut ipver = [0; 1];
stream.read(&mut ipver).ok()?;
addr = match ipver[0] {
0x01 => {
let mut octets = [0; 4];
stream.read(&mut octets).ok()?;
let mut port = [0; 2];
stream.read(&mut port).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()?;
let mut port = [0; 2];
stream.read(&mut port).ok()?;
let port = u16::from_be_bytes(port);
SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::from(octets), port, 0, 0))
}, _ => { return None },
};
},
_ => { }
addr = SocketAddr::from_str(&String::from_utf8(header).ok()?).ok()?;
},
IpForwarding::Modern => {
let mut ipver = [0; 1];
stream.read(&mut ipver).await.ok()?;
addr = match ipver[0] {
0x01 => {
let mut octets = [0; 4];
stream.read(&mut octets).await.ok()?;
let mut port = [0; 2];
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).await.ok()?;
let mut port = [0; 2];
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 },
};
},
_ => { }
}
let mut head = Vec::new();
{
let mut buf = [0; 1];
let mut counter = 0;
while let Ok(1) = stream.read(&mut buf).await {
let byte = buf[0];
head.push(byte);
counter = match (counter, byte) {
(0, b'\r') => 1,
(1, b'\n') => 2,
(2, b'\r') => 3,
(3, b'\n') => break,
_ => 0,
};
}
head.truncate(head.len() - 4);
}
if head.is_empty() { return None; }
let head_str = String::from_utf8(head.clone()).ok()?;
let head_str = head_str.trim_matches(char::from(0)).to_string();
let mut head_lines = head_str.split("\r\n");
let status = head_lines.next()?;
let status_seq: Vec<&str> = status.split(" ").collect();
let headers: Vec<(&str, &str)> = head_lines
.filter(|l| l.contains(": "))
.map(|l| l.split_once(": ").unwrap())
.collect();
let is_chunked = headers.iter()
.find(|o| o.0.to_lowercase() == "transfer-encoding")
.map(|o| o.1.split(",").map(|x| x.trim_matches(' ').to_string()).collect::<Vec<String>>())
.map(|o| o.contains(&"chunked".to_string()))
.unwrap_or(false);
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()?;
}
}
let mut conn: Connection = if conn.is_none() {
let mut host = String::new();
let mut keep_alive = false;
for (key, value) in &headers {
match key.to_lowercase().as_str() {
"host" => host = value.to_string(),
"connection" => keep_alive = *value == "keep-alive",
_ => {}
}
}
let site = config.read().await.get_site(&host)?.clone();
Connection {
stream: site.connect().await?,
config: site,
keep_alive,
host
}
} else {
conn?
};
let content_length = headers
.iter()
.filter(|(k, _)| k.to_lowercase() == "content-length")
.next()
.map(|o| o.1.parse().ok())
.flatten()
.unwrap_or(0usize);
let mut reqbuf: Vec<u8> = Vec::new();
if let Some(replace_host) = conn.config.replace_host.clone() {
let mut new_head = Vec::new();
let mut is_status = true;
for line in head_str.split("\r\n") {
if is_status {
new_head.append(&mut line.as_bytes().to_vec());
is_status = false;
} else {
new_head.append(&mut b"\r\n".to_vec());
let (key, _) = line.split_once(": ")?;
if key.to_lowercase() == "host" {
new_head.append(&mut key.as_bytes().to_vec());
new_head.append(&mut b": ".to_vec());
new_head.append(&mut replace_host.as_bytes().to_vec());
} else {
new_head.append(&mut line.as_bytes().to_vec());
}
}
}
head = new_head;
}
match &conn.config.ip_forwarding {
IpForwarding::Header(header) => {
reqbuf.append(&mut status.to_string().as_bytes().to_vec());
reqbuf.append(&mut b"\r\n".to_vec());
for (key, value) in String::from_utf8(head.clone()).ok()?
.split("\r\n")
.skip(1)
.filter_map(|o| o.split_once(": ")) {
if *key.to_lowercase() == header.to_lowercase() { continue }
reqbuf.append(&mut key.to_string().as_bytes().to_vec());
reqbuf.append(&mut b": ".to_vec());
reqbuf.append(&mut value.to_string().as_bytes().to_vec());
reqbuf.append(&mut b"\r\n".to_vec());
}
reqbuf.append(&mut header.as_bytes().to_vec());
reqbuf.append(&mut b": ".to_vec());
reqbuf.append(&mut addr.to_string().as_bytes().to_vec());
reqbuf.append(&mut b"\r\n\r\n".to_vec());
},
IpForwarding::Simple => {
reqbuf.append(&mut addr.to_string().as_bytes().to_vec());
reqbuf.push(b'\n');
reqbuf.append(&mut head.clone());
reqbuf.append(&mut b"\r\n\r\n".to_vec());
},
IpForwarding::Modern => {
reqbuf.push(if addr.is_ipv4() { 0x01 } else { 0x02 });
match addr.ip() {
IpAddr::V4(ip) => {
reqbuf.append(&mut ip.octets().to_vec());
}, IpAddr::V6(ip) => {
reqbuf.append(&mut ip.octets().to_vec());
}
}
reqbuf.append(&mut addr.port().to_be_bytes().to_vec());
reqbuf.append(&mut head.clone());
reqbuf.append(&mut b"\r\n\r\n".to_vec());
},
IpForwarding::None => {
reqbuf.append(&mut head.clone());
reqbuf.append(&mut b"\r\n\r\n".to_vec());
}
}
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).await {
if size == 0 { break }
read += size;
buf.truncate(size);
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).await {
let byte = buf[0];
length.push(byte);
counter = match (counter, byte) {
(0, b'\r') => 1,
(1, b'\n') => break,
_ => 0,
};
}
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).await.ok()?;
conn.stream.write_all(&data).await.ok()?;
if length == 0 {
break;
}
}
}
if conn.config.support_keep_alive {
let mut head = Vec::new();
{
let mut buf = [0; 1];
let mut counter = 0;
while let Ok(1) = 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).await.ok()?;
counter = match (counter, byte) {
(0, b'\r') => 1,
(1, b'\n') => 2,
@ -272,154 +434,46 @@ impl FlowgateServer {
if head.is_empty() { return None; }
let head_str = String::from_utf8(head.clone()).ok()?;
let head_str = head_str.trim_matches(char::from(0)).to_string();
let head_str = head_str.trim_matches(char::from(0));
let mut head_lines = head_str.split("\r\n");
let status = head_lines.next()?;
let status_seq: Vec<&str> = status.split(" ").collect();
let headers: Vec<(&str, &str)> = head_lines
let headers = head_str.split("\r\n")
.skip(1)
.filter(|l| l.contains(": "))
.map(|l| l.split_once(": ").unwrap())
.collect();
.map(|(k,v)| (k.to_lowercase(),v.to_string()))
.collect::<Vec<(String,String)>>();
let content_length = headers.iter()
.find(|(k, _)| k == "content-length")
.map(|o| o.1.parse().ok())
.flatten()
.unwrap_or(0usize);
let is_chunked = headers.iter()
.find(|o| o.0.to_lowercase() == "transfer-encoding")
.map(|o| o.1.split(",").map(|x| x.trim_matches(' ').to_string()).collect::<Vec<String>>())
.map(|o| o.contains(&"chunked".to_string()))
.unwrap_or(false);
if let IpForwarding::Header(header) = &config.read().ok()?.incoming_ip_forwarding {
if let Some(ip) = headers.iter().find(|o| o.0 == header).map(|o| o.1) {
addr = SocketAddr::from_str(ip).ok()?;
}
}
let mut conn: Connection = if conn.is_none() {
let mut host = String::new();
let mut keep_alive = false;
for (key, value) in &headers {
match key.to_lowercase().as_str() {
"host" => host = value.to_string(),
"connection" => keep_alive = *value == "keep-alive",
_ => {}
}
}
let site = config.read().ok()?.get_site(&host)?.clone();
Connection {
stream: site.connect()?,
config: site,
keep_alive,
host
}
} else {
conn?
};
let content_length = headers
.iter()
.filter(|(k, _)| k.to_lowercase() == "content-length")
.next()
.map(|o| o.1.parse().ok())
.flatten()
.unwrap_or(0usize);
let mut reqbuf: Vec<u8> = Vec::new();
if let Some(replace_host) = conn.config.replace_host.clone() {
let mut new_head = Vec::new();
let mut is_status = true;
for line in head_str.split("\r\n") {
if is_status {
new_head.append(&mut line.as_bytes().to_vec());
is_status = false;
} else {
new_head.append(&mut b"\r\n".to_vec());
let (key, _) = line.split_once(": ")?;
if key.to_lowercase() == "host" {
new_head.append(&mut key.as_bytes().to_vec());
new_head.append(&mut b": ".to_vec());
new_head.append(&mut replace_host.as_bytes().to_vec());
} else {
new_head.append(&mut line.as_bytes().to_vec());
}
}
}
head = new_head;
}
match &conn.config.ip_forwarding {
IpForwarding::Header(header) => {
reqbuf.append(&mut status.to_string().as_bytes().to_vec());
reqbuf.append(&mut b"\r\n".to_vec());
for (key, value) in String::from_utf8(head.clone()).ok()?
.split("\r\n")
.skip(1)
.filter_map(|o| o.split_once(": ")) {
if *key.to_lowercase() == header.to_lowercase() { continue }
reqbuf.append(&mut key.to_string().as_bytes().to_vec());
reqbuf.append(&mut b": ".to_vec());
reqbuf.append(&mut value.to_string().as_bytes().to_vec());
reqbuf.append(&mut b"\r\n".to_vec());
}
reqbuf.append(&mut header.as_bytes().to_vec());
reqbuf.append(&mut b": ".to_vec());
reqbuf.append(&mut addr.to_string().as_bytes().to_vec());
reqbuf.append(&mut b"\r\n\r\n".to_vec());
},
IpForwarding::Simple => {
reqbuf.append(&mut addr.to_string().as_bytes().to_vec());
reqbuf.push(b'\n');
reqbuf.append(&mut head.clone());
reqbuf.append(&mut b"\r\n\r\n".to_vec());
},
IpForwarding::Modern => {
reqbuf.push(if addr.is_ipv4() { 0x01 } else { 0x02 });
match addr.ip() {
IpAddr::V4(ip) => {
reqbuf.append(&mut ip.octets().to_vec());
}, IpAddr::V6(ip) => {
reqbuf.append(&mut ip.octets().to_vec());
}
}
reqbuf.append(&mut addr.port().to_be_bytes().to_vec());
reqbuf.append(&mut head.clone());
reqbuf.append(&mut b"\r\n\r\n".to_vec());
},
IpForwarding::None => {
reqbuf.append(&mut head.clone());
reqbuf.append(&mut b"\r\n\r\n".to_vec());
}
}
conn.stream.write_all(&reqbuf).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) = conn.stream.read(&mut buf).await {
if size == 0 { break }
read += size;
buf.truncate(size);
conn.stream.write_all(&buf).ok()?;
stream.write_all(&buf).await.ok()?;
buf = vec![0; 4096];
if read >= content_length { break }
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) = conn.stream.read(&mut buf).await {
let byte = buf[0];
length.push(byte);
@ -429,126 +483,32 @@ impl FlowgateServer {
_ => 0,
};
}
conn.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];
stream.read_exact(&mut data).ok()?;
conn.stream.read_exact(&mut data).await.ok()?;
conn.stream.write_all(&data).ok()?;
stream.write_all(&data).await.ok()?;
if length == 0 {
break;
}
}
}
if conn.config.support_keep_alive {
let mut head = Vec::new();
{
let mut buf = [0; 1];
let mut counter = 0;
while let Ok(1) = conn.stream.read(&mut buf) {
let byte = buf[0];
head.push(byte);
stream.write_all(&buf).ok()?;
counter = match (counter, byte) {
(0, b'\r') => 1,
(1, b'\n') => 2,
(2, b'\r') => 3,
(3, b'\n') => break,
_ => 0,
};
}
head.truncate(head.len() - 4);
}
if head.is_empty() { return None; }
let head_str = String::from_utf8(head.clone()).ok()?;
let head_str = head_str.trim_matches(char::from(0));
let headers = head_str.split("\r\n")
.skip(1)
.filter(|l| l.contains(": "))
.map(|l| l.split_once(": ").unwrap())
.map(|(k,v)| (k.to_lowercase(),v.to_string()))
.collect::<Vec<(String,String)>>();
let content_length = headers.iter()
.find(|(k, _)| k == "content-length")
.map(|o| o.1.parse().ok())
.flatten()
.unwrap_or(0usize);
let is_chunked = headers.iter()
.find(|o| o.0.to_lowercase() == "transfer-encoding")
.map(|o| o.1.split(",").map(|x| x.trim_matches(' ').to_string()).collect::<Vec<String>>())
.map(|o| o.contains(&"chunked".to_string()))
.unwrap_or(false);
if content_length > 0 {
let mut read = 0usize;
let mut buf = vec![0; 4096];
while let Ok(size) = conn.stream.read(&mut buf) {
if size == 0 { break }
read += size;
buf.truncate(size);
stream.write_all(&buf).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) = conn.stream.read(&mut buf) {
let byte = buf[0];
length.push(byte);
counter = match (counter, byte) {
(0, b'\r') => 1,
(1, b'\n') => break,
_ => 0,
};
}
stream.write_all(&length).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()?;
stream.write_all(&data).ok()?;
if length == 0 {
break;
}
}
}
} else {
let mut buf = vec![0;1024];
while let Ok(n) = conn.stream.read(&mut buf) {
if n == 0 { break }
buf.truncate(n);
stream.write_all(&buf).ok()?;
buf = vec![0;1024];
}
} else {
let mut buf = vec![0;1024];
while let Ok(n) = conn.stream.read(&mut buf).await {
if n == 0 { break }
buf.truncate(n);
stream.write_all(&buf).await.ok()?;
buf = vec![0;1024];
}
info!("{addr} > {} {}://{}{}", status_seq[0], if https { "https" } else { "http" }, conn.host, status_seq[1]);
Some(conn)
}
info!("{addr} > {} {}://{}{}", status_seq[0], if https { "https" } else { "http" }, conn.host, status_seq[1]);
Some(conn)
}

28
src/flowgate/ssl_cert.rs
View File

@ -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()
}
}

66
src/flowgate/tls.rs Executable file
View File

@ -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))
}

11
src/main.rs
View File

@ -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;
}