Skip to content

Commit

Permalink
optimize reader size retrieval performance
Browse files Browse the repository at this point in the history
  • Loading branch information
nikpivkin committed Feb 28, 2024
1 parent 5db9dc2 commit 557177b
Show file tree
Hide file tree
Showing 4 changed files with 56 additions and 19 deletions.
2 changes: 1 addition & 1 deletion pkg/iac/scanners/terraformplan/snapshot/scanner.go
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ func (s *Scanner) ScanFile(ctx context.Context, fsys fs.FS, filepath string) (sc
}

func (s *Scanner) Scan(ctx context.Context, reader io.Reader) (scan.Results, error) {
snap, err := readSnapshot(reader)
snap, err := parseSnapshot(reader)
if err != nil {
return nil, err
}
Expand Down
35 changes: 19 additions & 16 deletions pkg/iac/scanners/terraformplan/snapshot/snapshot.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package snapshot

import (
"archive/zip"
"bytes"
"encoding/json"
"errors"
"fmt"
Expand All @@ -14,6 +13,8 @@ import (
"strings"

"github.com/liamg/memoryfs"

iox "github.com/aquasecurity/trivy/pkg/x/io"
)

const (
Expand All @@ -34,34 +35,36 @@ type (
configSnapshotModuleManifest []configSnapshotModuleRecord
)

var errNoTerraformPlan = errors.New("no terraform plan file")

func IsPlanSnapshot(r io.Reader) bool {
if r == nil {
zr, err := readSnapshot(r)
if err != nil {
return false
}
return containsTfplanFile(zr)
}

buf, err := io.ReadAll(r)
if err != nil {
return false
func readSnapshot(r io.Reader) (*zip.Reader, error) {
if r == nil {
return nil, errors.New("reader is nil")
}

zr, err := zip.NewReader(bytes.NewReader(buf), int64(len(buf)))
rsa, size, err := iox.NewReadSeekerAtWithSize(r)
if err != nil {
return false
return nil, err
}

return containsTfplanFile(zr)
}

var errNoTerraformPlan = errors.New("no terraform plan file")

func readSnapshot(r io.Reader) (*snapshot, error) {
b, err := io.ReadAll(r)
zr, err := zip.NewReader(rsa, size)
if err != nil {
return nil, err
}

br := bytes.NewReader(b)
zr, err := zip.NewReader(br, int64(len(b)))
return zr, nil
}

func parseSnapshot(r io.Reader) (*snapshot, error) {
zr, err := readSnapshot(r)
if err != nil {
return nil, err
}
Expand Down
3 changes: 1 addition & 2 deletions pkg/iac/scanners/terraformplan/snapshot/snapshot_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ import (
)

func TestReadSnapshot(t *testing.T) {

tests := []struct {
name string
dir string
Expand Down Expand Up @@ -48,7 +47,7 @@ func TestReadSnapshot(t *testing.T) {
require.NoError(t, err)
defer f.Close()

snapshot, err := readSnapshot(f)
snapshot, err := parseSnapshot(f)
require.NoError(t, err)
require.NotNil(t, snapshot)

Expand Down
35 changes: 35 additions & 0 deletions pkg/x/io/io.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,41 @@ func NewReadSeekerAt(r io.Reader) (ReadSeekerAt, error) {
return bytes.NewReader(buff.Bytes()), nil
}

func NewReadSeekerAtWithSize(r io.Reader) (ReadSeekerAt, int64, error) {
rsa, err := NewReadSeekerAt(r)
if err != nil {
return nil, 0, err
}

br, ok := rsa.(*bytes.Reader)
if ok {
return rsa, br.Size(), nil
}

size, err := getSize(r)
if err != nil {
return nil, 0, xerrors.Errorf("get size error: %w", err)
}
return rsa, size, nil
}

func getSize(r io.Reader) (int64, error) {
s, ok := r.(io.Seeker)
if !ok {
return 0, xerrors.New("reader does not support seeking for size")
}

size, err := s.Seek(0, io.SeekEnd)
if err != nil {
return 0, xerrors.Errorf("seek error: %w", err)
}

if _, err = s.Seek(0, io.SeekStart); err != nil {
return 0, xerrors.Errorf("seek error: %w", err)
}
return size, nil
}

// NopCloser returns a ReadSeekCloserAt with a no-op Close method wrapping
// the provided Reader r.
func NopCloser(r ReadSeekerAt) ReadSeekCloserAt {
Expand Down

0 comments on commit 557177b

Please sign in to comment.