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