Merge pull request #6 from codeclysm/v3

Add Extractor struct

Author: matteosuppo

README.md

@@ -36,4 +36,28 @@ If you don't know which archive you're dealing with (life really is always a sur
 
 ```go
 extract.Archive(data, "/path/where/to/extract", nil)
+```
+
+If you need more control over how your files will be extracted you can use an Extractor.
+
+It Needs a FS object that implements the FS interface:
+
+```
+type FS interface {
+		Link(string, string) error
+		MkdirAll(string, os.FileMode) error
+		OpenFile(name string, flag int, perm os.FileMode) (*os.File, error)
+		Symlink(string, string) error
+	}
+```
+
+which contains only the required function to perform an extraction. This way it's easy to wrap os functions to 
+chroot the path, or scramble the files, or send an event for each operation or even reimplementing them for an in-memory store, I don't know.
+
+```go
+extractor := extract.Extractor{
+    FS: fs,
+}
+
+extractor.Archive(data, "path/where/to/extract", nil)
 ```

extract.go

@@ -27,22 +27,10 @@
 package extract
 
 import (
-	"archive/tar"
-	"archive/zip"
 	"bytes"
-	"compress/bzip2"
-	"compress/gzip"
 	"context"
 	"io"
-	"io/ioutil"
 	"os"
-	"path/filepath"
-	"strings"
-
-	filetype "github.com/h2non/filetype"
-	"github.com/h2non/filetype/types"
-
-	"github.com/juju/errors"
 )
 
 // Renamer is a function that can be used to rename the files when you're extracting
@@ -55,67 +43,31 @@ type Renamer func(string) string
 // handle the names of the files.
 // If the file is not an archive, an error is returned.
 func Archive(ctx context.Context, body io.Reader, location string, rename Renamer) error {
-	body, kind, err := match(body)
-	if err != nil {
-		errors.Annotatef(err, "Detect archive type")
+	extractor := Extractor{
+		FS: fs{},
 	}
 
-	switch kind.Extension {
-	case "zip":
-		return Zip(ctx, body, location, rename)
-	case "gz":
-		return Gz(ctx, body, location, rename)
-	case "bz2":
-		return Bz2(ctx, body, location, rename)
-	case "tar":
-		return Tar(ctx, body, location, rename)
-	default:
-		return errors.New("Not a supported archive")
-	}
+	return extractor.Archive(ctx, body, location, rename)
 }
 
 // Bz2 extracts a .bz2 or .tar.bz2 archived stream of data in the specified location.
 // It accepts a rename function to handle the names of the files (see the example)
 func Bz2(ctx context.Context, body io.Reader, location string, rename Renamer) error {
-	reader := bzip2.NewReader(body)
-
-	body, kind, err := match(reader)
-	if err != nil {
-		return errors.Annotatef(err, "extract bz2: detect")
-	}
-
-	if kind.Extension == "tar" {
-		return Tar(ctx, body, location, rename)
+	extractor := Extractor{
+		FS: fs{},
 	}
 
-	err = copy(ctx, location, 0666, body)
-	if err != nil {
-		return err
-	}
-	return nil
+	return extractor.Bz2(ctx, body, location, rename)
 }
 
 // Gz extracts a .gz or .tar.gz archived stream of data in the specified location.
 // It accepts a rename function to handle the names of the files (see the example)
 func Gz(ctx context.Context, body io.Reader, location string, rename Renamer) error {
-	reader, err := gzip.NewReader(body)
-	if err != nil {
-		return errors.Annotatef(err, "Gunzip")
-	}
-
-	body, kind, err := match(reader)
-	if err != nil {
-		return err
+	extractor := Extractor{
+		FS: fs{},
 	}
 
-	if kind.Extension == "tar" {
-		return Tar(ctx, body, location, rename)
-	}
-	err = copy(ctx, location, 0666, body)
-	if err != nil {
-		return err
-	}
-	return nil
+	return extractor.Gz(ctx, body, location, rename)
 }
 
 type file struct {
@@ -131,219 +83,37 @@ type link struct {
 // Tar extracts a .tar archived stream of data in the specified location.
 // It accepts a rename function to handle the names of the files (see the example)
 func Tar(ctx context.Context, body io.Reader, location string, rename Renamer) error {
-	files := []file{}
-	links := []link{}
-	symlinks := []link{}
-
-	// We make the first pass creating the directory structure, or we could end up
-	// attempting to create a file where there's no folder
-	tr := tar.NewReader(body)
-	for {
-		select {
-		case <-ctx.Done():
-			return errors.New("interrupted")
-		default:
-		}
-
-		header, err := tr.Next()
-		if err == io.EOF {
-			break
-		}
-
-		if err != nil {
-			return errors.Annotatef(err, "Read tar stream")
-		}
-
-		path := header.Name
-		if rename != nil {
-			path = rename(path)
-		}
-
-		if path == "" {
-			continue
-		}
-
-		path = filepath.Join(location, path)
-		info := header.FileInfo()
-
-		switch header.Typeflag {
-		case tar.TypeDir:
-			if err := os.MkdirAll(path, info.Mode()); err != nil {
-				return errors.Annotatef(err, "Create directory %s", path)
-			}
-		case tar.TypeReg, tar.TypeRegA:
-			var data bytes.Buffer
-			if _, err := copyCancel(ctx, &data, tr); err != nil {
-				return errors.Annotatef(err, "Read contents of file %s", path)
-			}
-			files = append(files, file{Path: path, Mode: info.Mode(), Data: data})
-		case tar.TypeLink:
-			name := header.Linkname
-			if rename != nil {
-				name = rename(name)
-			}
-
-			name = filepath.Join(location, name)
-			links = append(links, link{Path: path, Name: name})
-		case tar.TypeSymlink:
-			symlinks = append(symlinks, link{Path: path, Name: header.Linkname})
-		}
-	}
-
-	// Now we make another pass creating the files and links
-	for i := range files {
-		if err := copy(ctx, files[i].Path, files[i].Mode, &files[i].Data); err != nil {
-			return errors.Annotatef(err, "Create file %s", files[i].Path)
-		}
-	}
-
-	for i := range links {
-		select {
-		case <-ctx.Done():
-			return errors.New("interrupted")
-		default:
-		}
-		if err := os.Link(links[i].Name, links[i].Path); err != nil {
-			return errors.Annotatef(err, "Create link %s", links[i].Path)
-		}
+	extractor := Extractor{
+		FS: fs{},
 	}
 
-	for i := range symlinks {
-		select {
-		case <-ctx.Done():
-			return errors.New("interrupted")
-		default:
-		}
-		if err := os.Symlink(symlinks[i].Name, symlinks[i].Path); err != nil {
-			return errors.Annotatef(err, "Create link %s", symlinks[i].Path)
-		}
-	}
-	return nil
+	return extractor.Tar(ctx, body, location, rename)
 }
 
 // Zip extracts a .zip archived stream of data in the specified location.
 // It accepts a rename function to handle the names of the files (see the example).
 func Zip(ctx context.Context, body io.Reader, location string, rename Renamer) error {
-	// read the whole body into a buffer. Not sure this is the best way to do it
-	buffer := bytes.NewBuffer([]byte{})
-	copyCancel(ctx, buffer, body)
-
-	archive, err := zip.NewReader(bytes.NewReader(buffer.Bytes()), int64(buffer.Len()))
-	if err != nil {
-		return errors.Annotatef(err, "Read the zip file")
-	}
-
-	files := []file{}
-	links := []link{}
-
-	// We make the first pass creating the directory structure, or we could end up
-	// attempting to create a file where there's no folder
-	for _, header := range archive.File {
-		select {
-		case <-ctx.Done():
-			return errors.New("interrupted")
-		default:
-		}
-
-		path := header.Name
-
-		// Replace backslash with forward slash. There are archives in the wild made with
-		// buggy compressors that use backslash as path separator. The ZIP format explicitly
-		// denies the use of "\" so we just replace it with slash "/".
-		// Moreover it seems that folders are stored as "files" but with a final "\" in the
-		// filename... oh, well...
-		forceDir := strings.HasSuffix(path, "\\")
-		path = strings.Replace(path, "\\", "/", -1)
-
-		if rename != nil {
-			path = rename(path)
-		}
-
-		if path == "" {
-			continue
-		}
-
-		path = filepath.Join(location, path)
-		info := header.FileInfo()
-
-		switch {
-		case info.IsDir() || forceDir:
-			if err := os.MkdirAll(path, info.Mode()|os.ModeDir|100); err != nil {
-				return errors.Annotatef(err, "Create directory %s", path)
-			}
-		// We only check for symlinks because hard links aren't possible
-		case info.Mode()&os.ModeSymlink != 0:
-			f, err := header.Open()
-			if err != nil {
-				return errors.Annotatef(err, "Open link %s", path)
-			}
-			name, err := ioutil.ReadAll(f)
-			if err != nil {
-				return errors.Annotatef(err, "Read address of link %s", path)
-			}
-			links = append(links, link{Path: path, Name: string(name)})
-		default:
-			f, err := header.Open()
-			if err != nil {
-				return errors.Annotatef(err, "Open file %s", path)
-			}
-			var data bytes.Buffer
-			if _, err := copyCancel(ctx, &data, f); err != nil {
-				return errors.Annotatef(err, "Read contents of file %s", path)
-			}
-			files = append(files, file{Path: path, Mode: info.Mode(), Data: data})
-		}
+	extractor := Extractor{
+		FS: fs{},
 	}
 
-	// Now we make another pass creating the files and links
-	for i := range files {
-		if err := copy(ctx, files[i].Path, files[i].Mode, &files[i].Data); err != nil {
-			return errors.Annotatef(err, "Create file %s", files[i].Path)
-		}
-	}
+	return extractor.Zip(ctx, body, location, rename)
+}
 
-	for i := range links {
-		select {
-		case <-ctx.Done():
-			return errors.New("interrupted")
-		default:
-		}
-		if err := os.Symlink(links[i].Name, links[i].Path); err != nil {
-			return errors.Annotatef(err, "Create link %s", links[i].Path)
-		}
-	}
+type fs struct{}
 
-	return nil
+func (f fs) Link(oldname, newname string) error {
+	return os.Link(oldname, newname)
 }
 
-func copy(ctx context.Context, path string, mode os.FileMode, src io.Reader) error {
-	// We add the execution permission to be able to create files inside it
-	err := os.MkdirAll(filepath.Dir(path), mode|os.ModeDir|100)
-	if err != nil {
-		return err
-	}
-	file, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, mode)
-	if err != nil {
-		return err
-	}
-	defer file.Close()
-	_, err = copyCancel(ctx, file, src)
-	return err
+func (f fs) MkdirAll(path string, perm os.FileMode) error {
+	return os.MkdirAll(path, perm)
 }
 
-// match reads the first 512 bytes, calls types.Match and returns a reader
-// for the whole stream
-func match(r io.Reader) (io.Reader, types.Type, error) {
-	buffer := make([]byte, 512)
-
-	n, err := r.Read(buffer)
-	if err != nil && err != io.EOF {
-		return nil, types.Unknown, err
-	}
-
-	r = io.MultiReader(bytes.NewBuffer(buffer[:n]), r)
-
-	typ, err := filetype.Match(buffer)
+func (f fs) Symlink(oldname, newname string) error {
+	return os.Symlink(oldname, newname)
+}
 
-	return r, typ, err
+func (f fs) OpenFile(name string, flag int, perm os.FileMode) (*os.File, error) {
+	return os.OpenFile(name, flag, perm)
 }