@@ -2,6 +2,7 @@ package provisionerdserver_test
22
33import (
44 "context"
5+ crand "crypto/rand"
56 "database/sql"
67 "encoding/json"
78 "io"
@@ -25,6 +26,8 @@ import (
2526 "golang.org/x/xerrors"
2627 "google.golang.org/protobuf/types/known/timestamppb"
2728 "storj.io/drpc"
29+ "storj.io/drpc/drpcmux"
30+ "storj.io/drpc/drpcserver"
2831
2932 "cdr.dev/slog/v3"
3033 "cdr.dev/slog/v3/sloggers/slogtest"
@@ -52,6 +55,7 @@ import (
5255 "github.com/coder/coder/v2/coderd/usage/usagetypes"
5356 "github.com/coder/coder/v2/coderd/wspubsub"
5457 "github.com/coder/coder/v2/codersdk"
58+ "github.com/coder/coder/v2/codersdk/drpcsdk"
5559 "github.com/coder/coder/v2/provisionerd/proto"
5660 "github.com/coder/coder/v2/provisionersdk"
5761 sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
@@ -5420,3 +5424,141 @@ func newFakeUsageInserter() (*coderdtest.UsageInserter, *atomic.Pointer[usage.In
54205424 poitr .Store (& inserter )
54215425 return fake , poitr
54225426}
5427+
5428+ // serveProvisionerDaemon serves the provisioner daemon server over an
5429+ // in-memory pipe and returns a connected client, mirroring how coderd serves
5430+ // in-memory provisioner daemons. This exercises the real DRPC streaming path
5431+ // instead of a hand-rolled mock stream.
5432+ func serveProvisionerDaemon (t * testing.T , srv proto.DRPCProvisionerDaemonServer ) proto.DRPCProvisionerDaemonClient {
5433+ t .Helper ()
5434+ clientPipe , serverPipe := drpcsdk .MemTransportPipe ()
5435+ t .Cleanup (func () {
5436+ _ = clientPipe .Close ()
5437+ _ = serverPipe .Close ()
5438+ })
5439+ mux := drpcmux .New ()
5440+ require .NoError (t , proto .DRPCRegisterProvisionerDaemon (mux , srv ))
5441+ server := drpcserver .NewWithOptions (mux , drpcserver.Options {
5442+ Manager : drpcsdk .DefaultDRPCOptions (nil ),
5443+ })
5444+ ctx , cancel := context .WithCancel (context .Background ())
5445+ closed := make (chan struct {})
5446+ go func () {
5447+ defer close (closed )
5448+ _ = server .Serve (ctx , serverPipe )
5449+ }()
5450+ t .Cleanup (func () {
5451+ cancel ()
5452+ <- closed
5453+ })
5454+ return proto .NewDRPCProvisionerDaemonClient (clientPipe )
5455+ }
5456+
5457+ // insertModuleFile inserts a system-created (CreatedBy=uuid.Nil) tar file and
5458+ // links it as the cached module files of a template version in the given
5459+ // organization, returning the file.
5460+ func insertModuleFile (t * testing.T , db database.Store , orgID uuid.UUID , data []byte ) database.File {
5461+ t .Helper ()
5462+ ctx := testutil .Context (t , testutil .WaitShort )
5463+
5464+ user := dbgen .User (t , db , database.User {})
5465+ template := dbgen .Template (t , db , database.Template {
5466+ OrganizationID : orgID ,
5467+ CreatedBy : user .ID ,
5468+ })
5469+ jobID := uuid .New ()
5470+ version := dbgen .TemplateVersion (t , db , database.TemplateVersion {
5471+ OrganizationID : orgID ,
5472+ CreatedBy : user .ID ,
5473+ TemplateID : uuid.NullUUID {UUID : template .ID , Valid : true },
5474+ JobID : jobID ,
5475+ })
5476+ // Insert the file directly rather than via dbgen.File: the helper treats a
5477+ // zero CreatedBy as "unset" and replaces it with a random UUID, but module
5478+ // files must be system-created (CreatedBy=uuid.Nil) to match the handler's
5479+ // metadata check.
5480+ file , err := db .InsertFile (ctx , database.InsertFileParams {
5481+ ID : uuid .New (),
5482+ Hash : uuid .NewString (),
5483+ CreatedAt : dbtime .Now (),
5484+ CreatedBy : uuid .Nil ,
5485+ Mimetype : "application/x-tar" ,
5486+ Data : data ,
5487+ })
5488+ require .NoError (t , err )
5489+ err = db .InsertTemplateVersionTerraformValuesByJobID (ctx , database.InsertTemplateVersionTerraformValuesByJobIDParams {
5490+ JobID : version .JobID ,
5491+ CachedPlan : []byte ("{}" ),
5492+ CachedModuleFiles : uuid.NullUUID {UUID : file .ID , Valid : true },
5493+ UpdatedAt : dbtime .Now (),
5494+ })
5495+ require .NoError (t , err )
5496+ return file
5497+ }
5498+
5499+ // TestDownloadFile verifies that a provisioner daemon cannot download cached
5500+ // module archives belonging to other organizations (ANT-2026-22440), while
5501+ // still being able to download module files from its own organization.
5502+ func TestDownloadFile (t * testing.T ) {
5503+ t .Parallel ()
5504+
5505+ t .Run ("RejectsOtherOrgModuleFile" , func (t * testing.T ) {
5506+ t .Parallel ()
5507+
5508+ // The server is scoped to the default organization (org A).
5509+ srv , db , _ , daemon := setup (t , false , & overrides {
5510+ externalAuthConfigs : []* externalauth.Config {{}},
5511+ })
5512+ ctx := testutil .Context (t , testutil .WaitMedium )
5513+ client := serveProvisionerDaemon (t , srv )
5514+
5515+ // Create a module file belonging to a different organization (org B).
5516+ otherOrg := dbgen .Organization (t , db , database.Organization {})
5517+ require .NotEqual (t , daemon .OrganizationID , otherOrg .ID )
5518+
5519+ moduleData := make ([]byte , sdkproto .ChunkSize * 2 )
5520+ // crand.Read never returns an error as of Go 1.24.
5521+ _ , _ = crand .Read (moduleData )
5522+ file := insertModuleFile (t , db , otherOrg .ID , moduleData )
5523+
5524+ stream , err := client .DownloadFile (ctx , & proto.FileRequest {
5525+ FileId : file .ID .String (),
5526+ UploadType : sdkproto .DataUploadType_UPLOAD_TYPE_MODULE_FILES ,
5527+ })
5528+ require .NoError (t , err )
5529+
5530+ // The handler must reject the cross-org download with an error rather
5531+ // than streaming the file's contents.
5532+ _ , err = provisionersdk .HandleReceivingDataUpload (stream )
5533+ require .Error (t , err )
5534+ require .ErrorContains (t , err , "is not a modules file" )
5535+ })
5536+
5537+ t .Run ("AllowsSameOrgModuleFile" , func (t * testing.T ) {
5538+ t .Parallel ()
5539+
5540+ // The server is scoped to the default organization (org A).
5541+ srv , db , _ , daemon := setup (t , false , & overrides {
5542+ externalAuthConfigs : []* externalauth.Config {{}},
5543+ })
5544+ ctx := testutil .Context (t , testutil .WaitMedium )
5545+ client := serveProvisionerDaemon (t , srv )
5546+
5547+ moduleData := make ([]byte , sdkproto .ChunkSize * 2 + 512 )
5548+ // crand.Read never returns an error as of Go 1.24.
5549+ _ , _ = crand .Read (moduleData )
5550+ file := insertModuleFile (t , db , daemon .OrganizationID , moduleData )
5551+
5552+ stream , err := client .DownloadFile (ctx , & proto.FileRequest {
5553+ FileId : file .ID .String (),
5554+ UploadType : sdkproto .DataUploadType_UPLOAD_TYPE_MODULE_FILES ,
5555+ })
5556+ require .NoError (t , err )
5557+
5558+ builder , err := provisionersdk .HandleReceivingDataUpload (stream )
5559+ require .NoError (t , err )
5560+ data , err := builder .Complete ()
5561+ require .NoError (t , err )
5562+ require .Equal (t , moduleData , data )
5563+ })
5564+ }
0 commit comments