Add RunOnce and ScheduleRunOnce to SourceManager

This allows the manager to run an already initialized source.

pkg/context/context.go

@@ -26,9 +26,10 @@ type Context interface {
 	Logger() logr.Logger
 }
 
-// CancelFunc is a type alias to context.CancelFunc to allow use as if they are
-// the same types.
+// CancelFunc and CancelCauseFunc are type aliases to allow use as if they are
+// the same types as the standard library variants.
 type CancelFunc = context.CancelFunc
+type CancelCauseFunc = context.CancelCauseFunc
 
 // logCtx implements Context.
 type logCtx struct {
@@ -68,6 +69,16 @@ func WithCancel(parent Context) (Context, context.CancelFunc) {
 	return lCtx, cancel
 }
 
+// WithCancelCause returns context.WithCancelCause with the log object propagated.
+func WithCancelCause(parent Context) (Context, context.CancelCauseFunc) {
+	ctx, cancel := context.WithCancelCause(parent)
+	lCtx := logCtx{
+		log:     parent.Logger(),
+		Context: ctx,
+	}
+	return lCtx, cancel
+}
+
 // WithDeadline returns context.WithDeadline with the log object propagated and
 // the deadline added to the structured log values.
 func WithDeadline(parent Context, d time.Time) (Context, context.CancelFunc) {
@@ -79,6 +90,17 @@ func WithDeadline(parent Context, d time.Time) (Context, context.CancelFunc) {
 	return lCtx, cancel
 }
 
+// WithDeadlineCause returns context.WithDeadlineCause with the log object
+// propagated and the deadline added to the structured log values.
+func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, context.CancelFunc) {
+	ctx, cancel := context.WithDeadlineCause(parent, d, cause)
+	lCtx := logCtx{
+		log:     parent.Logger().WithValues("deadline", d),
+		Context: ctx,
+	}
+	return lCtx, cancel
+}
+
 // WithTimeout returns context.WithTimeout with the log object propagated and
 // the timeout added to the structured log values.
 func WithTimeout(parent Context, timeout time.Duration) (Context, context.CancelFunc) {
@@ -90,6 +112,22 @@ func WithTimeout(parent Context, timeout time.Duration) (Context, context.Cancel
 	return lCtx, cancel
 }
 
+// WithTimeoutCause returns context.WithTimeoutCause with the log object
+// propagated and the timeout added to the structured log values.
+func WithTimeoutCause(parent Context, timeout time.Duration, cause error) (Context, context.CancelFunc) {
+	ctx, cancel := context.WithTimeoutCause(parent, timeout, cause)
+	lCtx := logCtx{
+		log:     parent.Logger().WithValues("timeout", timeout),
+		Context: ctx,
+	}
+	return lCtx, cancel
+}
+
+// Cause returns the context.Cause of the context.
+func Cause(ctx context.Context) error {
+	return context.Cause(ctx)
+}
+
 // WithValue returns context.WithValue with the log object propagated and
 // the value added to the structured log values (if the key is a string).
 func WithValue(parent Context, key, val any) Context {

pkg/context/context_test.go

@@ -3,6 +3,7 @@ package context
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"strings"
 	"testing"
 	"time"
@@ -177,3 +178,10 @@ func TestRace(t *testing.T) {
 	cancel()
 	_ = ctx.Err()
 }
+
+func TestCause(t *testing.T) {
+	ctx, cancel := WithCancelCause(Background())
+	err := fmt.Errorf("oh no")
+	cancel(err)
+	assert.Equal(t, err, Cause(ctx))
+}

pkg/detectors/slackwebhook/slackwebhook.go

@@ -2,6 +2,7 @@ package slackwebhook
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"net/http"
 	"regexp"
@@ -12,14 +13,15 @@ import (
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
 )
 
-type Scanner struct{}
+type Scanner struct {
+	client *http.Client
+}
 
 // Ensure the Scanner satisfies the interface at compile time.
 var _ detectors.Detector = (*Scanner)(nil)
 
 var (
-	client = common.SaneHttpClient()
-
+	defaultClient = common.SaneHttpClient()
 	// Make sure that your group is surrounded in boundary characters such as below to reduce false positives.
 	keyPat = regexp.MustCompile(`(https://hooks\.slack\.com/services/T[A-Z0-9]+/B[A-Z0-9]+/[A-Za-z0-9]{23,25})`)
 )
@@ -48,6 +50,12 @@ func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (result
 		}
 
 		if verify {
+
+			client := s.client
+			if client == nil {
+				client = defaultClient
+			}
+
 			payload := strings.NewReader(`{"text": ""}`)
 			req, err := http.NewRequestWithContext(ctx, "POST", resMatch, payload)
 			if err != nil {
@@ -62,9 +70,18 @@ func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (result
 					continue
 				}
 				body := string(bodyBytes)
-				if (res.StatusCode >= 200 && res.StatusCode < 300) || (res.StatusCode == 400 && (strings.Contains(body, "no_text") || strings.Contains(body, "missing_text"))) {
+
+				defer res.Body.Close()
+
+				if res.StatusCode >= 200 && res.StatusCode < 300 || (res.StatusCode == 400 && (strings.Contains(body, "no_text") || strings.Contains(body, "missing_text"))) {
 					s1.Verified = true
+				} else if res.StatusCode == 401 || res.StatusCode == 403 {
+					// The secret is determinately not verified (nothing to do)
+				} else {
+					s1.VerificationError = fmt.Errorf("unexpected HTTP response status %d", res.StatusCode)
 				}
+			} else {
+				s1.VerificationError = err
 			}
 		}
 

pkg/detectors/slackwebhook/slackwebhook_test.go

@@ -6,10 +6,11 @@ package slackwebhook
 import (
 	"context"
 	"fmt"
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"testing"
 	"time"
 
-	"github.com/kylelemons/godebug/pretty"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
@@ -32,11 +33,12 @@ func TestSlackWebhook_FromChunk(t *testing.T) {
 		verify bool
 	}
 	tests := []struct {
-		name    string
-		s       Scanner
-		args    args
-		want    []detectors.Result
-		wantErr bool
+		name                string
+		s                   Scanner
+		args                args
+		want                []detectors.Result
+		wantErr             bool
+		wantVerificationErr bool
 	}{
 		{
 			name: "found, verified",
@@ -52,7 +54,8 @@ func TestSlackWebhook_FromChunk(t *testing.T) {
 					Verified:     true,
 				},
 			},
-			wantErr: false,
+			wantErr:             false,
+			wantVerificationErr: false,
 		},
 		{
 			name: "found, unverified",
@@ -68,7 +71,8 @@ func TestSlackWebhook_FromChunk(t *testing.T) {
 					Verified:     false,
 				},
 			},
-			wantErr: false,
+			wantErr:             false,
+			wantVerificationErr: false,
 		},
 		{
 			name: "not found",
@@ -78,14 +82,48 @@ func TestSlackWebhook_FromChunk(t *testing.T) {
 				data:   []byte("You cannot find the secret within"),
 				verify: true,
 			},
-			want:    nil,
-			wantErr: false,
+			want:                nil,
+			wantErr:             false,
+			wantVerificationErr: false,
+		},
+		{
+			name: "found, would be verified if not for timeout",
+			s:    Scanner{client: common.SaneHttpClientTimeOut(1 * time.Microsecond)},
+			args: args{
+				ctx:    context.Background(),
+				data:   []byte(fmt.Sprintf("You can find a slackwebhook secret %s within", secret)),
+				verify: true,
+			},
+			want: []detectors.Result{
+				{
+					DetectorType: detectorspb.DetectorType_SlackWebhook,
+					Verified:     false,
+				},
+			},
+			wantErr:             false,
+			wantVerificationErr: true,
+		},
+		{
+			name: "found, verified but unexpected api surface",
+			s:    Scanner{client: common.ConstantResponseHttpClient(404, "")},
+			args: args{
+				ctx:    context.Background(),
+				data:   []byte(fmt.Sprintf("You can find a slackwebhook secret %s within", secret)),
+				verify: true,
+			},
+			want: []detectors.Result{
+				{
+					DetectorType: detectorspb.DetectorType_SlackWebhook,
+					Verified:     false,
+				},
+			},
+			wantErr:             false,
+			wantVerificationErr: true,
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			s := Scanner{}
-			got, err := s.FromData(tt.args.ctx, tt.args.verify, tt.args.data)
+			got, err := tt.s.FromData(tt.args.ctx, tt.args.verify, tt.args.data)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("SlackWebhook.FromData() error = %v, wantErr %v", err, tt.wantErr)
 				return
@@ -94,9 +132,12 @@ func TestSlackWebhook_FromChunk(t *testing.T) {
 				if len(got[i].Raw) == 0 {
 					t.Fatalf("no raw secret present: \n %+v", got[i])
 				}
-				got[i].Raw = nil
+				if (got[i].VerificationError != nil) != tt.wantVerificationErr {
+					t.Fatalf("wantVerificationError = %v, verification error = %v", tt.wantVerificationErr, got[i].VerificationError)
+				}
 			}
-			if diff := pretty.Compare(got, tt.want); diff != "" {
+			ignoreOpts := cmpopts.IgnoreFields(detectors.Result{}, "Raw", "VerificationError")
+			if diff := cmp.Diff(got, tt.want, ignoreOpts); diff != "" {
 				t.Errorf("SlackWebhook.FromData() %s diff: (-got +want)\n%s", tt.name, diff)
 			}
 		})

pkg/sources/job_progress.go

@@ -71,11 +71,11 @@ func (r *JobProgressRef) Done() <-chan struct{} {
 // CancelRun requests that the job this is referencing is cancelled and stops
 // running. This method will have no effect if the job does not allow
 // cancellation.
-func (r *JobProgressRef) CancelRun() {
+func (r *JobProgressRef) CancelRun(cause error) {
 	if r.jobProgress == nil || r.jobProgress.jobCancel == nil {
 		return
 	}
-	r.jobProgress.jobCancel()
+	r.jobProgress.jobCancel(cause)
 }
 
 // Fatal is a wrapper around error to differentiate non-fatal errors from fatal
@@ -108,7 +108,7 @@ type JobProgress struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 	// Requests to cancel the job.
-	jobCancel context.CancelFunc
+	jobCancel context.CancelCauseFunc
 	// Metrics.
 	metrics     JobProgressMetrics
 	metricsLock sync.Mutex
@@ -150,7 +150,7 @@ func WithHooks(hooks ...JobProgressHook) func(*JobProgress) {
 }
 
 // WithCancel allows cancelling the job by the JobProgressRef.
-func WithCancel(cancel context.CancelFunc) func(*JobProgress) {
+func WithCancel(cancel context.CancelCauseFunc) func(*JobProgress) {
 	return func(jp *JobProgress) { jp.jobCancel = cancel }
 }
 

pkg/sources/s3/s3.go

@@ -104,22 +104,34 @@ func (s *Source) setMaxObjectSize(maxObjectSize int64) {
 	}
 }
 
-func (s *Source) newClient(region string) (*s3.S3, error) {
+func (s *Source) newClient(region, roleArn string) (*s3.S3, error) {
 	cfg := aws.NewConfig()
 	cfg.CredentialsChainVerboseErrors = aws.Bool(true)
 	cfg.Region = aws.String(region)
 
-	switch cred := s.conn.GetCredential().(type) {
-	case *sourcespb.S3_SessionToken:
-		cfg.Credentials = credentials.NewStaticCredentials(cred.SessionToken.Key, cred.SessionToken.Secret, cred.SessionToken.SessionToken)
-	case *sourcespb.S3_AccessKey:
-		cfg.Credentials = credentials.NewStaticCredentials(cred.AccessKey.Key, cred.AccessKey.Secret, "")
-	case *sourcespb.S3_Unauthenticated:
-		cfg.Credentials = credentials.AnonymousCredentials
-	case *sourcespb.S3_CloudEnvironment:
-		// Nothing needs to be done!
-	default:
-		return nil, errors.Errorf("invalid configuration given for %s source", s.name)
+	if roleArn == "" {
+		switch cred := s.conn.GetCredential().(type) {
+		case *sourcespb.S3_SessionToken:
+			cfg.Credentials = credentials.NewStaticCredentials(cred.SessionToken.Key, cred.SessionToken.Secret, cred.SessionToken.SessionToken)
+		case *sourcespb.S3_AccessKey:
+			cfg.Credentials = credentials.NewStaticCredentials(cred.AccessKey.Key, cred.AccessKey.Secret, "")
+		case *sourcespb.S3_Unauthenticated:
+			cfg.Credentials = credentials.AnonymousCredentials
+		case *sourcespb.S3_CloudEnvironment:
+			// Nothing needs to be done!
+		default:
+			return nil, errors.Errorf("invalid configuration given for %s source", s.name)
+		}
+	} else {
+		sess, err := session.NewSession(cfg)
+		if err != nil {
+			return nil, err
+		}
+
+		stsClient := sts.New(sess)
+		cfg.Credentials = stscreds.NewCredentialsWithClient(stsClient, roleArn, func(p *stscreds.AssumeRoleProvider) {
+			p.RoleSessionName = "trufflehog"
+		})
 	}
 
 	sess, err := session.NewSessionWithOptions(session.Options{
@@ -133,40 +145,6 @@ func (s *Source) newClient(region string) (*s3.S3, error) {
 	return s3.New(sess), nil
 }
 
-// Separate role assumption functionality into a different newClient function
-func (s *Source) newRoleClient(region string, roleArn string) (*s3.S3, error) {
-
-	cfg := aws.NewConfig()
-	cfg.CredentialsChainVerboseErrors = aws.Bool(true)
-	cfg.Region = aws.String(region)
-
-	var baseCredentials *credentials.Credentials
-
-	sess, err := session.NewSession(cfg)
-	if err != nil {
-		return nil, err
-	}
-
-	stsClient := sts.New(sess)
-	baseCredentials = stscreds.NewCredentialsWithClient(stsClient, roleArn, func(p *stscreds.AssumeRoleProvider) {
-		p.RoleSessionName = "trufflehog"
-	})
-
-	cfg.Credentials = baseCredentials
-
-	sess, err = session.NewSessionWithOptions(session.Options{
-		SharedConfigState: session.SharedConfigEnable,
-		Config:            *cfg,
-	})
-
-	if err != nil {
-		return nil, err
-	}
-
-	// Create and return the S3 client with the final configuration
-	return s3.New(sess), nil
-}
-
 // IAM identity needs s3:ListBuckets permission
 func (s *Source) getBucketsToScan(client *s3.S3) ([]string, error) {
 	if len(s.conn.Buckets) > 0 {
@@ -204,13 +182,10 @@ func (s *Source) scanBuckets(ctx context.Context, client *s3.S3, role string, bu
 
 		var regionalClient *s3.S3
 		if region != defaultAWSRegion {
-			if role != "" {
-				regionalClient, err = s.newRoleClient(region, role)
-			} else {
-				regionalClient, err = s.newClient(region)
-			}
+			regionalClient, err = s.newClient(region, role)
 			if err != nil {
 				s.log.Error(err, "could not make regional s3 client")
+				continue
 			}
 		} else {
 			regionalClient = client
@@ -238,24 +213,13 @@ func (s *Source) scanBuckets(ctx context.Context, client *s3.S3, role string, bu
 func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) error {
 	const defaultAWSRegion = "us-east-1"
 
-	if len(s.conn.Roles) > 0 {
-		for _, role := range s.conn.Roles {
-			client, err := s.newRoleClient(defaultAWSRegion, role)
-			if err != nil {
-				return errors.WrapPrefix(err, "could not create s3 client", 0)
-			}
+	roles := s.conn.Roles
+	if len(roles) == 0 {
+		roles = []string{""}
+	}
 
-			bucketsToScan, err := s.getBucketsToScan(client)
-			if err != nil {
-				return err
-			}
-
-			if err := s.scanBuckets(ctx, client, role, bucketsToScan, chunksChan); err != nil {
-				return err
-			}
-		}
-	} else {
-		client, err := s.newClient(defaultAWSRegion)
+	for _, role := range roles {
+		client, err := s.newClient(defaultAWSRegion, role)
 		if err != nil {
 			return errors.WrapPrefix(err, "could not create s3 client", 0)
 		}
@@ -265,7 +229,7 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
 			return err
 		}
 
-		if err := s.scanBuckets(ctx, client, "", bucketsToScan, chunksChan); err != nil {
+		if err := s.scanBuckets(ctx, client, role, bucketsToScan, chunksChan); err != nil {
 			return err
 		}
 	}

pkg/sources/source_manager.go

@@ -26,6 +26,7 @@ type SourceInitFunc func(ctx context.Context, jobID, sourceID int64) (Source, er
 // initialization.
 type sourceInfo struct {
 	initFunc SourceInitFunc
+	sourceID int64
 	name     string
 }
 
@@ -126,6 +127,7 @@ func (s *SourceManager) Enroll(ctx context.Context, name string, kind sourcespb.
 	}
 	s.handles[handleID] = sourceInfo{
 		initFunc: f,
+		sourceID: int64(handleID),
 		name:     name,
 	}
 	return handleID, nil
@@ -134,7 +136,18 @@ func (s *SourceManager) Enroll(ctx context.Context, name string, kind sourcespb.
 // Run blocks until a resource is available to run the source, then
 // synchronously runs it. The first fatal error, if any, will be returned.
 func (s *SourceManager) Run(ctx context.Context, handle handle) (JobProgressRef, error) {
-	progress, err := s.asyncRun(ctx, handle)
+	progress, err := s.ScheduleRun(ctx, handle)
+	if err != nil {
+		return progress, err
+	}
+	<-progress.Done()
+	return progress, progress.Snapshot().FatalError()
+}
+
+// RunOnce blocks until a resource is available to run the source, then
+// synchronously runs it. The first fatal error, if any, will be returned.
+func (s *SourceManager) RunOnce(ctx context.Context, name string, source Source) (JobProgressRef, error) {
+	progress, err := s.ScheduleRunOnce(ctx, name, source)
 	if err != nil {
 		return progress, err
 	}
@@ -149,15 +162,22 @@ func (s *SourceManager) ScheduleRun(ctx context.Context, handle handle) (JobProg
 	return s.asyncRun(ctx, handle)
 }
 
+// ScheduleRunOnce blocks until a resource is available to run the source, then
+// asynchronously runs it. Error information is stored and accessible via the
+// JobProgressRef as it becomes available.
+func (s *SourceManager) ScheduleRunOnce(ctx context.Context, name string, source Source) (JobProgressRef, error) {
+	return s.asyncRunOnce(ctx, name, source)
+}
+
 // asyncRun is a helper method to asynchronously run the Source. It calls out
 // to the API to get a job ID for this run, creates a JobProgress object, then
 // waits for an available goroutine to asynchronously run it.
 func (s *SourceManager) asyncRun(ctx context.Context, handle handle) (JobProgressRef, error) {
 	// Do preflight checks before waiting on the pool.
-	if err := s.preflightChecks(ctx, handle); err != nil {
+	if err := s.preflightChecks(ctx); err != nil {
 		return JobProgressRef{}, err
 	}
-	// Get the name. Should never fail due to preflight checks.
+	// Get the name.
 	sourceInfo, ok := s.getSourceInfo(handle)
 	if !ok {
 		return JobProgressRef{SourceID: int64(handle)}, fmt.Errorf("unrecognized handle")
@@ -170,7 +190,7 @@ func (s *SourceManager) asyncRun(ctx context.Context, handle handle) (JobProgres
 		return JobProgressRef{SourceID: int64(handle), SourceName: sourceName}, err
 	}
 	// Create a JobProgress object for tracking progress.
-	ctx, cancel := context.WithCancel(ctx)
+	ctx, cancel := context.WithCancelCause(ctx)
 	progress := NewJobProgress(jobID, int64(handle), sourceName, WithHooks(s.hooks...), WithCancel(cancel))
 	s.pool.Go(func() error {
 		atomic.AddInt32(&s.currentRunningCount, 1)
@@ -180,8 +200,37 @@ func (s *SourceManager) asyncRun(ctx context.Context, handle handle) (JobProgres
 			"source_manager_worker_id", common.RandomID(5),
 		)
 		defer common.Recover(ctx)
-		defer cancel()
-		return s.run(ctx, handle, jobID, progress)
+		defer cancel(nil)
+		return s.run(ctx, sourceInfo, jobID, progress)
+	})
+	return progress.Ref(), nil
+}
+
+// asyncRunOnce is a helper method to asynchronously run an already initialized
+// Source once. Unlike asyncRun, this method doesn't call out to the API for a
+// job ID since it's already available in the Source.
+func (s *SourceManager) asyncRunOnce(ctx context.Context, sourceName string, source Source) (JobProgressRef, error) {
+	if err := s.preflightChecks(ctx); err != nil {
+		return JobProgressRef{}, err
+	}
+	ctx = context.WithValue(ctx, "source_id", source.SourceID())
+	jobID := source.JobID()
+	// Create a JobProgress object for tracking progress.
+	ctx, cancel := context.WithCancelCause(ctx)
+	progress := NewJobProgress(jobID, source.SourceID(), sourceName, WithHooks(s.hooks...), WithCancel(cancel))
+	s.pool.Go(func() error {
+		ctx := context.WithValues(ctx,
+			"job_id", jobID,
+			"source_manager_worker_id", common.RandomID(5),
+		)
+		defer common.Recover(ctx)
+		defer cancel(nil)
+		sourceInfo := sourceInfo{
+			initFunc: func(context.Context, int64, int64) (Source, error) { return source, nil },
+			sourceID: source.SourceID(),
+			name:     sourceName,
+		}
+		return s.run(ctx, sourceInfo, jobID, progress)
 	})
 	return progress.Ref(), nil
 }
@@ -225,37 +274,31 @@ func (s *SourceManager) AvailableCapacity() int {
 	return s.poolLimit - int(runCount)
 }
 
-// preflightChecks is a helper method to check the Manager or the context isn't
-// done and that the handle is valid.
-func (s *SourceManager) preflightChecks(ctx context.Context, handle handle) error {
+// preflightChecks is a helper method to check the Manager or the context isn't done.
+func (s *SourceManager) preflightChecks(ctx context.Context) error {
 	// Check if the manager has been Waited.
 	if s.done {
 		return fmt.Errorf("manager is done")
 	}
-	// Check the handle is valid.
-	if _, ok := s.getSourceInfo(handle); !ok {
-		return fmt.Errorf("unrecognized handle")
-	}
 	return ctx.Err()
 }
 
 // run is a helper method to sychronously run the source. It does not check for
 // acquired resources. An error is returned if there was a fatal error during
 // the run. This information is also recorded in the JobProgress.
-func (s *SourceManager) run(ctx context.Context, handle handle, jobID int64, report *JobProgress) error {
+func (s *SourceManager) run(ctx context.Context, sourceInfo sourceInfo, jobID int64, report *JobProgress) error {
 	defer report.Finish()
 	report.Start(time.Now())
 	defer func() { report.End(time.Now()) }()
 
+	defer func() {
+		if err := context.Cause(ctx); err != nil {
+			report.ReportError(Fatal{err})
+		}
+	}()
+
 	// Initialize the source.
-	sourceInfo, ok := s.getSourceInfo(handle)
-	if !ok {
-		// Shouldn't happen due to preflight checks.
-		err := fmt.Errorf("unrecognized handle")
-		report.ReportError(Fatal{err})
-		return Fatal{err}
-	}
-	source, err := sourceInfo.initFunc(ctx, jobID, int64(handle))
+	source, err := sourceInfo.initFunc(ctx, jobID, sourceInfo.sourceID)
 	if err != nil {
 		report.ReportError(Fatal{err})
 		return Fatal{err}
@@ -267,14 +310,14 @@ func (s *SourceManager) run(ctx context.Context, handle handle, jobID int64, rep
 	)
 	// Check for the preferred method of tracking source units.
 	if enumChunker, ok := source.(SourceUnitEnumChunker); ok && s.useSourceUnits {
-		return s.runWithUnits(ctx, handle, enumChunker, report)
+		return s.runWithUnits(ctx, enumChunker, report)
 	}
-	return s.runWithoutUnits(ctx, handle, source, report)
+	return s.runWithoutUnits(ctx, source, report)
 }
 
 // runWithoutUnits is a helper method to run a Source. It has coarse-grained
 // job reporting.
-func (s *SourceManager) runWithoutUnits(ctx context.Context, handle handle, source Source, report *JobProgress) error {
+func (s *SourceManager) runWithoutUnits(ctx context.Context, source Source, report *JobProgress) error {
 	// Introspect on the chunks we get from the Chunks method.
 	ch := make(chan *Chunk, 1)
 	var wg sync.WaitGroup
@@ -304,7 +347,7 @@ func (s *SourceManager) runWithoutUnits(ctx context.Context, handle handle, sour
 // runWithUnits is a helper method to run a Source that is also a
 // SourceUnitEnumChunker. This allows better introspection of what is getting
 // scanned and any errors encountered.
-func (s *SourceManager) runWithUnits(ctx context.Context, handle handle, source SourceUnitEnumChunker, report *JobProgress) error {
+func (s *SourceManager) runWithUnits(ctx context.Context, source SourceUnitEnumChunker, report *JobProgress) error {
 	unitReporter := &mgrUnitReporter{
 		unitCh: make(chan SourceUnit, 1),
 		report: report,

pkg/sources/source_manager_test.go

@@ -14,8 +14,9 @@ import (
 
 // DummySource implements Source and is used for testing a SourceManager.
 type DummySource struct {
-	sourceID int64
-	jobID    int64
+	sourceID   int64
+	jobID      int64
+	initCalled bool
 	chunker
 }
 
@@ -23,6 +24,7 @@ func (d *DummySource) Type() sourcespb.SourceType { return 1337 }
 func (d *DummySource) SourceID() int64            { return d.sourceID }
 func (d *DummySource) JobID() int64               { return d.jobID }
 func (d *DummySource) Init(_ context.Context, _ string, jobID, sourceID int64, _ bool, _ *anypb.Any, _ int) error {
+	d.initCalled = true
 	d.sourceID = sourceID
 	d.jobID = jobID
 	return nil
@@ -312,10 +314,12 @@ func TestSourceManagerCancelRun(t *testing.T) {
 	ref, err := mgr.ScheduleRun(context.Background(), handle)
 	assert.NoError(t, err)
 
-	ref.CancelRun()
+	cancelErr := fmt.Errorf("abort! abort!")
+	ref.CancelRun(cancelErr)
 	<-ref.Done()
 	assert.Error(t, ref.Snapshot().FatalError())
 	assert.True(t, errors.Is(ref.Snapshot().FatalError(), returnedErr))
+	assert.True(t, errors.Is(ref.Snapshot().FatalErrors(), cancelErr))
 }
 
 func TestSourceManagerAvailableCapacity(t *testing.T) {
@@ -338,3 +342,21 @@ func TestSourceManagerAvailableCapacity(t *testing.T) {
 	<-ref.Done()      // Wait for the job to finish.
 	assert.Equal(t, 1337, mgr.AvailableCapacity())
 }
+
+func TestSourceManagerRunOnce(t *testing.T) {
+	mgr := NewManager()
+	var chunksCalled bool
+	source := &DummySource{sourceID: 123, jobID: 456, chunker: callbackChunker{
+		func(context.Context, chan *Chunk) error {
+			chunksCalled = true
+			return nil
+		},
+	}}
+	ref, err := mgr.RunOnce(context.Background(), "test", source)
+	assert.NoError(t, err)
+	assert.Equal(t, source.jobID, ref.JobID)
+	assert.Equal(t, source.sourceID, ref.SourceID)
+	assert.Equal(t, "test", ref.SourceName)
+	assert.True(t, chunksCalled)
+	assert.False(t, source.initCalled)
+}