Improve memory usage of software installers parsing (#23596)

This commit is contained in:
Martin Angers
2024-11-12 09:28:08 -05:00
committed by GitHub
parent eff1041b14
commit c3ebce9293
25 changed files with 550 additions and 133 deletions
@@ -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.
@@ -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
+9 -8
View File
@@ -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,
+31 -11
View File
@@ -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)
+11 -7
View File
@@ -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
}
+169
View File
@@ -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)
}
})
})
}
}
+4 -3
View File
@@ -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)
+14 -10
View File
@@ -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
}
+8 -6
View File
@@ -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,
+6 -3
View File
@@ -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)
+4 -7
View File
@@ -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)
+15 -12
View File
@@ -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
+9 -4
View File
@@ -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",
+24 -9
View File
@@ -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.
+6 -3
View File
@@ -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",
+28 -10
View File
@@ -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",
@@ -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",
+6 -3
View File
@@ -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",
+70 -2
View File
@@ -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
}
+67
View File
@@ -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))
}
+3 -4
View File
@@ -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
}
+3 -1
View File
@@ -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",
+3 -1
View File
@@ -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",
+15 -2
View File
@@ -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,
+4 -8
View File
@@ -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)