...
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	
    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	}
    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	// 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.
    41	func (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.
    50	func (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.
    56	func (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.
   187	func (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.
   202	func (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	
   249	func 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.
   276	func (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	
   283	func (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	
   291	type 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	
   301	func (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	
   309	func (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	
   358	type countWriter struct {
   359		w     io.Writer
   360		count int64
   361	}
   362	
   363	func (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	
   369	type nopCloser struct {
   370		io.Writer
   371	}
   372	
   373	func (w nopCloser) Close() error {
   374		return nil
   375	}
   376	
   377	type writeBuf []byte
   378	
   379	func (b *writeBuf) uint16(v uint16) {
   380		binary.LittleEndian.PutUint16(*b, v)
   381		*b = (*b)[2:]
   382	}
   383	
   384	func (b *writeBuf) uint32(v uint32) {
   385		binary.LittleEndian.PutUint32(*b, v)
   386		*b = (*b)[4:]
   387	}
   388	
   389	func (b *writeBuf) uint64(v uint64) {
   390		binary.LittleEndian.PutUint64(*b, v)
   391		*b = (*b)[8:]
   392	}
   393	

View as plain text