From a7dbe33c5a76e34e5d4bd01c6847c4c6b237fee2 Mon Sep 17 00:00:00 2001
From: Luke Curley <kixelated@gmail.com>
Date: Tue, 17 Oct 2023 15:36:06 +0900
Subject: [PATCH] Add a flag to manually specify roots.

Unfortunately cos-cloud kinda sucks and it's difficult to add roots.
---
 moq-relay/src/config.rs |  7 +++++++
 moq-relay/src/tls.rs    | 19 ++++++++++++++++---
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/moq-relay/src/config.rs b/moq-relay/src/config.rs
index 71ed02f6..a0a66460 100644
--- a/moq-relay/src/config.rs
+++ b/moq-relay/src/config.rs
@@ -24,6 +24,13 @@ pub struct Config {
 	#[arg(long)]
 	pub tls_key: Vec<path::PathBuf>,
 
+	/// Use the TLS root at this path, encoded as PEM.
+	///
+	/// This value can be provided multiple times for multiple roots.
+	/// If this is empty, system roots will be used instead
+	#[arg(long)]
+	pub tls_root: Vec<path::PathBuf>,
+
 	/// Optional: Use the moq-api via HTTP to store origin information.
 	#[arg(long)]
 	pub api: Option<Url>,
diff --git a/moq-relay/src/tls.rs b/moq-relay/src/tls.rs
index aa85aac5..14eeff24 100644
--- a/moq-relay/src/tls.rs
+++ b/moq-relay/src/tls.rs
@@ -38,9 +38,22 @@ impl Tls {
 		// Create a list of acceptable root certificates.
 		let mut roots = RootCertStore::empty();
 
-		// Add the platform's native root certificates.
-		for cert in rustls_native_certs::load_native_certs().context("could not load platform certs")? {
-			roots.add(&Certificate(cert.0)).context("failed to add root cert")?;
+		if config.tls_root.is_empty() {
+			// Add the platform's native root certificates.
+			for cert in rustls_native_certs::load_native_certs().context("could not load platform certs")? {
+				roots.add(&Certificate(cert.0)).context("failed to add root cert")?;
+			}
+		} else {
+			// Add the specified root certificates.
+			for root in &config.tls_root {
+				let root = fs::File::open(root).context("failed to open root cert file")?;
+				let mut root = io::BufReader::new(root);
+				let root = rustls_pemfile::certs(&mut root).context("failed to read root cert")?;
+				anyhow::ensure!(root.len() == 1, "expected a single root cert");
+				let root = Certificate(root[0].to_owned());
+
+				roots.add(&root).context("failed to add root cert")?;
+			}
 		}
 
 		let certs = Self {