Initial commit

vendor/github.com/hyperhq/hypercli/api/client/tarfile.go

@@ -0,0 +1,246 @@
+package client
+
+import (
+	"archive/tar"
+	"bytes"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+
+	"github.com/cheggaaa/pb"
+)
+
+const (
+	TARINFO_HEADER = iota
+	TARINFO_FILE
+	TARINFO_PAD
+	TARINFO_UPLOADED
+	TARINFO_FINISHED
+)
+
+type tarInfo struct {
+	info     os.FileInfo
+	relPath  string
+	linkName string
+	path     string
+	pad      int // number of zeros to pad at the end of the file entry
+
+	headerBuf *bytes.Buffer
+	pos       int64
+	state     int
+}
+
+type TarFile struct {
+	fileList  []*tarInfo
+	blockSize int
+	endPad    int
+	padding   []byte
+	closed    bool
+
+	source   string
+	progress *pb.ProgressBar
+}
+
+func (t *TarFile) writeHeader(p []byte, info *tarInfo) (int, error) {
+	if info.headerBuf == nil {
+		header, err := tar.FileInfoHeader(info.info, info.linkName)
+		if err != nil {
+			return 0, err
+		}
+		header.Name = info.relPath
+
+		buf := &bytes.Buffer{}
+		tarWriter := tar.NewWriter(buf)
+		defer tarWriter.Close()
+		err = tarWriter.WriteHeader(header)
+		if err != nil {
+			return 0, err
+		}
+
+		info.headerBuf = &bytes.Buffer{}
+		io.Copy(info.headerBuf, buf)
+	}
+
+	size, err := info.headerBuf.Read(p)
+	if err != nil && err != io.EOF {
+		return 0, err
+	}
+	if info.headerBuf.Len() == 0 {
+		info.headerBuf.Truncate(0)
+	}
+	return size, nil
+}
+
+func (t *TarFile) writeFile(p []byte, info *tarInfo) (int, error) {
+	var f *os.File
+	var err error
+
+	if f, err = os.Open(info.path); err != nil {
+		return 0, err
+	}
+	defer f.Close()
+	// resuming
+	if info.pos != 0 {
+		if _, err = f.Seek(info.pos, os.SEEK_SET); err != nil {
+			return 0, err
+		}
+	}
+	if ret, err := f.Read(p); err != nil && err != io.EOF {
+		return 0, err
+	} else {
+		return ret, nil
+	}
+}
+
+func (t *TarFile) writePad(p []byte, info *tarInfo) (int, error) {
+	buf := bytes.NewBuffer(t.padding[0:info.pad])
+	if ret, err := buf.Read(p); err != nil && err != io.EOF {
+		return 0, err
+	} else {
+		return ret, nil
+	}
+}
+
+// write the trailing zeros of a tar file
+func (t *TarFile) writeClose(p []byte) (n int, err error) {
+	size := t.endPad
+	buf := bytes.NewBuffer(make([]byte, size))
+	if ret, err := buf.Read(p); err != nil && err != io.EOF {
+		return 0, err
+	} else {
+		t.endPad -= ret
+	}
+	if t.endPad <= 0 {
+		if t.progress != nil {
+			t.progress.Finish()
+		}
+		t.closed = true
+		return size, io.EOF
+	}
+	return size - t.endPad, nil
+}
+
+func (t *TarFile) AllocBar(pool *pb.Pool) *pb.ProgressBar {
+	if t.progress == nil {
+		t.progress = pb.New(len(t.fileList)).Prefix(fmt.Sprintf("Sending %s", t.source))
+		pool.Add(t.progress)
+	}
+	return t.progress
+}
+
+func (t *TarFile) AddFile(info os.FileInfo, relPath, linkName, path string) {
+	t.fileList = append(t.fileList, &tarInfo{
+		info:     info,
+		relPath:  relPath,
+		linkName: linkName,
+		path:     path,
+		pad:      (t.blockSize - (int(info.Size()) % t.blockSize)) % t.blockSize,
+	})
+}
+
+func (t *TarFile) Close() error {
+	if !t.closed {
+		if t.progress != nil {
+			t.progress.Finish()
+		}
+	}
+	t.closed = true
+	return nil
+}
+
+func (t *TarFile) Read(p []byte) (n int, err error) {
+	var (
+		file   *tarInfo
+		idx    int
+		length = len(p)
+	)
+
+	if length == 0 {
+		return
+	}
+
+	if t.closed {
+		return 0, io.EOF
+	}
+
+	for idx, file = range t.fileList {
+		if file.state != TARINFO_FINISHED {
+			break
+		}
+		if idx == len(t.fileList)-1 {
+			file = nil
+			break
+		}
+	}
+
+	defer func() {
+		if err != nil && err != io.EOF {
+			t.Close()
+		}
+	}()
+
+	if file == nil {
+		return t.writeClose(p)
+	}
+
+	for n < length && file.state != TARINFO_FINISHED {
+		switch file.state {
+		case TARINFO_HEADER:
+			if ret, err := t.writeHeader(p, file); err != nil {
+				return 0, err
+			} else {
+				n += ret
+				p = p[n:]
+				if file.headerBuf.Len() == 0 {
+					if file.info.Mode().IsRegular() {
+						file.state = TARINFO_FILE
+					} else {
+						file.state = TARINFO_UPLOADED
+					}
+				}
+			}
+		case TARINFO_FILE:
+			if ret, err := t.writeFile(p, file); err != nil {
+				return 0, err
+			} else {
+				file.pos += int64(ret)
+				n += ret
+				p = p[ret:]
+				if file.pos >= file.info.Size() {
+					file.pos = 0
+					file.state = TARINFO_PAD
+				}
+			}
+		case TARINFO_PAD:
+			if file.pad == 0 {
+				file.state = TARINFO_UPLOADED
+			} else if ret, err := t.writePad(p, file); err != nil {
+				return 0, err
+			} else {
+				file.pos += int64(ret)
+				n += ret
+				p = p[ret:]
+				if file.pos >= int64(file.pad) {
+					file.state = TARINFO_UPLOADED
+				}
+			}
+		case TARINFO_UPLOADED:
+			file.state = TARINFO_FINISHED
+			if t.progress != nil {
+				t.progress.Increment()
+			}
+		}
+	}
+
+	return n, nil
+}
+
+func NewTarFile(path string, blockSize int) *TarFile {
+	return &TarFile{
+		blockSize: blockSize,
+		endPad:    blockSize * 2,
+		source:    filepath.Base(path),
+		padding:   make([]byte, blockSize),
+	}
+}