github.com/devseccon/trivy@v0.47.1-0.20231123133102-bd902a0bd996/internal/testutil/gzip.go (about)

     1  package testutil
     2  
     3  import (
     4  	"bytes"
     5  	"compress/gzip"
     6  	"io"
     7  	"os"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/require"
    11  )
    12  
    13  const (
    14  	max       = int64(10) << 30 // 10GB
    15  	blockSize = 4096
    16  )
    17  
    18  func DecompressGzip(t *testing.T, src, dst string) {
    19  	w, err := os.Create(dst)
    20  	require.NoError(t, err)
    21  	defer w.Close()
    22  
    23  	f, err := os.Open(src)
    24  	require.NoError(t, err)
    25  	defer f.Close()
    26  
    27  	gr, err := gzip.NewReader(f)
    28  	require.NoError(t, err)
    29  
    30  	_, err = io.CopyN(w, gr, max)
    31  	require.ErrorIs(t, err, io.EOF)
    32  }
    33  
    34  // DecompressSparseGzip decompresses a sparse gzip file for virtual machine image.
    35  func DecompressSparseGzip(t *testing.T, src, dst string) {
    36  	w, err := os.Create(dst)
    37  	require.NoError(t, err)
    38  	defer w.Close()
    39  
    40  	f, err := os.Open(src)
    41  	require.NoError(t, err)
    42  	defer f.Close()
    43  
    44  	gr, err := gzip.NewReader(f)
    45  	require.NoError(t, err)
    46  
    47  	buf := make([]byte, blockSize)
    48  	var size int
    49  	var written int64
    50  	for {
    51  		n, err := gr.Read(buf)
    52  		if n == 0 && err != nil {
    53  			if err == io.EOF {
    54  				break
    55  			}
    56  			require.NoError(t, err)
    57  		}
    58  
    59  		size += n
    60  		err = w.Truncate(int64(size))
    61  		require.NoError(t, err)
    62  
    63  		if !bytes.Equal(buf[:n], make([]byte, n)) {
    64  			wn, err := w.WriteAt(buf[:n], int64(size-n))
    65  			if err != nil {
    66  				if err == io.EOF {
    67  					break
    68  				}
    69  				require.NoError(t, err)
    70  			}
    71  			written += int64(wn)
    72  			if written > max {
    73  				require.Fail(t, "written size exceeds max")
    74  			}
    75  		}
    76  	}
    77  }