1// Copyright 2011 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package zip
6
7import (
8	"bytes"
9	"io"
10	"io/ioutil"
11	"math/rand"
12	"os"
13	"testing"
14)
15
16// TODO(adg): a more sophisticated test suite
17
18type WriteTest struct {
19	Name   string
20	Data   []byte
21	Method uint16
22	Mode   os.FileMode
23}
24
25var writeTests = []WriteTest{
26	{
27		Name:   "foo",
28		Data:   []byte("Rabbits, guinea pigs, gophers, marsupial rats, and quolls."),
29		Method: Store,
30		Mode:   0666,
31	},
32	{
33		Name:   "bar",
34		Data:   nil, // large data set in the test
35		Method: Deflate,
36		Mode:   0644,
37	},
38	{
39		Name:   "setuid",
40		Data:   []byte("setuid file"),
41		Method: Deflate,
42		Mode:   0755 | os.ModeSetuid,
43	},
44	{
45		Name:   "setgid",
46		Data:   []byte("setgid file"),
47		Method: Deflate,
48		Mode:   0755 | os.ModeSetgid,
49	},
50	{
51		Name:   "symlink",
52		Data:   []byte("../link/target"),
53		Method: Deflate,
54		Mode:   0755 | os.ModeSymlink,
55	},
56}
57
58func TestWriter(t *testing.T) {
59	largeData := make([]byte, 1<<17)
60	for i := range largeData {
61		largeData[i] = byte(rand.Int())
62	}
63	writeTests[1].Data = largeData
64	defer func() {
65		writeTests[1].Data = nil
66	}()
67
68	// write a zip file
69	buf := new(bytes.Buffer)
70	w := NewWriter(buf)
71
72	for _, wt := range writeTests {
73		testCreate(t, w, &wt)
74	}
75
76	if err := w.Close(); err != nil {
77		t.Fatal(err)
78	}
79
80	// read it back
81	r, err := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
82	if err != nil {
83		t.Fatal(err)
84	}
85	for i, wt := range writeTests {
86		testReadFile(t, r.File[i], &wt)
87	}
88}
89
90func TestWriterOffset(t *testing.T) {
91	largeData := make([]byte, 1<<17)
92	for i := range largeData {
93		largeData[i] = byte(rand.Int())
94	}
95	writeTests[1].Data = largeData
96	defer func() {
97		writeTests[1].Data = nil
98	}()
99
100	// write a zip file
101	buf := new(bytes.Buffer)
102	existingData := []byte{1, 2, 3, 1, 2, 3, 1, 2, 3}
103	n, _ := buf.Write(existingData)
104	w := NewWriter(buf)
105	w.SetOffset(int64(n))
106
107	for _, wt := range writeTests {
108		testCreate(t, w, &wt)
109	}
110
111	if err := w.Close(); err != nil {
112		t.Fatal(err)
113	}
114
115	// read it back
116	r, err := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
117	if err != nil {
118		t.Fatal(err)
119	}
120	for i, wt := range writeTests {
121		testReadFile(t, r.File[i], &wt)
122	}
123}
124
125func TestWriterFlush(t *testing.T) {
126	var buf bytes.Buffer
127	w := NewWriter(struct{ io.Writer }{&buf})
128	_, err := w.Create("foo")
129	if err != nil {
130		t.Fatal(err)
131	}
132	if buf.Len() > 0 {
133		t.Fatalf("Unexpected %d bytes already in buffer", buf.Len())
134	}
135	if err := w.Flush(); err != nil {
136		t.Fatal(err)
137	}
138	if buf.Len() == 0 {
139		t.Fatal("No bytes written after Flush")
140	}
141}
142
143func testCreate(t *testing.T, w *Writer, wt *WriteTest) {
144	header := &FileHeader{
145		Name:   wt.Name,
146		Method: wt.Method,
147	}
148	if wt.Mode != 0 {
149		header.SetMode(wt.Mode)
150	}
151	f, err := w.CreateHeader(header)
152	if err != nil {
153		t.Fatal(err)
154	}
155	_, err = f.Write(wt.Data)
156	if err != nil {
157		t.Fatal(err)
158	}
159}
160
161func testReadFile(t *testing.T, f *File, wt *WriteTest) {
162	if f.Name != wt.Name {
163		t.Fatalf("File name: got %q, want %q", f.Name, wt.Name)
164	}
165	testFileMode(t, wt.Name, f, wt.Mode)
166	rc, err := f.Open()
167	if err != nil {
168		t.Fatal("opening:", err)
169	}
170	b, err := ioutil.ReadAll(rc)
171	if err != nil {
172		t.Fatal("reading:", err)
173	}
174	err = rc.Close()
175	if err != nil {
176		t.Fatal("closing:", err)
177	}
178	if !bytes.Equal(b, wt.Data) {
179		t.Errorf("File contents %q, want %q", b, wt.Data)
180	}
181}
182
183func BenchmarkCompressedZipGarbage(b *testing.B) {
184	b.ReportAllocs()
185	var buf bytes.Buffer
186	bigBuf := bytes.Repeat([]byte("a"), 1<<20)
187	for i := 0; i <= b.N; i++ {
188		buf.Reset()
189		zw := NewWriter(&buf)
190		for j := 0; j < 3; j++ {
191			w, _ := zw.CreateHeader(&FileHeader{
192				Name:   "foo",
193				Method: Deflate,
194			})
195			w.Write(bigBuf)
196		}
197		zw.Close()
198		if i == 0 {
199			// Reset the timer after the first time through.
200			// This effectively discards the very large initial flate setup cost,
201			// as well as the initialization of bigBuf.
202			b.ResetTimer()
203		}
204	}
205}
206