Support large PAR executor messages

Author: embeaken

pkg/privateactionrunner/executor/BUILD.bazel

@@ -49,7 +49,10 @@ dd_agent_go_test(
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
         "@org_golang_google_grpc//:grpc",
+        "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//credentials",
         "@org_golang_google_grpc//credentials/insecure",
+        "@org_golang_google_grpc//status",
+        "@org_golang_google_protobuf//proto",
     ],
 )

pkg/privateactionrunner/executor/server.go

@@ -30,6 +30,12 @@ import (
 	pb "github.com/DataDog/datadog-agent/pkg/proto/pbgo/privateactionrunner/executor"
 )
 
+// maxMessageSize is the control<->executor protocol limit in bytes. Action
+// inputs and outputs can approach 15 MiB, so 20 MiB leaves protobuf headroom
+// while still bounding memory use. Keep this in sync with MAX_MESSAGE_SIZE in
+// par-control's executor client.
+const maxMessageSize = 20 * 1024 * 1024
+
 type actionExecutor interface {
 	PrepareTask(ctx context.Context, task *types.Task) (*runners.PreparedWorkflowTask, *types.Task, error)
 	RunPrepared(ctx context.Context, prepared *runners.PreparedWorkflowTask) (interface{}, error)
@@ -204,6 +210,12 @@ const idleCheckDivisor = 10
 // Serve serves the Executor on lis until ctx is cancelled, then stops gracefully
 // bounded by the drain timeout. Pass grpcOpts to secure the socket.
 func Serve(ctx context.Context, lis net.Listener, srv *Server, opts ServeOptions, grpcOpts ...grpc.ServerOption) error {
+	// Apply the protocol limits after caller-provided options so every executor
+	// endpoint accepts the same bounded action payload sizes.
+	grpcOpts = append(grpcOpts,
+		grpc.MaxRecvMsgSize(maxMessageSize),
+		grpc.MaxSendMsgSize(maxMessageSize),
+	)
 	grpcServer := grpc.NewServer(grpcOpts...)
 	pb.RegisterExecutorServer(grpcServer, srv)
 

pkg/privateactionrunner/executor/server_test.go

@@ -18,15 +18,19 @@ import (
 	"io"
 	"math/big"
 	"net"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/benbjohnson/clock"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/DataDog/datadog-agent/pkg/privateactionrunner/runners"
 	taskverifier "github.com/DataDog/datadog-agent/pkg/privateactionrunner/task-verifier"
@@ -95,6 +99,10 @@ func startTestServer(t *testing.T, srv *Server) pb.ExecutorClient {
 	conn, err := grpc.NewClient(
 		"passthrough:///"+socketPath,
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithDefaultCallOptions(
+			grpc.MaxCallRecvMsgSize(maxMessageSize*2),
+			grpc.MaxCallSendMsgSize(maxMessageSize*2),
+		),
 		grpc.WithContextDialer(func(dialCtx context.Context, _ string) (net.Conn, error) {
 			return Dial(dialCtx, socketPath, 2*time.Second)
 		}),
@@ -132,6 +140,95 @@ func runAction(t *testing.T, client pb.ExecutorClient, taskBytes []byte) *pb.Act
 	return result
 }
 
+func taskWithEncodedSize(t *testing.T, target int) []byte {
+	t.Helper()
+	const prefix = `{"data":{"id":"task","attributes":{"padding":"`
+	const suffix = `"}}}`
+	padding := target - len(prefix) - len(suffix)
+	require.Positive(t, padding)
+	for {
+		raw := []byte(prefix + strings.Repeat("a", padding) + suffix)
+		size := proto.Size(&pb.RunActionRequest{Task: raw})
+		if size == target {
+			return raw
+		}
+		padding += target - size
+		require.Positive(t, padding)
+	}
+}
+
+func outputWithEncodedSize(t *testing.T, target int) string {
+	t.Helper()
+	padding := target
+	for {
+		output := strings.Repeat("a", padding)
+		outputJSON, err := json.Marshal(output)
+		require.NoError(t, err)
+		size := proto.Size(&pb.RunActionResponse{
+			Event: &pb.RunActionResponse_Result{
+				Result: &pb.ActionResult{
+					Outcome: &pb.ActionResult_Output{Output: outputJSON},
+				},
+			},
+		})
+		if size == target {
+			return output
+		}
+		padding += target - size
+		require.Positive(t, padding)
+	}
+}
+
+func receiveRunActionError(client pb.ExecutorClient, task []byte) error {
+	stream, err := client.RunAction(context.Background(), &pb.RunActionRequest{Task: task})
+	if err != nil {
+		return err
+	}
+	_, err = stream.Recv()
+	return err
+}
+
+func TestServeEnforcesRequestMessageLimit(t *testing.T) {
+	fake := &fakeExecutor{
+		prepared: &runners.PreparedWorkflowTask{Task: &types.Task{}},
+		output:   map[string]interface{}{},
+	}
+	srv := NewServer(fake, "test-version", nil)
+	srv.SetReady(true)
+	client := startTestServer(t, srv)
+
+	below := taskWithEncodedSize(t, maxMessageSize-1)
+	result := runAction(t, client, below)
+	require.NotNil(t, result.GetOutput())
+
+	above := taskWithEncodedSize(t, maxMessageSize+1)
+	err := receiveRunActionError(client, above)
+	require.Error(t, err)
+	assert.Equal(t, codes.ResourceExhausted, status.Code(err))
+}
+
+func TestServeEnforcesResponseMessageLimit(t *testing.T) {
+	below := outputWithEncodedSize(t, maxMessageSize-1)
+	fake := &fakeExecutor{
+		prepared: &runners.PreparedWorkflowTask{Task: &types.Task{}},
+		output:   below,
+	}
+	srv := NewServer(fake, "test-version", nil)
+	srv.SetReady(true)
+	client := startTestServer(t, srv)
+
+	result := runAction(t, client, []byte(`{"data":{"id":"task-1"}}`))
+	belowJSON, err := json.Marshal(below)
+	require.NoError(t, err)
+	assert.Equal(t, belowJSON, result.GetOutput())
+
+	above := outputWithEncodedSize(t, maxMessageSize+1)
+	fake.output = above
+	err = receiveRunActionError(client, []byte(`{"data":{"id":"task-2"}}`))
+	require.Error(t, err)
+	assert.Equal(t, codes.ResourceExhausted, status.Code(err))
+}
+
 func TestServeSyncKeysSeedsAndReturnsCurrentSnapshot(t *testing.T) {
 	keysManager := &fakeKeysManager{snapshot: []taskverifier.SigningKey{{
 		ID:      "fresh",

pkg/privateactionrunner/par-control/src/executor.rs

@@ -12,9 +12,15 @@ use crate::proto::executor as pb;
 use crate::proto::executor::executor_client::ExecutorClient;
 use crate::transport;
 use anyhow::{Context, Result, bail};
+use prost::Message;
 use std::path::Path;
 use tonic::transport::Channel;
 
+/// Control<->executor protocol limit in bytes. Action inputs and outputs can
+/// approach 15 MiB, so 20 MiB leaves protobuf headroom while still bounding
+/// memory use. Keep this in sync with `maxMessageSize` in the Go executor.
+const MAX_MESSAGE_SIZE: usize = 20 * 1024 * 1024;
+
 /// Executor health snapshot used to gate dispatch.
 #[derive(Debug, Clone)]
 pub struct Health {
@@ -75,7 +81,9 @@ impl ExecutorDispatcher {
             None => transport::connect_lazy(socket),
         };
         ExecutorDispatcher {
-            client: ExecutorClient::new(channel),
+            client: ExecutorClient::new(channel)
+                .max_encoding_message_size(MAX_MESSAGE_SIZE)
+                .max_decoding_message_size(MAX_MESSAGE_SIZE),
         }
     }
 }
@@ -122,9 +130,16 @@ impl Dispatcher for ExecutorDispatcher {
     }
 
     async fn run_action(&self, raw: Vec<u8>) -> Result<Outcome> {
+        let request = pb::RunActionRequest { task: raw };
+        if request.encoded_len() > MAX_MESSAGE_SIZE {
+            return Err(tonic::Status::resource_exhausted(format!(
+                "RunAction request exceeds the {MAX_MESSAGE_SIZE}-byte protocol limit"
+            ))
+            .into());
+        }
         let mut client = self.client.clone();
         let mut stream = client
-            .run_action(pb::RunActionRequest { task: raw })
+            .run_action(request)
             .await
             .context("executor RunAction failed")?
             .into_inner();
@@ -152,3 +167,132 @@ impl Dispatcher for ExecutorDispatcher {
         }
     }
 }
+
+#[cfg(all(test, unix))]
+mod tests {
+    use super::*;
+    use crate::proto::executor::executor_server::{Executor, ExecutorServer};
+    use std::pin::Pin;
+    use tokio_stream::Stream;
+    use tokio_stream::wrappers::UnixListenerStream;
+    use tonic::{Request, Response, Status};
+
+    #[derive(Clone)]
+    struct FakeExecutor {
+        output: Vec<u8>,
+    }
+
+    #[tonic::async_trait]
+    impl Executor for FakeExecutor {
+        type RunActionStream =
+            Pin<Box<dyn Stream<Item = std::result::Result<pb::RunActionResponse, Status>> + Send>>;
+
+        async fn run_action(
+            &self,
+            _request: Request<pb::RunActionRequest>,
+        ) -> std::result::Result<Response<Self::RunActionStream>, Status> {
+            let response = pb::RunActionResponse {
+                event: Some(pb::run_action_response::Event::Result(pb::ActionResult {
+                    outcome: Some(pb::action_result::Outcome::Output(self.output.clone())),
+                })),
+            };
+            Ok(Response::new(Box::pin(tokio_stream::once(Ok(response)))))
+        }
+
+        async fn health(
+            &self,
+            _request: Request<pb::HealthRequest>,
+        ) -> std::result::Result<Response<pb::HealthResponse>, Status> {
+            Ok(Response::new(pb::HealthResponse::default()))
+        }
+
+        async fn sync_keys(
+            &self,
+            _request: Request<pb::SyncKeysRequest>,
+        ) -> std::result::Result<Response<pb::SyncKeysResponse>, Status> {
+            Ok(Response::new(pb::SyncKeysResponse::default()))
+        }
+    }
+
+    async fn test_dispatcher(output: Vec<u8>) -> (ExecutorDispatcher, tempfile::TempDir) {
+        let dir = tempfile::tempdir().expect("tempdir");
+        let socket = dir.path().join("executor.sock");
+        let listener = tokio::net::UnixListener::bind(&socket).expect("bind executor socket");
+        tokio::spawn(async move {
+            let service = ExecutorServer::new(FakeExecutor { output })
+                .max_decoding_message_size(MAX_MESSAGE_SIZE * 2)
+                .max_encoding_message_size(MAX_MESSAGE_SIZE * 2);
+            let _ = tonic::transport::Server::builder()
+                .add_service(service)
+                .serve_with_incoming(UnixListenerStream::new(listener))
+                .await;
+        });
+        (ExecutorDispatcher::new(&socket, None), dir)
+    }
+
+    fn request_payload_with_encoded_size(target: usize) -> Vec<u8> {
+        let mut payload = vec![b'a'; target];
+        loop {
+            let encoded = pb::RunActionRequest {
+                task: payload.clone(),
+            }
+            .encoded_len();
+            match encoded.cmp(&target) {
+                std::cmp::Ordering::Equal => return payload,
+                std::cmp::Ordering::Less => payload.resize(payload.len() + target - encoded, b'a'),
+                std::cmp::Ordering::Greater => payload.truncate(payload.len() - (encoded - target)),
+            }
+        }
+    }
+
+    fn response_output_with_encoded_size(target: usize) -> Vec<u8> {
+        let mut output = vec![b'a'; target];
+        loop {
+            let encoded = pb::RunActionResponse {
+                event: Some(pb::run_action_response::Event::Result(pb::ActionResult {
+                    outcome: Some(pb::action_result::Outcome::Output(output.clone())),
+                })),
+            }
+            .encoded_len();
+            match encoded.cmp(&target) {
+                std::cmp::Ordering::Equal => return output,
+                std::cmp::Ordering::Less => output.resize(output.len() + target - encoded, b'a'),
+                std::cmp::Ordering::Greater => output.truncate(output.len() - (encoded - target)),
+            }
+        }
+    }
+
+    fn assert_status(error: &anyhow::Error, expected: tonic::Code) {
+        let status = error
+            .downcast_ref::<Status>()
+            .unwrap_or_else(|| panic!("expected tonic status, got {error:#}"));
+        assert_eq!(status.code(), expected, "unexpected status: {status}");
+    }
+
+    #[tokio::test]
+    async fn enforces_request_encoding_limit() {
+        let (dispatcher, _dir) = test_dispatcher(Vec::new()).await;
+        let below = request_payload_with_encoded_size(MAX_MESSAGE_SIZE - 1);
+        assert!(dispatcher.run_action(below).await.is_ok());
+
+        let above = request_payload_with_encoded_size(MAX_MESSAGE_SIZE + 1);
+        let error = dispatcher.run_action(above).await.unwrap_err();
+        assert_status(&error, tonic::Code::ResourceExhausted);
+    }
+
+    #[tokio::test]
+    async fn enforces_response_decoding_limit() {
+        let below = response_output_with_encoded_size(MAX_MESSAGE_SIZE - 1);
+        let (dispatcher, _dir) = test_dispatcher(below.clone()).await;
+        match dispatcher.run_action(Vec::new()).await.unwrap() {
+            Outcome::Success { output_json } => assert_eq!(output_json, below),
+            Outcome::Failure { .. } => panic!("expected success"),
+        }
+
+        let above = response_output_with_encoded_size(MAX_MESSAGE_SIZE + 1);
+        let (dispatcher, _dir) = test_dispatcher(above).await;
+        let error = dispatcher.run_action(Vec::new()).await.unwrap_err();
+        // Tonic reports an oversized decoded response as OutOfRange.
+        assert_status(&error, tonic::Code::OutOfRange);
+    }
+}