-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.rs
134 lines (114 loc) · 4.63 KB
/
main.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
// axum mTLS example based on https://github.com/tokio-rs/axum/blob/main/examples/low-level-rustls/src/main.rs
use axum::{extract::ConnectInfo, routing::get, Router};
use axum::http::HeaderValue;
use futures_util::future::poll_fn;
use hyper::server::{
accept::Accept,
conn::{AddrIncoming, Http},
};
use rustls_pemfile::{certs, pkcs8_private_keys};
use std::{env, fs::File, io::BufReader, net::SocketAddr, pin::Pin, sync::Arc};
use axum::http::Request;
use axum::body::Body;
use axum::response::IntoResponse;
use tokio::net::TcpListener;
use tokio_rustls::{
rustls,
rustls::{Certificate, PrivateKey, ServerConfig},
TlsAcceptor,
};
use tower::MakeService;
use rustls_pemfile::Item;
#[tokio::main]
async fn main() {
// Set the RUST_LOG, if it hasn't been explicitly defined
if std::env::var_os("RUST_LOG").is_none() {
std::env::set_var("RUST_LOG", "rust_mtls_axum_example=debug")
}
tracing_subscriber::fmt::init();
let listen_addr = "127.0.0.1:3000";
let cert = "./tls/server.pem";
let key = "./tls/server-key-pkcs8.pem";
let mut ca = Some("./tls/ca.pem");
let mut args = env::args();
// skip self
args.next();
if let Some(flag) = args.next() {
if flag == "--no-mtls" {
ca = None
}
}
let rustls_config = build_rustls_server_config(cert, key, ca).await;
let acceptor = TlsAcceptor::from(rustls_config);
let listener = TcpListener::bind(listen_addr).await.unwrap();
let mut listener = AddrIncoming::from_listener(listener).unwrap();
let mut app = Router::new()
.route("/", get(handler))
.into_make_service_with_connect_info::<SocketAddr, _>();
tracing::info!("listening on https://{}", listen_addr);
loop {
let stream = poll_fn(|cx| Pin::new(&mut listener).poll_accept(cx))
.await
.unwrap()
.unwrap();
let acceptor = acceptor.clone();
let app = app.make_service(&stream).await.unwrap();
tokio::spawn(async move {
if let Ok(stream) = acceptor.accept(stream).await {
let _ = Http::new().serve_connection(stream, app).await;
}
});
}
}
async fn handler(ConnectInfo(addr): ConnectInfo<SocketAddr>, request: Request<Body>) -> impl IntoResponse {
let ConnectInfo(addr) = request.extensions().get::<ConnectInfo<SocketAddr>>().unwrap();
let empty_ua = HeaderValue::from_static("");
let user_agent = request.headers().get("User-Agent").or_else(||{Some(&empty_ua)}).unwrap().to_str().unwrap();
tracing::info!("remote_addr={} method={} uri={} user_agent={}", addr, request.method(), request.uri(), user_agent);
addr.to_string()
}
async fn build_rustls_server_config(cert: &str, key: &str, ca: Option<&str>) -> Arc<ServerConfig> {
let cert = tokio::fs::read(cert).await.unwrap();
let key = tokio::fs::read(key).await.unwrap();
// get pem from file
let cert = rustls_pemfile::certs(&mut cert.as_ref()).unwrap();
let key = match rustls_pemfile::read_one(&mut key.as_ref()).unwrap() {
Some(Item::RSAKey(key)) | Some(Item::PKCS8Key(key)) => key,
// rustls only support PKCS8, does not support ECC private key
_ => panic!("private key invalid or not supported"),
};
let cert = cert.into_iter().map(rustls::Certificate).collect();
let key = rustls::PrivateKey(key);
let config_builder = rustls::ServerConfig::builder().with_safe_defaults();
let mut server_config = match ca {
None => {
tracing::info!("mTLS disabled");
config_builder
.with_no_client_auth()
.with_single_cert(cert, key)
.expect("bad certificate/key")
},
Some(ca) => {
tracing::info!("mTLS enabled, ca cert path={}", ca);
let ca = tokio::fs::read(ca).await.unwrap();
if let Some(Item::X509Certificate(ca)) =
rustls_pemfile::read_one(&mut ca.as_ref()).unwrap()
{
let mut root_cert_store = rustls::RootCertStore::empty();
root_cert_store
.add(&rustls::Certificate(ca))
.expect("bad ca cert");
config_builder
.with_client_cert_verifier(rustls::server::AllowAnyAuthenticatedClient::new(
root_cert_store,
))
.with_single_cert(cert, key)
.expect("bad certificate/key")
} else {
panic!("invalid root ca cert")
}
}
};
server_config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
Arc::new(server_config)
}