blob: 3a9292e380705573fdd0249602cb5ab83177b42e [file] [log] [blame]
Dan Willemsen25a4e072016-08-05 16:34:03 -07001// 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 "bufio"
9 "encoding/binary"
10 "errors"
11 "hash"
12 "hash/crc32"
13 "io"
14)
15
16// TODO(adg): support zip file comments
17
18// Writer implements a zip file writer.
19type Writer struct {
20 cw *countWriter
21 dir []*header
22 last *fileWriter
23 closed bool
24 compressors map[uint16]Compressor
25}
26
27type header struct {
28 *FileHeader
29 offset uint64
30}
31
32// NewWriter returns a new Writer writing a zip file to w.
33func NewWriter(w io.Writer) *Writer {
34 return &Writer{cw: &countWriter{w: bufio.NewWriter(w)}}
35}
36
37// SetOffset sets the offset of the beginning of the zip data within the
38// underlying writer. It should be used when the zip data is appended to an
39// existing file, such as a binary executable.
40// It must be called before any data is written.
41func (w *Writer) SetOffset(n int64) {
42 if w.cw.count != 0 {
43 panic("zip: SetOffset called after data was written")
44 }
45 w.cw.count = n
46}
47
48// Flush flushes any buffered data to the underlying writer.
49// Calling Flush is not normally necessary; calling Close is sufficient.
50func (w *Writer) Flush() error {
51 return w.cw.w.(*bufio.Writer).Flush()
52}
53
54// Close finishes writing the zip file by writing the central directory.
55// It does not (and cannot) close the underlying writer.
56func (w *Writer) Close() error {
57 if w.last != nil && !w.last.closed {
58 if err := w.last.close(); err != nil {
59 return err
60 }
61 w.last = nil
62 }
63 if w.closed {
64 return errors.New("zip: writer closed twice")
65 }
66 w.closed = true
67
68 // write central directory
69 start := w.cw.count
70 for _, h := range w.dir {
71 var buf [directoryHeaderLen]byte
72 b := writeBuf(buf[:])
73 b.uint32(uint32(directoryHeaderSignature))
74 b.uint16(h.CreatorVersion)
75 b.uint16(h.ReaderVersion)
76 b.uint16(h.Flags)
77 b.uint16(h.Method)
78 b.uint16(h.ModifiedTime)
79 b.uint16(h.ModifiedDate)
80 b.uint32(h.CRC32)
81 if h.isZip64() || h.offset >= uint32max {
82 // the file needs a zip64 header. store maxint in both
83 // 32 bit size fields (and offset later) to signal that the
84 // zip64 extra header should be used.
85 b.uint32(uint32max) // compressed size
86 b.uint32(uint32max) // uncompressed size
87
88 // append a zip64 extra block to Extra
89 var buf [28]byte // 2x uint16 + 3x uint64
90 eb := writeBuf(buf[:])
91 eb.uint16(zip64ExtraId)
92 eb.uint16(24) // size = 3x uint64
93 eb.uint64(h.UncompressedSize64)
94 eb.uint64(h.CompressedSize64)
95 eb.uint64(h.offset)
96 h.Extra = append(h.Extra, buf[:]...)
97 } else {
98 b.uint32(h.CompressedSize)
99 b.uint32(h.UncompressedSize)
100 }
101 b.uint16(uint16(len(h.Name)))
102 b.uint16(uint16(len(h.Extra)))
103 b.uint16(uint16(len(h.Comment)))
104 b = b[4:] // skip disk number start and internal file attr (2x uint16)
105 b.uint32(h.ExternalAttrs)
106 if h.offset > uint32max {
107 b.uint32(uint32max)
108 } else {
109 b.uint32(uint32(h.offset))
110 }
111 if _, err := w.cw.Write(buf[:]); err != nil {
112 return err
113 }
114 if _, err := io.WriteString(w.cw, h.Name); err != nil {
115 return err
116 }
117 if _, err := w.cw.Write(h.Extra); err != nil {
118 return err
119 }
120 if _, err := io.WriteString(w.cw, h.Comment); err != nil {
121 return err
122 }
123 }
124 end := w.cw.count
125
126 records := uint64(len(w.dir))
127 size := uint64(end - start)
128 offset := uint64(start)
129
130 if records > uint16max || size > uint32max || offset > uint32max {
131 var buf [directory64EndLen + directory64LocLen]byte
132 b := writeBuf(buf[:])
133
134 // zip64 end of central directory record
135 b.uint32(directory64EndSignature)
136 b.uint64(directory64EndLen - 12) // length minus signature (uint32) and length fields (uint64)
137 b.uint16(zipVersion45) // version made by
138 b.uint16(zipVersion45) // version needed to extract
139 b.uint32(0) // number of this disk
140 b.uint32(0) // number of the disk with the start of the central directory
141 b.uint64(records) // total number of entries in the central directory on this disk
142 b.uint64(records) // total number of entries in the central directory
143 b.uint64(size) // size of the central directory
144 b.uint64(offset) // offset of start of central directory with respect to the starting disk number
145
146 // zip64 end of central directory locator
147 b.uint32(directory64LocSignature)
148 b.uint32(0) // number of the disk with the start of the zip64 end of central directory
149 b.uint64(uint64(end)) // relative offset of the zip64 end of central directory record
150 b.uint32(1) // total number of disks
151
152 if _, err := w.cw.Write(buf[:]); err != nil {
153 return err
154 }
155
156 // store max values in the regular end record to signal that
157 // that the zip64 values should be used instead
158 records = uint16max
159 size = uint32max
160 offset = uint32max
161 }
162
163 // write end record
164 var buf [directoryEndLen]byte
165 b := writeBuf(buf[:])
166 b.uint32(uint32(directoryEndSignature))
167 b = b[4:] // skip over disk number and first disk number (2x uint16)
168 b.uint16(uint16(records)) // number of entries this disk
169 b.uint16(uint16(records)) // number of entries total
170 b.uint32(uint32(size)) // size of directory
171 b.uint32(uint32(offset)) // start of directory
172 // skipped size of comment (always zero)
173 if _, err := w.cw.Write(buf[:]); err != nil {
174 return err
175 }
176
177 return w.cw.w.(*bufio.Writer).Flush()
178}
179
180// Create adds a file to the zip file using the provided name.
181// It returns a Writer to which the file contents should be written.
182// The name must be a relative path: it must not start with a drive
183// letter (e.g. C:) or leading slash, and only forward slashes are
184// allowed.
185// The file's contents must be written to the io.Writer before the next
186// call to Create, CreateHeader, or Close.
187func (w *Writer) Create(name string) (io.Writer, error) {
188 header := &FileHeader{
189 Name: name,
190 Method: Deflate,
191 }
192 return w.CreateHeader(header)
193}
194
195// CreateHeader adds a file to the zip file using the provided FileHeader
196// for the file metadata.
197// It returns a Writer to which the file contents should be written.
198//
199// The file's contents must be written to the io.Writer before the next
200// call to Create, CreateHeader, or Close. The provided FileHeader fh
201// must not be modified after a call to CreateHeader.
202func (w *Writer) CreateHeader(fh *FileHeader) (io.Writer, error) {
203 if w.last != nil && !w.last.closed {
204 if err := w.last.close(); err != nil {
205 return nil, err
206 }
207 }
208 if len(w.dir) > 0 && w.dir[len(w.dir)-1].FileHeader == fh {
209 // See https://golang.org/issue/11144 confusion.
210 return nil, errors.New("archive/zip: invalid duplicate FileHeader")
211 }
212
213 fh.Flags |= 0x8 // we will write a data descriptor
214
215 fh.CreatorVersion = fh.CreatorVersion&0xff00 | zipVersion20 // preserve compatibility byte
216 fh.ReaderVersion = zipVersion20
217
218 fw := &fileWriter{
219 zipw: w.cw,
220 compCount: &countWriter{w: w.cw},
221 crc32: crc32.NewIEEE(),
222 }
223 comp := w.compressor(fh.Method)
224 if comp == nil {
225 return nil, ErrAlgorithm
226 }
227 var err error
228 fw.comp, err = comp(fw.compCount)
229 if err != nil {
230 return nil, err
231 }
232 fw.rawCount = &countWriter{w: fw.comp}
233
234 h := &header{
235 FileHeader: fh,
236 offset: uint64(w.cw.count),
237 }
238 w.dir = append(w.dir, h)
239 fw.header = h
240
241 if err := writeHeader(w.cw, fh); err != nil {
242 return nil, err
243 }
244
245 w.last = fw
246 return fw, nil
247}
248
249func writeHeader(w io.Writer, h *FileHeader) error {
250 var buf [fileHeaderLen]byte
251 b := writeBuf(buf[:])
252 b.uint32(uint32(fileHeaderSignature))
253 b.uint16(h.ReaderVersion)
254 b.uint16(h.Flags)
255 b.uint16(h.Method)
256 b.uint16(h.ModifiedTime)
257 b.uint16(h.ModifiedDate)
258 b.uint32(0) // since we are writing a data descriptor crc32,
259 b.uint32(0) // compressed size,
260 b.uint32(0) // and uncompressed size should be zero
261 b.uint16(uint16(len(h.Name)))
262 b.uint16(uint16(len(h.Extra)))
263 if _, err := w.Write(buf[:]); err != nil {
264 return err
265 }
266 if _, err := io.WriteString(w, h.Name); err != nil {
267 return err
268 }
269 _, err := w.Write(h.Extra)
270 return err
271}
272
273// RegisterCompressor registers or overrides a custom compressor for a specific
274// method ID. If a compressor for a given method is not found, Writer will
275// default to looking up the compressor at the package level.
276func (w *Writer) RegisterCompressor(method uint16, comp Compressor) {
277 if w.compressors == nil {
278 w.compressors = make(map[uint16]Compressor)
279 }
280 w.compressors[method] = comp
281}
282
283func (w *Writer) compressor(method uint16) Compressor {
284 comp := w.compressors[method]
285 if comp == nil {
286 comp = compressor(method)
287 }
288 return comp
289}
290
291type fileWriter struct {
292 *header
293 zipw io.Writer
294 rawCount *countWriter
295 comp io.WriteCloser
296 compCount *countWriter
297 crc32 hash.Hash32
298 closed bool
299}
300
301func (w *fileWriter) Write(p []byte) (int, error) {
302 if w.closed {
303 return 0, errors.New("zip: write to closed file")
304 }
305 w.crc32.Write(p)
306 return w.rawCount.Write(p)
307}
308
309func (w *fileWriter) close() error {
310 if w.closed {
311 return errors.New("zip: file closed twice")
312 }
313 w.closed = true
314 if err := w.comp.Close(); err != nil {
315 return err
316 }
317
318 // update FileHeader
319 fh := w.header.FileHeader
320 fh.CRC32 = w.crc32.Sum32()
321 fh.CompressedSize64 = uint64(w.compCount.count)
322 fh.UncompressedSize64 = uint64(w.rawCount.count)
323
324 if fh.isZip64() {
325 fh.CompressedSize = uint32max
326 fh.UncompressedSize = uint32max
327 fh.ReaderVersion = zipVersion45 // requires 4.5 - File uses ZIP64 format extensions
328 } else {
329 fh.CompressedSize = uint32(fh.CompressedSize64)
330 fh.UncompressedSize = uint32(fh.UncompressedSize64)
331 }
332
333 // Write data descriptor. This is more complicated than one would
334 // think, see e.g. comments in zipfile.c:putextended() and
335 // http://bugs.sun.com/bugdatabase/view_bug.do?bug_id=7073588.
336 // The approach here is to write 8 byte sizes if needed without
337 // adding a zip64 extra in the local header (too late anyway).
338 var buf []byte
339 if fh.isZip64() {
340 buf = make([]byte, dataDescriptor64Len)
341 } else {
342 buf = make([]byte, dataDescriptorLen)
343 }
344 b := writeBuf(buf)
345 b.uint32(dataDescriptorSignature) // de-facto standard, required by OS X
346 b.uint32(fh.CRC32)
347 if fh.isZip64() {
348 b.uint64(fh.CompressedSize64)
349 b.uint64(fh.UncompressedSize64)
350 } else {
351 b.uint32(fh.CompressedSize)
352 b.uint32(fh.UncompressedSize)
353 }
354 _, err := w.zipw.Write(buf)
355 return err
356}
357
358type countWriter struct {
359 w io.Writer
360 count int64
361}
362
363func (w *countWriter) Write(p []byte) (int, error) {
364 n, err := w.w.Write(p)
365 w.count += int64(n)
366 return n, err
367}
368
369type nopCloser struct {
370 io.Writer
371}
372
373func (w nopCloser) Close() error {
374 return nil
375}
376
377type writeBuf []byte
378
379func (b *writeBuf) uint16(v uint16) {
380 binary.LittleEndian.PutUint16(*b, v)
381 *b = (*b)[2:]
382}
383
384func (b *writeBuf) uint32(v uint32) {
385 binary.LittleEndian.PutUint32(*b, v)
386 *b = (*b)[4:]
387}
388
389func (b *writeBuf) uint64(v uint64) {
390 binary.LittleEndian.PutUint64(*b, v)
391 *b = (*b)[8:]
392}