aprl/hai
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

Compare commits

base: aprl:a8948c817373220087570a666eb4151833edd548
Branches Tags
aprl:master
..
compare: aprl:58c3ec763a6ece7ee73605792b3aae2807e9fb31
Branches Tags
aprl:master

No commits in common. "a8948c817373220087570a666eb4151833edd548" and "58c3ec763a6ece7ee73605792b3aae2807e9fb31" have entirely different histories.
a8948c8173 ... 58c3ec763a

5 changed files with 36 additions and 321 deletions
Show stats Expand all files Collapse all files

47
Cargo.lock generated
View file

@ -92,7 +92,7 @@ dependencies = [
"hashbrown 0.15.2",
"paste",
"static_assertions",
"windows 0.58.0",
"windows",
"windows-core 0.58.0",
]
@ -1874,10 +1874,8 @@ dependencies = [
"quinn",
"rcgen",
"rustls 0.23.23",
"rustls-pemfile 2.2.0",
"shadow-rs",
"tokio",
"tracing",
"trust-dns-resolver",
]
@ -1943,13 +1941,13 @@ dependencies = [
[[package]]
name = "hostname"
version = "0.4.0"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f9c7c7c8ac16c798734b8a24560c1362120597c40d5e1459f09498f8f6c8f2ba"
checksum = "3c731c3e10504cc8ed35cfe2f1db4c9274c3d35fa486e3b31df46f068ef3e867"
dependencies = [
"cfg-if",
"libc",
"windows 0.52.0",
"match_cfg",
"winapi",
]
[[package]]
@ -2153,7 +2151,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b77d01e822461baa8409e156015a1d91735549f0f2c17691bd2d996bef238f7f"
dependencies = [
"byteorder-lite",
"quick-error",
"quick-error 2.0.1",
]
[[package]]
@ -2454,6 +2452,12 @@ dependencies = [
"libc",
]
[[package]]
name = "match_cfg"
version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ffbee8634e0d45d258acb448e7eaab3fce7a0a467395d4d9f228e3c1f01fb2e4"
[[package]]
name = "memchr"
version = "2.7.4"
@ -3210,6 +3214,12 @@ version = "1.0.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "afbdc74edc00b6f6a218ca6a5364d6226a259d4b8ea1af4a0ea063f27e179f4d"
[[package]]
name = "quick-error"
version = "1.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0"
[[package]]
name = "quick-error"
version = "2.0.1"
@ -3418,11 +3428,12 @@ checksum = "19b30a45b0cd0bcca8037f3d0dc3421eaf95327a17cad11964fb8179b4fc4832"
[[package]]
name = "resolv-conf"
version = "0.7.1"
version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "48375394603e3dd4b2d64371f7148fd8c7baa2680e28741f2cb8d23b59e3d4c4"
checksum = "52e44394d2086d010551b14b53b1f24e31647570cd1deb0379e2c21b329aba00"
dependencies = [
"hostname",
"quick-error 1.2.3",
]
[[package]]
@ -4959,7 +4970,7 @@ dependencies = [
"wasm-bindgen",
"web-sys",
"wgpu-types",
"windows 0.58.0",
"windows",
]
[[package]]
@ -4988,9 +4999,9 @@ dependencies = [
[[package]]
name = "widestring"
version = "1.2.0"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dd7cf3379ca1aac9eea11fba24fd7e315d621f8dfe35c8d7d2be8b793726e07d"
checksum = "7219d36b6eac893fa81e84ebe06485e7dcbb616177469b142df14f1f4deb1311"
[[package]]
name = "winapi"
@ -5023,16 +5034,6 @@ version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
[[package]]
name = "windows"
version = "0.52.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e48a53791691ab099e5e2ad123536d0fff50652600abaf43bbf952894110d0be"
dependencies = [
"windows-core 0.52.0",
"windows-targets 0.52.6",
]
[[package]]
name = "windows"
version = "0.58.0"

2
Cargo.toml
View file

@ -22,8 +22,6 @@ rcgen = { version = "0.13" }
rustls = { version = "0.23", features = ["prefer-post-quantum"] }
shadow-rs = "0.38"
tokio = { version = "1.43", features = ["full"] }
rustls-pemfile = "2.2.0"
tracing = { version = "0.1.41", features = ["log-always"] }
trust-dns-resolver = { version = "0.23.2", features = ["tokio", "tokio-rustls", "rustls", "dns-over-rustls"] }
[build-dependencies]

196
src/main.rs
View file

@ -11,26 +11,18 @@ mod pman;
mod quic;
use std::fs::File;
use std::io::Write;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::path::{Path, PathBuf};
use std::net::SocketAddr;
use std::str::FromStr;
use bunt::println;
use clap::{Parser, Subcommand};
use config::{Config, FileFormat, Source};
use log::{debug, error, info};
use log::{debug, info};
use pman::{init_process_manager, ProcessCommand, ProcessManager};
use shadow_rs::shadow;
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};
use anyhow::{anyhow, bail, Context, Error};
use quinn::crypto::rustls::QuicServerConfig;
use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer};
use tokio::{signal, sync::mpsc::Sender};
use tracing::{info_span, Instrument};
use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
use trust_dns_resolver::{Resolver, TokioAsyncResolver};
use crate::quic::server::handle_connection;
shadow!(build);
@ -61,30 +53,6 @@ enum Commands {
#[command(about = "Start client as GUI")]
GuiClient,
Devtest,
#[command(about = "Start server")]
Server {
/// file to log TLS keys to for debugging
#[clap(long = "keylog")]
keylog: bool,
/// TLS private key in PEM format
#[clap(short = 'k', long = "key", requires = "cert")]
key: Option<PathBuf>,
/// TLS certificate in PEM format
#[clap(short = 'c', long = "cert", requires = "key")]
cert: Option<PathBuf>,
/// Enable stateless retries
#[clap(long = "stateless-retry")]
stateless_retry: bool,
/// Address to listen on
#[clap(long = "listen", default_value = "[::1]:4433")]
listen: SocketAddr,
/// Client address to block
#[clap(long = "block")]
block: Option<SocketAddr>,
/// Maximum number of concurrent connections to allow
#[clap(long = "connection-limit")]
connection_limit: Option<usize>,
}
}
fn config() -> &'static Config {
@ -124,7 +92,6 @@ async fn main() -> anyhow::Result<()> {
match args.command {
Commands::CliClient => {
let custom_certs_fut = read_opt_server_cert();
let host = config().get_string("remote_host")?;
let port = config().get::<u16>("remote_port")?;
let resolver = TokioAsyncResolver::tokio(
@ -133,52 +100,10 @@ async fn main() -> anyhow::Result<()> {
);
let response = resolver.lookup_ip(&host).await?;
let ip_addr = response.iter().next().expect("no addresses returned!");
info!("Remote host {} resolved to {:?}", host, ip_addr);
let server_sock = SocketAddr::new(ip_addr, port);
let client_sock = if server_sock.is_ipv6() {
SocketAddr::new(IpAddr::V6(Ipv6Addr::new(0,0,0,0,0,0,0,1)), 0)
} else {
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127,0,0,1)), 0)
};
info!("Client address: {:?}", client_sock);
let custom_certs = custom_certs_fut.await?;
let client_endpoint = quic::client::make_client_endpoint(client_sock.clone(), custom_certs)?;
let client = client_endpoint.connect(server_sock, &host)?.await?;
let ip_sock = SocketAddr::new(ip_addr, port);
let client_endpoint = quic::client::make_client_endpoint(ip_sock.clone(), None).unwrap();
let client = client_endpoint.connect(ip_sock, &host)?.await?;
println!("[client] connected: addr={}", client.remote_address());
loop {
let stream = client.open_bi().await;
let (mut send, mut recv) = match stream {
Err(quinn::ConnectionError::ApplicationClosed { .. }) => {
info!("connection closed");
return Ok(());
}
Err(e) => {
return bail!(e);
}
Ok(s) => s,
};
send.write_all("ping".as_bytes()).await.expect("TODO: panic message");
send.finish()?;
let response_start = Instant::now();
let resp = recv
.read_to_end(usize::MAX)
.await
.map_err(|e| anyhow!("failed to read response: {}", e))?;
let duration = response_start.elapsed();
eprintln!(
"response received in {:?} - {} KiB/s",
duration,
resp.len() as f32 / (duration_secs(&duration) * 1024.0)
);
std::io::stdout().write_all(&resp).unwrap();
std::io::stdout().flush().unwrap();
client.close(0u32.into(), b"done");
// Give the server a fair chance to receive the close packet
client_endpoint.wait_idle().await;
}
}
Commands::Shadow => {
@ -206,120 +131,9 @@ async fn main() -> anyhow::Result<()> {
}
}
}
Commands::Server {
keylog, key, cert, stateless_retry, listen, block, connection_limit
} => {
let (certs, key) = if let (Some(key_path), Some(cert_path)) = (&key, &cert) {
let key = std::fs::read(key_path).context("failed to read private key")?;
let key = if key_path.extension().is_some_and(|x| x == "der") {
PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(key))
} else {
rustls_pemfile::private_key(&mut &*key)
.context("malformed PKCS #1 private key")?
.ok_or_else(|| anyhow::Error::msg("no private keys found"))?
};
let cert_chain = std::fs::read(cert_path).context("failed to read certificate chain")?;
let cert_chain = if cert_path.extension().is_some_and(|x| x == "der") {
vec![CertificateDer::from(cert_chain)]
} else {
rustls_pemfile::certs(&mut &*cert_chain)
.collect::<Result<_, _>>()
.context("invalid PEM-encoded certificate")?
};
(cert_chain, key)
} else {
let default_dir = dirs::data_dir().unwrap();
let path = default_dir.join("hai-server");
let cert_path = path.join("cert.der");
let key_path = path.join("key.der");
let (cert, key) = match std::fs::read(&cert_path).and_then(|x| Ok((x, std::fs::read(&key_path)?))) {
Ok((cert, key)) => (
CertificateDer::from(cert),
PrivateKeyDer::try_from(key).map_err(anyhow::Error::msg)?,
),
Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
info!("generating self-signed certificate");
let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
let key = PrivatePkcs8KeyDer::from(cert.key_pair.serialize_der());
let cert = cert.cert.into();
std::fs::create_dir_all(path).context("failed to create certificate directory")?;
std::fs::write(&cert_path, &cert).context("failed to write certificate")?;
std::fs::write(&key_path, key.secret_pkcs8_der())
.context("failed to write private key")?;
(cert, key.into())
}
Err(e) => {
bail!("failed to read certificate: {}", e);
}
};
(vec![cert], key)
};
let mut server_crypto = rustls::ServerConfig::builder()
//TODO change to auth
.with_no_client_auth()
.with_single_cert(certs, key)?;
if keylog {
server_crypto.key_log = Arc::new(rustls::KeyLogFile::new());
}
let mut server_config =
quinn::ServerConfig::with_crypto(Arc::new(QuicServerConfig::try_from(server_crypto)?));
let transport_config = Arc::get_mut(&mut server_config.transport).unwrap();
transport_config.max_concurrent_uni_streams(0_u8.into());
let endpoint = quinn::Endpoint::server(server_config, listen)?;
eprintln!("listening on {}", endpoint.local_addr().unwrap());
while let Some(conn) = endpoint.accept().await {
if connection_limit
.is_some_and(|n| endpoint.open_connections() >= n)
{
info!("refusing due to open connection limit");
conn.refuse();
} else if Some(conn.remote_address()) == block {
info!("refusing blocked client IP address");
conn.refuse();
} else if stateless_retry && !conn.remote_address_validated() {
info!("requiring connection to validate its address");
conn.retry().unwrap();
} else {
info!("accepting connection");
let fut = handle_connection(conn);
tokio::spawn(async move {
if let Err(e) = fut.await {
error!("connection failed: {reason}", reason = e.to_string())
}
});
}
}
}
};
//handling anything here for gui wont work
println!("exit");
Ok(())
}
async fn read_opt_server_cert() -> anyhow::Result<Option<Vec<u8>>> {
let cert_path_res = config().get_string("remote_cert");
if cert_path_res.is_err() {
return Ok(None);
}
let cert_path = cert_path_res?;
let res = tokio::fs::read(cert_path).await;
if let Err(e) = res {
error!("failed to read certificate: {reason}", reason = e.to_string());
Err(Error::from(e))
} else {
info!("read certificate successfully read");
Ok(Some(res?))
}
}
fn duration_secs(x: &Duration) -> f32 {
x.as_secs() as f32 + x.subsec_nanos() as f32 * 1e-9
}

