Prioritize H2 via ALPN. (#30)

Author: UwUDev

Cargo.toml

@@ -41,6 +41,7 @@ hex = "0.4.3"
 nom = "8.0.0"
 
 tokio-rustls = { version = "0.26", default-features = false }
+rustls-pemfile = "2.0"
 futures-util = { version = "0.3.31", default-features = false }
 httlib-hpack = "0.1.3"
 pin-project-lite = "0.2.9"

src/server/certificate.rs

@@ -1,10 +1,94 @@
+use std::{io, path::Path, sync::Arc};
+
+use axum_server::tls_rustls::RustlsConfig;
 use rcgen::{
     date_time_ymd, BasicConstraints, CertificateParams, DistinguishedName, DnType, IsCa, KeyPair,
     KeyUsagePurpose, SanType,
 };
+use rustls_pemfile::Item;
+use tokio_rustls::rustls::{
+    pki_types::{CertificateDer, PrivateKeyDer},
+    ServerConfig,
+};
+
+// Load TLS configuration from self-signed PEM files
+pub async fn config_self_signed() -> crate::Result<RustlsConfig> {
+    let (cert, key) = get_self_signed_cert()?;
+    let cert = rustls_pemfile::certs(&mut cert.as_ref())
+        .map(|it| it.map(|it| it.to_vec()))
+        .collect::<Result<Vec<_>, _>>()?;
+
+    // Check the entire PEM file for the key in case it is not first section
+    let mut key_vec: Vec<Vec<u8>> = rustls_pemfile::read_all(&mut key.as_ref())
+        .filter_map(|i| match i.ok()? {
+            Item::Sec1Key(key) => Some(key.secret_sec1_der().to_vec()),
+            Item::Pkcs1Key(key) => Some(key.secret_pkcs1_der().to_vec()),
+            Item::Pkcs8Key(key) => Some(key.secret_pkcs8_der().to_vec()),
+            _ => None,
+        })
+        .collect();
+
+    // Make sure file contains only one key
+    if key_vec.len() != 1 {
+        return Err(io::Error::other("private key format not supported").into());
+    }
+
+    let cert = cert.into_iter().map(CertificateDer::from).collect();
+    let key = PrivateKeyDer::try_from(
+        key_vec
+            .pop()
+            .expect("private key should be present in the file"),
+    )
+    .map_err(io::Error::other)?;
+
+    Ok(config_from_der(cert, key)?)
+}
+
+/// Load TLS configuration from PEM files
+pub async fn config_from_pem_chain_file(
+    cert: impl AsRef<Path>,
+    chain: impl AsRef<Path>,
+) -> crate::Result<RustlsConfig> {
+    let cert = tokio::fs::read(cert.as_ref()).await?;
+    let cert = rustls_pemfile::certs(&mut cert.as_ref())
+        .map(|it| it.map(|it| CertificateDer::from(it.to_vec())))
+        .collect::<Result<Vec<_>, _>>()?;
+
+    let key = tokio::fs::read(chain.as_ref()).await?;
+    let key_cert: PrivateKeyDer = match rustls_pemfile::read_one(&mut key.as_ref())?
+        .ok_or_else(|| io::Error::other("could not parse pem file"))?
+    {
+        Item::Pkcs8Key(key) => Ok(key.into()),
+        Item::Sec1Key(key) => Ok(key.into()),
+        Item::Pkcs1Key(key) => Ok(key.into()),
+        x => Err(io::Error::other(format!(
+            "invalid certificate format, received: {x:?}"
+        ))),
+    }?;
+
+    Ok(config_from_der(cert, key_cert)?)
+}
+
+fn config_from_der(
+    cert_chain: Vec<CertificateDer<'static>>,
+    key_der: PrivateKeyDer<'static>,
+) -> io::Result<RustlsConfig> {
+    let mut config = ServerConfig::builder()
+        .with_no_client_auth()
+        .with_single_cert(cert_chain, key_der)
+        .map_err(io::Error::other)?;
+
+    config.alpn_protocols = vec![
+        b"h2".to_vec(),
+        b"http/1.1".to_vec(),
+        b"http/1.0".to_vec(),
+        b"http/0.9".to_vec(),
+    ];
+
+    Ok(RustlsConfig::from_config(Arc::new(config)))
+}
 
-/// Get self-signed certificate and key.
-pub fn get_self_signed_cert() -> crate::Result<(Vec<u8>, Vec<u8>)> {
+fn get_self_signed_cert() -> crate::Result<(Vec<u8>, Vec<u8>)> {
     let temp_dir = std::env::temp_dir().join(env!("CARGO_PKG_NAME"));
     if !temp_dir.exists() {
         tracing::info!("Creating temp cert directory: {}", temp_dir.display());
@@ -25,7 +109,6 @@ pub fn get_self_signed_cert() -> crate::Result<(Vec<u8>, Vec<u8>)> {
     Ok((cert, key))
 }
 
-/// Generate self-signed certificate and key.
 fn generate_self_signed() -> crate::Result<(Vec<u8>, Vec<u8>)> {
     let mut params = CertificateParams::default();
     params.not_before = date_time_ymd(1975, 1, 1);

src/server/mod.rs

@@ -13,7 +13,7 @@ use axum::{
     Extension, Router,
 };
 use axum_extra::response::ErasedJson;
-use axum_server::{tls_rustls::RustlsConfig, Handle};
+use axum_server::Handle;
 use tower::limit::ConcurrencyLimitLayer;
 use tower_http::{
     cors::{AllowHeaders, AllowMethods, AllowOrigin, CorsLayer},
@@ -74,17 +74,20 @@ pub async fn run(args: Args) -> Result<()> {
     // Spawn a task to gracefully shutdown server.
     tokio::spawn(signal::graceful_shutdown(handle.clone()));
 
-    // Load TLS configuration
-    let tls_config = match (args.tls_cert.as_ref(), args.tls_key.as_ref()) {
-        (Some(cert), Some(key)) => RustlsConfig::from_pem_chain_file(cert, key).await,
+    // Load TLS configuration with HTTP/2 ALPN preference
+    let config = match (args.tls_cert.as_ref(), args.tls_key.as_ref()) {
+        (Some(cert_path), Some(key_path)) => {
+            // Load TLS configuration from PEM files
+            certificate::config_from_pem_chain_file(cert_path, key_path).await?
+        }
         _ => {
-            let (cert, key) = certificate::get_self_signed_cert()?;
-            RustlsConfig::from_pem(cert, key).await
+            // Generate self-signed certificate configuration
+            certificate::config_self_signed().await?
         }
-    }?;
+    };
 
     // Use TLS configuration to create a secure server
-    let mut server = axum_server::bind_rustls(args.bind, tls_config);
+    let mut server = axum_server::bind_rustls(args.bind, config);
     server
         .http_builder()
         .http2()