github.com/kubeshop/testkube@v1.17.23/cmd/tcl/testworkflow-toolkit/artifacts/tar_stream.go (about)

     1  // Copyright 2024 Testkube.
     2  //
     3  // Licensed as a Testkube Pro file under the Testkube Community
     4  // License (the "License"); you may not use this file except in compliance with
     5  // the License. You may obtain a copy of the License at
     6  //
     7  //	https://github.com/kubeshop/testkube/blob/main/licenses/TCL.txt
     8  
     9  package artifacts
    10  
    11  import (
    12  	"archive/tar"
    13  	"compress/gzip"
    14  	"fmt"
    15  	"io"
    16  	"io/fs"
    17  	"sync"
    18  )
    19  
    20  type tarStream struct {
    21  	reader io.ReadCloser
    22  	writer io.WriteCloser
    23  	gzip   io.WriteCloser
    24  	tar    *tar.Writer
    25  	mu     *sync.Mutex
    26  	wg     sync.WaitGroup
    27  }
    28  
    29  func NewTarStream() *tarStream {
    30  	reader, writer := io.Pipe()
    31  	gzip := gzip.NewWriter(writer)
    32  	tar := tar.NewWriter(gzip)
    33  	return &tarStream{
    34  		reader: reader,
    35  		writer: writer,
    36  		gzip:   gzip,
    37  		tar:    tar,
    38  		mu:     &sync.Mutex{},
    39  	}
    40  }
    41  
    42  func (t *tarStream) Add(path string, file fs.File, stat fs.FileInfo) error {
    43  	t.wg.Add(1)
    44  	t.mu.Lock()
    45  	defer t.mu.Unlock()
    46  	defer t.wg.Done()
    47  
    48  	// Write file header
    49  	name := stat.Name()
    50  	header, err := tar.FileInfoHeader(stat, name)
    51  	if err != nil {
    52  		return err
    53  	}
    54  	header.Name = path
    55  	err = t.tar.WriteHeader(header)
    56  	if err != nil {
    57  		return err
    58  	}
    59  	_, err = io.Copy(t.tar, file)
    60  	return err
    61  }
    62  
    63  func (t *tarStream) Read(p []byte) (n int, err error) {
    64  	return t.reader.Read(p)
    65  }
    66  
    67  func (t *tarStream) Done() chan struct{} {
    68  	ch := make(chan struct{})
    69  	go func() {
    70  		t.wg.Wait()
    71  		close(ch)
    72  	}()
    73  	return ch
    74  }
    75  
    76  func (t *tarStream) Close() (err error) {
    77  	err = t.tar.Close()
    78  	if err != nil {
    79  		_ = t.gzip.Close()
    80  		_ = t.writer.Close()
    81  		return fmt.Errorf("closing tar: tar: %v", err)
    82  	}
    83  	err = t.gzip.Close()
    84  	if err != nil {
    85  		_ = t.writer.Close()
    86  		return fmt.Errorf("closing tar: gzip: %v", err)
    87  	}
    88  	err = t.writer.Close()
    89  	if err != nil {
    90  		return fmt.Errorf("closing tar: pipe: %v", err)
    91  	}
    92  	return nil
    93  }