Merge pull request #4 from jonsyu1/jyu/archive

Stream archive contents from Reader to Writer

archive/archive.go

@@ -0,0 +1,198 @@
+package archive
+
+import (
+	"archive/tar"
+	"archive/zip"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"time"
+
+	"github.com/jonsyu1/seelog/archive/gzip"
+)
+
+// Reader is the interface for reading files from an archive.
+type Reader interface {
+	NextFile() (name string, err error)
+	io.Reader
+}
+
+// ReadCloser is the interface that groups Reader with the Close method.
+type ReadCloser interface {
+	Reader
+	io.Closer
+}
+
+// Writer is the interface for writing files to an archived format.
+type Writer interface {
+	NextFile(name string, fi os.FileInfo) error
+	io.Writer
+}
+
+// WriteCloser is the interface that groups Writer with the Close method.
+type WriteCloser interface {
+	Writer
+	io.Closer
+}
+
+type nopCloser struct{ Reader }
+
+func (nopCloser) Close() error { return nil }
+
+// NopCloser returns a ReadCloser with a no-op Close method wrapping the
+// provided Reader r.
+func NopCloser(r Reader) ReadCloser {
+	return nopCloser{r}
+}
+
+// Copy copies from src to dest until either EOF is reached on src or an error
+// occurs.
+//
+// When the archive format of src matches that of dst, Copy streams the files
+// directly into dst. Otherwise, copy buffers the contents to disk to compute
+// headers before writing to dst.
+func Copy(dst Writer, src Reader) error {
+	switch src := src.(type) {
+	case tarReader:
+		if dst, ok := dst.(tarWriter); ok {
+			return copyTar(dst, src)
+		}
+	case zipReader:
+		if dst, ok := dst.(zipWriter); ok {
+			return copyZip(dst, src)
+		}
+	// Switch on concrete type because gzip has no special methods
+	case *gzip.Reader:
+		if dst, ok := dst.(*gzip.Writer); ok {
+			_, err := io.Copy(dst, src)
+			return err
+		}
+	}
+
+	return copyBuffer(dst, src)
+}
+
+func copyBuffer(dst Writer, src Reader) (err error) {
+	const defaultFileMode = 0666
+
+	buf, err := ioutil.TempFile("", "archive_copy_buffer")
+	if err != nil {
+		return err
+	}
+	defer os.Remove(buf.Name()) // Do not care about failure removing temp
+	defer buf.Close()           // Do not care about failure closing temp
+	for {
+		// Handle the next file
+		name, err := src.NextFile()
+		switch err {
+		case io.EOF: // Done copying
+			return nil
+		default: // Failed to write: bail out
+			return err
+		case nil: // Proceed below
+		}
+
+		// Buffer the file
+		if _, err := io.Copy(buf, src); err != nil {
+			return fmt.Errorf("buffer to disk: %v", err)
+		}
+
+		// Seek to the start of the file for full file copy
+		if _, err := buf.Seek(0, os.SEEK_SET); err != nil {
+			return err
+		}
+
+		// Set desired file permissions
+		if err := os.Chmod(buf.Name(), defaultFileMode); err != nil {
+			return err
+		}
+		fi, err := buf.Stat()
+		if err != nil {
+			return err
+		}
+
+		// Write the buffered file
+		if err := dst.NextFile(name, fi); err != nil {
+			return err
+		}
+		if _, err := io.Copy(dst, buf); err != nil {
+			return fmt.Errorf("copy to dst: %v", err)
+		}
+		if err := buf.Truncate(0); err != nil {
+			return err
+		}
+		if _, err := buf.Seek(0, os.SEEK_SET); err != nil {
+			return err
+		}
+	}
+}
+
+type tarReader interface {
+	Next() (*tar.Header, error)
+	io.Reader
+}
+
+type tarWriter interface {
+	WriteHeader(hdr *tar.Header) error
+	io.Writer
+}
+
+type zipReader interface {
+	Files() []*zip.File
+}
+
+type zipWriter interface {
+	CreateHeader(fh *zip.FileHeader) (io.Writer, error)
+}
+
+func copyTar(w tarWriter, r tarReader) error {
+	for {
+		hdr, err := r.Next()
+		switch err {
+		case io.EOF:
+			return nil
+		default: // Handle error
+			return err
+		case nil: // Proceed below
+		}
+
+		info := hdr.FileInfo()
+		// Skip directories
+		if info.IsDir() {
+			continue
+		}
+		if err := w.WriteHeader(hdr); err != nil {
+			return err
+		}
+		if _, err := io.Copy(w, r); err != nil {
+			return err
+		}
+	}
+}
+
+func copyZip(zw zipWriter, r zipReader) error {
+	for _, f := range r.Files() {
+		if err := copyZipFile(zw, f); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func copyZipFile(zw zipWriter, f *zip.File) error {
+	rc, err := f.Open()
+	if err != nil {
+		return err
+	}
+	defer rc.Close() // Read-only
+
+	hdr := f.FileHeader
+	hdr.SetModTime(time.Now())
+	w, err := zw.CreateHeader(&hdr)
+	if err != nil {
+		return err
+	}
+	_, err = io.Copy(w, rc)
+	return err
+}

archive/archive_test.go

@@ -0,0 +1,178 @@
+package archive_test
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"testing"
+
+	"github.com/jonsyu1/seelog/archive"
+	"github.com/jonsyu1/seelog/archive/gzip"
+	"github.com/jonsyu1/seelog/archive/tar"
+	"github.com/jonsyu1/seelog/archive/zip"
+	"github.com/jonsyu1/seelog/io/iotest"
+)
+
+const (
+	gzipType = "gzip"
+	tarType  = "tar"
+	zipType  = "zip"
+)
+
+var types = []string{gzipType, tarType, zipType}
+
+type file struct {
+	name     string
+	contents []byte
+}
+
+var (
+	oneFile = []file{
+		{
+			name:     "file1",
+			contents: []byte("This is a single log."),
+		},
+	}
+	twoFiles = []file{
+		{
+			name:     "file1",
+			contents: []byte("This is a log."),
+		},
+		{
+			name:     "file2",
+			contents: []byte("This is another log."),
+		},
+	}
+)
+
+type testCase struct {
+	srcType, dstType string
+	in               []file
+}
+
+func copyTests() map[string]testCase {
+	// types X types X files
+	tests := make(map[string]testCase, len(types)*len(types)*2)
+	for _, srct := range types {
+		for _, dstt := range types {
+			tests[fmt.Sprintf("%s to %s: one file", srct, dstt)] = testCase{
+				srcType: srct,
+				dstType: dstt,
+				in:      oneFile,
+			}
+			// gzip does not handle more than one file
+			if srct != gzipType && dstt != gzipType {
+				tests[fmt.Sprintf("%s to %s: two files", srct, dstt)] = testCase{
+					srcType: srct,
+					dstType: dstt,
+					in:      twoFiles,
+				}
+			}
+		}
+	}
+	return tests
+}
+
+func TestCopy(t *testing.T) {
+	srcb, dstb := new(bytes.Buffer), new(bytes.Buffer)
+	for tname, tt := range copyTests() {
+		// Reset buffers between tests
+		srcb.Reset()
+		dstb.Reset()
+
+		// Last file name (needed for gzip.NewReader)
+		var fname string
+
+		// Seed the src
+		srcw := writer(t, tname, srcb, tt.srcType)
+		for _, f := range tt.in {
+			srcw.NextFile(f.name, iotest.FileInfo(t, f.contents))
+			mustCopy(t, tname, srcw, bytes.NewReader(f.contents))
+			fname = f.name
+		}
+		mustClose(t, tname, srcw)
+
+		// Perform the copy
+		srcr := reader(t, tname, srcb, tt.srcType, fname)
+		dstw := writer(t, tname, dstb, tt.dstType)
+		if err := archive.Copy(dstw, srcr); err != nil {
+			t.Fatalf("%s: %v", tname, err)
+		}
+		srcr.Close() // Read-only
+		mustClose(t, tname, dstw)
+
+		// Read back dst to confirm our expectations
+		dstr := reader(t, tname, dstb, tt.dstType, fname)
+		for _, want := range tt.in {
+			buf := new(bytes.Buffer)
+			name, err := dstr.NextFile()
+			if err != nil {
+				t.Fatalf("%s: %v", tname, err)
+			}
+			mustCopy(t, tname, buf, dstr)
+			got := file{
+				name:     name,
+				contents: buf.Bytes(),
+			}
+
+			switch {
+			case got.name != want.name:
+				t.Errorf("%s: got file %q but want file %q",
+					tname, got.name, want.name)
+
+			case !bytes.Equal(got.contents, want.contents):
+				t.Errorf("%s: mismatched contents in %q: got %q but want %q",
+					tname, got.name, got.contents, want.contents)
+			}
+		}
+		dstr.Close()
+	}
+}
+
+func writer(t *testing.T, tname string, w io.Writer, atype string) archive.WriteCloser {
+	switch atype {
+	case gzipType:
+		return gzip.NewWriter(w)
+	case tarType:
+		return tar.NewWriter(w)
+	case zipType:
+		return zip.NewWriter(w)
+	}
+	t.Fatalf("%s: unrecognized archive type: %s", tname, atype)
+	panic("execution continued after (*testing.T).Fatalf")
+}
+
+func reader(t *testing.T, tname string, buf *bytes.Buffer, atype string, fname string) archive.ReadCloser {
+	switch atype {
+	case gzipType:
+		gr, err := gzip.NewReader(buf, fname)
+		if err != nil {
+			t.Fatalf("%s: %v", tname, err)
+		}
+		return gr
+	case tarType:
+		return archive.NopCloser(tar.NewReader(buf))
+	case zipType:
+		zr, err := zip.NewReader(
+			bytes.NewReader(buf.Bytes()),
+			int64(buf.Len()))
+		if err != nil {
+			t.Fatalf("%s: new zip reader: %v", tname, err)
+		}
+		return archive.NopCloser(zr)
+	}
+	t.Fatalf("%s: unrecognized archive type: %s", tname, atype)
+	panic("execution continued after (*testing.T).Fatalf")
+}
+
+func mustCopy(t *testing.T, tname string, dst io.Writer, src io.Reader) {
+	if _, err := io.Copy(dst, src); err != nil {
+		t.Fatalf("%s: copy: %v", tname, err)
+	}
+}
+
+func mustClose(t *testing.T, tname string, c io.Closer) {
+	if err := c.Close(); err != nil {
+		t.Fatalf("%s: close: %v", tname, err)
+	}
+}