...

Source file src/compress/gzip/gzip.go

Documentation: compress/gzip

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

View as plain text