[feature] gRPC File Download API

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-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-checksum"), 1)
+		assert.Equal(t, f.Hash, metadata.Get("sha3-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,19 @@ 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) {}
+    rpc DownloadFile(DownloadFileRequest) returns (stream DownloadFileResponse);
 }