Add TLS options

googleforgames · Feb 3, 2025 · 2c3b0d7 · 2c3b0d7
1 parent 43fb309
commit 2c3b0d7
Show file tree

Hide file tree

Showing 4 changed files with 100 additions and 8 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -235,7 +235,7 @@ tokio = { version = "1.41.1", features = [
     "parking_lot",
     "tracing",
 ] }
-tonic = "0.12"
+tonic = { version = "0.12", features = ["tls"] }
 tower = "0.5"
 tracing = "0.1.40"
 tracing-futures = { version = "0.2.5", features = ["futures-03"] }

diff --git a/crates/xds/src/server.rs b/crates/xds/src/server.rs
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-use std::{io, sync::Arc, time::Duration};
+use std::{sync::Arc, time::Duration};
 
 use futures::{Stream, TryFutureExt};
 use tokio_stream::StreamExt;
@@ -34,6 +34,30 @@ use crate::{
     net::TcpListener,
 };
 
+#[derive(Clone)]
+pub struct TlsIdentity {
+    identity: tonic::transport::Identity,
+}
+
+impl TlsIdentity {
+    pub fn from_raw(cert: &[u8], key: &[u8]) -> Self {
+        Self {
+            identity: tonic::transport::Identity::from_pem(cert, key),
+        }
+    }
+
+    pub fn from_files(cert: &std::path::Path, key: &std::path::Path) -> eyre::Result<Self> {
+        use eyre::WrapErr as _;
+        let cert = std::fs::read(cert)
+            .with_context(|| format!("failed to read PEM certificate from {cert:?}"))?;
+        let key = std::fs::read(key).with_context(|| format!("failed to read key from {key:?}"))?;
+
+        Ok(Self {
+            identity: tonic::transport::Identity::from_pem(cert, key),
+        })
+    }
+}
+
 const VERSION_INFO: &str = "9";
 
 pub struct ControlPlane<C> {
@@ -69,7 +93,8 @@ impl<C: crate::config::Configuration> ControlPlane<C> {
     pub fn management_server(
         mut self,
         listener: TcpListener,
-    ) -> io::Result<impl std::future::Future<Output = crate::Result<()>>> {
+        tls: Option<TlsIdentity>,
+    ) -> eyre::Result<impl std::future::Future<Output = crate::Result<()>>> {
         self.is_relay = false;
         tokio::spawn({
             let this = self.clone();
@@ -78,7 +103,15 @@ impl<C: crate::config::Configuration> ControlPlane<C> {
 
         let server = AggregatedDiscoveryServiceServer::new(self)
             .max_encoding_message_size(crate::config::max_grpc_message_size());
-        let server = tonic::transport::Server::builder().add_service(server);
+        let builder = tonic::transport::Server::builder();
+
+        let mut builder = if let Some(tls) = tls {
+            builder.tls_config(tonic::transport::ServerTlsConfig::new().identity(tls.identity))?
+        } else {
+            builder
+        };
+
+        let server = builder.add_service(server);
         tracing::info!("serving management server on port `{}`", listener.port());
         Ok(server
             .serve_with_incoming(listener.into_stream()?)
@@ -88,7 +121,8 @@ impl<C: crate::config::Configuration> ControlPlane<C> {
     pub fn relay_server(
         mut self,
         listener: TcpListener,
-    ) -> io::Result<impl std::future::Future<Output = crate::Result<()>>> {
+        tls: Option<TlsIdentity>,
+    ) -> eyre::Result<impl std::future::Future<Output = crate::Result<()>>> {
         self.is_relay = true;
         tokio::spawn({
             let this = self.clone();
@@ -97,7 +131,15 @@ impl<C: crate::config::Configuration> ControlPlane<C> {
 
         let server = AggregatedControlPlaneDiscoveryServiceServer::new(self)
             .max_encoding_message_size(crate::config::max_grpc_message_size());
-        let server = tonic::transport::Server::builder().add_service(server);
+        let builder = tonic::transport::Server::builder();
+
+        let mut builder = if let Some(tls) = tls {
+            builder.tls_config(tonic::transport::ServerTlsConfig::new().identity(tls.identity))?
+        } else {
+            builder
+        };
+
+        let server = builder.add_service(server);
         tracing::info!("serving relay server on port `{}`", listener.port());
         Ok(server
             .serve_with_incoming(listener.into_stream()?)

diff --git a/src/cli/service.rs b/src/cli/service.rs
@@ -80,6 +80,36 @@ pub struct Service {
         default_value_t = 7800
     )]
     xds_port: u16,
+    /// A PEM encoded certificate, if supplied, applies to the mds and xds service(s)
+    #[clap(
+        long = "service.tls.cert",
+        env = "QUILKIN_SERVICE_TLS_CERT",
+        requires("service.tls.key")
+    )]
+    tls_cert: Option<Vec<u8>>,
+    /// The private key for the cert
+    #[clap(
+        long = "service.tls.key",
+        env = "QUILKIN_SERVICE_TLS_KEY",
+        requires("service.tls.cert")
+    )]
+    tls_key: Option<Vec<u8>>,
+    /// Path to a PEM encoded certificate, if supplied, applies to the mds and xds service(s)
+    #[clap(
+        long = "service.tls.cert-path",
+        env = "QUILKIN_SERVICE_TLS_CERT_PATH",
+        requires("service.tls.key-path"),
+        conflicts_with("service.tls.cert")
+    )]
+    tls_cert_path: Option<std::path::PathBuf>,
+    /// Path to the private key for the cert
+    #[clap(
+        long = "service.tls.key-path",
+        env = "QUILKIN_SERVICE_TLS_KEY_PATH",
+        requires("service.tls.cert-path"),
+        conflicts_with("service.tls.key")
+    )]
+    tls_key_path: Option<std::path::PathBuf>,
 }
 
 impl Default for Service {
@@ -97,6 +127,10 @@ impl Default for Service {
             xds_enabled: <_>::default(),
             xds_port: 7800,
             xdp: <_>::default(),
+            tls_cert: None,
+            tls_key: None,
+            tls_cert_path: None,
+            tls_key_path: None,
         }
     }
 }
@@ -171,6 +205,20 @@ impl Service {
             || self.mds_enabled
     }
 
+    fn tls_identity(&self) -> crate::Result<Option<quilkin_xds::server::TlsIdentity>> {
+        if let Some((cert, key)) = self.tls_cert.as_ref().zip(self.tls_key.as_ref()) {
+            Ok(Some(quilkin_xds::server::TlsIdentity::from_raw(cert, key)))
+        } else if let Some((certp, keyp)) =
+            self.tls_cert_path.as_ref().zip(self.tls_key_path.as_ref())
+        {
+            Ok(Some(quilkin_xds::server::TlsIdentity::from_files(
+                certp, keyp,
+            )?))
+        } else {
+            Ok(None)
+        }
+    }
+
     /// The main entrypoint for listening network servers. When called will
     /// spawn any and all enabled services, if successful returning a future
     /// that can be await to wait on services to be cancelled.
@@ -260,7 +308,7 @@ impl Service {
                     config.clone(),
                     crate::components::admin::IDLE_REQUEST_INTERVAL,
                 )
-                .relay_server(listener)?,
+                .relay_server(listener, self.tls_identity()?)?,
             )
             .map_err(From::from)
             .and_then(std::future::ready),
@@ -286,7 +334,7 @@ impl Service {
                     config.clone(),
                     crate::components::admin::IDLE_REQUEST_INTERVAL,
                 )
-                .management_server(listener)?,
+                .management_server(listener, self.tls_identity()?)?,
             )
             .map_err(From::from)
             .and_then(std::future::ready),