123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252 |
- // Copyright 2010 The Go Authors. All rights reserved.
- // Use of this source code is governed by a BSD-style
- // license that can be found in the LICENSE file.
- package gzip
- import (
- "errors"
- "fmt"
- "io"
- "github.com/klauspost/compress/flate"
- "github.com/klauspost/crc32"
- )
- // These constants are copied from the flate package, so that code that imports
- // "compress/gzip" does not also have to import "compress/flate".
- const (
- NoCompression = flate.NoCompression
- BestSpeed = flate.BestSpeed
- BestCompression = flate.BestCompression
- DefaultCompression = flate.DefaultCompression
- ConstantCompression = flate.ConstantCompression
- HuffmanOnly = flate.HuffmanOnly
- )
- // A Writer is an io.WriteCloser.
- // Writes to a Writer are compressed and written to w.
- type Writer struct {
- Header // written at first call to Write, Flush, or Close
- w io.Writer
- level int
- wroteHeader bool
- compressor *flate.Writer
- digest uint32 // CRC-32, IEEE polynomial (section 8)
- size uint32 // Uncompressed size (section 2.3.1)
- closed bool
- buf [10]byte
- err error
- }
- // NewWriter returns a new Writer.
- // Writes to the returned writer are compressed and written to w.
- //
- // It is the caller's responsibility to call Close on the WriteCloser when done.
- // Writes may be buffered and not flushed until Close.
- //
- // Callers that wish to set the fields in Writer.Header must do so before
- // the first call to Write, Flush, or Close.
- func NewWriter(w io.Writer) *Writer {
- z, _ := NewWriterLevel(w, DefaultCompression)
- return z
- }
- // NewWriterLevel is like NewWriter but specifies the compression level instead
- // of assuming DefaultCompression.
- //
- // The compression level can be DefaultCompression, NoCompression, or any
- // integer value between BestSpeed and BestCompression inclusive. The error
- // returned will be nil if the level is valid.
- func NewWriterLevel(w io.Writer, level int) (*Writer, error) {
- if level < HuffmanOnly || level > BestCompression {
- return nil, fmt.Errorf("gzip: invalid compression level: %d", level)
- }
- z := new(Writer)
- z.init(w, level)
- return z, nil
- }
- func (z *Writer) init(w io.Writer, level int) {
- compressor := z.compressor
- if compressor != nil {
- compressor.Reset(w)
- }
- *z = Writer{
- Header: Header{
- OS: 255, // unknown
- },
- w: w,
- level: level,
- compressor: compressor,
- }
- }
- // Reset discards the Writer z's state and makes it equivalent to the
- // result of its original state from NewWriter or NewWriterLevel, but
- // writing to w instead. This permits reusing a Writer rather than
- // allocating a new one.
- func (z *Writer) Reset(w io.Writer) {
- z.init(w, z.level)
- }
- // writeBytes writes a length-prefixed byte slice to z.w.
- func (z *Writer) writeBytes(b []byte) error {
- if len(b) > 0xffff {
- return errors.New("gzip.Write: Extra data is too large")
- }
- le.PutUint16(z.buf[:2], uint16(len(b)))
- _, err := z.w.Write(z.buf[:2])
- if err != nil {
- return err
- }
- _, err = z.w.Write(b)
- return err
- }
- // writeString writes a UTF-8 string s in GZIP's format to z.w.
- // GZIP (RFC 1952) specifies that strings are NUL-terminated ISO 8859-1 (Latin-1).
- func (z *Writer) writeString(s string) (err error) {
- // GZIP stores Latin-1 strings; error if non-Latin-1; convert if non-ASCII.
- needconv := false
- for _, v := range s {
- if v == 0 || v > 0xff {
- return errors.New("gzip.Write: non-Latin-1 header string")
- }
- if v > 0x7f {
- needconv = true
- }
- }
- if needconv {
- b := make([]byte, 0, len(s))
- for _, v := range s {
- b = append(b, byte(v))
- }
- _, err = z.w.Write(b)
- } else {
- _, err = io.WriteString(z.w, s)
- }
- if err != nil {
- return err
- }
- // GZIP strings are NUL-terminated.
- z.buf[0] = 0
- _, err = z.w.Write(z.buf[:1])
- return err
- }
- // Write writes a compressed form of p to the underlying io.Writer. The
- // compressed bytes are not necessarily flushed until the Writer is closed.
- func (z *Writer) Write(p []byte) (int, error) {
- if z.err != nil {
- return 0, z.err
- }
- var n int
- // Write the GZIP header lazily.
- if !z.wroteHeader {
- z.wroteHeader = true
- z.buf[0] = gzipID1
- z.buf[1] = gzipID2
- z.buf[2] = gzipDeflate
- z.buf[3] = 0
- if z.Extra != nil {
- z.buf[3] |= 0x04
- }
- if z.Name != "" {
- z.buf[3] |= 0x08
- }
- if z.Comment != "" {
- z.buf[3] |= 0x10
- }
- le.PutUint32(z.buf[4:8], uint32(z.ModTime.Unix()))
- if z.level == BestCompression {
- z.buf[8] = 2
- } else if z.level == BestSpeed {
- z.buf[8] = 4
- } else {
- z.buf[8] = 0
- }
- z.buf[9] = z.OS
- n, z.err = z.w.Write(z.buf[:10])
- if z.err != nil {
- return n, z.err
- }
- if z.Extra != nil {
- z.err = z.writeBytes(z.Extra)
- if z.err != nil {
- return n, z.err
- }
- }
- if z.Name != "" {
- z.err = z.writeString(z.Name)
- if z.err != nil {
- return n, z.err
- }
- }
- if z.Comment != "" {
- z.err = z.writeString(z.Comment)
- if z.err != nil {
- return n, z.err
- }
- }
- if z.compressor == nil {
- z.compressor, _ = flate.NewWriter(z.w, z.level)
- }
- }
- z.size += uint32(len(p))
- z.digest = crc32.Update(z.digest, crc32.IEEETable, p)
- n, z.err = z.compressor.Write(p)
- return n, z.err
- }
- // Flush flushes any pending compressed data to the underlying writer.
- //
- // It is useful mainly in compressed network protocols, to ensure that
- // a remote reader has enough data to reconstruct a packet. Flush does
- // not return until the data has been written. If the underlying
- // writer returns an error, Flush returns that error.
- //
- // In the terminology of the zlib library, Flush is equivalent to Z_SYNC_FLUSH.
- func (z *Writer) Flush() error {
- if z.err != nil {
- return z.err
- }
- if z.closed {
- return nil
- }
- if !z.wroteHeader {
- z.Write(nil)
- if z.err != nil {
- return z.err
- }
- }
- z.err = z.compressor.Flush()
- return z.err
- }
- // Close closes the Writer, flushing any unwritten data to the underlying
- // io.Writer, but does not close the underlying io.Writer.
- func (z *Writer) Close() error {
- if z.err != nil {
- return z.err
- }
- if z.closed {
- return nil
- }
- z.closed = true
- if !z.wroteHeader {
- z.Write(nil)
- if z.err != nil {
- return z.err
- }
- }
- z.err = z.compressor.Close()
- if z.err != nil {
- return z.err
- }
- le.PutUint32(z.buf[:4], z.digest)
- le.PutUint32(z.buf[4:8], z.size)
- _, z.err = z.w.Write(z.buf[:8])
- return z.err
- }
|