assemble.go

Documentation: github.com/vbatts/tar-split/tar/asm

     1  package asm
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"hash"
     7  	"hash/crc64"
     8  	"io"
     9  	"sync"
    10  
    11  	"github.com/vbatts/tar-split/tar/storage"
    12  )
    13  
    14  // NewOutputTarStream returns an io.ReadCloser that is an assembled tar archive
    15  // stream.
    16  //
    17  // It takes a storage.FileGetter, for mapping the file payloads that are to be read in,
    18  // and a storage.Unpacker, which has access to the rawbytes and file order
    19  // metadata. With the combination of these two items, a precise assembled Tar
    20  // archive is possible.
    21  func NewOutputTarStream(fg storage.FileGetter, up storage.Unpacker) io.ReadCloser {
    22  	// ... Since these are interfaces, this is possible, so let's not have a nil pointer
    23  	if fg == nil || up == nil {
    24  		return nil
    25  	}
    26  	pr, pw := io.Pipe()
    27  	go func() {
    28  		err := WriteOutputTarStream(fg, up, pw)
    29  		if err != nil {
    30  			pw.CloseWithError(err)
    31  		} else {
    32  			pw.Close()
    33  		}
    34  	}()
    35  	return pr
    36  }
    37  
    38  // WriteOutputTarStream writes assembled tar archive to a writer.
    39  func WriteOutputTarStream(fg storage.FileGetter, up storage.Unpacker, w io.Writer) error {
    40  	// ... Since these are interfaces, this is possible, so let's not have a nil pointer
    41  	if fg == nil || up == nil {
    42  		return nil
    43  	}
    44  	var copyBuffer []byte
    45  	var crcHash hash.Hash
    46  	var crcSum []byte
    47  	var multiWriter io.Writer
    48  	for {
    49  		entry, err := up.Next()
    50  		if err != nil {
    51  			if err == io.EOF {
    52  				return nil
    53  			}
    54  			return err
    55  		}
    56  		switch entry.Type {
    57  		case storage.SegmentType:
    58  			if _, err := w.Write(entry.Payload); err != nil {
    59  				return err
    60  			}
    61  		case storage.FileType:
    62  			if entry.Size == 0 {
    63  				continue
    64  			}
    65  			fh, err := fg.Get(entry.GetName())
    66  			if err != nil {
    67  				return err
    68  			}
    69  			if crcHash == nil {
    70  				crcHash = crc64.New(storage.CRCTable)
    71  				crcSum = make([]byte, 8)
    72  				multiWriter = io.MultiWriter(w, crcHash)
    73  				copyBuffer = byteBufferPool.Get().([]byte)
    74  				// TODO once we have some benchmark or memory profile then we can experiment with using *bytes.Buffer
    75  				//nolint:staticcheck // SA6002 not going to do a pointer here
    76  				defer byteBufferPool.Put(copyBuffer)
    77  			} else {
    78  				crcHash.Reset()
    79  			}
    80  
    81  			if _, err := copyWithBuffer(multiWriter, fh, copyBuffer); err != nil {
    82  				fh.Close()
    83  				return err
    84  			}
    85  
    86  			if !bytes.Equal(crcHash.Sum(crcSum[:0]), entry.Payload) {
    87  				// I would rather this be a comparable ErrInvalidChecksum or such,
    88  				// but since it's coming through the PipeReader, the context of
    89  				// _which_ file would be lost...
    90  				fh.Close()
    91  				return fmt.Errorf("file integrity checksum failed for %q", entry.GetName())
    92  			}
    93  			fh.Close()
    94  		}
    95  	}
    96  }
    97  
    98  var byteBufferPool = &sync.Pool{
    99  	New: func() interface{} {
   100  		return make([]byte, 32*1024)
   101  	},
   102  }
   103  
   104  // copyWithBuffer is taken from stdlib io.Copy implementation
   105  // https://github.com/golang/go/blob/go1.5.1/src/io/io.go#L367
   106  func copyWithBuffer(dst io.Writer, src io.Reader, buf []byte) (written int64, err error) {
   107  	for {
   108  		nr, er := src.Read(buf)
   109  		if nr > 0 {
   110  			nw, ew := dst.Write(buf[0:nr])
   111  			if nw > 0 {
   112  				written += int64(nw)
   113  			}
   114  			if ew != nil {
   115  				err = ew
   116  				break
   117  			}
   118  			if nr != nw {
   119  				err = io.ErrShortWrite
   120  				break
   121  			}
   122  		}
   123  		if er == io.EOF {
   124  			break
   125  		}
   126  		if er != nil {
   127  			err = er
   128  			break
   129  		}
   130  	}
   131  	return written, err
   132  }
   133
View as plain text