libgo/go/compress/gzip/gzip.go

   1 // Copyright 2010 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 gzip
   6
   7 import (
   8         "compress/flate"
   9         "errors"
  10         "fmt"
  11         "hash"
  12         "hash/crc32"
  13         "io"
  14 )
  15
  16 // These constants are copied from the flate package, so that code that imports
  17 // "compress/gzip" does not also have to import "compress/flate".
  18 const (
  19         NoCompression      = flate.NoCompression
  20         BestSpeed          = flate.BestSpeed
  21         BestCompression    = flate.BestCompression
  22         DefaultCompression = flate.DefaultCompression
  23 )
  24
  25 // A Writer is an io.WriteCloser that satisfies writes by compressing data written
  26 // to its wrapped io.Writer.
  27 type Writer struct {
  28         Header
  29         w          io.Writer
  30         level      int
  31         compressor *flate.Writer
  32         digest     hash.Hash32
  33         size       uint32
  34         closed     bool
  35         buf        [10]byte
  36         err        error
  37 }
  38
  39 // NewWriter creates a new Writer that satisfies writes by compressing data
  40 // written to w.
  41 //
  42 // It is the caller's responsibility to call Close on the WriteCloser when done.
  43 // Writes may be buffered and not flushed until Close.
  44 //
  45 // Callers that wish to set the fields in Writer.Header must do so before
  46 // the first call to Write or Close. The Comment and Name header fields are
  47 // UTF-8 strings in Go, but the underlying format requires NUL-terminated ISO
  48 // 8859-1 (Latin-1). NUL or non-Latin-1 runes in those strings will lead to an
  49 // error on Write.
  50 func NewWriter(w io.Writer) *Writer {
  51         z, _ := NewWriterLevel(w, DefaultCompression)
  52         return z
  53 }
  54
  55 // NewWriterLevel is like NewWriter but specifies the compression level instead
  56 // of assuming DefaultCompression.
  57 //
  58 // The compression level can be DefaultCompression, NoCompression, or any
  59 // integer value between BestSpeed and BestCompression inclusive. The error
  60 // returned will be nil if the level is valid.
  61 func NewWriterLevel(w io.Writer, level int) (*Writer, error) {
  62         if level < DefaultCompression || level > BestCompression {
  63                 return nil, fmt.Errorf("gzip: invalid compression level: %d", level)
  64         }
  65         return &Writer{
  66                 Header: Header{
  67                         OS: 255, // unknown
  68                 },
  69                 w:      w,
  70                 level:  level,
  71                 digest: crc32.NewIEEE(),
  72         }, nil
  73 }
  74
  75 // GZIP (RFC 1952) is little-endian, unlike ZLIB (RFC 1950).
  76 func put2(p []byte, v uint16) {
  77         p[0] = uint8(v >> 0)
  78         p[1] = uint8(v >> 8)
  79 }
  80
  81 func put4(p []byte, v uint32) {
  82         p[0] = uint8(v >> 0)
  83         p[1] = uint8(v >> 8)
  84         p[2] = uint8(v >> 16)
  85         p[3] = uint8(v >> 24)
  86 }
  87
  88 // writeBytes writes a length-prefixed byte slice to z.w.
  89 func (z *Writer) writeBytes(b []byte) error {
  90         if len(b) > 0xffff {
  91                 return errors.New("gzip.Write: Extra data is too large")
  92         }
  93         put2(z.buf[0:2], uint16(len(b)))
  94         _, err := z.w.Write(z.buf[0:2])
  95         if err != nil {
  96                 return err
  97         }
  98         _, err = z.w.Write(b)
  99         return err
 100 }
 101
 102 // writeString writes a UTF-8 string s in GZIP's format to z.w.
 103 // GZIP (RFC 1952) specifies that strings are NUL-terminated ISO 8859-1 (Latin-1).
 104 func (z *Writer) writeString(s string) (err error) {
 105         // GZIP stores Latin-1 strings; error if non-Latin-1; convert if non-ASCII.
 106         needconv := false
 107         for _, v := range s {
 108                 if v == 0 || v > 0xff {
 109                         return errors.New("gzip.Write: non-Latin-1 header string")
 110                 }
 111                 if v > 0x7f {
 112                         needconv = true
 113                 }
 114         }
 115         if needconv {
 116                 b := make([]byte, 0, len(s))
 117                 for _, v := range s {
 118                         b = append(b, byte(v))
 119                 }
 120                 _, err = z.w.Write(b)
 121         } else {
 122                 _, err = io.WriteString(z.w, s)
 123         }
 124         if err != nil {
 125                 return err
 126         }
 127         // GZIP strings are NUL-terminated.
 128         z.buf[0] = 0
 129         _, err = z.w.Write(z.buf[0:1])
 130         return err
 131 }
 132
 133 // Write writes a compressed form of p to the underlying io.Writer. The
 134 // compressed bytes are not necessarily flushed until the Writer is closed.
 135 func (z *Writer) Write(p []byte) (int, error) {
 136         if z.err != nil {
 137                 return 0, z.err
 138         }
 139         var n int
 140         // Write the GZIP header lazily.
 141         if z.compressor == nil {
 142                 z.buf[0] = gzipID1
 143                 z.buf[1] = gzipID2
 144                 z.buf[2] = gzipDeflate
 145                 z.buf[3] = 0
 146                 if z.Extra != nil {
 147                         z.buf[3] |= 0x04
 148                 }
 149                 if z.Name != "" {
 150                         z.buf[3] |= 0x08
 151                 }
 152                 if z.Comment != "" {
 153                         z.buf[3] |= 0x10
 154                 }
 155                 put4(z.buf[4:8], uint32(z.ModTime.Unix()))
 156                 if z.level == BestCompression {
 157                         z.buf[8] = 2
 158                 } else if z.level == BestSpeed {
 159                         z.buf[8] = 4
 160                 } else {
 161                         z.buf[8] = 0
 162                 }
 163                 z.buf[9] = z.OS
 164                 n, z.err = z.w.Write(z.buf[0:10])
 165                 if z.err != nil {
 166                         return n, z.err
 167                 }
 168                 if z.Extra != nil {
 169                         z.err = z.writeBytes(z.Extra)
 170                         if z.err != nil {
 171                                 return n, z.err
 172                         }
 173                 }
 174                 if z.Name != "" {
 175                         z.err = z.writeString(z.Name)
 176                         if z.err != nil {
 177                                 return n, z.err
 178                         }
 179                 }
 180                 if z.Comment != "" {
 181                         z.err = z.writeString(z.Comment)
 182                         if z.err != nil {
 183                                 return n, z.err
 184                         }
 185                 }
 186                 z.compressor, _ = flate.NewWriter(z.w, z.level)
 187         }
 188         z.size += uint32(len(p))
 189         z.digest.Write(p)
 190         n, z.err = z.compressor.Write(p)
 191         return n, z.err
 192 }
 193
 194 // Flush flushes any pending compressed data to the underlying writer.
 195 //
 196 // It is useful mainly in compressed network protocols, to ensure that
 197 // a remote reader has enough data to reconstruct a packet. Flush does
 198 // not return until the data has been written. If the underlying
 199 // writer returns an error, Flush returns that error.
 200 //
 201 // In the terminology of the zlib library, Flush is equivalent to Z_SYNC_FLUSH.
 202 func (z *Writer) Flush() error {
 203         if z.err != nil {
 204                 return z.err
 205         }
 206         if z.closed {
 207                 return nil
 208         }
 209         if z.compressor == nil {
 210                 z.Write(nil)
 211         }
 212         z.err = z.compressor.Flush()
 213         return z.err
 214 }
 215
 216 // Close closes the Writer. It does not close the underlying io.Writer.
 217 func (z *Writer) Close() error {
 218         if z.err != nil {
 219                 return z.err
 220         }
 221         if z.closed {
 222                 return nil
 223         }
 224         z.closed = true
 225         if z.compressor == nil {
 226                 z.Write(nil)
 227                 if z.err != nil {
 228                         return z.err
 229                 }
 230         }
 231         z.err = z.compressor.Close()
 232         if z.err != nil {
 233                 return z.err
 234         }
 235         put4(z.buf[0:4], z.digest.Sum32())
 236         put4(z.buf[4:8], z.size)
 237         _, z.err = z.w.Write(z.buf[0:8])
 238         return z.err
 239 }