Skip to content

Commit da597c4

Browse files
knqyf263knqyf263
authored
refactor: propagate time through context values (aquasecurity#5858)
Signed-off-by: knqyf263 <[email protected]>
1 parent 1607eee commit da597c4

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

77 files changed

+384
-546
lines changed

‎integration/client_server_test.go

-4
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ import (
1717
"github.com/stretchr/testify/require"
1818
testcontainers "github.com/testcontainers/testcontainers-go"
1919

20-
"github.com/aquasecurity/trivy/pkg/clock"
2120
"github.com/aquasecurity/trivy/pkg/report"
2221
"github.com/aquasecurity/trivy/pkg/uuid"
2322
)
@@ -364,8 +363,6 @@ func TestClientServerWithFormat(t *testing.T) {
364363
}
365364

366365
fakeTime := time.Date(2021, 8, 25, 12, 20, 30, 5, time.UTC)
367-
clock.SetFakeTime(t, fakeTime)
368-
369366
report.CustomTemplateFuncMap = map[string]interface{}{
370367
"now": func() time.Time {
371368
return fakeTime
@@ -428,7 +425,6 @@ func TestClientServerWithCycloneDX(t *testing.T) {
428425
addr, cacheDir := setup(t, setupOptions{})
429426
for _, tt := range tests {
430427
t.Run(tt.name, func(t *testing.T) {
431-
clock.SetFakeTime(t, time.Date(2021, 8, 25, 12, 20, 30, 5, time.UTC))
432428
uuid.SetFakeUUID(t, "3ff14136-e09f-4df9-80ea-%012d")
433429

434430
osArgs, outputFile := setupClient(t, tt.args, addr, cacheDir, tt.golden)

‎integration/integration_test.go

+6-5
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"encoding/json"
88
"flag"
99
"fmt"
10+
"github.com/aquasecurity/trivy/pkg/clock"
1011
"io"
1112
"net"
1213
"os"
@@ -27,7 +28,6 @@ import (
2728

2829
"github.com/aquasecurity/trivy-db/pkg/db"
2930
"github.com/aquasecurity/trivy-db/pkg/metadata"
30-
"github.com/aquasecurity/trivy/pkg/clock"
3131
"github.com/aquasecurity/trivy/pkg/commands"
3232
"github.com/aquasecurity/trivy/pkg/dbtest"
3333
"github.com/aquasecurity/trivy/pkg/types"
@@ -44,8 +44,6 @@ func initDB(t *testing.T) string {
4444
entries, err := os.ReadDir(fixtureDir)
4545
require.NoError(t, err)
4646

47-
clock.SetFakeTime(t, time.Date(2021, 8, 25, 12, 20, 30, 5, time.UTC))
48-
4947
var fixtures []string
5048
for _, entry := range entries {
5149
if entry.IsDir() {
@@ -193,13 +191,16 @@ func readSpdxJson(t *testing.T, filePath string) *spdx.Document {
193191
}
194192

195193
func execute(osArgs []string) error {
194+
// Set a fake time
195+
ctx := clock.With(context.Background(), time.Date(2021, 8, 25, 12, 20, 30, 5, time.UTC))
196+
196197
// Setup CLI App
197198
app := commands.NewApp()
198199
app.SetOut(io.Discard)
200+
app.SetArgs(osArgs)
199201

200202
// Run Trivy
201-
app.SetArgs(osArgs)
202-
return app.Execute()
203+
return app.ExecuteContext(ctx)
203204
}
204205

205206
func compareReports(t *testing.T, wantFile, gotFile string, override func(*types.Report)) {

‎integration/repo_test.go

+8-9
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,13 @@ package integration
44

55
import (
66
"fmt"
7+
"github.com/stretchr/testify/assert"
8+
"github.com/stretchr/testify/require"
79
"os"
810
"path/filepath"
911
"strings"
1012
"testing"
11-
"time"
12-
13-
"github.com/stretchr/testify/assert"
14-
"github.com/stretchr/testify/require"
1513

16-
"github.com/aquasecurity/trivy/pkg/clock"
1714
ftypes "github.com/aquasecurity/trivy/pkg/fanal/types"
1815
"github.com/aquasecurity/trivy/pkg/types"
1916
"github.com/aquasecurity/trivy/pkg/uuid"
@@ -416,12 +413,15 @@ func TestRepository(t *testing.T) {
416413

417414
osArgs := []string{
418415
"-q",
419-
"--cache-dir", cacheDir,
416+
"--cache-dir",
417+
cacheDir,
420418
command,
421419
"--skip-db-update",
422420
"--skip-policy-update",
423-
"--format", string(format),
424-
"--parallel", fmt.Sprint(tt.args.parallel),
421+
"--format",
422+
string(format),
423+
"--parallel",
424+
fmt.Sprint(tt.args.parallel),
425425
"--offline-scan",
426426
}
427427

@@ -499,7 +499,6 @@ func TestRepository(t *testing.T) {
499499
osArgs = append(osArgs, "--output", outputFile)
500500
osArgs = append(osArgs, tt.args.input)
501501

502-
clock.SetFakeTime(t, time.Date(2021, 8, 25, 12, 20, 30, 5, time.UTC))
503502
uuid.SetFakeUUID(t, "3ff14136-e09f-4df9-80ea-%012d")
504503

505504
// Run "trivy repo"

‎integration/testdata/debian-stretch.json.golden

+1-2
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,7 @@
66
"Metadata": {
77
"OS": {
88
"Family": "debian",
9-
"Name": "9.9",
10-
"EOSL": true
9+
"Name": "9.9"
1110
},
1211
"ImageID": "sha256:f26939cc87ef44a6fc554eedd0a976ab30b5bc2769d65d2e986b6c5f1fd4053d",
1312
"DiffIDs": [

‎integration/testdata/distroless-base.json.golden

+1-2
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,7 @@
66
"Metadata": {
77
"OS": {
88
"Family": "debian",
9-
"Name": "9.9",
10-
"EOSL": true
9+
"Name": "9.9"
1110
},
1211
"ImageID": "sha256:7f04a8d247173b1f2546d22913af637bbab4e7411e00ae6207da8d94c445750d",
1312
"DiffIDs": [

‎integration/testdata/distroless-python27.json.golden

+1-2
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,7 @@
66
"Metadata": {
77
"OS": {
88
"Family": "debian",
9-
"Name": "9.9",
10-
"EOSL": true
9+
"Name": "9.9"
1110
},
1211
"ImageID": "sha256:6fcac2cc8a710f21577b5bbd534e0bfc841c0cca569b57182ba19054696cddda",
1312
"DiffIDs": [

‎integration/testdata/ubuntu-1804-ignore-unfixed.json.golden

+1-2
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,7 @@
66
"Metadata": {
77
"OS": {
88
"Family": "ubuntu",
9-
"Name": "18.04",
10-
"EOSL": true
9+
"Name": "18.04"
1110
},
1211
"ImageID": "sha256:a2a15febcdf362f6115e801d37b5e60d6faaeedcb9896155e5fe9d754025be12",
1312
"DiffIDs": [

‎integration/testdata/ubuntu-1804.json.golden

+1-2
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,7 @@
66
"Metadata": {
77
"OS": {
88
"Family": "ubuntu",
9-
"Name": "18.04",
10-
"EOSL": true
9+
"Name": "18.04"
1110
},
1211
"ImageID": "sha256:a2a15febcdf362f6115e801d37b5e60d6faaeedcb9896155e5fe9d754025be12",
1312
"DiffIDs": [

‎pkg/clock/clock.go

+20-10
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,33 @@
11
package clock
22

33
import (
4-
"testing"
4+
"context"
55
"time"
66

77
"k8s.io/utils/clock"
88
clocktesting "k8s.io/utils/clock/testing"
99
)
1010

11-
var c clock.Clock = clock.RealClock{}
11+
// clockKey is the context key for clock. It is unexported to prevent collisions with context keys defined in
12+
// other packages.
13+
type clockKey struct{}
1214

13-
// SetFakeTime sets a fake time for testing.
14-
func SetFakeTime(t *testing.T, fakeTime time.Time) {
15-
c = clocktesting.NewFakeClock(fakeTime)
16-
t.Cleanup(func() {
17-
c = clock.RealClock{}
18-
})
15+
// With returns a new context with the given time.
16+
func With(ctx context.Context, t time.Time) context.Context {
17+
c := clocktesting.NewFakeClock(t)
18+
return context.WithValue(ctx, clockKey{}, c)
1919
}
2020

21-
func Now() time.Time {
22-
return c.Now()
21+
// Now returns the current time.
22+
func Now(ctx context.Context) time.Time {
23+
return Clock(ctx).Now()
24+
}
25+
26+
// Clock returns the clock from the context.
27+
func Clock(ctx context.Context) clock.Clock {
28+
t, ok := ctx.Value(clockKey{}).(clock.Clock)
29+
if !ok {
30+
return clock.RealClock{}
31+
}
32+
return t
2333
}

‎pkg/cloud/aws/commands/run_test.go

+22-9
Original file line numberDiff line numberDiff line change
@@ -1068,8 +1068,11 @@ Summary Report for compliance: my-custom-spec
10681068
MisconfOptions: flag.MisconfOptions{IncludeNonFailures: true},
10691069
},
10701070
cacheContent: "testdata/s3andcloudtrailcache.json",
1071-
allServices: []string{"s3", "cloudtrail"},
1072-
want: expectedS3AndCloudTrailResult,
1071+
allServices: []string{
1072+
"s3",
1073+
"cloudtrail",
1074+
},
1075+
want: expectedS3AndCloudTrailResult,
10731076
},
10741077
{
10751078
name: "skip certain services and include specific services",
@@ -1087,7 +1090,10 @@ Summary Report for compliance: my-custom-spec
10871090
MisconfOptions: flag.MisconfOptions{IncludeNonFailures: true},
10881091
},
10891092
cacheContent: "testdata/s3andcloudtrailcache.json",
1090-
allServices: []string{"s3", "cloudtrail"},
1093+
allServices: []string{
1094+
"s3",
1095+
"cloudtrail",
1096+
},
10911097
// we skip cloudtrail but still expect results from it as it is cached
10921098
want: expectedS3AndCloudTrailResult,
10931099
},
@@ -1096,16 +1102,23 @@ Summary Report for compliance: my-custom-spec
10961102
options: flag.Options{
10971103
RegoOptions: flag.RegoOptions{SkipPolicyUpdate: true},
10981104
AWSOptions: flag.AWSOptions{
1099-
Region: "us-east-1",
1100-
SkipServices: []string{"cloudtrail", "iam"},
1101-
Account: "12345678",
1105+
Region: "us-east-1",
1106+
SkipServices: []string{
1107+
"cloudtrail",
1108+
"iam",
1109+
},
1110+
Account: "12345678",
11021111
},
11031112
CloudOptions: flag.CloudOptions{
11041113
MaxCacheAge: time.Hour * 24 * 365 * 100,
11051114
},
11061115
MisconfOptions: flag.MisconfOptions{IncludeNonFailures: true},
11071116
},
1108-
allServices: []string{"s3", "cloudtrail", "iam"},
1117+
allServices: []string{
1118+
"s3",
1119+
"cloudtrail",
1120+
"iam",
1121+
},
11091122
cacheContent: "testdata/s3onlycache.json",
11101123
want: expectedS3ScanResult,
11111124
},
@@ -1129,7 +1142,7 @@ Summary Report for compliance: my-custom-spec
11291142
},
11301143
}
11311144

1132-
clock.SetFakeTime(t, time.Date(2021, 8, 25, 12, 20, 30, 5, time.UTC))
1145+
ctx := clock.With(context.Background(), time.Date(2021, 8, 25, 12, 20, 30, 5, time.UTC))
11331146
for _, test := range tests {
11341147
t.Run(test.name, func(t *testing.T) {
11351148
if test.allServices != nil {
@@ -1179,7 +1192,7 @@ Summary Report for compliance: my-custom-spec
11791192
require.NoError(t, os.WriteFile(cacheFile, cacheData, 0600))
11801193
}
11811194

1182-
err := Run(context.Background(), test.options)
1195+
err := Run(ctx, test.options)
11831196
if test.expectErr {
11841197
assert.Error(t, err)
11851198
return

‎pkg/cloud/report/report.go

+4-4
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ func Write(ctx context.Context, rep *Report, opt flag.Options, fromCache bool) e
6767
defer cleanup()
6868

6969
if opt.Compliance.Spec.ID != "" {
70-
return writeCompliance(rep, opt, output)
70+
return writeCompliance(ctx, rep, opt, output)
7171
}
7272

7373
var filtered []types.Result
@@ -93,7 +93,7 @@ func Write(ctx context.Context, rep *Report, opt flag.Options, fromCache bool) e
9393
})
9494

9595
base := types.Report{
96-
CreatedAt: clock.Now(),
96+
CreatedAt: clock.Now(ctx),
9797
ArtifactName: rep.AccountID,
9898
ArtifactType: ftypes.ArtifactAWSAccount,
9999
Results: filtered,
@@ -139,7 +139,7 @@ func Write(ctx context.Context, rep *Report, opt flag.Options, fromCache bool) e
139139
}
140140
}
141141

142-
func writeCompliance(rep *Report, opt flag.Options, output io.Writer) error {
142+
func writeCompliance(ctx context.Context, rep *Report, opt flag.Options, output io.Writer) error {
143143
var crr []types.Results
144144
for _, r := range rep.Results {
145145
crr = append(crr, r.Results)
@@ -150,7 +150,7 @@ func writeCompliance(rep *Report, opt flag.Options, output io.Writer) error {
150150
return xerrors.Errorf("compliance report build error: %w", err)
151151
}
152152

153-
return cr.Write(complianceReport, cr.Option{
153+
return cr.Write(ctx, complianceReport, cr.Option{
154154
Format: opt.Format,
155155
Report: opt.ReportFormat,
156156
Output: output,

‎pkg/cloud/report/service_test.go

+11-4
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,10 @@ This scan report was loaded from cached results. If you'd like to run a fresh sc
8888
},
8989
},
9090
AWSOptions: flag.AWSOptions{
91-
Services: []string{"s3", "ec2"},
91+
Services: []string{
92+
"s3",
93+
"ec2",
94+
},
9295
},
9396
},
9497
fromCache: false,
@@ -117,7 +120,11 @@ Scan Overview for AWS Account
117120
},
118121
},
119122
AWSOptions: flag.AWSOptions{
120-
Services: []string{"ec2", "s3", "iam"},
123+
Services: []string{
124+
"ec2",
125+
"s3",
126+
"iam",
127+
},
121128
},
122129
},
123130
fromCache: false,
@@ -310,7 +317,7 @@ Scan Overview for AWS Account
310317
}`,
311318
},
312319
}
313-
clock.SetFakeTime(t, time.Date(2021, 8, 25, 12, 20, 30, 5, time.UTC))
320+
ctx := clock.With(context.Background(), time.Date(2021, 8, 25, 12, 20, 30, 5, time.UTC))
314321
for _, tt := range tests {
315322
t.Run(tt.name, func(t *testing.T) {
316323
report := New(
@@ -323,7 +330,7 @@ Scan Overview for AWS Account
323330

324331
output := bytes.NewBuffer(nil)
325332
tt.options.SetOutputWriter(output)
326-
require.NoError(t, Write(context.Background(), report, tt.options, tt.fromCache))
333+
require.NoError(t, Write(ctx, report, tt.options, tt.fromCache))
327334

328335
assert.Equal(t, "AWS", report.Provider)
329336
assert.Equal(t, tt.options.AWSOptions.Account, report.AccountID)

‎pkg/commands/server/run.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -59,5 +59,5 @@ func Run(ctx context.Context, opts flag.Options) (err error) {
5959

6060
server := rpcServer.NewServer(opts.AppVersion, opts.Listen, opts.CacheDir, opts.Token, opts.TokenHeader,
6161
opts.DBRepository, opts.RegistryOpts())
62-
return server.ListenAndServe(cache, opts.SkipDBUpdate)
62+
return server.ListenAndServe(ctx, cache, opts.SkipDBUpdate)
6363
}

‎pkg/compliance/report/report.go

+4-3
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package report
22

33
import (
4+
"context"
45
"io"
56

67
"golang.org/x/xerrors"
@@ -63,8 +64,8 @@ type Writer interface {
6364
Write(ComplianceReport) error
6465
}
6566

66-
// Write writes the results in the give format
67-
func Write(report *ComplianceReport, option Option) error {
67+
// Write writes the results in the given format
68+
func Write(ctx context.Context, report *ComplianceReport, option Option) error {
6869
switch option.Format {
6970
case types.FormatJSON:
7071
jwriter := JSONWriter{
@@ -79,7 +80,7 @@ func Write(report *ComplianceReport, option Option) error {
7980
Report: option.Report,
8081
Severities: option.Severities,
8182
}
82-
err := complianceWriter.Write(report)
83+
err := complianceWriter.Write(ctx, report)
8384
if err != nil {
8485
return err
8586
}

0 commit comments

Comments
 (0)