tokio-rustls/examples/server.rs
2023-05-31 17:09:52 +02:00

105 lines
3.1 KiB
Rust

use argh::FromArgs;
use rustls_pemfile::{certs, rsa_private_keys};
use std::fs::File;
use std::io::{self, BufReader};
use std::net::ToSocketAddrs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::io::{copy, sink, split, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio_rustls::rustls::{self, Certificate, PrivateKey};
use tokio_rustls::TlsAcceptor;
/// Tokio Rustls server example
#[derive(FromArgs)]
struct Options {
/// bind addr
#[argh(positional)]
addr: String,
/// cert file
#[argh(option, short = 'c')]
cert: PathBuf,
/// key file
#[argh(option, short = 'k')]
key: PathBuf,
/// echo mode
#[argh(switch, short = 'e')]
echo_mode: bool,
}
fn load_certs(path: &Path) -> io::Result<Vec<Certificate>> {
certs(&mut BufReader::new(File::open(path)?))
.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid cert"))
.map(|mut certs| certs.drain(..).map(Certificate).collect())
}
fn load_keys(path: &Path) -> io::Result<Vec<PrivateKey>> {
rsa_private_keys(&mut BufReader::new(File::open(path)?))
.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid key"))
.map(|mut keys| keys.drain(..).map(PrivateKey).collect())
}
#[tokio::main]
async fn main() -> io::Result<()> {
let options: Options = argh::from_env();
let addr = options
.addr
.to_socket_addrs()?
.next()
.ok_or_else(|| io::Error::from(io::ErrorKind::AddrNotAvailable))?;
let certs = load_certs(&options.cert)?;
let mut keys = load_keys(&options.key)?;
let flag_echo = options.echo_mode;
let config = rustls::ServerConfig::builder()
.with_safe_defaults()
.with_no_client_auth()
.with_single_cert(certs, keys.remove(0))
.map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?;
let acceptor = TlsAcceptor::from(Arc::new(config));
let listener = TcpListener::bind(&addr).await?;
loop {
let (stream, peer_addr) = listener.accept().await?;
let acceptor = acceptor.clone();
let fut = async move {
let mut stream = acceptor.accept(stream).await?;
if flag_echo {
let (mut reader, mut writer) = split(stream);
let n = copy(&mut reader, &mut writer).await?;
writer.flush().await?;
println!("Echo: {} - {}", peer_addr, n);
} else {
let mut output = sink();
stream
.write_all(
&b"HTTP/1.0 200 ok\r\n\
Connection: close\r\n\
Content-length: 12\r\n\
\r\n\
Hello world!"[..],
)
.await?;
stream.shutdown().await?;
copy(&mut stream, &mut output).await?;
println!("Hello: {}", peer_addr);
}
Ok(()) as io::Result<()>
};
tokio::spawn(async move {
if let Err(err) = fut.await {
eprintln!("{:?}", err);
}
});
}
}