[feature] gRPC File Download API

docs/_docs/dev-guide/tavern.md

@@ -141,17 +141,6 @@ terraform apply -var="gcp_project=<PROJECT_ID>" -var="oauth_client_id=<OAUTH_CLI
 
 ## User Interface
 
-## CDN API
-
-### Uploading Files
-
-* File uploads require 2 form parameters: `fileName` and `fileContent`
-* A successful response returns JSON with the following content: `{"data":{"file":{"id":<FILE_ID>}}}`
-
-### Downloading Files
-
-* TODO (CDN is not yet added to Tavern)
-
 ## GraphQL API
 
 ### Playground
@@ -296,9 +285,12 @@ query get_task_res {
 
 Tavern also supports a gRPC API for agents to claim tasks and report execution output. This API is defined by our c2.proto spec and is still under active development.
 
-### YAML Test Reference (gRPC)
+### Downloading Files
+
+You may download files from Tavern utilizing the `DownloadFile` gRPC method. This method streams responses, each of which will contain a chunk of the desired file. We rely on the [ordering guarantees](https://grpc.io/docs/languages/go/basics/#defining-the-service) provided by gRPC to ensure the file is assembled correctly. This API also sets two headers to ensure the integrity of files:
 
-Still under development.
+* `sha3-256-checksum`: Set to the SHA3 hash of the entire file.
+* `file-size`: Set to the number of bytes contained by the file.
 
 ## Performance Profiling
 

implants/lib/c2/src/c2.rs

@@ -143,6 +143,18 @@ pub struct ReportTaskOutputRequest {
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ReportTaskOutputResponse {}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct DownloadFileRequest {
+    #[prost(string, tag = "1")]
+    pub name: ::prost::alloc::string::String,
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct DownloadFileResponse {
+    #[prost(bytes = "vec", tag = "1")]
+    pub chunk: ::prost::alloc::vec::Vec<u8>,
+}
 /// Generated client implementations.
 pub mod c2_client {
     #![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)]
@@ -272,5 +284,35 @@ pub mod c2_client {
             req.extensions_mut().insert(GrpcMethod::new("c2.C2", "ReportTaskOutput"));
             self.inner.unary(req, path, codec).await
         }
+        ///
+        /// Download a file from the server, returning one or more chunks of data.
+        /// The maximum size of these chunks is determined by the server.
+        /// The server should reply with two headers:
+        ///   - "sha3-256-checksum": A SHA3-256 digest of the entire file contents.
+        ///   - "file-size": The number of bytes contained by the file.
+        ///
+        /// If no associated file can be found, a NotFound status error is returned.
+        pub async fn download_file(
+            &mut self,
+            request: impl tonic::IntoRequest<super::DownloadFileRequest>,
+        ) -> std::result::Result<
+            tonic::Response<tonic::codec::Streaming<super::DownloadFileResponse>>,
+            tonic::Status,
+        > {
+            self.inner
+                .ready()
+                .await
+                .map_err(|e| {
+                    tonic::Status::new(
+                        tonic::Code::Unknown,
+                        format!("Service was not ready: {}", e.into()),
+                    )
+                })?;
+            let codec = tonic::codec::ProstCodec::default();
+            let path = http::uri::PathAndQuery::from_static("/c2.C2/DownloadFile");
+            let mut req = request.into_request();
+            req.extensions_mut().insert(GrpcMethod::new("c2.C2", "DownloadFile"));
+            self.inner.server_streaming(req, path, codec).await
+        }
     }
 }

tavern/internal/c2/api_download_file.go

@@ -0,0 +1,64 @@
+package c2
+
+import (
+	"bytes"
+	"fmt"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+	"realm.pub/tavern/internal/c2/c2pb"
+	"realm.pub/tavern/internal/ent"
+	"realm.pub/tavern/internal/ent/file"
+)
+
+func (srv *Server) DownloadFile(req *c2pb.DownloadFileRequest, stream c2pb.C2_DownloadFileServer) error {
+	ctx := stream.Context()
+
+	// Load File
+	name := req.GetName()
+	f, err := srv.graph.File.Query().
+		Where(file.Name(name)).
+		Only(ctx)
+	if ent.IsNotFound(err) {
+		return status.Errorf(codes.NotFound, "%v", err)
+	}
+	if err != nil {
+		return status.Errorf(codes.Internal, "failed to query file (%q): %v", name, err)
+	}
+
+	// Set Header Metadata
+	stream.SetHeader(metadata.Pairs(
+		"sha3-256-checksum", f.Hash,
+		"file-size", fmt.Sprintf("%d", f.Size),
+	))
+
+	// Send File Chunks
+	buf := bytes.NewBuffer(f.Content)
+	for {
+		// Check Empty Buffer
+		if buf.Len() < 1 {
+			return nil
+		}
+
+		// Determine Chunk Size
+		chunkLen := srv.MaxFileChunkSize
+		if uint64(buf.Len()) < chunkLen {
+			chunkLen = uint64(buf.Len())
+		}
+
+		// Read Chunk
+		chunk := make([]byte, chunkLen)
+		if _, err := buf.Read(chunk); err != nil {
+			return status.Errorf(codes.Internal, "failed to read file content: %v", err)
+		}
+
+		// Send Chunk
+		sendErr := stream.Send(&c2pb.DownloadFileResponse{
+			Chunk: chunk,
+		})
+		if sendErr != nil {
+			return status.Errorf(codes.Internal, "failed to send file content: %v", err)
+		}
+	}
+}

tavern/internal/c2/api_download_file_test.go

@@ -0,0 +1,113 @@
+package c2_test
+
+import (
+	"bytes"
+	"context"
+	"crypto/rand"
+	"errors"
+	"fmt"
+	"io"
+	"testing"
+
+	_ "github.com/mattn/go-sqlite3"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"realm.pub/tavern/internal/c2/c2pb"
+	"realm.pub/tavern/internal/c2/c2test"
+)
+
+func TestDownloadFile(t *testing.T) {
+	// Setup Dependencies
+	ctx := context.Background()
+	client, graph, close := c2test.New(t)
+	defer close()
+
+	// Test Cases
+	type testCase struct {
+		name     string
+		fileName string
+		fileSize int
+		req      *c2pb.DownloadFileRequest
+		wantCode codes.Code
+	}
+	tests := []testCase{
+		{
+			name:     "Small_File",
+			fileName: "small_file",
+			fileSize: 100,
+			req:      &c2pb.DownloadFileRequest{Name: "small_file"},
+			wantCode: codes.OK,
+		},
+		{
+			name:     "Large_File",
+			fileName: "large_file",
+			fileSize: 1024 * 1024 * 10, // 10 MB
+			req:      &c2pb.DownloadFileRequest{Name: "large_file"},
+			wantCode: codes.OK,
+		},
+		{
+			name:     "File Not Found",
+			fileName: "n/a",
+			req:      &c2pb.DownloadFileRequest{Name: "this_file_does_not_exist"},
+			wantCode: codes.NotFound,
+		},
+	}
+
+	testHandler := func(t *testing.T, tc testCase) {
+		// Generate Random Content
+		data := make([]byte, tc.fileSize)
+		_, err := rand.Read(data)
+		require.NoError(t, err)
+
+		// Create File
+		f := graph.File.Create().
+			SetName(tc.fileName).
+			SetContent(data).
+			SaveX(ctx)
+
+		// Send Request
+		fileClient, err := client.DownloadFile(ctx, tc.req)
+		require.NoError(t, err)
+
+		// Read All Chunks
+		var buf bytes.Buffer
+		for {
+			// Receive Chunk
+			resp, err := fileClient.Recv()
+			if errors.Is(err, io.EOF) {
+				break
+			}
+
+			// Check Status
+			require.Equal(t, tc.wantCode.String(), status.Code(err).String())
+			if status.Code(err) != codes.OK {
+				// Do not continue if we expected error code
+				return
+			}
+
+			// Write Chunk
+			_, err = buf.Write(resp.Chunk)
+			require.NoError(t, err)
+		}
+
+		// Assert Content
+		assert.Equal(t, f.Content, buf.Bytes())
+
+		// Assert Headers
+		metadata, err := fileClient.Header()
+		require.NoError(t, err)
+		require.Len(t, metadata.Get("sha3-256-checksum"), 1)
+		assert.Equal(t, f.Hash, metadata.Get("sha3-256-checksum")[0])
+		require.Len(t, metadata.Get("file-size"), 1)
+		assert.Equal(t, fmt.Sprintf("%d", f.Size), metadata.Get("file-size")[0])
+	}
+
+	// Run Tests
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			testHandler(t, tc)
+		})
+	}
+}

tavern/internal/c2/c2.proto

@@ -84,11 +84,30 @@ message ReportTaskOutputRequest {
 
 message ReportTaskOutputResponse {}
 
+message DownloadFileRequest {
+    string name = 1;
+}
+
+message DownloadFileResponse {
+    bytes chunk = 1;
+}
+
 /*
  * Service
  */
 
 service C2 {
     rpc ClaimTasks(ClaimTasksRequest) returns (ClaimTasksResponse) {}
     rpc ReportTaskOutput(ReportTaskOutputRequest) returns (ReportTaskOutputResponse) {}
+
+    /*
+     * Download a file from the server, returning one or more chunks of data.
+     * The maximum size of these chunks is determined by the server.
+     * The server should reply with two headers:
+     *   - "sha3-256-checksum": A SHA3-256 digest of the entire file contents.
+     *   - "file-size": The number of bytes contained by the file.
+     *
+     * If no associated file can be found, a NotFound status error is returned.
+     */
+    rpc DownloadFile(DownloadFileRequest) returns (stream DownloadFileResponse);
 }