39
src/quic/client.rs
View file

@ -1,8 +1,6 @@
use std::error::Error;
use std::net::SocketAddr;
use std::sync::Arc;
use anyhow::anyhow;
use log::info;
use quinn::{ClientConfig, Endpoint};
use rustls::pki_types::CertificateDer;
@ -12,11 +10,13 @@ use rustls::pki_types::CertificateDer;
///
/// - server_certs: a list of trusted certificates in DER format.
fn configure_client(
server_certs: Option<Vec<u8>>,
) -> anyhow::Result<ClientConfig> {
server_certs: Option<&[&[u8]]>,
) -> Result<ClientConfig, Box<dyn Error + Send + Sync + 'static>> {
if let Some(server_certs) = server_certs {
let mut certs = rustls::RootCertStore::empty();
certs.add(CertificateDer::from(server_certs.as_slice()))?;
for cert in server_certs {
certs.add(CertificateDer::from(*cert))?;
}
Ok(ClientConfig::with_root_certificates(Arc::new(certs))?)
} else {
@ -32,35 +32,10 @@ fn configure_client(
#[allow(unused)]
pub fn make_client_endpoint(
bind_addr: SocketAddr,
server_certs: Option<Vec<u8>>,
) -> anyhow::Result<Endpoint> {
server_certs: Option<&[&[u8]]>,
) -> Result<Endpoint, Box<dyn Error + Send + Sync + 'static>> {
let client_cfg = configure_client(server_certs)?;
let mut endpoint = Endpoint::client(bind_addr)?;
endpoint.set_default_client_config(client_cfg);
Ok(endpoint)
}
pub(crate) async fn handle_request(
(mut send, mut recv): (quinn::SendStream, quinn::RecvStream),
) -> anyhow::Result<()> {
let req = recv
.read_to_end(64 * 1024)
.await
.map_err(|e| anyhow!("failed reading request: {}", e))?;
let mut escaped = String::new();
for &x in &req[..] {
let part = std::ascii::escape_default(x).collect::<Vec<_>>();
escaped.push_str(std::str::from_utf8(&part)?);
}
tracing::info!(content = %escaped);
// Execute the request
let resp = "ping {}";
// Write the response
send.write_all(&resp.as_bytes())
.await
.map_err(|e| anyhow!("failed to send response: {}", e))?;
// Gracefully terminate the stream
send.finish()?;
info!("complete");
Ok(())
}

73
src/quic/server.rs
View file

@ -1,13 +1,9 @@
use std::error::Error;
use std::fmt::Debug;
use std::net::SocketAddr;
use std::sync::Arc;
use anyhow::anyhow;
use log::{error, info};
use quinn::{Endpoint, ServerConfig};
use rustls::pki_types::{CertificateDer, PrivatePkcs8KeyDer};
use rustls::pki_types::pem::PemObject;
use tracing::{info_span, Instrument};
/// Constructs a QUIC endpoint configured to listen for incoming connections on a certain address
/// and port.
@ -49,73 +45,4 @@ fn configure_server(
transport_config.max_concurrent_uni_streams(0_u8.into());
Ok((server_config, cert_der))
}
pub async fn handle_connection(conn: quinn::Incoming) -> anyhow::Result<()> {
let connection = conn.await?;
let span = info_span!(
"connection",
remote = %connection.remote_address(),
protocol = %connection
.handshake_data()
.unwrap()
.downcast::<quinn::crypto::rustls::HandshakeData>().unwrap()
.protocol
.map_or_else(|| "<none>".into(), |x| String::from_utf8_lossy(&x).into_owned())
);
async {
info!("established");
// Each stream initiated by the client constitutes a new request.
loop {
let stream = connection.accept_bi().await;
let stream = match stream {
Err(quinn::ConnectionError::ApplicationClosed { .. }) => {
info!("connection closed");
return Ok(());
}
Err(e) => {
return Err(e);
}
Ok(s) => s,
};
let fut = handle_request(stream);
tokio::spawn(
async move {
if let Err(e) = fut.await {
error!("failed: {reason}", reason = e.to_string());
}
}
.instrument(info_span!("request")),
);
}
}
.instrument(span)
.await?;
Ok(())
}
pub(crate) async fn handle_request(
(mut send, mut recv): (quinn::SendStream, quinn::RecvStream),
) -> anyhow::Result<()> {
let req = recv
.read_to_end(64 * 1024)
.await
.map_err(|e| anyhow!("failed reading request: {}", e))?;
let mut escaped = String::new();
for &x in &req[..] {
let part = std::ascii::escape_default(x).collect::<Vec<_>>();
escaped.push_str(std::str::from_utf8(&part)?);
}
tracing::info!(content = %escaped);
// Execute the request
let resp = "Hello World {}";
// Write the response
send.write_all(&resp.as_bytes())
.await
.map_err(|e| anyhow!("failed to send response: {}", e))?;
// Gracefully terminate the stream
send.finish()?;
info!("complete");
Ok(())
}