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

View as plain text