1use std::io;
13use std::net::SocketAddr;
14use std::path::{Path, PathBuf};
15use std::pin::Pin;
16use std::sync::{Arc, Mutex, RwLock};
17use std::task::{Context, Poll};
18use std::time::{Duration, SystemTime};
19
20use rustls::pki_types::pem::PemObject;
21use rustls::pki_types::{CertificateDer, PrivateKeyDer};
22use rustls::server::WebPkiClientVerifier;
23use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
24use tokio::net::{TcpListener, TcpStream};
25use tokio::sync::mpsc;
26use tokio_rustls::TlsAcceptor;
27use tokio_rustls::server::TlsStream;
28use tracing::{debug, info, warn};
29
30pub use crate::api::TlsSettings;
31
32pub const RELOAD_INTERVAL: Duration = Duration::from_secs(10);
34const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
36const ACCEPT_QUEUE: usize = 256;
38pub const MAX_HANDSHAKES: usize = 1024;
44
45#[derive(Debug)]
49#[non_exhaustive]
50pub enum Error {
51 Certificate {
53 path: PathBuf,
55 source: BoxError,
57 },
58 NoCertificate {
60 path: PathBuf,
62 },
63 PrivateKey {
65 path: PathBuf,
67 source: BoxError,
69 },
70 ClientCa {
72 path: PathBuf,
74 source: BoxError,
76 },
77 CertificateAndKey {
80 cert: PathBuf,
82 key: PathBuf,
84 source: BoxError,
86 },
87 Setup(BoxError),
89}
90
91impl std::fmt::Display for Error {
92 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
93 match self {
94 Self::Certificate { path, source } | Self::ClientCa { path, source } => {
95 write!(f, "{}: {source}", path.display())
96 }
97 Self::NoCertificate { path } => write!(f, "{}: no PEM certificate", path.display()),
98 Self::PrivateKey { path, source } => {
99 write!(f, "{}: no PEM private key ({source})", path.display())
100 }
101 Self::CertificateAndKey { cert, key, source } => {
102 write!(f, "{} and {}: {source}", cert.display(), key.display())
103 }
104 Self::Setup(source) => write!(f, "TLS setup: {source}"),
105 }
106 }
107}
108
109impl std::error::Error for Error {
110 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
111 match self {
112 Self::Certificate { source, .. }
113 | Self::PrivateKey { source, .. }
114 | Self::ClientCa { source, .. }
115 | Self::CertificateAndKey { source, .. }
116 | Self::Setup(source) => Some(&**source),
117 Self::NoCertificate { .. } => None,
118 }
119 }
120}
121
122pub type BoxError = Box<dyn std::error::Error + Send + Sync>;
124
125#[derive(Debug, Clone, Copy, PartialEq)]
127pub enum Protocols {
128 Http,
130 Grpc,
132}
133
134pub struct TlsConfigs {
136 settings: TlsSettings,
137 http: RwLock<Arc<rustls::ServerConfig>>,
138 grpc: RwLock<Arc<rustls::ServerConfig>>,
139 loaded: Mutex<Option<SystemTime>>,
141}
142
143impl TlsConfigs {
144 pub fn load(settings: TlsSettings) -> Result<Arc<Self>, Error> {
146 let (http, grpc) = build(&settings)?;
147 let loaded = newest_change(&settings);
148 Ok(Arc::new(Self {
149 settings,
150 http: RwLock::new(http),
151 grpc: RwLock::new(grpc),
152 loaded: Mutex::new(loaded),
153 }))
154 }
155
156 pub fn mutual(&self) -> bool {
158 self.settings.client_ca.is_some()
159 }
160
161 fn current(&self, protocols: Protocols) -> Arc<rustls::ServerConfig> {
162 let config = match protocols {
163 Protocols::Http => &self.http,
164 Protocols::Grpc => &self.grpc,
165 };
166 config.read().unwrap_or_else(|e| e.into_inner()).clone()
167 }
168
169 pub fn reload_if_changed(&self) -> Result<bool, Error> {
172 let mut loaded = self.loaded.lock().unwrap_or_else(|e| e.into_inner());
175 let changed = newest_change(&self.settings);
176 if changed == *loaded {
177 return Ok(false);
178 }
179 *loaded = changed;
181 let (http, grpc) = build(&self.settings)?;
182 *self.http.write().unwrap_or_else(|e| e.into_inner()) = http;
183 *self.grpc.write().unwrap_or_else(|e| e.into_inner()) = grpc;
184 Ok(true)
185 }
186
187 pub fn watch(self: &Arc<Self>) {
189 let configs = Arc::downgrade(self);
190 tokio::spawn(async move {
191 let mut interval = tokio::time::interval(RELOAD_INTERVAL);
192 interval.tick().await;
193 loop {
194 interval.tick().await;
195 let Some(configs) = configs.upgrade() else {
196 return;
197 };
198 match configs.reload_if_changed() {
199 Ok(true) => info!("🔐 Reloaded the TLS certificate"),
200 Ok(false) => {}
201 Err(e) => warn!("Keeping the previous TLS certificate: {e}"),
202 }
203 }
204 });
205 }
206}
207
208fn newest_change(settings: &TlsSettings) -> Option<SystemTime> {
210 let mut newest = None;
211 for path in [&settings.cert, &settings.key]
212 .into_iter()
213 .chain(settings.client_ca.as_ref())
214 {
215 let modified = std::fs::metadata(path).and_then(|m| m.modified()).ok()?;
216 newest = newest.max(Some(modified));
217 }
218 newest
219}
220
221fn build(
222 settings: &TlsSettings,
223) -> Result<(Arc<rustls::ServerConfig>, Arc<rustls::ServerConfig>), Error> {
224 let certs = read_certs(&settings.cert)?;
225 let key = PrivateKeyDer::from_pem_file(&settings.key).map_err(|e| Error::PrivateKey {
226 path: settings.key.clone(),
227 source: e.into(),
228 })?;
229 let provider = Arc::new(rustls_pure::provider());
230 let builder = rustls::ServerConfig::builder_with_provider(provider.clone())
231 .with_protocol_versions(&[&rustls::version::TLS13])
232 .map_err(|e| Error::Setup(e.into()))?;
233 let builder = match &settings.client_ca {
234 None => builder.with_no_client_auth(),
235 Some(ca) => {
236 let client_ca = |e: BoxError| Error::ClientCa {
237 path: ca.clone(),
238 source: e,
239 };
240 let mut roots = rustls::RootCertStore::empty();
241 for cert in read_certs(ca)? {
242 roots.add(cert).map_err(|e| client_ca(e.into()))?;
243 }
244 let verifier = WebPkiClientVerifier::builder_with_provider(Arc::new(roots), provider)
245 .build()
246 .map_err(|e| client_ca(e.into()))?;
247 builder.with_client_cert_verifier(verifier)
248 }
249 };
250 let config = builder
251 .with_single_cert(certs, key)
252 .map_err(|e| Error::CertificateAndKey {
253 cert: settings.cert.clone(),
254 key: settings.key.clone(),
255 source: e.into(),
256 })?;
257 let mut http = config.clone();
258 http.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
259 let mut grpc = config;
260 grpc.alpn_protocols = vec![b"h2".to_vec()];
261 Ok((Arc::new(http), Arc::new(grpc)))
262}
263
264fn read_certs(path: &Path) -> Result<Vec<CertificateDer<'static>>, Error> {
265 let certs = CertificateDer::pem_file_iter(path)
266 .and_then(|certs| certs.collect::<Result<Vec<_>, _>>())
267 .map_err(|e| Error::Certificate {
268 path: path.to_path_buf(),
269 source: e.into(),
270 })?;
271 if certs.is_empty() {
272 return Err(Error::NoCertificate {
273 path: path.to_path_buf(),
274 });
275 }
276 Ok(certs)
277}
278
279pub struct TlsConnection {
281 stream: TlsStream<TcpStream>,
282 local: Option<SocketAddr>,
283 peer: SocketAddr,
284}
285
286impl TlsConnection {
287 pub fn peer(&self) -> SocketAddr {
289 self.peer
290 }
291
292 pub fn local(&self) -> Option<SocketAddr> {
294 self.local
295 }
296}
297
298impl AsyncRead for TlsConnection {
299 fn poll_read(
300 mut self: Pin<&mut Self>,
301 cx: &mut Context<'_>,
302 buf: &mut ReadBuf<'_>,
303 ) -> Poll<io::Result<()>> {
304 Pin::new(&mut self.stream).poll_read(cx, buf)
305 }
306}
307
308impl AsyncWrite for TlsConnection {
309 fn poll_write(
310 mut self: Pin<&mut Self>,
311 cx: &mut Context<'_>,
312 buf: &[u8],
313 ) -> Poll<io::Result<usize>> {
314 Pin::new(&mut self.stream).poll_write(cx, buf)
315 }
316
317 fn poll_write_vectored(
318 mut self: Pin<&mut Self>,
319 cx: &mut Context<'_>,
320 bufs: &[io::IoSlice<'_>],
321 ) -> Poll<io::Result<usize>> {
322 Pin::new(&mut self.stream).poll_write_vectored(cx, bufs)
323 }
324
325 fn is_write_vectored(&self) -> bool {
326 self.stream.is_write_vectored()
327 }
328
329 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
330 Pin::new(&mut self.stream).poll_flush(cx)
331 }
332
333 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
334 Pin::new(&mut self.stream).poll_shutdown(cx)
335 }
336}
337
338#[cfg(feature = "grpc")]
339impl tonic::transport::server::Connected for TlsConnection {
340 type ConnectInfo = tonic::transport::server::TcpConnectInfo;
341
342 fn connect_info(&self) -> Self::ConnectInfo {
343 tonic::transport::server::TcpConnectInfo {
344 local_addr: self.local(),
345 remote_addr: Some(self.peer()),
346 }
347 }
348}
349
350pub fn accept(
354 listener: TcpListener,
355 configs: Arc<TlsConfigs>,
356 protocols: Protocols,
357) -> mpsc::Receiver<TlsConnection> {
358 accept_bounded(listener, configs, protocols, MAX_HANDSHAKES)
359}
360
361fn accept_bounded(
362 listener: TcpListener,
363 configs: Arc<TlsConfigs>,
364 protocols: Protocols,
365 max_handshakes: usize,
366) -> mpsc::Receiver<TlsConnection> {
367 let (tx, rx) = mpsc::channel(ACCEPT_QUEUE);
368 let local = listener.local_addr().ok();
369 let handshakes = Arc::new(tokio::sync::Semaphore::new(max_handshakes));
370 tokio::spawn(async move {
371 loop {
372 let handshake = tokio::select! {
374 permit = handshakes.clone().acquire_owned() => match permit {
375 Ok(permit) => permit,
376 Err(_) => return,
377 },
378 () = tx.closed() => return,
379 };
380 let (tcp, peer) = tokio::select! {
381 accepted = listener.accept() => match accepted {
382 Ok(accepted) => accepted,
383 Err(e) => {
384 warn!("Accepting a connection failed: {e}");
386 tokio::time::sleep(Duration::from_millis(50)).await;
387 continue;
388 }
389 },
390 () = tx.closed() => return,
391 };
392 let _ = tcp.set_nodelay(true);
393 let acceptor = TlsAcceptor::from(configs.current(protocols));
394 let tx = tx.clone();
395 tokio::spawn(async move {
396 let handshaken =
397 tokio::time::timeout(HANDSHAKE_TIMEOUT, acceptor.accept(tcp)).await;
398 drop(handshake);
399 match handshaken {
400 Ok(Ok(stream)) => {
401 let _ = tx
402 .send(TlsConnection {
403 stream,
404 local,
405 peer,
406 })
407 .await;
408 }
409 Ok(Err(e)) => debug!("TLS handshake with {peer} failed: {e}"),
410 Err(_) => debug!("TLS handshake with {peer} timed out"),
411 }
412 });
413 }
414 });
415 rx
416}
417
418pub struct TlsListener {
420 connections: mpsc::Receiver<TlsConnection>,
421 local: SocketAddr,
422}
423
424impl TlsListener {
425 pub fn new(listener: TcpListener, configs: Arc<TlsConfigs>) -> io::Result<Self> {
427 let local = listener.local_addr()?;
428 Ok(Self {
429 connections: accept(listener, configs, Protocols::Http),
430 local,
431 })
432 }
433}
434
435impl axum::serve::Listener for TlsListener {
436 type Io = TlsConnection;
437 type Addr = SocketAddr;
438
439 async fn accept(&mut self) -> (Self::Io, Self::Addr) {
440 match self.connections.recv().await {
441 Some(connection) => {
442 let peer = connection.peer;
443 (connection, peer)
444 }
445 None => std::future::pending().await,
447 }
448 }
449
450 fn local_addr(&self) -> io::Result<Self::Addr> {
451 Ok(self.local)
452 }
453}
454
455pub fn describe(settings: &TlsSettings) -> String {
457 let mut text = format!("{}", settings.cert.display());
458 if let Some(ca) = &settings.client_ca {
459 text.push_str(&format!(", clients verified against {}", ca.display()));
460 }
461 text
462}
463
464#[cfg(test)]
465pub(crate) mod tests {
466 use super::*;
467 use crate::api::tests::router;
468 use crate::api::{ClientAddr, ServerConfig};
469 use crate::model::{DummyModel, Model};
470 use std::path::PathBuf;
471 use tokio::io::{AsyncReadExt, AsyncWriteExt};
472
473 const FIXTURES: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/tls");
474
475 fn fixture(name: &str) -> PathBuf {
476 Path::new(FIXTURES).join(name)
477 }
478
479 fn fixture_cert(name: &str) -> CertificateDer<'static> {
482 match read_certs(&fixture(name)) {
483 Ok(mut certs) if !certs.is_empty() => certs.remove(0),
484 _ => panic!("no certificate in the fixture {name}"),
485 }
486 }
487
488 pub(crate) fn settings(server: &str, client_ca: bool) -> TlsSettings {
489 TlsSettings {
490 cert: fixture(&format!("{server}.crt")),
491 key: fixture(&format!("{server}.key")),
492 client_ca: client_ca.then(|| fixture("ca.crt")),
493 }
494 }
495
496 async fn serve(configs: Arc<TlsConfigs>) -> SocketAddr {
498 let model = Model::new(
499 Box::new(DummyModel::new("test_model".into(), vec![4], vec![2])),
500 "dummy",
501 0,
502 );
503 let app = router(model, &ServerConfig::default()).unwrap();
504 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
505 let addr = listener.local_addr().unwrap();
506 let listener = TlsListener::new(listener, configs).unwrap();
507 tokio::spawn(async move {
508 axum::serve(
509 listener,
510 app.into_make_service_with_connect_info::<ClientAddr>(),
511 )
512 .await
513 });
514 addr
515 }
516
517 fn client(certificate: Option<&str>) -> tokio_rustls::TlsConnector {
519 tokio_rustls::TlsConnector::from(Arc::new(client_config(certificate, b"http/1.1")))
520 }
521
522 pub(crate) fn client_config(certificate: Option<&str>, alpn: &[u8]) -> rustls::ClientConfig {
524 let mut roots = rustls::RootCertStore::empty();
525 roots.add(fixture_cert("ca.crt")).unwrap();
526 let builder =
527 rustls::ClientConfig::builder_with_provider(Arc::new(rustls_pure::provider()))
528 .with_protocol_versions(&[&rustls::version::TLS13])
529 .unwrap()
530 .with_root_certificates(roots);
531 let mut config = match certificate {
532 None => builder.with_no_client_auth(),
533 Some(name) => builder
534 .with_client_auth_cert(
535 read_certs(&fixture(&format!("{name}.crt"))).unwrap(),
536 PrivateKeyDer::from_pem_file(fixture(&format!("{name}.key"))).unwrap(),
537 )
538 .unwrap(),
539 };
540 config.alpn_protocols = vec![alpn.to_vec()];
541 config
542 }
543
544 async fn get_healthz(
546 addr: SocketAddr,
547 connector: &tokio_rustls::TlsConnector,
548 ) -> io::Result<(String, CertificateDer<'static>)> {
549 let tcp = TcpStream::connect(addr).await?;
550 let name = rustls::pki_types::ServerName::try_from("localhost").unwrap();
551 let mut tls = connector.connect(name, tcp).await?;
552 let certificate = tls
553 .get_ref()
554 .1
555 .peer_certificates()
556 .and_then(|chain| chain.first())
557 .cloned()
558 .ok_or_else(|| io::Error::other("the server sent no certificate"))?;
559 tls.write_all(b"GET /healthz HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n")
560 .await?;
561 let mut response = String::new();
562 tls.read_to_string(&mut response).await?;
563 let status = response.lines().next().unwrap_or_default().to_string();
564 Ok((status, certificate))
565 }
566
567 #[tokio::test]
568 async fn https_with_each_key_type() {
569 for server in ["server-p256", "server-p384", "server-ed25519"] {
570 let addr = serve(TlsConfigs::load(settings(server, false)).unwrap()).await;
571 let (status, certificate) = get_healthz(addr, &client(None)).await.unwrap();
572 assert_eq!(status, "HTTP/1.1 200 OK", "{server}");
573 let expected = fixture_cert(&format!("{server}.crt"));
574 assert!(
575 certificate == expected,
576 "{server}: another certificate was served"
577 );
578 }
579 }
580
581 #[tokio::test]
582 async fn handshakes_in_flight_are_bounded() {
583 let configs = TlsConfigs::load(settings("server-p256", false)).unwrap();
584 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
585 let addr = listener.local_addr().unwrap();
586 let mut connections = accept_bounded(listener, configs, Protocols::Http, 1);
587 let silent = TcpStream::connect(addr).await.unwrap();
589 let handshake = tokio::spawn(async move {
590 let tcp = TcpStream::connect(addr).await?;
591 let name = rustls::pki_types::ServerName::try_from("localhost").unwrap();
592 client(None).connect(name, tcp).await
593 });
594 tokio::time::sleep(Duration::from_millis(100)).await;
596 assert!(!handshake.is_finished());
597 drop(silent);
599 let handshaken = tokio::time::timeout(Duration::from_secs(10), handshake)
600 .await
601 .expect("the handshake goes through")
602 .unwrap();
603 assert!(handshaken.is_ok(), "{handshaken:?}");
604 assert!(connections.recv().await.is_some());
605 }
606
607 #[tokio::test]
608 async fn plain_http_gets_no_answer() {
609 let addr = serve(TlsConfigs::load(settings("server-p256", false)).unwrap()).await;
610 let mut tcp = TcpStream::connect(addr).await.unwrap();
611 tcp.write_all(b"GET /healthz HTTP/1.1\r\nHost: localhost\r\n\r\n")
612 .await
613 .unwrap();
614 let mut response = Vec::new();
615 let _ = tcp.read_to_end(&mut response).await;
616 assert!(!String::from_utf8_lossy(&response).contains("200 OK"));
617 }
618
619 #[tokio::test]
620 async fn mutual_tls_needs_a_certificate_from_the_ca() {
621 let configs = TlsConfigs::load(settings("server-p256", true)).unwrap();
622 assert!(configs.mutual());
623 let addr = serve(configs).await;
624 for client_name in ["client", "client-p384", "client-ed25519"] {
626 let (status, _) = get_healthz(addr, &client(Some(client_name))).await.unwrap();
627 assert_eq!(status, "HTTP/1.1 200 OK", "{client_name}");
628 }
629 for (client_name, presented) in [("no certificate", None), ("stranger", Some("stranger"))] {
633 let status = get_healthz(addr, &client(presented))
634 .await
635 .map(|(status, _)| status);
636 assert!(
637 !matches!(&status, Ok(status) if status.contains("200")),
638 "{client_name}: {status:?}"
639 );
640 }
641 }
642
643 #[tokio::test]
644 async fn certificates_are_reloaded_and_bad_files_keep_the_old_one() {
645 let dir = tempfile::tempdir().unwrap();
646 let (cert, key) = (dir.path().join("tls.crt"), dir.path().join("tls.key"));
647 let mut stamp = SystemTime::now();
649 let mut touch = |path: &Path| {
650 stamp += Duration::from_secs(60);
651 let file = std::fs::File::options().write(true).open(path).unwrap();
652 file.set_modified(stamp).unwrap();
653 };
654 let mut install = |server: &str| {
655 std::fs::copy(fixture(&format!("{server}.crt")), &cert).unwrap();
656 std::fs::copy(fixture(&format!("{server}.key")), &key).unwrap();
657 touch(&cert);
658 touch(&key);
659 };
660 install("server-p256");
661 let configs = TlsConfigs::load(TlsSettings {
662 cert: cert.clone(),
663 key: key.clone(),
664 client_ca: None,
665 })
666 .unwrap();
667 let addr = serve(configs.clone()).await;
668 let p256 = fixture_cert("server-p256.crt");
669 let p384 = fixture_cert("server-p384.crt");
670 let served = || async { get_healthz(addr, &client(None)).await.unwrap().1 };
671 assert!(served().await == p256, "the first certificate is served");
672
673 assert!(matches!(configs.reload_if_changed(), Ok(false)));
674 install("server-p384");
675 assert!(matches!(configs.reload_if_changed(), Ok(true)));
676 assert!(served().await == p384, "the renewed certificate is served");
677
678 std::fs::write(&key, "not a key").unwrap();
680 let file = std::fs::File::options().write(true).open(&key).unwrap();
681 file.set_modified(SystemTime::now() + Duration::from_secs(3600))
682 .unwrap();
683 assert!(matches!(
684 configs.reload_if_changed(),
685 Err(Error::PrivateKey { .. })
686 ));
687 assert!(matches!(configs.reload_if_changed(), Ok(false)));
688 assert!(served().await == p384, "the renewed certificate stays");
689 }
690
691 #[test]
692 fn key_formats() {
693 let load = |key: &str| {
694 build(&TlsSettings {
695 cert: fixture("server-p256.crt"),
696 key: fixture(key),
697 client_ca: None,
698 })
699 };
700 assert!(load("server-p256.key").is_ok());
701 assert!(load("server-p256-sec1.key").is_ok());
702 let error = build(&settings("server-rsa", false)).err().unwrap();
703 assert!(matches!(error, Error::CertificateAndKey { .. }), "{error}");
704 let error = error.to_string();
705 assert!(error.contains("RSA keys are not supported"), "{error}");
706 assert!(error.contains("ECDSA"), "{error}");
707 let error = load("server-p521.key").err().unwrap().to_string();
708 assert!(
710 error.contains("unsupported EC key curve 1.3.132.0.35"),
711 "{error}"
712 );
713
714 let missing = load("nope.key").err().unwrap().to_string();
715 assert!(missing.contains("nope.key"), "{missing}");
716 }
717}