2 files changed, 48 insertions, 16 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
diff --git a/archive/unarchive.go b/archive/unarchive.go
index 5089370..d283d4e 100644
--- a/archive/unarchive.go
+++ b/archive/unarchive.go
@@ -4,29 +4,43 @@ import (
 	"archive/tar"
 	fp "path/filepath"
 
+	"compress/gzip"
+
 	"io"
 	"os"
 )
 
 type Unarchive struct {
 	Input     string
-	inputFile *os.File
 	tarReader *tar.Reader
 }
 
 // Extracts a tar arcive to the current working directory
 // This will overwrite everything. So be careful
-func (u *Unarchive) UnTar() error {
-	var err error
-	u.inputFile, err = os.Open(u.Input)
+func (u *Unarchive) UnTar(compressionType string) error {
+	f, err := os.Open(u.Input)
 	if err != nil {
 		return err
 	}
-	defer u.inputFile.Close()
-	u.tarReader = tar.NewReader(u.inputFile)
+	defer f.Close()
+
+	var treader *tar.Reader
+
+	switch compressionType {
+	case "gz":
+		gzreader, err := gzip.NewReader(f)
+		if err != nil {
+			return err
+		}
+		// Read from the gzip reader instead of the file
+		treader = tar.NewReader(gzreader)
+	default:
+		// Read from the file directly
+		treader = tar.NewReader(f)
+	}
 
 	for {
-		hdr, err := u.tarReader.Next()
+		hdr, err := treader.Next()
 		if err == io.EOF {
 			// We've reached the end! Whoee
 			break
@@ -54,7 +68,7 @@ func (u *Unarchive) UnTar() error {
 			return err
 		}
 		defer f.Close()
-		if _, err := io.Copy(f, u.tarReader); err != nil {
+		if _, err := io.Copy(f, treader); err != nil {
 			return err
 		}