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