From 557177b20a8aacc96a7e2b9e38ad8075389cbf0d Mon Sep 17 00:00:00 2001
From: nikpivkin <nikita.pivkin@smartforce.io>
Date: Wed, 28 Feb 2024 15:09:55 +0700
Subject: [PATCH] optimize reader size retrieval performance

---
 .../terraformplan/snapshot/scanner.go         |  2 +-
 .../terraformplan/snapshot/snapshot.go        | 35 ++++++++++---------
 .../terraformplan/snapshot/snapshot_test.go   |  3 +-
 pkg/x/io/io.go                                | 35 +++++++++++++++++++
 4 files changed, 56 insertions(+), 19 deletions(-)

diff --git a/pkg/iac/scanners/terraformplan/snapshot/scanner.go b/pkg/iac/scanners/terraformplan/snapshot/scanner.go
index 525716f63379..e2a8d2807fa1 100644
--- a/pkg/iac/scanners/terraformplan/snapshot/scanner.go
+++ b/pkg/iac/scanners/terraformplan/snapshot/scanner.go
@@ -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
 	}
diff --git a/pkg/iac/scanners/terraformplan/snapshot/snapshot.go b/pkg/iac/scanners/terraformplan/snapshot/snapshot.go
index 7fc5d7c791fb..e32e47790c32 100644
--- a/pkg/iac/scanners/terraformplan/snapshot/snapshot.go
+++ b/pkg/iac/scanners/terraformplan/snapshot/snapshot.go
@@ -2,7 +2,6 @@ package snapshot
 
 import (
 	"archive/zip"
-	"bytes"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -14,6 +13,8 @@ import (
 	"strings"
 
 	"github.com/liamg/memoryfs"
+
+	iox "github.com/aquasecurity/trivy/pkg/x/io"
 )
 
 const (
@@ -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
 	}
diff --git a/pkg/iac/scanners/terraformplan/snapshot/snapshot_test.go b/pkg/iac/scanners/terraformplan/snapshot/snapshot_test.go
index 0da3d3d8f7c6..22f26e5e7b6b 100644
--- a/pkg/iac/scanners/terraformplan/snapshot/snapshot_test.go
+++ b/pkg/iac/scanners/terraformplan/snapshot/snapshot_test.go
@@ -14,7 +14,6 @@ import (
 )
 
 func TestReadSnapshot(t *testing.T) {
-
 	tests := []struct {
 		name          string
 		dir           string
@@ -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)
 
diff --git a/pkg/x/io/io.go b/pkg/x/io/io.go
index c4ba4bfe4c59..d7d584fc6349 100644
--- a/pkg/x/io/io.go
+++ b/pkg/x/io/io.go
@@ -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 {