...
Run Format

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

View as plain text