1package archive
2
3import (
4	"archive/tar"
5	"archive/zip"
6	"fmt"
7	"io"
8	"io/ioutil"
9	"os"
10	"time"
11
12	"github.com/cihub/seelog/archive/gzip"
13)
14
15// Reader is the interface for reading files from an archive.
16type Reader interface {
17	NextFile() (name string, err error)
18	io.Reader
19}
20
21// ReadCloser is the interface that groups Reader with the Close method.
22type ReadCloser interface {
23	Reader
24	io.Closer
25}
26
27// Writer is the interface for writing files to an archived format.
28type Writer interface {
29	NextFile(name string, fi os.FileInfo) error
30	io.Writer
31}
32
33// WriteCloser is the interface that groups Writer with the Close method.
34type WriteCloser interface {
35	Writer
36	io.Closer
37}
38
39type nopCloser struct{ Reader }
40
41func (nopCloser) Close() error { return nil }
42
43// NopCloser returns a ReadCloser with a no-op Close method wrapping the
44// provided Reader r.
45func NopCloser(r Reader) ReadCloser {
46	return nopCloser{r}
47}
48
49// Copy copies from src to dest until either EOF is reached on src or an error
50// occurs.
51//
52// When the archive format of src matches that of dst, Copy streams the files
53// directly into dst. Otherwise, copy buffers the contents to disk to compute
54// headers before writing to dst.
55func Copy(dst Writer, src Reader) error {
56	switch src := src.(type) {
57	case tarReader:
58		if dst, ok := dst.(tarWriter); ok {
59			return copyTar(dst, src)
60		}
61	case zipReader:
62		if dst, ok := dst.(zipWriter); ok {
63			return copyZip(dst, src)
64		}
65	// Switch on concrete type because gzip has no special methods
66	case *gzip.Reader:
67		if dst, ok := dst.(*gzip.Writer); ok {
68			_, err := io.Copy(dst, src)
69			return err
70		}
71	}
72
73	return copyBuffer(dst, src)
74}
75
76func copyBuffer(dst Writer, src Reader) (err error) {
77	const defaultFileMode = 0666
78
79	buf, err := ioutil.TempFile("", "archive_copy_buffer")
80	if err != nil {
81		return err
82	}
83	defer os.Remove(buf.Name()) // Do not care about failure removing temp
84	defer buf.Close()           // Do not care about failure closing temp
85	for {
86		// Handle the next file
87		name, err := src.NextFile()
88		switch err {
89		case io.EOF: // Done copying
90			return nil
91		default: // Failed to write: bail out
92			return err
93		case nil: // Proceed below
94		}
95
96		// Buffer the file
97		if _, err := io.Copy(buf, src); err != nil {
98			return fmt.Errorf("buffer to disk: %v", err)
99		}
100
101		// Seek to the start of the file for full file copy
102		if _, err := buf.Seek(0, os.SEEK_SET); err != nil {
103			return err
104		}
105
106		// Set desired file permissions
107		if err := os.Chmod(buf.Name(), defaultFileMode); err != nil {
108			return err
109		}
110		fi, err := buf.Stat()
111		if err != nil {
112			return err
113		}
114
115		// Write the buffered file
116		if err := dst.NextFile(name, fi); err != nil {
117			return err
118		}
119		if _, err := io.Copy(dst, buf); err != nil {
120			return fmt.Errorf("copy to dst: %v", err)
121		}
122		if err := buf.Truncate(0); err != nil {
123			return err
124		}
125		if _, err := buf.Seek(0, os.SEEK_SET); err != nil {
126			return err
127		}
128	}
129}
130
131type tarReader interface {
132	Next() (*tar.Header, error)
133	io.Reader
134}
135
136type tarWriter interface {
137	WriteHeader(hdr *tar.Header) error
138	io.Writer
139}
140
141type zipReader interface {
142	Files() []*zip.File
143}
144
145type zipWriter interface {
146	CreateHeader(fh *zip.FileHeader) (io.Writer, error)
147}
148
149func copyTar(w tarWriter, r tarReader) error {
150	for {
151		hdr, err := r.Next()
152		switch err {
153		case io.EOF:
154			return nil
155		default: // Handle error
156			return err
157		case nil: // Proceed below
158		}
159
160		info := hdr.FileInfo()
161		// Skip directories
162		if info.IsDir() {
163			continue
164		}
165		if err := w.WriteHeader(hdr); err != nil {
166			return err
167		}
168		if _, err := io.Copy(w, r); err != nil {
169			return err
170		}
171	}
172}
173
174func copyZip(zw zipWriter, r zipReader) error {
175	for _, f := range r.Files() {
176		if err := copyZipFile(zw, f); err != nil {
177			return err
178		}
179	}
180	return nil
181}
182
183func copyZipFile(zw zipWriter, f *zip.File) error {
184	rc, err := f.Open()
185	if err != nil {
186		return err
187	}
188	defer rc.Close() // Read-only
189
190	hdr := f.FileHeader
191	hdr.SetModTime(time.Now())
192	w, err := zw.CreateHeader(&hdr)
193	if err != nil {
194		return err
195	}
196	_, err = io.Copy(w, rc)
197	return err
198}
199