github.com/arieschain/arieschain@v0.0.0-20191023063405-37c074544356/internal/build/archive.go (about)

     1  package build
     2  
     3  import (
     4  	"archive/tar"
     5  	"archive/zip"
     6  	"compress/gzip"
     7  	"fmt"
     8  	"io"
     9  	"os"
    10  	"path/filepath"
    11  	"strings"
    12  )
    13  
    14  type Archive interface {
    15  	// Directory adds a new directory entry to the archive and sets the
    16  	// directory for subsequent calls to Header.
    17  	Directory(name string) error
    18  
    19  	// Header adds a new file to the archive. The file is added to the directory
    20  	// set by Directory. The content of the file must be written to the returned
    21  	// writer.
    22  	Header(os.FileInfo) (io.Writer, error)
    23  
    24  	// Close flushes the archive and closes the underlying file.
    25  	Close() error
    26  }
    27  
    28  func NewArchive(file *os.File) (Archive, string) {
    29  	switch {
    30  	case strings.HasSuffix(file.Name(), ".zip"):
    31  		return NewZipArchive(file), strings.TrimSuffix(file.Name(), ".zip")
    32  	case strings.HasSuffix(file.Name(), ".tar.gz"):
    33  		return NewTarballArchive(file), strings.TrimSuffix(file.Name(), ".tar.gz")
    34  	default:
    35  		return nil, ""
    36  	}
    37  }
    38  
    39  // AddFile appends an existing file to an archive.
    40  func AddFile(a Archive, file string) error {
    41  	fd, err := os.Open(file)
    42  	if err != nil {
    43  		return err
    44  	}
    45  	defer fd.Close()
    46  	fi, err := fd.Stat()
    47  	if err != nil {
    48  		return err
    49  	}
    50  	w, err := a.Header(fi)
    51  	if err != nil {
    52  		return err
    53  	}
    54  	if _, err := io.Copy(w, fd); err != nil {
    55  		return err
    56  	}
    57  	return nil
    58  }
    59  
    60  // WriteArchive creates an archive containing the given files.
    61  func WriteArchive(name string, files []string) (err error) {
    62  	archfd, err := os.Create(name)
    63  	if err != nil {
    64  		return err
    65  	}
    66  
    67  	defer func() {
    68  		archfd.Close()
    69  		// Remove the half-written archive on failure.
    70  		if err != nil {
    71  			os.Remove(name)
    72  		}
    73  	}()
    74  	archive, basename := NewArchive(archfd)
    75  	if archive == nil {
    76  		return fmt.Errorf("unknown archive extension")
    77  	}
    78  	fmt.Println(name)
    79  	if err := archive.Directory(basename); err != nil {
    80  		return err
    81  	}
    82  	for _, file := range files {
    83  		fmt.Println("   +", filepath.Base(file))
    84  		if err := AddFile(archive, file); err != nil {
    85  			return err
    86  		}
    87  	}
    88  	return archive.Close()
    89  }
    90  
    91  type ZipArchive struct {
    92  	dir  string
    93  	zipw *zip.Writer
    94  	file io.Closer
    95  }
    96  
    97  func NewZipArchive(w io.WriteCloser) Archive {
    98  	return &ZipArchive{"", zip.NewWriter(w), w}
    99  }
   100  
   101  func (a *ZipArchive) Directory(name string) error {
   102  	a.dir = name + "/"
   103  	return nil
   104  }
   105  
   106  func (a *ZipArchive) Header(fi os.FileInfo) (io.Writer, error) {
   107  	head, err := zip.FileInfoHeader(fi)
   108  	if err != nil {
   109  		return nil, fmt.Errorf("can't make zip header: %v", err)
   110  	}
   111  	head.Name = a.dir + head.Name
   112  	head.Method = zip.Deflate
   113  	w, err := a.zipw.CreateHeader(head)
   114  	if err != nil {
   115  		return nil, fmt.Errorf("can't add zip header: %v", err)
   116  	}
   117  	return w, nil
   118  }
   119  
   120  func (a *ZipArchive) Close() error {
   121  	if err := a.zipw.Close(); err != nil {
   122  		return err
   123  	}
   124  	return a.file.Close()
   125  }
   126  
   127  type TarballArchive struct {
   128  	dir  string
   129  	tarw *tar.Writer
   130  	gzw  *gzip.Writer
   131  	file io.Closer
   132  }
   133  
   134  func NewTarballArchive(w io.WriteCloser) Archive {
   135  	gzw := gzip.NewWriter(w)
   136  	tarw := tar.NewWriter(gzw)
   137  	return &TarballArchive{"", tarw, gzw, w}
   138  }
   139  
   140  func (a *TarballArchive) Directory(name string) error {
   141  	a.dir = name + "/"
   142  	return a.tarw.WriteHeader(&tar.Header{
   143  		Name:     a.dir,
   144  		Mode:     0755,
   145  		Typeflag: tar.TypeDir,
   146  	})
   147  }
   148  
   149  func (a *TarballArchive) Header(fi os.FileInfo) (io.Writer, error) {
   150  	head, err := tar.FileInfoHeader(fi, "")
   151  	if err != nil {
   152  		return nil, fmt.Errorf("can't make tar header: %v", err)
   153  	}
   154  	head.Name = a.dir + head.Name
   155  	if err := a.tarw.WriteHeader(head); err != nil {
   156  		return nil, fmt.Errorf("can't add tar header: %v", err)
   157  	}
   158  	return a.tarw, nil
   159  }
   160  
   161  func (a *TarballArchive) Close() error {
   162  	if err := a.tarw.Close(); err != nil {
   163  		return err
   164  	}
   165  	if err := a.gzw.Close(); err != nil {
   166  		return err
   167  	}
   168  	return a.file.Close()
   169  }