From c3ebce929371b60d7fa0a56df7c1e6a2289f749b Mon Sep 17 00:00:00 2001 From: Martin Angers Date: Tue, 12 Nov 2024 09:28:08 -0500 Subject: [PATCH] Improve memory usage of software installers parsing (#23596) --- ...5-improve-memory-usage-software-installers | 1 + .../installer_cache/installer-cache.go | 7 +- ee/server/service/maintained_apps.go | 17 +- ee/server/service/software_installers.go | 42 +++-- pkg/file/file.go | 18 +- pkg/file/file_bench_test.go | 169 ++++++++++++++++++ pkg/file/file_test.go | 7 +- pkg/file/msi.go | 24 +-- pkg/file/pe.go | 14 +- pkg/file/pe_test.go | 9 +- pkg/file/rpm_test.go | 11 +- pkg/file/xar.go | 27 +-- server/datastore/mysql/activities_test.go | 13 +- server/datastore/mysql/policies_test.go | 33 +++- server/datastore/mysql/scripts_test.go | 9 +- .../datastore/mysql/setup_experience_test.go | 38 ++-- .../mysql/software_installers_test.go | 52 ++++-- server/datastore/mysql/software_test.go | 9 +- server/fleet/software_installer.go | 72 +++++++- server/fleet/software_installer_test.go | 67 +++++++ server/mdm/maintainedapps/installers.go | 7 +- server/service/integration_core_test.go | 4 +- server/service/integration_mdm_test.go | 4 +- server/service/software_installers.go | 17 +- server/service/testing_client.go | 12 +- 25 files changed, 550 insertions(+), 133 deletions(-) create mode 100644 changes/20595-improve-memory-usage-software-installers create mode 100644 pkg/file/file_bench_test.go create mode 100644 server/fleet/software_installer_test.go diff --git a/changes/20595-improve-memory-usage-software-installers b/changes/20595-improve-memory-usage-software-installers new file mode 100644 index 0000000000..7e15f3b935 --- /dev/null +++ b/changes/20595-improve-memory-usage-software-installers @@ -0,0 +1 @@ +* Improved memory usage of the Fleet server when uploading a large software installer file. Note that the installer will now use (temporary) disk space and sufficient storage space is required. diff --git a/cmd/osquery-perf/installer_cache/installer-cache.go b/cmd/osquery-perf/installer_cache/installer-cache.go index 57994de5f6..2581528fab 100644 --- a/cmd/osquery-perf/installer_cache/installer-cache.go +++ b/cmd/osquery-perf/installer_cache/installer-cache.go @@ -6,6 +6,7 @@ import ( "sync" "github.com/fleetdm/fleet/v4/pkg/file" + "github.com/fleetdm/fleet/v4/server/fleet" "github.com/fleetdm/fleet/v4/server/service" ) @@ -51,13 +52,13 @@ func populateMetadata(orbitClient *service.OrbitClient, installerID uint) (*file return nil, err } // Figure out what we're actually installing here and add it to software inventory - f, err := os.Open(path) + tfr, err := fleet.NewKeepFileReader(path) if err != nil { log.Println("open installer:", err) return nil, err } - defer f.Close() - item, err := file.ExtractInstallerMetadata(f) + defer tfr.Close() + item, err := file.ExtractInstallerMetadata(tfr) if err != nil { log.Println("extract installer metadata:", err) return nil, err diff --git a/ee/server/service/maintained_apps.go b/ee/server/service/maintained_apps.go index 2cc057308f..76798a16be 100644 --- a/ee/server/service/maintained_apps.go +++ b/ee/server/service/maintained_apps.go @@ -1,10 +1,10 @@ package service import ( - "bytes" "context" "crypto/sha256" "encoding/hex" + "io" "os" "path/filepath" "time" @@ -48,10 +48,11 @@ func (svc *Service) AddFleetMaintainedApp( } client := fleethttp.NewClient(fleethttp.WithTimeout(timeout)) - installerBytes, filename, err := maintainedapps.DownloadInstaller(ctx, app.InstallerURL, client) + installerTFR, filename, err := maintainedapps.DownloadInstaller(ctx, app.InstallerURL, client) if err != nil { return ctxerr.Wrap(ctx, err, "downloading app installer") } + defer installerTFR.Close() extension, err := maintainedapps.ExtensionForBundleIdentifier(app.BundleIdentifier) if err != nil { @@ -63,15 +64,16 @@ func (svc *Service) AddFleetMaintainedApp( // hash shouldn't be checked. if app.SHA256 != noCheckHash { h := sha256.New() - _, err = h.Write(installerBytes) - if err != nil { - return ctxerr.Wrap(ctx, err, "generating SHA256 of maintained app installer") - } + _, _ = io.Copy(h, installerTFR) // writes to a Hash can never fail gotHash := hex.EncodeToString(h.Sum(nil)) if gotHash != app.SHA256 { return ctxerr.New(ctx, "mismatch in maintained app SHA256 hash") } + + if err := installerTFR.Rewind(); err != nil { + return ctxerr.Wrap(ctx, err, "rewind installer reader") + } } // Fall back to the filename if we weren't able to extract a filename from the installer response @@ -95,9 +97,8 @@ func (svc *Service) AddFleetMaintainedApp( uninstallScript = app.UninstallScript } - installerReader := bytes.NewReader(installerBytes) payload := &fleet.UploadSoftwareInstallerPayload{ - InstallerFile: installerReader, + InstallerFile: installerTFR, Title: app.Name, UserID: vc.UserID(), TeamID: teamID, diff --git a/ee/server/service/software_installers.go b/ee/server/service/software_installers.go index d7984a9cd9..b5e1e2a943 100644 --- a/ee/server/service/software_installers.go +++ b/ee/server/service/software_installers.go @@ -1,13 +1,11 @@ package service import ( - "bytes" "context" "encoding/hex" "encoding/json" "errors" "fmt" - "io" "mime" "net/http" "net/url" @@ -1090,7 +1088,7 @@ func (svc *Service) addMetadataToSoftwarePayload(ctx context.Context, payload *f payload.Extension = meta.Extension // reset the reader (it was consumed to extract metadata) - if _, err := payload.InstallerFile.Seek(0, 0); err != nil { + if err := payload.InstallerFile.Rewind(); err != nil { return "", ctxerr.Wrap(ctx, err, "resetting installer file reader") } @@ -1230,7 +1228,7 @@ func (svc *Service) softwareBatchUpload( } }(time.Now()) - downloadURLFn := func(ctx context.Context, url string) (http.Header, []byte, error) { + downloadURLFn := func(ctx context.Context, url string) (http.Header, *fleet.TempFileReader, error) { client := fleethttp.NewClient() client.Transport = fleethttp.NewSizeLimitTransport(fleet.MaxSoftwareInstallerSize) @@ -1268,7 +1266,7 @@ func (svc *Service) softwareBatchUpload( ) } - bodyBytes, err := io.ReadAll(resp.Body) + tfr, err := fleet.NewTempFileReader(resp.Body, nil) if err != nil { // the max size error can be received either at client.Do or here when // reading the body if it's caught via a limited body reader. @@ -1282,7 +1280,7 @@ func (svc *Service) softwareBatchUpload( return nil, nil, fmt.Errorf("reading installer %q contents: %w", url, err) } - return resp.Header, bodyBytes, nil + return resp.Header, tfr, nil } var g errgroup.Group @@ -1295,18 +1293,22 @@ func (svc *Service) softwareBatchUpload( i, p := i, p g.Go(func() error { - headers, bodyBytes, err := downloadURLFn(ctx, p.URL) + headers, tfr, err := downloadURLFn(ctx, p.URL) if err != nil { return err } + // NOTE: cannot defer tfr.Close() here because the reader needs to be + // available after the goroutine completes. Instead, all temp file + // readers will have their Close deferred after the join/wait of + // goroutines. installer := &fleet.UploadSoftwareInstallerPayload{ TeamID: teamID, InstallScript: p.InstallScript, PreInstallQuery: p.PreInstallQuery, PostInstallScript: p.PostInstallScript, UninstallScript: p.UninstallScript, - InstallerFile: bytes.NewReader(bodyBytes), + InstallerFile: tfr, SelfService: p.SelfService, UserID: userID, URL: p.URL, @@ -1326,6 +1328,7 @@ func (svc *Service) softwareBatchUpload( ext, err := svc.addMetadataToSoftwarePayload(ctx, installer) if err != nil { + _ = tfr.Close() // closing the temp file here since it will not be available after the goroutine completes return err } @@ -1352,9 +1355,18 @@ func (svc *Service) softwareBatchUpload( }) } - if err := g.Wait(); err != nil { + waitErr := g.Wait() + + // defer close for any valid temp file reader + for _, payload := range installers { + if payload != nil && payload.InstallerFile != nil { + defer payload.InstallerFile.Close() + } + } + + if waitErr != nil { // NOTE: intentionally not wrapping to avoid polluting user errors. - batchErr = err + batchErr = waitErr return } @@ -1563,7 +1575,15 @@ func UninstallSoftwareMigration( return ctxerr.Wrap(ctx, err, "getting installer from store") } - meta, err := file.ExtractInstallerMetadata(installer) + tfr, err := fleet.NewTempFileReader(installer, nil) + _ = installer.Close() + if err != nil { + level.Warn(logger).Log("msg", "extracting metadata from installer", "software_installer_id", id, "storage_id", storageID, "err", + err) + continue + } + meta, err := file.ExtractInstallerMetadata(tfr) + _ = tfr.Close() // best-effort closing and deleting of temp file if err != nil { level.Warn(logger).Log("msg", "extracting metadata from installer", "software_installer_id", id, "storage_id", storageID, "err", err) diff --git a/pkg/file/file.go b/pkg/file/file.go index 45774a3e9d..d1246c2b6e 100644 --- a/pkg/file/file.go +++ b/pkg/file/file.go @@ -15,6 +15,7 @@ import ( "strings" "github.com/fleetdm/fleet/v4/pkg/secure" + "github.com/fleetdm/fleet/v4/server/fleet" ) var ErrUnsupportedType = errors.New("unsupported file type") @@ -31,25 +32,28 @@ type InstallerMetadata struct { // ExtractInstallerMetadata extracts the software name and version from the // installer file and returns them along with the sha256 hash of the bytes. The // format of the installer is determined based on the magic bytes of the content. -func ExtractInstallerMetadata(r io.Reader) (*InstallerMetadata, error) { - br := bufio.NewReader(r) +func ExtractInstallerMetadata(tfr *fleet.TempFileReader) (*InstallerMetadata, error) { + br := bufio.NewReader(tfr) extension, err := typeFromBytes(br) if err != nil { return nil, err } + if err := tfr.Rewind(); err != nil { + return nil, err + } var meta *InstallerMetadata switch extension { case "deb": - meta, err = ExtractDebMetadata(br) + meta, err = ExtractDebMetadata(tfr) case "rpm": - meta, err = ExtractRPMMetadata(br) + meta, err = ExtractRPMMetadata(tfr) case "exe": - meta, err = ExtractPEMetadata(br) + meta, err = ExtractPEMetadata(tfr) case "pkg": - meta, err = ExtractXARMetadata(br) + meta, err = ExtractXARMetadata(tfr) case "msi": - meta, err = ExtractMSIMetadata(br) + meta, err = ExtractMSIMetadata(tfr) default: return nil, ErrUnsupportedType } diff --git a/pkg/file/file_bench_test.go b/pkg/file/file_bench_test.go new file mode 100644 index 0000000000..19cd1dfe48 --- /dev/null +++ b/pkg/file/file_bench_test.go @@ -0,0 +1,169 @@ +package file_test + +import ( + "encoding/hex" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/fleetdm/fleet/v4/pkg/file" + "github.com/fleetdm/fleet/v4/server/fleet" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Results @b81f69d16220524866fc90e9260a0af0d2aeb94c before any change: +// +// $ GO_TEST_EXTRA_FLAGS="--timeout 20m" FLEET_INTEGRATION_TESTS_DISABLE_LOG=1 REDIS_TEST=1 MYSQL_TEST=1 MINIO_STORAGE_TEST=1 go test ./pkg/file -run zzz -bench . -benchmem | prettybench +// goos: linux +// goarch: amd64 +// pkg: github.com/fleetdm/fleet/v4/pkg/file +// cpu: Intel(R) Core(TM) i7-10510U CPU @ 1.80GHz +// PASS +// benchmark iter time/iter bytes alloc allocs +// --------- ---- --------- ----------- ------ +// BenchmarkExtractInstallerMetadata/.exe/file_size:_39712_kb-8 4 266.94 ms/op 257662412 B/op 251387 allocs/op +// BenchmarkExtractInstallerMetadata/Box.app.pkg/file_size:_67219_kb-8 3 464.89 ms/op 393721768 B/op 3962 allocs/op +// BenchmarkExtractInstallerMetadata/Fleet_osquery.msi/file_size:_43775_kb-8 4 277.60 ms/op 252408264 B/op 3737 allocs/op +// BenchmarkExtractInstallerMetadata/Go_Programming_Language_amd64_go1.22.2.msi/file_size:_61680_kb-8 3 513.12 ms/op 402892376 B/op 161092 allocs/op +// BenchmarkExtractInstallerMetadata/Go.pkg/file_size:_69628_kb-8 3 472.77 ms/op 393635344 B/op 1778 allocs/op +// BenchmarkExtractInstallerMetadata/Go.pkg#01/file_size:_66444_kb-8 3 465.80 ms/op 393635776 B/op 1785 allocs/op +// BenchmarkExtractInstallerMetadata/NordVPN.app.pkg/file_size:_155592_kb-8 1 1011.86 ms/op 961921824 B/op 1919 allocs/op +// BenchmarkExtractInstallerMetadata/Notion_3.11.1.exe/file_size:_77768_kb-8 2 528.20 ms/op 492055496 B/op 567 allocs/op +// BenchmarkExtractInstallerMetadata/Python.pkg/file_size:_44601_kb-8 4 291.94 ms/op 251876720 B/op 5834 allocs/op +// BenchmarkExtractInstallerMetadata/TeamViewer.app.pkg/file_size:_93051_kb-8 2 594.36 ms/op 492383088 B/op 6823 allocs/op +// BenchmarkExtractInstallerMetadata/Vim.exe/file_size:_10704_kb-8 15 117.45 ms/op 65504394 B/op 640 allocs/op +// BenchmarkExtractInstallerMetadata/Visual_Studio_Code.exe/file_size:_97156_kb-8 2 637.27 ms/op 615259264 B/op 637 allocs/op +// BenchmarkExtractInstallerMetadata/code.deb/file_size:_99278_kb-8 3 379.23 ms/op 8455728 B/op 116 allocs/op +// BenchmarkExtractInstallerMetadata/code.rpm/file_size:_138886_kb-8 2 556.37 ms/op 3397024 B/op 11274 allocs/op +// BenchmarkExtractInstallerMetadata/fleet-osquery.deb/file_size:_79581_kb-8 4 308.36 ms/op 59696 B/op 90 allocs/op +// BenchmarkExtractInstallerMetadata/htop.deb/file_size:_90_kb-8 822 1.96 ms/op 8446331 B/op 110 allocs/op +// BenchmarkExtractInstallerMetadata/ruby.deb/file_size:_11_kb-8 649 1.66 ms/op 8448424 B/op 122 allocs/op +// ok github.com/fleetdm/fleet/v4/pkg/file 36.644s + +// Results @0c700ca40e5d3602b6206f12232c4c123b6c4ee9 with the use of TempFileReader but not change otherwise: +// $ GO_TEST_EXTRA_FLAGS="--timeout 20m" FLEET_INTEGRATION_TESTS_DISABLE_LOG=1 REDIS_TEST=1 MYSQL_TEST=1 MINIO_STORAGE_TEST=1 go test ./pkg/file -run zzz -bench . -benchmem | prettybench +// goos: linux +// goarch: amd64 +// pkg: github.com/fleetdm/fleet/v4/pkg/file +// cpu: Intel(R) Core(TM) i7-10510U CPU @ 1.80GHz +// PASS +// benchmark iter time/iter bytes alloc allocs +// --------- ---- --------- ----------- ------ +// BenchmarkExtractInstallerMetadata/.exe/file_size:_39712_kb-8 4 315.43 ms/op 257661592 B/op 251389 allocs/op +// BenchmarkExtractInstallerMetadata/Box.app.pkg/file_size:_67219_kb-8 3 418.06 ms/op 393721613 B/op 3962 allocs/op +// BenchmarkExtractInstallerMetadata/Fleet_osquery.msi/file_size:_43775_kb-8 4 296.01 ms/op 252408208 B/op 3737 allocs/op +// BenchmarkExtractInstallerMetadata/Go_Programming_Language_amd64_go1.22.2.msi/file_size:_61680_kb-8 3 475.20 ms/op 402892136 B/op 161092 allocs/op +// BenchmarkExtractInstallerMetadata/Go.pkg/file_size:_69628_kb-8 2 508.10 ms/op 393635328 B/op 1779 allocs/op +// BenchmarkExtractInstallerMetadata/Go.pkg#01/file_size:_66444_kb-8 2 533.04 ms/op 393635672 B/op 1785 allocs/op +// BenchmarkExtractInstallerMetadata/NordVPN.app.pkg/file_size:_155592_kb-8 1 1363.48 ms/op 961921792 B/op 1920 allocs/op +// BenchmarkExtractInstallerMetadata/Notion_3.11.1.exe/file_size:_77768_kb-8 2 621.95 ms/op 492055376 B/op 566 allocs/op +// BenchmarkExtractInstallerMetadata/Python.pkg/file_size:_44601_kb-8 2 513.91 ms/op 251876472 B/op 5832 allocs/op +// BenchmarkExtractInstallerMetadata/TeamViewer.app.pkg/file_size:_93051_kb-8 2 573.12 ms/op 492382960 B/op 6823 allocs/op +// BenchmarkExtractInstallerMetadata/Vim.exe/file_size:_10704_kb-8 18 107.35 ms/op 65504316 B/op 640 allocs/op +// BenchmarkExtractInstallerMetadata/Visual_Studio_Code.exe/file_size:_97156_kb-8 2 606.29 ms/op 615259376 B/op 639 allocs/op +// BenchmarkExtractInstallerMetadata/code.deb/file_size:_99278_kb-8 3 397.47 ms/op 8447426 B/op 114 allocs/op +// BenchmarkExtractInstallerMetadata/code.rpm/file_size:_138886_kb-8 2 594.10 ms/op 3396936 B/op 11274 allocs/op +// BenchmarkExtractInstallerMetadata/fleet-osquery.deb/file_size:_79581_kb-8 3 335.98 ms/op 60384 B/op 90 allocs/op +// BenchmarkExtractInstallerMetadata/htop.deb/file_size:_90_kb-8 732 3.16 ms/op 8446791 B/op 110 allocs/op +// BenchmarkExtractInstallerMetadata/ruby.deb/file_size:_11_kb-8 578 3.48 ms/op 8449575 B/op 122 allocs/op +// ok github.com/fleetdm/fleet/v4/pkg/file 37.775s + +// Results @64321f8d241bba9233a1de21845ac0c7a6f4dda6 with the .exe improvements (read from disk with mmap) - massively +// better memory usage (only exe benchmarks show): +// +// $ GO_TEST_EXTRA_FLAGS="--timeout 20m" FLEET_INTEGRATION_TESTS_DISABLE_LOG=1 REDIS_TEST=1 MYSQL_TEST=1 MINIO_STORAGE_TEST=1 go test ./pkg/file -run zzz -bench . -benchmem | prettybench +// goos: linux +// goarch: amd64 +// pkg: github.com/fleetdm/fleet/v4/pkg/file +// cpu: Intel(R) Core(TM) i7-10510U CPU @ 1.80GHz +// PASS +// benchmark iter time/iter bytes alloc allocs +// --------- ---- --------- ----------- ------ +// BenchmarkExtractInstallerMetadata/.exe/file_size:_39712_kb-8 6 208.14 ms/op 6135304 B/op 251345 allocs/op +// BenchmarkExtractInstallerMetadata/Notion_3.11.1.exe/file_size:_77768_kb-8 3 337.30 ms/op 61258 B/op 521 allocs/op +// BenchmarkExtractInstallerMetadata/Vim.exe/file_size:_10704_kb-8 22 47.32 ms/op 67321 B/op 604 allocs/op +// BenchmarkExtractInstallerMetadata/Visual_Studio_Code.exe/file_size:_97156_kb-8 3 421.23 ms/op 65573 B/op 591 allocs/op +// ok github.com/fleetdm/fleet/v4/pkg/file 35.887s + +// Results @e5ad9300701f0aa1f7b40efffdb6944988038dc7 with the .pkg improvements +// - massively better memory usage (only pkg benchmarks shown): +// +// $ GO_TEST_EXTRA_FLAGS="--timeout 20m" FLEET_INTEGRATION_TESTS_DISABLE_LOG=1 REDIS_TEST=1 MYSQL_TEST=1 MINIO_STORAGE_TEST=1 go test ./pkg/file -run zzz -bench . -benchmem | prettybench +// goos: linux +// goarch: amd64 +// pkg: github.com/fleetdm/fleet/v4/pkg/file +// cpu: Intel(R) Core(TM) i7-10510U CPU @ 1.80GHz +// PASS +// benchmark iter time/iter bytes alloc allocs +// --------- ---- --------- ----------- ------ +// BenchmarkExtractInstallerMetadata/Box.app.pkg/file_size:_67219_kb-8 4 319.98 ms/op 285844 B/op 3915 allocs/op +// BenchmarkExtractInstallerMetadata/Go.pkg/file_size:_69628_kb-8 4 280.80 ms/op 199696 B/op 1734 allocs/op +// BenchmarkExtractInstallerMetadata/Go.pkg#01/file_size:_66444_kb-8 4 286.77 ms/op 199824 B/op 1738 allocs/op +// BenchmarkExtractInstallerMetadata/NordVPN.app.pkg/file_size:_155592_kb-8 2 713.27 ms/op 223144 B/op 1866 allocs/op +// BenchmarkExtractInstallerMetadata/Python.pkg/file_size:_44601_kb-8 6 247.57 ms/op 350997 B/op 5791 allocs/op +// BenchmarkExtractInstallerMetadata/TeamViewer.app.pkg/file_size:_93051_kb-8 2 586.15 ms/op 389312 B/op 6776 allocs/op +// ok github.com/fleetdm/fleet/v4/pkg/file 39.536s + +// Results @532daf10bebe7c432a2b5e6c3822639c5937dc29 with the .msi improvements +// - massively better memory usage (only msi benchmarks shown): +// $ GO_TEST_EXTRA_FLAGS="--timeout 20m" FLEET_INTEGRATION_TESTS_DISABLE_LOG=1 REDIS_TEST=1 MYSQL_TEST=1 MINIO_STORAGE_TEST=1 go test ./pkg/file -run zzz -bench . -benchmem | prettybench +// goos: linux +// goarch: amd64 +// pkg: github.com/fleetdm/fleet/v4/pkg/file +// cpu: Intel(R) Core(TM) i7-10510U CPU @ 1.80GHz +// PASS +// benchmark iter time/iter bytes alloc allocs +// --------- ---- --------- ----------- ------ +// BenchmarkExtractInstallerMetadata/Fleet_osquery.msi/file_size:_43775_kb-8 6 191.69 ms/op 879274 B/op 3752 allocs/op +// BenchmarkExtractInstallerMetadata/Go_Programming_Language_amd64_go1.22.2.msi/file_size:_61680_kb-8 4 305.72 ms/op 8430244 B/op 161054 allocs/op +// ok github.com/fleetdm/fleet/v4/pkg/file 32.193s + +func BenchmarkExtractInstallerMetadata(b *testing.B) { + dents, err := os.ReadDir(filepath.Join("testdata", "installers")) + if err != nil { + b.Fatal(err) + } + + for _, dent := range dents { + if !dent.Type().IsRegular() || strings.HasPrefix(dent.Name(), ".") { + continue + } + parts := strings.Split(strings.TrimSuffix(dent.Name(), filepath.Ext(dent.Name())), "$") + if len(parts) < 4 { + b.Fatalf("invalid filename, expected at least 4 sections, got %d: %s", len(parts), dent.Name()) + } + wantName, wantVersion, wantHash, wantBundleIdentifier := parts[0], parts[1], parts[2], parts[3] + wantExtension := strings.TrimPrefix(filepath.Ext(dent.Name()), ".") + + b.Run(wantName+"."+wantExtension, func(b *testing.B) { + + b.ResetTimer() + b.ReportAllocs() + info, err := dent.Info() + require.NoError(b, err) + + b.Run(fmt.Sprintf("file size: %d kb", info.Size()/1024), func(b *testing.B) { + // the goal of this benchmark is not so much accuracy of time performance, but + // memory usage, so it doesn't matter that the file is read from disk on each + // iteration. + for i := 0; i < b.N; i++ { + tfr, err := fleet.NewKeepFileReader(filepath.Join("testdata", "installers", dent.Name())) + require.NoError(b, err) + + meta, err := file.ExtractInstallerMetadata(tfr) + require.NoError(b, err) + tfr.Close() + + assert.Equal(b, wantName, meta.Name) + assert.Equal(b, wantVersion, meta.Version) + assert.Equal(b, wantHash, hex.EncodeToString(meta.SHASum)) + assert.Equal(b, wantExtension, meta.Extension) + assert.Equal(b, wantBundleIdentifier, meta.BundleIdentifier) + } + }) + }) + } +} diff --git a/pkg/file/file_test.go b/pkg/file/file_test.go index 7e520c41db..d59f09d95b 100644 --- a/pkg/file/file_test.go +++ b/pkg/file/file_test.go @@ -9,6 +9,7 @@ import ( "testing" "github.com/fleetdm/fleet/v4/pkg/file" + "github.com/fleetdm/fleet/v4/server/fleet" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -127,11 +128,11 @@ func TestExtractInstallerMetadata(t *testing.T) { wantName, wantVersion, wantHash, wantBundleIdentifier := parts[0], parts[1], parts[2], parts[3] wantExtension := strings.TrimPrefix(filepath.Ext(dent.Name()), ".") - f, err := os.Open(filepath.Join("testdata", "installers", dent.Name())) + tfr, err := fleet.NewKeepFileReader(filepath.Join("testdata", "installers", dent.Name())) require.NoError(t, err) - defer f.Close() + defer tfr.Close() - meta, err := file.ExtractInstallerMetadata(f) + meta, err := file.ExtractInstallerMetadata(tfr) require.NoError(t, err) assert.Equal(t, wantName, meta.Name) assert.Equal(t, wantVersion, meta.Version) diff --git a/pkg/file/msi.go b/pkg/file/msi.go index 78ac33f86d..29538ae5d7 100644 --- a/pkg/file/msi.go +++ b/pkg/file/msi.go @@ -9,19 +9,20 @@ import ( "io" "strings" + "github.com/fleetdm/fleet/v4/server/fleet" "github.com/sassoftware/relic/v8/lib/comdoc" ) -func ExtractMSIMetadata(r io.Reader) (*InstallerMetadata, error) { +func ExtractMSIMetadata(tfr *fleet.TempFileReader) (*InstallerMetadata, error) { + // compute its hash h := sha256.New() - r = io.TeeReader(r, h) - b, err := io.ReadAll(r) - if err != nil { - return nil, fmt.Errorf("failed to read all content: %w", err) + _, _ = io.Copy(h, tfr) // writes to a hash cannot fail + + if err := tfr.Rewind(); err != nil { + return nil, err } - rr := bytes.NewReader(b) - c, err := comdoc.ReadFile(rr) + c, err := comdoc.ReadFile(tfr) if err != nil { return nil, fmt.Errorf("reading msi file: %w", err) } @@ -232,6 +233,7 @@ func decodeStrings(dataReader, poolReader io.Reader) ([]string, error) { } var stringEntry entry var stringTable []string + var buf bytes.Buffer for { err := binary.Read(poolReader, binary.LittleEndian, &stringEntry) if err != nil { @@ -240,11 +242,13 @@ func decodeStrings(dataReader, poolReader io.Reader) ([]string, error) { } return nil, fmt.Errorf("failed to read pool entry: %w", err) } - buf := make([]byte, stringEntry.Size) - if _, err := io.ReadFull(dataReader, buf); err != nil { + buf.Reset() + buf.Grow(int(stringEntry.Size)) + _, err = io.CopyN(&buf, dataReader, int64(stringEntry.Size)) + if err != nil { return nil, fmt.Errorf("failed to read string data: %w", err) } - stringTable = append(stringTable, string(buf)) + stringTable = append(stringTable, buf.String()) } return stringTable, nil } diff --git a/pkg/file/pe.go b/pkg/file/pe.go index 70d2596ec0..27e872221d 100644 --- a/pkg/file/pe.go +++ b/pkg/file/pe.go @@ -6,22 +6,24 @@ import ( "io" "strings" + "github.com/fleetdm/fleet/v4/server/fleet" "github.com/saferwall/pe" ) // ExtractPEMetadata extracts the name and version metadata from a .exe file in // the Portable Executable (PE) format. -func ExtractPEMetadata(r io.Reader) (*InstallerMetadata, error) { +func ExtractPEMetadata(tfr *fleet.TempFileReader) (*InstallerMetadata, error) { + // compute its hash h := sha256.New() - r = io.TeeReader(r, h) - b, err := io.ReadAll(r) - if err != nil { - return nil, fmt.Errorf("failed to read all content: %w", err) + _, _ = io.Copy(h, tfr) // writes to a hash cannot fail + + if err := tfr.Rewind(); err != nil { + return nil, err } // cannot use the "Fast" option, we need the data directories for the // resources to be available. - pep, err := pe.NewBytes(b, &pe.Options{ + pep, err := pe.New(tfr.Name(), &pe.Options{ OmitExportDirectory: true, OmitImportDirectory: true, OmitExceptionDirectory: true, diff --git a/pkg/file/pe_test.go b/pkg/file/pe_test.go index 0bfd0776aa..c8c8da4b0d 100644 --- a/pkg/file/pe_test.go +++ b/pkg/file/pe_test.go @@ -1,18 +1,21 @@ package file import ( - "os" "testing" + "github.com/fleetdm/fleet/v4/server/fleet" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestExtractPEMetadata(t *testing.T) { t.Parallel() - file, err := os.Open("testdata/software-installers/hello-world-installer.exe") + + tfr, err := fleet.NewKeepFileReader("testdata/software-installers/hello-world-installer.exe") require.NoError(t, err) - meta, err := ExtractPEMetadata(file) + defer tfr.Close() + + meta, err := ExtractPEMetadata(tfr) require.NoError(t, err) require.NotNil(t, meta) assert.Equal(t, "Hello world", meta.Name) diff --git a/pkg/file/rpm_test.go b/pkg/file/rpm_test.go index d9d3127250..a8aae27d70 100644 --- a/pkg/file/rpm_test.go +++ b/pkg/file/rpm_test.go @@ -8,6 +8,7 @@ import ( "testing" "github.com/fleetdm/fleet/v4/orbit/pkg/constant" + "github.com/fleetdm/fleet/v4/server/fleet" "github.com/goreleaser/nfpm/v2" "github.com/goreleaser/nfpm/v2/files" "github.com/goreleaser/nfpm/v2/rpm" @@ -62,14 +63,10 @@ func TestExtractRPMMetadata(t *testing.T) { // Using ExtractInstallerMetadata for broader testing (for a file // with rpm extension it will call ExtractRPMMetadata). // - f, err := os.Open(rpmPath) + tfr, err := fleet.NewKeepFileReader(rpmPath) require.NoError(t, err) - t.Cleanup(func() { - f.Close() - }) - m, err := ExtractInstallerMetadata(f) - require.NoError(t, err) - err = f.Close() + t.Cleanup(func() { tfr.Close() }) + m, err := ExtractInstallerMetadata(tfr) require.NoError(t, err) require.Empty(t, m.BundleIdentifier) require.Equal(t, "rpm", m.Extension) diff --git a/pkg/file/xar.go b/pkg/file/xar.go index 89d285dc4f..31e3773b2b 100644 --- a/pkg/file/xar.go +++ b/pkg/file/xar.go @@ -159,27 +159,28 @@ type distributionApp struct { // ExtractXARMetadata extracts the name and version metadata from a .pkg file // in the XAR format. -func ExtractXARMetadata(r io.Reader) (*InstallerMetadata, error) { +func ExtractXARMetadata(tfr *fleet.TempFileReader) (*InstallerMetadata, error) { var hdr xarHeader h := sha256.New() - r = io.TeeReader(r, h) - b, err := io.ReadAll(r) - if err != nil { - return nil, fmt.Errorf("failed to read all content: %w", err) + size, _ := io.Copy(h, tfr) // writes to a hash cannot fail + + if err := tfr.Rewind(); err != nil { + return nil, fmt.Errorf("rewind reader: %w", err) } - rr := bytes.NewReader(b) - if err := binary.Read(rr, binary.BigEndian, &hdr); err != nil { + // read the file header + if err := binary.Read(tfr, binary.BigEndian, &hdr); err != nil { return nil, fmt.Errorf("decode xar header: %w", err) } - zr, err := zlib.NewReader(io.LimitReader(rr, hdr.CompressedSize)) + zr, err := zlib.NewReader(io.LimitReader(tfr, hdr.CompressedSize)) if err != nil { return nil, fmt.Errorf("create zlib reader: %w", err) } defer zr.Close() + // decode the TOC data (in XML inside the zlib-compressed data) var root xmlXar decoder := xml.NewDecoder(zr) decoder.Strict = false @@ -187,12 +188,13 @@ func ExtractXARMetadata(r io.Reader) (*InstallerMetadata, error) { return nil, fmt.Errorf("decode xar xml: %w", err) } + // look for the distribution file, with the metadata information heapOffset := xarHeaderSize + hdr.CompressedSize var packageInfoFile *xmlFile for _, f := range root.TOC.Files { switch f.Name { case "Distribution": - contents, err := readCompressedFile(rr, heapOffset, int64(len(b)), f) + contents, err := readCompressedFile(tfr, heapOffset, size, f) if err != nil { return nil, err } @@ -203,6 +205,7 @@ func ExtractXARMetadata(r io.Reader) (*InstallerMetadata, error) { } meta.SHASum = h.Sum(nil) return meta, nil + case "PackageInfo": // If Distribution archive was not found, we will use the top-level PackageInfo archive packageInfoFile = f @@ -210,7 +213,7 @@ func ExtractXARMetadata(r io.Reader) (*InstallerMetadata, error) { } if packageInfoFile != nil { - contents, err := readCompressedFile(rr, heapOffset, int64(len(b)), packageInfoFile) + contents, err := readCompressedFile(tfr, heapOffset, size, packageInfoFile) if err != nil { return nil, err } @@ -226,9 +229,9 @@ func ExtractXARMetadata(r io.Reader) (*InstallerMetadata, error) { return &InstallerMetadata{SHASum: h.Sum(nil)}, nil } -func readCompressedFile(rr *bytes.Reader, heapOffset int64, sectionLength int64, f *xmlFile) ([]byte, error) { +func readCompressedFile(rat io.ReaderAt, heapOffset int64, sectionLength int64, f *xmlFile) ([]byte, error) { var fileReader io.Reader - heapReader := io.NewSectionReader(rr, heapOffset, sectionLength-heapOffset) + heapReader := io.NewSectionReader(rat, heapOffset, sectionLength-heapOffset) fileReader = io.NewSectionReader(heapReader, f.Data.Offset, f.Data.Length) // the distribution file can be compressed differently than the TOC, the diff --git a/server/datastore/mysql/activities_test.go b/server/datastore/mysql/activities_test.go index 0cdeeac692..2e19fddccf 100644 --- a/server/datastore/mysql/activities_test.go +++ b/server/datastore/mysql/activities_test.go @@ -394,10 +394,11 @@ func testListHostUpcomingActivities(t *testing.T, ds *Datastore) { require.NoError(t, err) // create a couple of software installers - installer := strings.NewReader("echo") + installer1, err := fleet.NewTempFileReader(strings.NewReader("echo"), t.TempDir) + require.NoError(t, err) sw1, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "install foo", - InstallerFile: installer, + InstallerFile: installer1, StorageID: uuid.NewString(), Filename: "foo.pkg", Title: "foo", @@ -406,9 +407,11 @@ func testListHostUpcomingActivities(t *testing.T, ds *Datastore) { UserID: u.ID, }) require.NoError(t, err) + installer2, err := fleet.NewTempFileReader(strings.NewReader("echo"), t.TempDir) + require.NoError(t, err) sw2, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "install bar", - InstallerFile: installer, + InstallerFile: installer2, StorageID: uuid.NewString(), Filename: "bar.pkg", Title: "bar", @@ -417,9 +420,11 @@ func testListHostUpcomingActivities(t *testing.T, ds *Datastore) { UserID: u.ID, }) require.NoError(t, err) + installer3, err := fleet.NewTempFileReader(strings.NewReader("echo"), t.TempDir) + require.NoError(t, err) sw3, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "install to delete", - InstallerFile: installer, + InstallerFile: installer3, StorageID: uuid.NewString(), Filename: "todelete.pkg", Title: "todelete", diff --git a/server/datastore/mysql/policies_test.go b/server/datastore/mysql/policies_test.go index 7eb2d82ee3..d3626231ac 100644 --- a/server/datastore/mysql/policies_test.go +++ b/server/datastore/mysql/policies_test.go @@ -1,7 +1,6 @@ package mysql import ( - "bytes" "context" "crypto/md5" //nolint:gosec // (only used for tests) "encoding/hex" @@ -1235,11 +1234,13 @@ func testPoliciesByID(t *testing.T, ds *Datastore) { host1 := newTestHostWithPlatform(t, ds, "host1", "darwin", nil) // Associate an installer to policy2 + installer, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID, err := ds.MatchOrCreateSoftwareInstaller(context.Background(), &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello", PreInstallQuery: "SELECT 1", PostInstallScript: "world", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: installer, StorageID: "storage1", Filename: "file1", Title: "file1", @@ -3981,11 +3982,13 @@ func testTeamPoliciesWithInstaller(t *testing.T, ds *Datastore) { }) require.NoError(t, err) // Create and associate an installer to p2. + installer, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID, err := ds.MatchOrCreateSoftwareInstaller(context.Background(), &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello", PreInstallQuery: "SELECT 1", PostInstallScript: "world", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: installer, StorageID: "storage1", Filename: "file1", Title: "file1", @@ -4017,11 +4020,13 @@ func testTeamPoliciesWithInstaller(t *testing.T, ds *Datastore) { require.Equal(t, installerID, *p2.SoftwareInstallerID) // Policy p4 in "No team" with associated installer. + installer1, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) noTeamInstallerID, err := ds.MatchOrCreateSoftwareInstaller(context.Background(), &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello", PreInstallQuery: "SELECT 1", PostInstallScript: "world", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: installer1, StorageID: "storage1", Filename: "file1", Title: "file1", @@ -4236,11 +4241,13 @@ func testApplyPolicySpecWithInstallers(t *testing.T, ds *Datastore) { host1Team1 := newHost("host1Team1", &team1.ID, "darwin") + tfr1, err := fleet.NewTempFileReader(strings.NewReader("hello1"), t.TempDir) + require.NoError(t, err) installer1ID, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello", PreInstallQuery: "SELECT 1;", PostInstallScript: "world1", - InstallerFile: bytes.NewReader([]byte("hello1")), + InstallerFile: tfr1, StorageID: "storage1", Filename: "file1", Title: "file1", @@ -4253,11 +4260,13 @@ func testApplyPolicySpecWithInstallers(t *testing.T, ds *Datastore) { installer1, err := ds.GetSoftwareInstallerMetadataByID(ctx, installer1ID) require.NoError(t, err) require.NotNil(t, installer1.TitleID) + tfr2, err := fleet.NewTempFileReader(strings.NewReader("hello2"), t.TempDir) + require.NoError(t, err) installer2ID, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello2", PreInstallQuery: "SELECT 2;", PostInstallScript: "world2", - InstallerFile: bytes.NewReader([]byte("hello2")), + InstallerFile: tfr2, StorageID: "storage2", Filename: "file2", Title: "file2", @@ -4270,11 +4279,13 @@ func testApplyPolicySpecWithInstallers(t *testing.T, ds *Datastore) { installer2, err := ds.GetSoftwareInstallerMetadataByID(ctx, installer2ID) require.NoError(t, err) require.NotNil(t, installer2.TitleID) + tfr3, err := fleet.NewTempFileReader(strings.NewReader("hello3"), t.TempDir) + require.NoError(t, err) installer3ID, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello3", PreInstallQuery: "SELECT 3;", PostInstallScript: "world3", - InstallerFile: bytes.NewReader([]byte("hello3")), + InstallerFile: tfr3, StorageID: "storage3", Filename: "file3", Title: "file3", @@ -4288,11 +4299,13 @@ func testApplyPolicySpecWithInstallers(t *testing.T, ds *Datastore) { require.NoError(t, err) require.NotNil(t, installer3.TitleID) // Another installer on team1 to test changing installers. + tfr5, err := fleet.NewTempFileReader(strings.NewReader("hello5"), t.TempDir) + require.NoError(t, err) installer5ID, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello5", PreInstallQuery: "SELECT 5;", PostInstallScript: "world5", - InstallerFile: bytes.NewReader([]byte("hello5")), + InstallerFile: tfr5, StorageID: "storage5", Filename: "file5", Title: "file5", @@ -4477,11 +4490,13 @@ func testApplyPolicySpecWithInstallers(t *testing.T, ds *Datastore) { require.Nil(t, team2Policies[0].SoftwareInstallerID) // Apply team policies associated to two installers (again, with two installers with the same title). + tfr4, err := fleet.NewTempFileReader(strings.NewReader("hello3"), t.TempDir) + require.NoError(t, err) installer4ID, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello3", PreInstallQuery: "SELECT 3;", PostInstallScript: "world3", - InstallerFile: bytes.NewReader([]byte("hello3")), + InstallerFile: tfr4, StorageID: "storage3", Filename: "file1", Title: "file1", // same title as installer1. diff --git a/server/datastore/mysql/scripts_test.go b/server/datastore/mysql/scripts_test.go index c01584f8a6..a476e72870 100644 --- a/server/datastore/mysql/scripts_test.go +++ b/server/datastore/mysql/scripts_test.go @@ -1,7 +1,6 @@ package mysql import ( - "bytes" "context" _ "embed" "fmt" @@ -1284,12 +1283,14 @@ func testCleanupUnusedScriptContents(t *testing.T, ds *Datastore) { require.NoError(t, err) // create a software install that references scripts + tfr1, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) swi, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "install-script", UninstallScript: "uninstall-script", PreInstallQuery: "SELECT 1", PostInstallScript: "post-install-script", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr1, StorageID: "storage1", Filename: "file1", Title: "file1", @@ -1343,11 +1344,13 @@ func testCleanupUnusedScriptContents(t *testing.T, ds *Datastore) { require.Equal(t, md5ChecksumScriptContent(res.ScriptContents), sc[0].Checksum) // create a software install without a post-install script + tfr2, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) swi, err = ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ PreInstallQuery: "SELECT 1", InstallScript: "install-script", UninstallScript: "uninstall-script", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr2, StorageID: "storage1", Filename: "file1", Title: "file1", diff --git a/server/datastore/mysql/setup_experience_test.go b/server/datastore/mysql/setup_experience_test.go index 5d56b71ffc..fa48a4107c 100644 --- a/server/datastore/mysql/setup_experience_test.go +++ b/server/datastore/mysql/setup_experience_test.go @@ -1,9 +1,9 @@ package mysql import ( - "bytes" "context" "database/sql" + "strings" "testing" "time" @@ -53,12 +53,14 @@ func testEnqueueSetupExperienceItems(t *testing.T, ds *Datastore) { user1 := test.NewUser(t, ds, "Alice", "alice@example.com", true) + tfr1, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID1, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello", PreInstallQuery: "SELECT 1", PostInstallScript: "world", UninstallScript: "goodbye", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr1, StorageID: "storage1", Filename: "file1", Title: "Software1", @@ -70,11 +72,13 @@ func testEnqueueSetupExperienceItems(t *testing.T, ds *Datastore) { }) require.NoError(t, err) + tfr2, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID2, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "banana", PreInstallQuery: "SELECT 3", PostInstallScript: "apple", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr2, StorageID: "storage3", Filename: "file3", Title: "Software2", @@ -317,12 +321,14 @@ func testGetSetupExperienceTitles(t *testing.T, ds *Datastore) { user1 := test.NewUser(t, ds, "Alice", "alice@example.com", true) + tfr1, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID1, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello", PreInstallQuery: "SELECT 1", PostInstallScript: "world", UninstallScript: "goodbye", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr1, StorageID: "storage1", Filename: "file1", Title: "file1", @@ -334,11 +340,13 @@ func testGetSetupExperienceTitles(t *testing.T, ds *Datastore) { }) require.NoError(t, err) + tfr3, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID3, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "banana", PreInstallQuery: "SELECT 3", PostInstallScript: "apple", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr3, StorageID: "storage3", Filename: "file3", Title: "file3", @@ -351,11 +359,13 @@ func testGetSetupExperienceTitles(t *testing.T, ds *Datastore) { }) require.NoError(t, err) + tfr4, err := fleet.NewTempFileReader(strings.NewReader("hello2"), t.TempDir) + require.NoError(t, err) installerID4, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "pear", PreInstallQuery: "SELECT 4", PostInstallScript: "apple", - InstallerFile: bytes.NewReader([]byte("hello2")), + InstallerFile: tfr4, StorageID: "storage3", Filename: "file4", Title: "file4", @@ -445,12 +455,14 @@ func testSetSetupExperienceTitles(t *testing.T, ds *Datastore) { user1 := test.NewUser(t, ds, "Alice", "alice@example.com", true) + tfr1, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID1, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello", PreInstallQuery: "SELECT 1", PostInstallScript: "world", UninstallScript: "goodbye", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr1, StorageID: "storage1", Filename: "file1", Title: "file1", @@ -463,11 +475,13 @@ func testSetSetupExperienceTitles(t *testing.T, ds *Datastore) { _ = installerID1 require.NoError(t, err) + tfr2, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID2, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "world", PreInstallQuery: "SELECT 2", PostInstallScript: "hello", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr2, StorageID: "storage2", Filename: "file2", Title: "file2", @@ -480,11 +494,13 @@ func testSetSetupExperienceTitles(t *testing.T, ds *Datastore) { _ = installerID2 require.NoError(t, err) + tfr3, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID3, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "banana", PreInstallQuery: "SELECT 3", PostInstallScript: "apple", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr3, StorageID: "storage3", Filename: "file3", Title: "file3", @@ -498,11 +514,13 @@ func testSetSetupExperienceTitles(t *testing.T, ds *Datastore) { _ = installerID3 require.NoError(t, err) + tfr4, err := fleet.NewTempFileReader(strings.NewReader("hello2"), t.TempDir) + require.NoError(t, err) installerID4, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "pear", PreInstallQuery: "SELECT 4", PostInstallScript: "apple", - InstallerFile: bytes.NewReader([]byte("hello2")), + InstallerFile: tfr4, StorageID: "storage3", Filename: "file4", Title: "file4", diff --git a/server/datastore/mysql/software_installers_test.go b/server/datastore/mysql/software_installers_test.go index 648e5f16a3..ec1d5dbf32 100644 --- a/server/datastore/mysql/software_installers_test.go +++ b/server/datastore/mysql/software_installers_test.go @@ -5,6 +5,7 @@ import ( "context" "os" "path/filepath" + "strings" "testing" "time" @@ -52,12 +53,14 @@ func testListPendingSoftwareInstalls(t *testing.T, ds *Datastore) { host2 := test.NewHost(t, ds, "host2", "2", "host2key", "host2uuid", time.Now()) user1 := test.NewUser(t, ds, "Alice", "alice@example.com", true) + tfr1, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID1, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello", PreInstallQuery: "SELECT 1", PostInstallScript: "world", UninstallScript: "goodbye", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr1, StorageID: "storage1", Filename: "file1", Title: "file1", @@ -67,11 +70,13 @@ func testListPendingSoftwareInstalls(t *testing.T, ds *Datastore) { }) require.NoError(t, err) + tfr2, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID2, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "world", PreInstallQuery: "SELECT 2", PostInstallScript: "hello", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr2, StorageID: "storage2", Filename: "file2", Title: "file2", @@ -81,11 +86,13 @@ func testListPendingSoftwareInstalls(t *testing.T, ds *Datastore) { }) require.NoError(t, err) + tfr3, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID3, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "banana", PreInstallQuery: "SELECT 3", PostInstallScript: "apple", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr3, StorageID: "storage3", Filename: "file3", Title: "file3", @@ -617,11 +624,14 @@ func testCleanupUnusedSoftwareInstallers(t *testing.T, ds *Datastore) { ins0File := bytes.NewReader([]byte("installer0")) err = store.Put(ctx, ins0, ins0File) require.NoError(t, err) + _, _ = ins0File.Seek(0, 0) + tfr0, err := fleet.NewTempFileReader(ins0File, t.TempDir) + require.NoError(t, err) assertExisting([]string{ins0}) swi, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "install", - InstallerFile: ins0File, + InstallerFile: tfr0, StorageID: ins0, Filename: "installer0", Title: "ins0", @@ -695,9 +705,11 @@ func testBatchSetSoftwareInstallers(t *testing.T, ds *Datastore) { // add a single installer ins0 := "installer0" ins0File := bytes.NewReader([]byte("installer0")) + tfr0, err := fleet.NewTempFileReader(ins0File, t.TempDir) + require.NoError(t, err) err = ds.BatchSetSoftwareInstallers(ctx, &team.ID, []*fleet.UploadSoftwareInstallerPayload{{ InstallScript: "install", - InstallerFile: ins0File, + InstallerFile: tfr0, StorageID: ins0, Filename: "installer0", Title: "ins0", @@ -724,10 +736,12 @@ func testBatchSetSoftwareInstallers(t *testing.T, ds *Datastore) { // mark ins0 as install_during_setup ins1 := "installer1" ins1File := bytes.NewReader([]byte("installer1")) + tfr1, err := fleet.NewTempFileReader(ins1File, t.TempDir) + require.NoError(t, err) err = ds.BatchSetSoftwareInstallers(ctx, &team.ID, []*fleet.UploadSoftwareInstallerPayload{ { InstallScript: "install", - InstallerFile: ins0File, + InstallerFile: tfr0, StorageID: ins0, Filename: ins0, Title: ins0, @@ -742,7 +756,7 @@ func testBatchSetSoftwareInstallers(t *testing.T, ds *Datastore) { { InstallScript: "install", PostInstallScript: "post-install", - InstallerFile: ins1File, + InstallerFile: tfr1, StorageID: ins1, Filename: ins1, Title: ins1, @@ -776,7 +790,7 @@ func testBatchSetSoftwareInstallers(t *testing.T, ds *Datastore) { { InstallScript: "install", PostInstallScript: "post-install", - InstallerFile: ins1File, + InstallerFile: tfr1, StorageID: ins1, Filename: ins1, Title: ins1, @@ -794,7 +808,7 @@ func testBatchSetSoftwareInstallers(t *testing.T, ds *Datastore) { err = ds.BatchSetSoftwareInstallers(ctx, &team.ID, []*fleet.UploadSoftwareInstallerPayload{ { InstallScript: "install", - InstallerFile: ins0File, + InstallerFile: tfr0, StorageID: ins0, Filename: ins0, Title: ins0, @@ -809,7 +823,7 @@ func testBatchSetSoftwareInstallers(t *testing.T, ds *Datastore) { { InstallScript: "install", PostInstallScript: "post-install", - InstallerFile: ins1File, + InstallerFile: tfr1, StorageID: ins1, Filename: ins1, Title: ins1, @@ -827,7 +841,7 @@ func testBatchSetSoftwareInstallers(t *testing.T, ds *Datastore) { err = ds.BatchSetSoftwareInstallers(ctx, &team.ID, []*fleet.UploadSoftwareInstallerPayload{ { InstallScript: "install", - InstallerFile: ins0File, + InstallerFile: tfr0, StorageID: ins0, Filename: ins0, Title: ins0, @@ -842,7 +856,7 @@ func testBatchSetSoftwareInstallers(t *testing.T, ds *Datastore) { { InstallScript: "install", PostInstallScript: "post-install", - InstallerFile: ins1File, + InstallerFile: tfr1, StorageID: ins1, Filename: ins1, Title: ins1, @@ -861,7 +875,7 @@ func testBatchSetSoftwareInstallers(t *testing.T, ds *Datastore) { { InstallScript: "install", PostInstallScript: "post-install", - InstallerFile: ins1File, + InstallerFile: tfr1, StorageID: ins1, Filename: ins1, Title: ins1, @@ -1061,13 +1075,16 @@ func testDeleteSoftwareInstallers(t *testing.T, ds *Datastore) { ins0File := bytes.NewReader([]byte("installer0")) err = store.Put(ctx, ins0, ins0File) require.NoError(t, err) + _, _ = ins0File.Seek(0, 0) + tfr0, err := fleet.NewTempFileReader(ins0File, t.TempDir) + require.NoError(t, err) team1, err := ds.NewTeam(ctx, &fleet.Team{Name: "team1"}) require.NoError(t, err) softwareInstallerID, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "install", - InstallerFile: ins0File, + InstallerFile: tfr0, StorageID: ins0, Filename: "installer.pkg", Title: "ins0", @@ -1136,10 +1153,13 @@ func testGetHostLastInstallData(t *testing.T, ds *Datastore) { ins0File := bytes.NewReader([]byte("installer0")) err = store.Put(ctx, ins0, ins0File) require.NoError(t, err) + _, _ = ins0File.Seek(0, 0) + tfr0, err := fleet.NewTempFileReader(ins0File, t.TempDir) + require.NoError(t, err) softwareInstallerID1, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "install", - InstallerFile: ins0File, + InstallerFile: tfr0, StorageID: ins0, Filename: "installer.pkg", Title: "ins1", @@ -1151,7 +1171,7 @@ func testGetHostLastInstallData(t *testing.T, ds *Datastore) { require.NoError(t, err) softwareInstallerID2, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "install2", - InstallerFile: ins0File, + InstallerFile: tfr0, StorageID: ins0, Filename: "installer2.pkg", Title: "ins2", diff --git a/server/datastore/mysql/software_test.go b/server/datastore/mysql/software_test.go index c7463789f6..4e99fe4b6a 100644 --- a/server/datastore/mysql/software_test.go +++ b/server/datastore/mysql/software_test.go @@ -1,7 +1,6 @@ package mysql import ( - "bytes" "context" "database/sql" "encoding/hex" @@ -4600,9 +4599,11 @@ func testListHostSoftwareInstallThenTransferTeam(t *testing.T, ds *Datastore) { require.NoError(t, err) // create a software installer for team 1 + tfr1, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerTm1, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr1, StorageID: "storage1", Filename: "file1", Title: "file1", @@ -4710,9 +4711,11 @@ func testListHostSoftwareInstallThenDeleteInstallers(t *testing.T, ds *Datastore require.NoError(t, err) // create a software installer for team 1 + tfr1, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerTm1, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr1, StorageID: "storage1", Filename: "file1", Title: "file1", diff --git a/server/fleet/software_installer.go b/server/fleet/software_installer.go index d715e1ddc5..aae130902a 100644 --- a/server/fleet/software_installer.go +++ b/server/fleet/software_installer.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "io" + "os" "path/filepath" "strings" "time" @@ -312,7 +313,7 @@ type UploadSoftwareInstallerPayload struct { InstallScript string PreInstallQuery string PostInstallScript string - InstallerFile io.ReadSeeker // TODO: maybe pull this out of the payload and only pass it to methods that need it (e.g., won't be needed when storing metadata in the database) + InstallerFile *TempFileReader // TODO: maybe pull this out of the payload and only pass it to methods that need it (e.g., won't be needed when storing metadata in the database) StorageID string Filename string Title string @@ -338,7 +339,7 @@ type UpdateSoftwareInstallerPayload struct { // used for authorization and persisted as author UserID uint // optional; used for pulling metadata + persisting new installer package to file system - InstallerFile io.ReadSeeker + InstallerFile *TempFileReader // update the installer with these fields (*not* PATCH semantics at that point; while the // associated endpoint is a PATCH, the entire row will be updated to these values, including // blanks, so make sure they're set from either user input or the existing installer row @@ -533,3 +534,70 @@ type SoftwareInstallerTokenMetadata struct { } const SoftwareInstallerURLMaxLength = 255 + +// TempFileReader is an io.Reader with all extra io interfaces supported by a +// file on disk reader (e.g. io.ReaderAt, io.Seeker, etc.). When created with +// NewTempFileReader, it is backed by a temporary file on disk, and that file +// is deleted when Close is called. +type TempFileReader struct { + *os.File + keepFile bool +} + +// Rewind seeks to the beginning of the file so the next read will read from +// the start of the bytes. +func (r *TempFileReader) Rewind() error { + if _, err := r.Seek(0, io.SeekStart); err != nil { + return err + } + return nil +} + +// Close closes the TempFileReader and deletes the underlying temp file unless +// it was instructed not to do so at creation time. +func (r *TempFileReader) Close() error { + cerr := r.File.Close() + var rerr error + if !r.keepFile { + rerr = os.Remove(r.File.Name()) + } + if cerr != nil { + return cerr + } + return rerr +} + +// NewKeepFileReader creates a TempFileReader from a file path and keeps the +// file on Close, instead of deleting it. +func NewKeepFileReader(filename string) (*TempFileReader, error) { + f, err := os.Open(filename) + if err != nil { + return nil, err + } + return &TempFileReader{File: f, keepFile: true}, nil +} + +// NewTempFileReader creates a temp file to store the data from the provided +// reader and returns a TempFileReader that reads from that temp file, deleting +// it on close. +func NewTempFileReader(from io.Reader, tempDirFn func() string) (*TempFileReader, error) { + if tempDirFn == nil { + tempDirFn = os.TempDir + } + + tempFile, err := os.CreateTemp(tempDirFn(), "fleet-temp-file-*") + if err != nil { + return nil, err + } + tfr := &TempFileReader{File: tempFile} + + if _, err := io.Copy(tempFile, from); err != nil { + _ = tfr.Close() // best-effort close/delete + return nil, err + } + if err := tfr.Rewind(); err != nil { + _ = tfr.Close() // best-effort close/delete + return nil, err + } + return tfr, nil +} diff --git a/server/fleet/software_installer_test.go b/server/fleet/software_installer_test.go new file mode 100644 index 0000000000..e26cb22332 --- /dev/null +++ b/server/fleet/software_installer_test.go @@ -0,0 +1,67 @@ +package fleet + +import ( + "io" + "os" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestTempFileReader(t *testing.T) { + content1And2 := "Hello, World!" + tfr1, err := NewTempFileReader(strings.NewReader(content1And2), t.TempDir) + require.NoError(t, err) + tfr2, err := NewTempFileReader(strings.NewReader(content1And2), t.TempDir) + require.NoError(t, err) + + content3 := "Hello, Temp!" + keepFile, err := os.CreateTemp(t.TempDir(), "test") + require.NoError(t, err) + _, err = io.Copy(keepFile, strings.NewReader(content3)) + require.NoError(t, err) + err = keepFile.Close() + require.NoError(t, err) + tfr3, err := NewKeepFileReader(keepFile.Name()) + require.NoError(t, err) + + b, err := io.ReadAll(tfr1) + require.NoError(t, err) + require.Equal(t, content1And2, string(b)) + b, err = io.ReadAll(tfr2) + require.NoError(t, err) + require.Equal(t, content1And2, string(b)) + + // rewind and read again gets the same content + err = tfr1.Rewind() + require.NoError(t, err) + b, err = io.ReadAll(tfr1) + require.NoError(t, err) + require.Equal(t, content1And2, string(b)) + + // tfr2 is at EOF, so it reads nothing + b, err = io.ReadAll(tfr2) + require.NoError(t, err) + require.Equal(t, "", string(b)) + + b, err = io.ReadAll(tfr3) + require.NoError(t, err) + require.Equal(t, content3, string(b)) + + // closing deletes the file + err = tfr1.Close() + require.NoError(t, err) + _, err = os.Stat(tfr1.Name()) + require.True(t, os.IsNotExist(err)) + + // tfr2 still exists + _, err = os.Stat(tfr2.Name()) + require.False(t, os.IsNotExist(err)) + + // tfr3 still exists even after Close + err = tfr3.Close() + require.NoError(t, err) + _, err = os.Stat(tfr3.Name()) + require.False(t, os.IsNotExist(err)) +} diff --git a/server/mdm/maintainedapps/installers.go b/server/mdm/maintainedapps/installers.go index e9b6e145ce..c98ff2d4c1 100644 --- a/server/mdm/maintainedapps/installers.go +++ b/server/mdm/maintainedapps/installers.go @@ -3,7 +3,6 @@ package maintainedapps import ( "context" "fmt" - "io" "mime" "net/http" "net/url" @@ -18,7 +17,7 @@ import ( const InstallerTimeout = 15 * time.Minute // DownloadInstaller downloads the maintained app installer located at the given URL. -func DownloadInstaller(ctx context.Context, installerURL string, client *http.Client) ([]byte, string, error) { +func DownloadInstaller(ctx context.Context, installerURL string, client *http.Client) (*fleet.TempFileReader, string, error) { // validate the URL before doing the request _, err := url.ParseRequestURI(installerURL) if err != nil { @@ -70,10 +69,10 @@ func DownloadInstaller(ctx context.Context, installerURL string, client *http.Cl filename = path.Base(resp.Request.URL.Path) } - bodyBytes, err := io.ReadAll(resp.Body) + tfr, err := fleet.NewTempFileReader(resp.Body, nil) if err != nil { return nil, "", ctxerr.Wrapf(ctx, err, "reading installer %q contents", installerURL) } - return bodyBytes, filename, nil + return tfr, filename, nil } diff --git a/server/service/integration_core_test.go b/server/service/integration_core_test.go index 1c7a73f6d0..2e08b4ef9d 100644 --- a/server/service/integration_core_test.go +++ b/server/service/integration_core_test.go @@ -11697,9 +11697,11 @@ func (s *integrationTestSuite) TestListHostUpcomingActivities() { h1E := hsr.ExecutionID // create a software installation request + tfr1, err := fleet.NewTempFileReader(strings.NewReader("echo"), t.TempDir) + require.NoError(t, err) sw1, err := s.ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "install foo", - InstallerFile: strings.NewReader("echo"), + InstallerFile: tfr1, StorageID: uuid.NewString(), Filename: "foo.pkg", Title: "foo", diff --git a/server/service/integration_mdm_test.go b/server/service/integration_mdm_test.go index 987b22bddf..0e51ef7cc1 100644 --- a/server/service/integration_mdm_test.go +++ b/server/service/integration_mdm_test.go @@ -11824,12 +11824,14 @@ func (s *integrationMDMTestSuite) TestSetupExperience() { team1, err := ds.NewTeam(ctx, &fleet.Team{Name: "team1"}) require.NoError(t, err) + tfr1, err := fleet.NewTempFileReader(strings.NewReader("hello"), t.TempDir) + require.NoError(t, err) installerID1, err := ds.MatchOrCreateSoftwareInstaller(ctx, &fleet.UploadSoftwareInstallerPayload{ InstallScript: "hello", PreInstallQuery: "SELECT 1", PostInstallScript: "world", UninstallScript: "goodbye", - InstallerFile: bytes.NewReader([]byte("hello")), + InstallerFile: tfr1, StorageID: "storage1", Filename: "file1", Title: "file1", diff --git a/server/service/software_installers.go b/server/service/software_installers.go index 37f0036feb..0045ed962b 100644 --- a/server/service/software_installers.go +++ b/server/service/software_installers.go @@ -152,7 +152,14 @@ func updateSoftwareInstallerEndpoint(ctx context.Context, request interface{}, s return uploadSoftwareInstallerResponse{Err: err}, nil } defer ff.Close() - payload.InstallerFile = ff + + tfr, err := fleet.NewTempFileReader(ff, nil) + if err != nil { + return uploadSoftwareInstallerResponse{Err: err}, nil + } + defer tfr.Close() + + payload.InstallerFile = tfr payload.Filename = req.File.Filename } @@ -267,12 +274,18 @@ func uploadSoftwareInstallerEndpoint(ctx context.Context, request interface{}, s } defer ff.Close() + tfr, err := fleet.NewTempFileReader(ff, nil) + if err != nil { + return uploadSoftwareInstallerResponse{Err: err}, nil + } + defer tfr.Close() + payload := &fleet.UploadSoftwareInstallerPayload{ TeamID: req.TeamID, InstallScript: req.InstallScript, PreInstallQuery: req.PreInstallQuery, PostInstallScript: req.PostInstallScript, - InstallerFile: ff, + InstallerFile: tfr, Filename: req.File.Filename, SelfService: req.SelfService, UninstallScript: req.UninstallScript, diff --git a/server/service/testing_client.go b/server/service/testing_client.go index e2da93ed0e..291b671a3e 100644 --- a/server/service/testing_client.go +++ b/server/service/testing_client.go @@ -539,16 +539,12 @@ func (ts *withServer) uploadSoftwareInstaller( expectedError string, ) { t.Helper() - openFile := func(name string) *os.File { - f, err := os.Open(filepath.Join("testdata", "software-installers", name)) - require.NoError(t, err) - return f - } - f := openFile(payload.Filename) - defer f.Close() + tfr, err := fleet.NewKeepFileReader(filepath.Join("testdata", "software-installers", payload.Filename)) + require.NoError(t, err) + defer tfr.Close() - payload.InstallerFile = f + payload.InstallerFile = tfr var b bytes.Buffer w := multipart.NewWriter(&b)