Accept compressed RPCs in Oak Functions trusted app

project-oak · Jan 23, 2024 · a399a1b · a399a1b
1 parent 5641a13
commit a399a1b
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 4 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/oak_functions_containers_app/Cargo.toml b/oak_functions_containers_app/Cargo.toml
@@ -33,7 +33,7 @@ opentelemetry-otlp = { version = "*", default-features = false, features = [
 prost = "*"
 tokio = { version = "*", features = ["rt-multi-thread", "macros", "sync"] }
 tokio-stream = { version = "*", features = ["net"] }
-tonic = { workspace = true }
+tonic = { workspace = true, features = ["gzip"] }
 tower = { version = "*", features = ["load-shed"] }
 tower-http = { version = "*", features = ["trace"] }
 tracing = "*"

diff --git a/oak_functions_containers_app/src/lib.rs b/oak_functions_containers_app/src/lib.rs
@@ -41,6 +41,7 @@ use opentelemetry::{
 use prost::Message;
 use tokio::net::TcpListener;
 use tokio_stream::{wrappers::TcpListenerStream, StreamExt};
+use tonic::codec::CompressionEncoding;
 use tracing::Span;
 
 use crate::proto::oak::functions::oak_functions_server::{OakFunctions, OakFunctionsServer};
@@ -407,7 +408,8 @@ pub async fn serve<G: AsyncEncryptionKeyHandle + Send + Sync + 'static>(
                 encryption_key_handle,
                 Some(Arc::new(OtelObserver::new(meter))),
             ))
-            .max_decoding_message_size(MAX_DECODING_MESSAGE_SIZE),
+            .max_decoding_message_size(MAX_DECODING_MESSAGE_SIZE)
+            .accept_compressed(CompressionEncoding::Gzip),
         )
         .serve_with_incoming(TcpListenerStream::new(listener))
         .await

diff --git a/oak_functions_containers_app/tests/integration_test.rs b/oak_functions_containers_app/tests/integration_test.rs
@@ -34,7 +34,7 @@ use oak_functions_containers_app::serve;
 use oak_functions_service::proto::oak::functions::InitializeRequest;
 use opentelemetry::metrics::{noop::NoopMeterProvider, MeterProvider};
 use tokio::net::TcpListener;
-use tonic::transport::Endpoint;
+use tonic::{codec::CompressionEncoding, transport::Endpoint};
 
 use crate::proto::oak::functions::oak_functions_client::OakFunctionsClient;
 
@@ -60,7 +60,7 @@ async fn test_lookup() {
             .connect()
             .await
             .expect("couldn't connect to trusted app");
-        OakFunctionsClient::new(channel)
+        OakFunctionsClient::new(channel).send_compressed(CompressionEncoding::Gzip)
     };
 
     let _ = oak_functions_client