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 uint32
    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			b.uint32(h.CompressedSize)
    66			b.uint32(h.UncompressedSize)
    67			b.uint16(uint16(len(h.Name)))
    68			b.uint16(uint16(len(h.Extra)))
    69			b.uint16(uint16(len(h.Comment)))
    70			b = b[4:] // skip disk number start and internal file attr (2x uint16)
    71			b.uint32(h.ExternalAttrs)
    72			b.uint32(h.offset)
    73			if _, err := w.cw.Write(buf[:]); err != nil {
    74				return err
    75			}
    76			if _, err := io.WriteString(w.cw, h.Name); err != nil {
    77				return err
    78			}
    79			if _, err := w.cw.Write(h.Extra); err != nil {
    80				return err
    81			}
    82			if _, err := io.WriteString(w.cw, h.Comment); err != nil {
    83				return err
    84			}
    85		}
    86		end := w.cw.count
    87	
    88		// write end record
    89		var buf [directoryEndLen]byte
    90		b := writeBuf(buf[:])
    91		b.uint32(uint32(directoryEndSignature))
    92		b = b[4:]                     // skip over disk number and first disk number (2x uint16)
    93		b.uint16(uint16(len(w.dir)))  // number of entries this disk
    94		b.uint16(uint16(len(w.dir)))  // number of entries total
    95		b.uint32(uint32(end - start)) // size of directory
    96		b.uint32(uint32(start))       // start of directory
    97		// skipped size of comment (always zero)
    98		if _, err := w.cw.Write(buf[:]); err != nil {
    99			return err
   100		}
   101	
   102		return w.cw.w.(*bufio.Writer).Flush()
   103	}
   104	
   105	// Create adds a file to the zip file using the provided name.
   106	// It returns a Writer to which the file contents should be written.
   107	// The file's contents must be written to the io.Writer before the next
   108	// call to Create, CreateHeader, or Close.
   109	func (w *Writer) Create(name string) (io.Writer, error) {
   110		header := &FileHeader{
   111			Name:   name,
   112			Method: Deflate,
   113		}
   114		return w.CreateHeader(header)
   115	}
   116	
   117	// CreateHeader adds a file to the zip file using the provided FileHeader
   118	// for the file metadata. 
   119	// It returns a Writer to which the file contents should be written.
   120	// The file's contents must be written to the io.Writer before the next
   121	// call to Create, CreateHeader, or Close.
   122	func (w *Writer) CreateHeader(fh *FileHeader) (io.Writer, error) {
   123		if w.last != nil && !w.last.closed {
   124			if err := w.last.close(); err != nil {
   125				return nil, err
   126			}
   127		}
   128	
   129		fh.Flags |= 0x8 // we will write a data descriptor
   130		fh.CreatorVersion = fh.CreatorVersion&0xff00 | 0x14
   131		fh.ReaderVersion = 0x14
   132	
   133		fw := &fileWriter{
   134			zipw:      w.cw,
   135			compCount: &countWriter{w: w.cw},
   136			crc32:     crc32.NewIEEE(),
   137		}
   138		switch fh.Method {
   139		case Store:
   140			fw.comp = nopCloser{fw.compCount}
   141		case Deflate:
   142			var err error
   143			fw.comp, err = flate.NewWriter(fw.compCount, 5)
   144			if err != nil {
   145				return nil, err
   146			}
   147		default:
   148			return nil, ErrAlgorithm
   149		}
   150		fw.rawCount = &countWriter{w: fw.comp}
   151	
   152		h := &header{
   153			FileHeader: fh,
   154			offset:     uint32(w.cw.count),
   155		}
   156		w.dir = append(w.dir, h)
   157		fw.header = h
   158	
   159		if err := writeHeader(w.cw, fh); err != nil {
   160			return nil, err
   161		}
   162	
   163		w.last = fw
   164		return fw, nil
   165	}
   166	
   167	func writeHeader(w io.Writer, h *FileHeader) error {
   168		var buf [fileHeaderLen]byte
   169		b := writeBuf(buf[:])
   170		b.uint32(uint32(fileHeaderSignature))
   171		b.uint16(h.ReaderVersion)
   172		b.uint16(h.Flags)
   173		b.uint16(h.Method)
   174		b.uint16(h.ModifiedTime)
   175		b.uint16(h.ModifiedDate)
   176		b.uint32(h.CRC32)
   177		b.uint32(h.CompressedSize)
   178		b.uint32(h.UncompressedSize)
   179		b.uint16(uint16(len(h.Name)))
   180		b.uint16(uint16(len(h.Extra)))
   181		if _, err := w.Write(buf[:]); err != nil {
   182			return err
   183		}
   184		if _, err := io.WriteString(w, h.Name); err != nil {
   185			return err
   186		}
   187		_, err := w.Write(h.Extra)
   188		return err
   189	}
   190	
   191	type fileWriter struct {
   192		*header
   193		zipw      io.Writer
   194		rawCount  *countWriter
   195		comp      io.WriteCloser
   196		compCount *countWriter
   197		crc32     hash.Hash32
   198		closed    bool
   199	}
   200	
   201	func (w *fileWriter) Write(p []byte) (int, error) {
   202		if w.closed {
   203			return 0, errors.New("zip: write to closed file")
   204		}
   205		w.crc32.Write(p)
   206		return w.rawCount.Write(p)
   207	}
   208	
   209	func (w *fileWriter) close() error {
   210		if w.closed {
   211			return errors.New("zip: file closed twice")
   212		}
   213		w.closed = true
   214		if err := w.comp.Close(); err != nil {
   215			return err
   216		}
   217	
   218		// update FileHeader
   219		fh := w.header.FileHeader
   220		fh.CRC32 = w.crc32.Sum32()
   221		fh.CompressedSize = uint32(w.compCount.count)
   222		fh.UncompressedSize = uint32(w.rawCount.count)
   223	
   224		// write data descriptor
   225		var buf [dataDescriptorLen]byte
   226		b := writeBuf(buf[:])
   227		b.uint32(dataDescriptorSignature) // de-facto standard, required by OS X
   228		b.uint32(fh.CRC32)
   229		b.uint32(fh.CompressedSize)
   230		b.uint32(fh.UncompressedSize)
   231		_, err := w.zipw.Write(buf[:])
   232		return err
   233	}
   234	
   235	type countWriter struct {
   236		w     io.Writer
   237		count int64
   238	}
   239	
   240	func (w *countWriter) Write(p []byte) (int, error) {
   241		n, err := w.w.Write(p)
   242		w.count += int64(n)
   243		return n, err
   244	}
   245	
   246	type nopCloser struct {
   247		io.Writer
   248	}
   249	
   250	func (w nopCloser) Close() error {
   251		return nil
   252	}
   253	
   254	type writeBuf []byte
   255	
   256	func (b *writeBuf) uint16(v uint16) {
   257		binary.LittleEndian.PutUint16(*b, v)
   258		*b = (*b)[2:]
   259	}
   260	
   261	func (b *writeBuf) uint32(v uint32) {
   262		binary.LittleEndian.PutUint32(*b, v)
   263		*b = (*b)[4:]
   264	}