From ea1ea8aa0cb4c8a57b003b4303f88ae9183a640f Mon Sep 17 00:00:00 2001
From: XAMPPRocky <erin.power@embark-studios.com>
Date: Wed, 22 Nov 2023 21:29:35 +0100
Subject: [PATCH] Increase gRPC message size

---
 src/net/xds/client.rs | 10 ++++++++--
 src/net/xds/server.rs |  6 ++++--
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/src/net/xds/client.rs b/src/net/xds/client.rs
index 692dc2b025..b34622a2bb 100644
--- a/src/net/xds/client.rs
+++ b/src/net/xds/client.rs
@@ -71,7 +71,10 @@ impl ServiceClient for AdsGrpcClient {
     async fn connect(
         endpoint: tonic::transport::Endpoint,
     ) -> Result<Self, tonic::transport::Error> {
-        AdsGrpcClient::connect(endpoint).await
+        Ok(AdsGrpcClient::connect(endpoint)
+            .await?
+            .max_decoding_message_size(256 * 1024 * 1024)
+            .max_encoding_message_size(256 * 1024 * 1024))
     }
 
     async fn stream_requests<S: tonic::IntoStreamingRequest<Message = Self::Request> + Send>(
@@ -90,7 +93,10 @@ impl ServiceClient for MdsGrpcClient {
     async fn connect(
         endpoint: tonic::transport::Endpoint,
     ) -> Result<Self, tonic::transport::Error> {
-        MdsGrpcClient::connect(endpoint).await
+        Ok(MdsGrpcClient::connect(endpoint)
+            .await?
+            .max_decoding_message_size(256 * 1024 * 1024)
+            .max_encoding_message_size(256 * 1024 * 1024))
     }
 
     async fn stream_requests<S: tonic::IntoStreamingRequest<Message = Self::Request> + Send>(
diff --git a/src/net/xds/server.rs b/src/net/xds/server.rs
index 8494b35904..ecb9e0dc87 100644
--- a/src/net/xds/server.rs
+++ b/src/net/xds/server.rs
@@ -44,7 +44,8 @@ pub fn spawn(
     mode: Admin,
     config: std::sync::Arc<crate::Config>,
 ) -> impl std::future::Future<Output = crate::Result<()>> {
-    let server = AggregatedDiscoveryServiceServer::new(ControlPlane::from_arc(config, mode));
+    let server = AggregatedDiscoveryServiceServer::new(ControlPlane::from_arc(config, mode))
+            .max_encoding_message_size(256 * 1024 * 1024);
     let server = tonic::transport::Server::builder().add_service(server);
     tracing::info!("serving management server on port `{port}`");
     server
@@ -58,7 +59,8 @@ pub(crate) fn control_plane_discovery_server(
     config: Arc<Config>,
 ) -> impl std::future::Future<Output = crate::Result<()>> {
     let server =
-        AggregatedControlPlaneDiscoveryServiceServer::new(ControlPlane::from_arc(config, mode));
+        AggregatedControlPlaneDiscoveryServiceServer::new(ControlPlane::from_arc(config, mode))
+            .max_encoding_message_size(256 * 1024 * 1024);
     let server = tonic::transport::Server::builder().add_service(server);
     tracing::info!("serving relay server on port `{port}`");
     server