From 8701ae69dda4dad853ae532e195a26bf72412737 Mon Sep 17 00:00:00 2001
From: Mike Dilger <mike@mikedilger.com>
Date: Thu, 26 Oct 2023 19:41:10 +1300
Subject: [PATCH] TLS

---
 .gitignore               |  1 +
 sample/sample.config.ron |  2 +
 src/config.rs            |  4 ++
 src/error.rs             |  7 ++++
 src/main.rs              | 79 ++++++++++++++++++++++++++++++----------
 5 files changed, 74 insertions(+), 19 deletions(-)

diff --git a/.gitignore b/.gitignore
index ea8c4bf..c5d3465 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
 /target
+/tls/
diff --git a/sample/sample.config.ron b/sample/sample.config.ron
index 21b6b07..4a2ab8e 100644
--- a/sample/sample.config.ron
+++ b/sample/sample.config.ron
@@ -2,6 +2,8 @@ Config(
     data_directory: "./sample",
     ip_address: "127.0.0.1",
     port: 8080,
+    certchain_pem_path: "tls/fullchain.pem",
+    key_pem_path: "tls/privkey.pem",
     name: None,
     description: None,
     public_key_hex: None,
diff --git a/src/config.rs b/src/config.rs
index c79ea52..643381c 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -5,6 +5,8 @@ pub struct Config {
     pub data_directory: String,
     pub ip_address: String,
     pub port: u16,
+    pub certchain_pem_path: String,
+    pub key_pem_path: String,
     pub name: Option<String>,
     pub description: Option<String>,
     pub public_key_hex: Option<String>,
@@ -16,6 +18,8 @@ impl Default for Config {
             data_directory: "/tmp".to_string(),
             ip_address: "127.0.0.1".to_string(),
             port: 80,
+            certchain_pem_path: "./tls/fullchain.pem".to_string(),
+            key_pem_path: "./tls/privkey.pem".to_string(),
             name: None,
             description: None,
             public_key_hex: None,
diff --git a/src/error.rs b/src/error.rs
index e5e68f0..00d1d1a 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -67,6 +67,13 @@ pub enum Error {
     #[error("LMDB: {0}")]
     Lmdb(#[from] heed::Error),
 
+    #[error("Private Key Not Found")]
+    NoPrivateKey,
+
+    // Rustls
+    #[error("TLS: {0}")]
+    Rustls(#[from] tokio_rustls::rustls::Error),
+
     // Filter is underspecified
     #[error("Filter is underspecified. Scrapers are not allowed")]
     Scraper,
diff --git a/src/main.rs b/src/main.rs
index 44432f7..e3cc93a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -12,11 +12,14 @@ use crate::error::Error;
 use crate::globals::GLOBALS;
 use crate::store::Store;
 use hyper::{Body, Request, Response};
+use rustls::{Certificate, PrivateKey};
 use std::env;
 use std::error::Error as StdError;
-use std::fs::OpenOptions;
-use std::io::Read;
+use std::fs::{File, OpenOptions};
+use std::io::{BufReader, Read};
+use std::sync::Arc;
 use tokio::net::TcpListener;
+use tokio_rustls::{rustls, TlsAcceptor};
 
 #[tokio::main]
 async fn main() -> Result<(), Error> {
@@ -41,6 +44,35 @@ async fn main() -> Result<(), Error> {
     let store = Store::new(&config.data_directory)?;
     let _ = GLOBALS.store.set(store);
 
+    // TLS setup
+    let tls_acceptor = {
+        let certs: Vec<Certificate> =
+            rustls_pemfile::certs(&mut BufReader::new(File::open(&config.certchain_pem_path)?))?
+                .drain(..)
+                .map(Certificate)
+                .collect();
+
+        let mut keys: Vec<PrivateKey> = rustls_pemfile::pkcs8_private_keys(&mut BufReader::new(
+            File::open(&config.key_pem_path)?,
+        ))?
+        .drain(..)
+        .rev()
+        .map(PrivateKey)
+        .collect();
+
+        let key = match keys.pop() {
+            Some(k) => k,
+            None => return Err(Error::NoPrivateKey),
+        };
+
+        let tls_config = rustls::ServerConfig::builder()
+            .with_safe_defaults()
+            .with_no_client_auth()
+            .with_single_cert(certs, key)?;
+
+        TlsAcceptor::from(Arc::new(tls_config))
+    };
+
     // Bind listener to port
     let listener = TcpListener::bind((&*config.ip_address, config.port)).await?;
     log::info!("Running on {}:{}", config.ip_address, config.port);
@@ -56,24 +88,33 @@ async fn main() -> Result<(), Error> {
     loop {
         let (tcp_stream, peer_addr) = listener.accept().await?;
 
-        let connection =
-            http_server.serve_connection(tcp_stream, hyper::service::service_fn(handle_request));
-
+        let acceptor = tls_acceptor.clone();
+        let http_server_clone = http_server.clone();
         tokio::spawn(async move {
-            // If our service exits with an error, log the error
-            if let Err(he) = connection.await {
-                if let Some(src) = he.source() {
-                    if &*format!("{}", src) == "Transport endpoint is not connected (os error 107)"
-                    {
-                        // do nothing
-                    } else {
-                        // Print in detail
-                        log::info!("{:?}", src);
-                    }
-                } else {
-                    // Print in less detail
-                    let e: Error = he.into();
-                    log::info!("{}", e);
+            match acceptor.accept(tcp_stream).await {
+                Err(e) => log::error!("{}", e),
+                Ok(tls_stream) => {
+                    let connection = http_server_clone
+                        .serve_connection(tls_stream, hyper::service::service_fn(handle_request));
+                    tokio::spawn(async move {
+                        // If our service exits with an error, log the error
+                        if let Err(he) = connection.await {
+                            if let Some(src) = he.source() {
+                                if &*format!("{}", src)
+                                    == "Transport endpoint is not connected (os error 107)"
+                                {
+                                    // do nothing
+                                } else {
+                                    // Print in detail
+                                    log::info!("{:?}", src);
+                                }
+                            } else {
+                                // Print in less detail
+                                let e: Error = he.into();
+                                log::info!("{}", e);
+                            }
+                        }
+                    });
                 }
             }
         });