From 9abeec89bcb451e48277f2b67e7fd2830db0c8bc Mon Sep 17 00:00:00 2001
From: Edward Shen <code@eddie.sh>
Date: Sat, 24 Apr 2021 00:56:58 -0400
Subject: [PATCH] use arc_swap for tls

---
 Cargo.lock   |  7 +++++++
 Cargo.toml   |  1 +
 src/main.rs  |  3 ++-
 src/ping.rs  | 15 +++++++++++----
 src/state.rs | 28 +++++++++++++++++++---------
 5 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4d716ef..378497e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -234,6 +234,12 @@ dependencies = [
  "memchr",
 ]
 
+[[package]]
+name = "arc-swap"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d4d7d63395147b81a9e570bcc6243aaf71c017bd666d4909cfef0085bdda8d73"
+
 [[package]]
 name = "arrayvec"
 version = "0.5.2"
@@ -1019,6 +1025,7 @@ name = "mangadex-home"
 version = "0.3.0"
 dependencies = [
  "actix-web",
+ "arc-swap",
  "async-trait",
  "base64",
  "bincode",
diff --git a/Cargo.toml b/Cargo.toml
index 200f105..0a280f4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,6 +15,7 @@ debug = 1
 
 [dependencies]
 actix-web = { version = "4.0.0-beta.4", features = [ "rustls" ] }
+arc-swap = "1"
 async-trait = "0.1"
 base64 = "0.13"
 bincode = "1"
diff --git a/src/main.rs b/src/main.rs
index 01640ee..c6cdde1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -25,6 +25,7 @@ use stop::send_stop;
 use thiserror::Error;
 
 use crate::cache::MemoryLruCache;
+use crate::state::DynamicServerCert;
 
 mod cache;
 mod config;
@@ -88,7 +89,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     // What's nice is that Rustls only supports TLS 1.2 and 1.3.
     let mut tls_config = ServerConfig::new(NoClientAuth::new());
-    tls_config.cert_resolver = data_0.clone();
+    tls_config.cert_resolver = Arc::new(DynamicServerCert);
 
     //
     // At this point, the server is ready to start, and starts the necessary
diff --git a/src/ping.rs b/src/ping.rs
index 613528f..476040f 100644
--- a/src/ping.rs
+++ b/src/ping.rs
@@ -11,9 +11,9 @@ use serde::{Deserialize, Serialize};
 use sodiumoxide::crypto::box_::PrecomputedKey;
 use url::Url;
 
-use crate::config::CliArgs;
-use crate::state::PREVIOUSLY_PAUSED;
+use crate::state::{PREVIOUSLY_PAUSED, TLS_CERTS, TLS_SIGNING_KEY};
 use crate::{client_api_version, state::PREVIOUSLY_COMPROMISED};
+use crate::{config::CliArgs, state::TLS_PREVIOUSLY_CREATED};
 use crate::{config::VALIDATE_TOKENS, state::RwLockServerState};
 
 pub const CONTROL_CENTER_PING_URL: &str = "https://api.mangadex.network/ping";
@@ -42,7 +42,9 @@ impl<'a> Request<'a> {
             build_version: client_api_version!()
                 .parse()
                 .expect("to parse the build version"),
-            tls_created_at: Some(state.0.read().tls_config.created_at.clone()),
+            tls_created_at: TLS_PREVIOUSLY_CREATED
+                .get()
+                .map(|v| v.load().as_ref().to_owned()),
         }
     }
 }
@@ -182,7 +184,12 @@ pub async fn update_server_state(secret: &str, cli: &CliArgs, data: &mut Arc<RwL
                 }
 
                 if let Some(tls) = resp.tls {
-                    write_guard.tls_config = tls;
+                    TLS_PREVIOUSLY_CREATED
+                        .get()
+                        .unwrap()
+                        .swap(Arc::new(tls.created_at));
+                    TLS_SIGNING_KEY.get().unwrap().swap(tls.priv_key);
+                    TLS_CERTS.get().unwrap().swap(Arc::new(tls.certs));
                 }
 
                 let previously_compromised = PREVIOUSLY_COMPROMISED.load(Ordering::Acquire);
diff --git a/src/state.rs b/src/state.rs
index dfe5bc2..3b184c2 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -1,11 +1,13 @@
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Arc;
 
 use crate::config::{CliArgs, SEND_SERVER_VERSION, VALIDATE_TOKENS};
-use crate::ping::{Request, Response, Tls, CONTROL_CENTER_PING_URL};
+use crate::ping::{Request, Response, CONTROL_CENTER_PING_URL};
+use arc_swap::ArcSwap;
 use log::{error, info, warn};
+use once_cell::sync::OnceCell;
 use parking_lot::RwLock;
-use rustls::sign::CertifiedKey;
+use rustls::sign::{CertifiedKey, SigningKey};
+use rustls::Certificate;
 use rustls::{ClientHello, ResolvesServerCert};
 use sodiumoxide::crypto::box_::PrecomputedKey;
 use thiserror::Error;
@@ -14,7 +16,6 @@ use url::Url;
 pub struct ServerState {
     pub precomputed_key: PrecomputedKey,
     pub image_server: Url,
-    pub tls_config: Tls,
     pub url: Url,
     pub url_overridden: bool,
 }
@@ -22,6 +23,10 @@ pub struct ServerState {
 pub static PREVIOUSLY_PAUSED: AtomicBool = AtomicBool::new(false);
 pub static PREVIOUSLY_COMPROMISED: AtomicBool = AtomicBool::new(false);
 
+pub static TLS_PREVIOUSLY_CREATED: OnceCell<ArcSwap<String>> = OnceCell::new();
+pub static TLS_SIGNING_KEY: OnceCell<ArcSwap<Box<dyn SigningKey>>> = OnceCell::new();
+pub static TLS_CERTS: OnceCell<ArcSwap<Vec<Certificate>>> = OnceCell::new();
+
 #[derive(Error, Debug)]
 pub enum ServerInitError {
     #[error(transparent)]
@@ -96,10 +101,14 @@ impl ServerState {
                         VALIDATE_TOKENS.store(resp.force_tokens, Ordering::Release);
                     }
 
+                    let tls = resp.tls.unwrap();
+                    TLS_PREVIOUSLY_CREATED.set(ArcSwap::from_pointee(tls.created_at));
+                    TLS_SIGNING_KEY.set(ArcSwap::new(tls.priv_key));
+                    TLS_CERTS.set(ArcSwap::from_pointee(tls.certs));
+
                     Ok(Self {
                         precomputed_key: key,
                         image_server: resp.image_server,
-                        tls_config: resp.tls.unwrap(),
                         url: resp.url,
                         url_overridden: config.override_upstream.is_some(),
                     })
@@ -125,14 +134,15 @@ impl ServerState {
 
 pub struct RwLockServerState(pub RwLock<ServerState>);
 
-impl ResolvesServerCert for RwLockServerState {
+pub struct DynamicServerCert;
+
+impl ResolvesServerCert for DynamicServerCert {
     fn resolve(&self, _: ClientHello) -> Option<CertifiedKey> {
         // TODO: wait for actix-web to use a new version of rustls so we can
         // remove cloning the certs all the time
-        let read_guard = self.0.read();
         Some(CertifiedKey {
-            cert: read_guard.tls_config.certs.clone(),
-            key: Arc::clone(&read_guard.tls_config.priv_key),
+            cert: TLS_CERTS.get().unwrap().load().as_ref().clone(),
+            key: TLS_SIGNING_KEY.get().unwrap().load_full(),
             ocsp: None,
             sct_list: None,
         })