Run Format

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

View as plain text