feat: use proto streams to increase maximum module files payload (coder#18268)

This PR implements protobuf streaming to handle large module files by:
1. **Streaming large payloads**: When module files exceed the 4MB limit,
they're streamed in chunks using a new UploadFile RPC method
2. **Database storage**: Streamed files are stored in the database and
referenced by hash for deduplication
3. **Backward compatibility**: Small module files continue using the
existing direct payload method

Author: Emyrk

cli/testdata/coder_provisioner_list_--output_json.golden

@@ -7,7 +7,7 @@
     "last_seen_at": "====[timestamp]=====",
     "name": "test-daemon",
     "version": "v0.0.0-devel",
-    "api_version": "1.6",
+    "api_version": "1.7",
     "provisioners": [
       "echo"
     ],

coderd/database/dbauthz/dbauthz.go

@@ -171,7 +171,7 @@ var (
 				DisplayName: "Provisioner Daemon",
 				Site: rbac.Permissions(map[string][]policy.Action{
 					rbac.ResourceProvisionerJobs.Type: {policy.ActionRead, policy.ActionUpdate, policy.ActionCreate},
-					rbac.ResourceFile.Type:            {policy.ActionRead},
+					rbac.ResourceFile.Type:            {policy.ActionCreate, policy.ActionRead},
 					rbac.ResourceSystem.Type:          {policy.WildcardSymbol},
 					rbac.ResourceTemplate.Type:        {policy.ActionRead, policy.ActionUpdate},
 					// Unsure why provisionerd needs update and read personal

coderd/database/dbmem/dbmem.go

@@ -8743,6 +8743,12 @@ func (q *FakeQuerier) InsertFile(_ context.Context, arg database.InsertFileParam
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
 
+	if slices.ContainsFunc(q.files, func(file database.File) bool {
+		return file.CreatedBy == arg.CreatedBy && file.Hash == arg.Hash
+	}) {
+		return database.File{}, newUniqueConstraintError(database.UniqueFilesHashCreatedByKey)
+	}
+
 	//nolint:gosimple
 	file := database.File{
 		ID:        arg.ID,

coderd/provisionerdserver/provisionerdserver.go

@@ -773,7 +773,7 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
 	case database.ProvisionerStorageMethodFile:
 		file, err := s.Database.GetFileByID(ctx, job.FileID)
 		if err != nil {
-			return nil, failJob(fmt.Sprintf("get file by hash: %s", err))
+			return nil, failJob(fmt.Sprintf("get file by id: %s", err))
 		}
 		protoJob.TemplateSourceArchive = file.Data
 	default:
@@ -1321,6 +1321,104 @@ func (s *server) prepareForNotifyWorkspaceManualBuildFailed(ctx context.Context,
 	return templateAdmins, template, templateVersion, workspaceOwner, nil
 }
 
+func (s *server) UploadFile(stream proto.DRPCProvisionerDaemon_UploadFileStream) error {
+	var file *sdkproto.DataBuilder
+	// Always terminate the stream with an empty response.
+	defer stream.SendAndClose(&proto.Empty{})
+
+UploadFileStream:
+	for {
+		msg, err := stream.Recv()
+		if err != nil {
+			return xerrors.Errorf("receive complete job with files: %w", err)
+		}
+
+		switch typed := msg.Type.(type) {
+		case *proto.UploadFileRequest_DataUpload:
+			if file != nil {
+				return xerrors.New("unexpected file upload while waiting for file completion")
+			}
+
+			file, err = sdkproto.NewDataBuilder(&sdkproto.DataUpload{
+				UploadType: typed.DataUpload.UploadType,
+				DataHash:   typed.DataUpload.DataHash,
+				FileSize:   typed.DataUpload.FileSize,
+				Chunks:     typed.DataUpload.Chunks,
+			})
+			if err != nil {
+				return xerrors.Errorf("unable to create file upload: %w", err)
+			}
+
+			if file.IsDone() {
+				// If a file is 0 bytes, we can consider it done immediately.
+				// This should never really happen in practice, but we handle it gracefully.
+				break UploadFileStream
+			}
+		case *proto.UploadFileRequest_ChunkPiece:
+			if file == nil {
+				return xerrors.New("unexpected chunk piece while waiting for file upload")
+			}
+
+			done, err := file.Add(&sdkproto.ChunkPiece{
+				Data:         typed.ChunkPiece.Data,
+				FullDataHash: typed.ChunkPiece.FullDataHash,
+				PieceIndex:   typed.ChunkPiece.PieceIndex,
+			})
+			if err != nil {
+				return xerrors.Errorf("unable to add chunk piece: %w", err)
+			}
+
+			if done {
+				break UploadFileStream
+			}
+		}
+	}
+
+	fileData, err := file.Complete()
+	if err != nil {
+		return xerrors.Errorf("complete file upload: %w", err)
+	}
+
+	// Just rehash the data to be sure it is correct.
+	hashBytes := sha256.Sum256(fileData)
+	hash := hex.EncodeToString(hashBytes[:])
+
+	var insert database.InsertFileParams
+
+	switch file.Type {
+	case sdkproto.DataUploadType_UPLOAD_TYPE_MODULE_FILES:
+		insert = database.InsertFileParams{
+			ID:        uuid.New(),
+			Hash:      hash,
+			CreatedAt: dbtime.Now(),
+			CreatedBy: uuid.Nil,
+			Mimetype:  tarMimeType,
+			Data:      fileData,
+		}
+	default:
+		return xerrors.Errorf("unsupported file upload type: %s", file.Type)
+	}
+
+	//nolint:gocritic // Provisionerd actor
+	_, err = s.Database.InsertFile(dbauthz.AsProvisionerd(s.lifecycleCtx), insert)
+	if err != nil {
+		// Duplicated files already exist in the database, so we can ignore this error.
+		if !database.IsUniqueViolation(err, database.UniqueFilesHashCreatedByKey) {
+			return xerrors.Errorf("insert file: %w", err)
+		}
+	}
+
+	s.Logger.Info(s.lifecycleCtx, "file uploaded to database",
+		slog.F("type", file.Type.String()),
+		slog.F("hash", hash),
+		slog.F("size", len(fileData)),
+		// new_insert indicates whether the file was newly inserted or already existed.
+		slog.F("new_insert", err == nil),
+	)
+
+	return nil
+}
+
 // CompleteJob is triggered by a provision daemon to mark a provisioner job as completed.
 func (s *server) CompleteJob(ctx context.Context, completed *proto.CompletedJob) (*proto.Empty, error) {
 	ctx, span := s.startTrace(ctx, tracing.FuncName())
@@ -1606,6 +1704,20 @@ func (s *server) completeTemplateImportJob(ctx context.Context, job database.Pro
 				}
 			}
 
+			if len(jobType.TemplateImport.ModuleFilesHash) > 0 {
+				hashString := hex.EncodeToString(jobType.TemplateImport.ModuleFilesHash)
+				//nolint:gocritic // Acting as provisioner
+				file, err := db.GetFileByHashAndCreator(dbauthz.AsProvisionerd(ctx), database.GetFileByHashAndCreatorParams{Hash: hashString, CreatedBy: uuid.Nil})
+				if err != nil {
+					return xerrors.Errorf("get file by hash, it should have been uploaded: %w", err)
+				}
+
+				fileID = uuid.NullUUID{
+					Valid: true,
+					UUID:  file.ID,
+				}
+			}
+
 			err = db.InsertTemplateVersionTerraformValuesByJobID(ctx, database.InsertTemplateVersionTerraformValuesByJobIDParams{
 				JobID:               jobID,
 				UpdatedAt:           now,

coderd/provisionerdserver/upload_file_test.go

@@ -0,0 +1,191 @@
+package provisionerdserver_test
+
+import (
+	"context"
+	crand "crypto/rand"
+	"fmt"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
+	"storj.io/drpc"
+
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/externalauth"
+	"github.com/coder/coder/v2/codersdk/drpcsdk"
+	proto "github.com/coder/coder/v2/provisionerd/proto"
+	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
+	"github.com/coder/coder/v2/testutil"
+)
+
+// TestUploadFileLargeModuleFiles tests the UploadFile RPC with large module files
+func TestUploadFileLargeModuleFiles(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitMedium)
+
+	// Create server
+	server, db, _, _ := setup(t, false, &overrides{
+		externalAuthConfigs: []*externalauth.Config{{}},
+	})
+
+	testSizes := []int{
+		0,                             // Empty file
+		512,                           // A small file
+		drpcsdk.MaxMessageSize + 1024, // Just over the limit
+		drpcsdk.MaxMessageSize * 2,    // 2x the limit
+		sdkproto.ChunkSize*3 + 512,    // Multiple chunks with partial last
+	}
+
+	for _, size := range testSizes {
+		t.Run(fmt.Sprintf("size_%d_bytes", size), func(t *testing.T) {
+			t.Parallel()
+
+			// Generate test module files data
+			moduleData := make([]byte, size)
+			_, err := crand.Read(moduleData)
+			require.NoError(t, err)
+
+			// Convert to upload format
+			upload, chunks := sdkproto.BytesToDataUpload(sdkproto.DataUploadType_UPLOAD_TYPE_MODULE_FILES, moduleData)
+
+			stream := newMockUploadStream(upload, chunks...)
+
+			// Execute upload
+			err = server.UploadFile(stream)
+			require.NoError(t, err)
+
+			// Upload should be done
+			require.True(t, stream.isDone(), "stream should be done after upload")
+
+			// Verify file was stored in database
+			hashString := fmt.Sprintf("%x", upload.DataHash)
+			file, err := db.GetFileByHashAndCreator(ctx, database.GetFileByHashAndCreatorParams{
+				Hash:      hashString,
+				CreatedBy: uuid.Nil, // Provisionerd creates with Nil UUID
+			})
+			require.NoError(t, err)
+			require.Equal(t, hashString, file.Hash)
+			require.Equal(t, moduleData, file.Data)
+			require.Equal(t, "application/x-tar", file.Mimetype)
+
+			// Try to upload it again, and it should still be successful
+			stream = newMockUploadStream(upload, chunks...)
+			err = server.UploadFile(stream)
+			require.NoError(t, err, "re-upload should succeed without error")
+			require.True(t, stream.isDone(), "stream should be done after re-upload")
+		})
+	}
+}
+
+// TestUploadFileErrorScenarios tests various error conditions in file upload
+func TestUploadFileErrorScenarios(t *testing.T) {
+	t.Parallel()
+
+	//nolint:dogsled
+	server, _, _, _ := setup(t, false, &overrides{
+		externalAuthConfigs: []*externalauth.Config{{}},
+	})
+
+	// Generate test data
+	moduleData := make([]byte, sdkproto.ChunkSize*2)
+	_, err := crand.Read(moduleData)
+	require.NoError(t, err)
+
+	upload, chunks := sdkproto.BytesToDataUpload(sdkproto.DataUploadType_UPLOAD_TYPE_MODULE_FILES, moduleData)
+
+	t.Run("chunk_before_upload", func(t *testing.T) {
+		t.Parallel()
+
+		stream := newMockUploadStream(nil, chunks[0])
+
+		err := server.UploadFile(stream)
+		require.ErrorContains(t, err, "unexpected chunk piece while waiting for file upload")
+		require.True(t, stream.isDone(), "stream should be done after error")
+	})
+
+	t.Run("duplicate_upload", func(t *testing.T) {
+		t.Parallel()
+
+		stream := &mockUploadStream{
+			done:     make(chan struct{}),
+			messages: make(chan *proto.UploadFileRequest, 2),
+		}
+
+		up := &proto.UploadFileRequest{Type: &proto.UploadFileRequest_DataUpload{DataUpload: upload}}
+
+		// Send it twice
+		stream.messages <- up
+		stream.messages <- up
+
+		err := server.UploadFile(stream)
+		require.ErrorContains(t, err, "unexpected file upload while waiting for file completion")
+		require.True(t, stream.isDone(), "stream should be done after error")
+	})
+
+	t.Run("unsupported_upload_type", func(t *testing.T) {
+		t.Parallel()
+
+		//nolint:govet // Ignore lock copy
+		cpy := *upload
+		cpy.UploadType = sdkproto.DataUploadType_UPLOAD_TYPE_UNKNOWN // Set to an unsupported type
+		stream := newMockUploadStream(&cpy, chunks...)
+
+		err := server.UploadFile(stream)
+		require.ErrorContains(t, err, "unsupported file upload type")
+		require.True(t, stream.isDone(), "stream should be done after error")
+	})
+}
+
+type mockUploadStream struct {
+	done     chan struct{}
+	messages chan *proto.UploadFileRequest
+}
+
+func (m mockUploadStream) SendAndClose(empty *proto.Empty) error {
+	close(m.done)
+	return nil
+}
+
+func (m mockUploadStream) Recv() (*proto.UploadFileRequest, error) {
+	msg, ok := <-m.messages
+	if !ok {
+		return nil, xerrors.New("no more messages to receive")
+	}
+	return msg, nil
+}
+func (*mockUploadStream) Context() context.Context { panic(errUnimplemented) }
+func (*mockUploadStream) MsgSend(msg drpc.Message, enc drpc.Encoding) error {
+	panic(errUnimplemented)
+}
+
+func (*mockUploadStream) MsgRecv(msg drpc.Message, enc drpc.Encoding) error {
+	panic(errUnimplemented)
+}
+func (*mockUploadStream) CloseSend() error { panic(errUnimplemented) }
+func (*mockUploadStream) Close() error     { panic(errUnimplemented) }
+func (m *mockUploadStream) isDone() bool {
+	select {
+	case <-m.done:
+		return true
+	default:
+		return false
+	}
+}
+
+func newMockUploadStream(up *sdkproto.DataUpload, chunks ...*sdkproto.ChunkPiece) *mockUploadStream {
+	stream := &mockUploadStream{
+		done:     make(chan struct{}),
+		messages: make(chan *proto.UploadFileRequest, 1+len(chunks)),
+	}
+	if up != nil {
+		stream.messages <- &proto.UploadFileRequest{Type: &proto.UploadFileRequest_DataUpload{DataUpload: up}}
+	}
+
+	for _, chunk := range chunks {
+		stream.messages <- &proto.UploadFileRequest{Type: &proto.UploadFileRequest_ChunkPiece{ChunkPiece: chunk}}
+	}
+	close(stream.messages)
+	return stream
+}