github.com/kubeshop/testkube@v1.17.23/pkg/archive/tarball.go (about)

     1  package archive
     2  
     3  import (
     4  	"archive/tar"
     5  	"bytes"
     6  	"compress/gzip"
     7  	"io"
     8  	"path/filepath"
     9  	"strings"
    10  
    11  	"github.com/pkg/errors"
    12  )
    13  
    14  type Tarball struct{}
    15  
    16  func NewTarballService() *Tarball {
    17  	return &Tarball{}
    18  }
    19  
    20  func (t *Tarball) Extract(in io.Reader) ([]*File, error) {
    21  	tarReader, err := GetTarballReader(in)
    22  	if err != nil {
    23  		return nil, errors.Wrapf(err, "error creating tarball reader")
    24  	}
    25  
    26  	var files []*File
    27  	for true {
    28  		header, err := tarReader.Next()
    29  
    30  		if errors.Is(err, io.EOF) {
    31  			break
    32  		}
    33  		if err != nil {
    34  			return nil, errors.Wrapf(err, "error reading next tarball header")
    35  		}
    36  
    37  		switch header.Typeflag {
    38  		case tar.TypeDir:
    39  			// do nothing
    40  		case tar.TypeReg:
    41  			sanitizedFilepath := t.sanitizeFilepath(header.Name)
    42  			file := &File{Name: sanitizedFilepath, Size: header.Size, Mode: header.Mode, ModTime: header.ModTime, Data: new(bytes.Buffer)}
    43  			if _, err := io.Copy(file.Data, tarReader); err != nil {
    44  				return nil, errors.Wrapf(err, "error copying file %s data to tarball", file.Name)
    45  			}
    46  			files = append(files, file)
    47  		default:
    48  			return nil, errors.Errorf("uknown %v type in tarball %s", header.Typeflag, header.Name)
    49  		}
    50  	}
    51  	return files, nil
    52  }
    53  
    54  func (t *Tarball) sanitizeFilepath(path string) string {
    55  	cleaned := filepath.Clean(path)
    56  	if strings.HasPrefix(cleaned, "..") {
    57  		cleaned = strings.TrimPrefix(cleaned, "..")
    58  	}
    59  	if strings.HasPrefix(cleaned, "/") {
    60  		cleaned = strings.TrimPrefix(cleaned, "/")
    61  	}
    62  	return cleaned
    63  }
    64  
    65  func (t *Tarball) Create(out io.Writer, files []*File) error {
    66  	gzipWriter := gzip.NewWriter(out)
    67  	defer gzipWriter.Close()
    68  
    69  	tarWriter := tar.NewWriter(gzipWriter)
    70  	defer tarWriter.Close()
    71  
    72  	for _, file := range files {
    73  		if err := t.addFileToTarWriter(file, tarWriter); err != nil {
    74  			return errors.Wrapf(err, "error adding file %s to tarball", file.Name)
    75  		}
    76  	}
    77  
    78  	return nil
    79  }
    80  
    81  func (t *Tarball) addFileToTarWriter(file *File, tarWriter *tar.Writer) error {
    82  	tarHeader := &tar.Header{Name: file.Name, Mode: file.Mode, Size: file.Size, ModTime: file.ModTime}
    83  	if err := tarWriter.WriteHeader(tarHeader); err != nil {
    84  		return errors.Wrapf(err, "error writing header for file %s in tarball", file.Name)
    85  	}
    86  
    87  	_, err := tarWriter.Write(file.Data.Bytes())
    88  	if err != nil {
    89  		return errors.Wrapf(err, "error copying file %s data to tarball", file.Name)
    90  	}
    91  
    92  	return nil
    93  }
    94  
    95  func GetTarballReader(gzipStream io.Reader) (*tar.Reader, error) {
    96  	uncompressedStream, err := gzip.NewReader(gzipStream)
    97  	if err != nil {
    98  		return nil, errors.Wrapf(err, "error creating gzip reader")
    99  	}
   100  
   101  	return tar.NewReader(uncompressedStream), nil
   102  }
   103  
   104  var _ Archive = (*Tarball)(nil)