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