Add support for gzipped archives

1 files changed, 26 insertions, 8 deletions

diff --git a/archive/archive.go b/archive/archive.go
index 39ff9f8..6941b64 100644
--- a/archive/archive.go
+++ b/archive/archive.go
@@ -4,6 +4,8 @@ import (
 	"archive/tar"
 	"path/filepath"
 
+	"compress/gzip"
+
 	"io"
 	"os"
 )
@@ -12,10 +14,9 @@ type Archive struct {
 	Output string
 	Input  []string
 	file   *os.File
-	writer *tar.Writer
 }
 
-func (a *Archive) Tar() error {
+func (a *Archive) Tar(compressionType string) error {
 	var err error
 	if a.file, err = os.Create(a.Output); err != nil {
 		return err
@@ -23,11 +24,28 @@ func (a *Archive) Tar() error {
 
 	defer a.file.Close()
 
-	a.writer = tar.NewWriter(a.file)
-	defer a.writer.Close()
+	var twriter *tar.Writer
+
+	// Set the compression type... if any
+	switch compressionType {
+	case "gz":
+		gzwriter, err := gzip.NewWriterLevel(a.file, gzip.BestSpeed)
+		if err != nil {
+			return err
+		}
+		defer gzwriter.Close()
+		// Write to the gzip writer
+		twriter = tar.NewWriter(gzwriter)
+	default:
+		// Write directly to the file
+		twriter = tar.NewWriter(a.file)
+	}
+
+	// Close off the tar writer when we're done
+	defer twriter.Close()
 
 	for _, v := range a.Input {
-		if err := filepath.Walk(v, a.tarWalkfn()); err != nil {
+		if err := filepath.Walk(v, tarWalkfn(twriter)); err != nil {
 			return err
 		}
 	}
@@ -36,7 +54,7 @@ func (a *Archive) Tar() error {
 
 }
 
-func (a *Archive) tarWalkfn() filepath.WalkFunc {
+func tarWalkfn(writer *tar.Writer) filepath.WalkFunc {
 	// This is an interesting trick to get around scoping issues
 	return func(path string, info os.FileInfo, err error) error {
 		if err != nil {
@@ -62,12 +80,12 @@ func (a *Archive) tarWalkfn() filepath.WalkFunc {
 			ModTime: info.ModTime(),
 		}
 
-		err = a.writer.WriteHeader(h)
+		err = writer.WriteHeader(h)
 		if err != nil {
 			return err
 		}
 
-		_, err = io.Copy(a.writer, f)
+		_, err = io.Copy(writer, f)
 		if err != nil {
 			// TODO: Figure out how to add more useful information to
 			// These errors