← Back to branch summary

~nskaggs/+junk/xenial-test

« back to all changes in this revision

Viewing changes to src/github.com/juju/zip/writer.go

  • Committer: Nicholas Skaggs
  • Date: 2016-10-24 20:56:05 UTC
  • Revision ID: nicholas.skaggs@canonical.com-20161024205605-z8lta0uvuhtxwzwl
Initi with beta15

Show diffs side-by-side

added added

removed removed

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