From d0bf7d0a81cdcdf57e8dc4a8d8f93b4f2b4a617f Mon Sep 17 00:00:00 2001
From: Liam Vovk <63673978+vovkman@users.noreply.github.com>
Date: Mon, 18 Sep 2023 21:35:13 -0700
Subject: [PATCH] geyser: add TLS config to gRPC server (#183)

---
 CHANGELOG.md                          |  2 ++
 yellowstone-grpc-geyser/src/config.rs |  9 +++++++++
 yellowstone-grpc-geyser/src/grpc.rs   | 21 ++++++++++++++++++---
 yellowstone-grpc-geyser/src/plugin.rs |  1 +
 4 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1e69be2d..cbadee16 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,8 @@ The minor version will be incremented upon a breaking change and the patch versi
 
 ### Features
 
+- geyser: add optional TLS to gRPC server config ([#183](https://github.com/rpcpool/yellowstone-grpc/pull/183)).
+
 ### Fixes
 
 ### Breaking
diff --git a/yellowstone-grpc-geyser/src/config.rs b/yellowstone-grpc-geyser/src/config.rs
index 87c1dd71..e2e603e0 100644
--- a/yellowstone-grpc-geyser/src/config.rs
+++ b/yellowstone-grpc-geyser/src/config.rs
@@ -62,6 +62,8 @@ impl ConfigLog {
 pub struct ConfigGrpc {
     /// Address of Grpc service.
     pub address: SocketAddr,
+    /// TLS config
+    pub tls_config: Option<ConfigGrpcServerTls>,
     /// Capacity of the channel per connection
     #[serde(
         default = "ConfigGrpc::channel_capacity_default",
@@ -92,6 +94,13 @@ impl ConfigGrpc {
     }
 }
 
+#[derive(Debug, Clone, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct ConfigGrpcServerTls {
+    pub cert_path: String,
+    pub key_path: String,
+}
+
 #[derive(Debug, Default, Clone, Deserialize)]
 #[serde(deny_unknown_fields)]
 pub struct ConfigGrpcFilters {
diff --git a/yellowstone-grpc-geyser/src/grpc.rs b/yellowstone-grpc-geyser/src/grpc.rs
index ad576d7b..97cf800a 100644
--- a/yellowstone-grpc-geyser/src/grpc.rs
+++ b/yellowstone-grpc-geyser/src/grpc.rs
@@ -37,13 +37,17 @@ use {
         },
     },
     tokio::{
+        fs,
         sync::{broadcast, mpsc, Notify, RwLock, Semaphore},
         time::{sleep, Duration, Instant},
     },
     tokio_stream::wrappers::ReceiverStream,
     tonic::{
         codec::CompressionEncoding,
-        transport::server::{Server, TcpIncoming},
+        transport::{
+            server::{Server, TcpIncoming},
+            Identity, ServerTlsConfig,
+        },
         Request, Response, Result as TonicResult, Status, Streaming,
     },
     tonic_health::server::health_reporter,
@@ -679,7 +683,7 @@ pub struct GrpcService {
 }
 
 impl GrpcService {
-    pub fn create(
+    pub async fn create(
         config: ConfigGrpc,
         block_fail_action: ConfigBlockFailAction,
     ) -> Result<
@@ -705,6 +709,17 @@ impl GrpcService {
         // Messages to clients combined by commitment
         let (broadcast_tx, _) = broadcast::channel(config.channel_capacity);
 
+        // gRPC server builder with optional TLS
+        let mut server_builder = Server::builder();
+        if let Some(tls_config) = &config.tls_config {
+            let (cert, key) = tokio::try_join!(
+                fs::read(&tls_config.cert_path),
+                fs::read(&tls_config.key_path)
+            )?;
+            server_builder = server_builder
+                .tls_config(ServerTlsConfig::new().identity(Identity::from_pem(cert, key)))?;
+        }
+
         // Create Server
         let service = GeyserServer::new(Self {
             config,
@@ -732,7 +747,7 @@ impl GrpcService {
             let (mut health_reporter, health_service) = health_reporter();
             health_reporter.set_serving::<GeyserServer<Self>>().await;
 
-            Server::builder()
+            server_builder
                 .http2_keepalive_interval(Some(Duration::from_secs(5)))
                 .add_service(health_service)
                 .add_service(service)
diff --git a/yellowstone-grpc-geyser/src/plugin.rs b/yellowstone-grpc-geyser/src/plugin.rs
index a7c02c17..3296f551 100644
--- a/yellowstone-grpc-geyser/src/plugin.rs
+++ b/yellowstone-grpc-geyser/src/plugin.rs
@@ -81,6 +81,7 @@ impl GeyserPlugin for Plugin {
         let (grpc_channel, grpc_shutdown, prometheus) = runtime.block_on(async move {
             let (grpc_channel, grpc_shutdown) =
                 GrpcService::create(config.grpc, config.block_fail_action)
+                    .await
                     .map_err(|error| GeyserPluginError::Custom(error))?;
             let prometheus = PrometheusService::new(config.prometheus)
                 .map_err(|error| GeyserPluginError::Custom(Box::new(error)))?;