writer.go

     1  // Copyright 2009 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 zlib
     6  
     7  import (
     8  	"compress/flate"
     9  	"encoding/binary"
    10  	"fmt"
    11  	"hash"
    12  	"hash/adler32"
    13  	"io"
    14  )
    15  
    16  // These constants are copied from the [flate] package, so that code that imports
    17  // [compress/zlib] 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 takes data written to it and writes the compressed
    27  // form of that data to an underlying writer (see [NewWriter]).
    28  type Writer struct {
    29  	w           io.Writer
    30  	level       int
    31  	dict        []byte
    32  	compressor  *flate.Writer
    33  	digest      hash.Hash32
    34  	err         error
    35  	scratch     [4]byte
    36  	wroteHeader bool
    37  }
    38  
    39  // NewWriter creates a new [Writer].
    40  // Writes to the returned Writer are compressed and written to w.
    41  //
    42  // It is the caller's responsibility to call Close on the Writer when done.
    43  // Writes may be buffered and not flushed until Close.
    44  //
    45  // Note that the exact bytes written to w are not covered by the Go 1
    46  // compatibility promise. Callers, including tests, should not depend on the
    47  // exact written bytes.
    48  func NewWriter(w io.Writer) *Writer {
    49  	z, _ := NewWriterLevelDict(w, DefaultCompression, nil)
    50  	return z
    51  }
    52  
    53  // NewWriterLevel is like [NewWriter] but specifies the compression level instead
    54  // of assuming [DefaultCompression].
    55  //
    56  // The compression level can be [DefaultCompression], [NoCompression], [HuffmanOnly]
    57  // or any integer value between [BestSpeed] and [BestCompression] inclusive.
    58  // The error returned will be nil if the level is valid.
    59  //
    60  // Note that the exact bytes written to w are not covered by the Go 1
    61  // compatibility promise. Callers, including tests, should not depend on the
    62  // exact written bytes.
    63  func NewWriterLevel(w io.Writer, level int) (*Writer, error) {
    64  	return NewWriterLevelDict(w, level, nil)
    65  }
    66  
    67  // NewWriterLevelDict is like [NewWriterLevel] but specifies a dictionary to
    68  // compress with.
    69  //
    70  // The dictionary may be nil. If not, its contents should not be modified until
    71  // the Writer is closed.
    72  //
    73  // Note that the exact bytes written to w are not covered by the Go 1
    74  // compatibility promise. Callers, including tests, should not depend on the
    75  // exact written bytes.
    76  func NewWriterLevelDict(w io.Writer, level int, dict []byte) (*Writer, error) {
    77  	if level < HuffmanOnly || level > BestCompression {
    78  		return nil, fmt.Errorf("zlib: invalid compression level: %d", level)
    79  	}
    80  	return &Writer{
    81  		w:     w,
    82  		level: level,
    83  		dict:  dict,
    84  	}, nil
    85  }
    86  
    87  // Reset clears the state of the [Writer] z such that it is equivalent to its
    88  // initial state from [NewWriterLevel] or [NewWriterLevelDict], but instead writing
    89  // to w.
    90  func (z *Writer) Reset(w io.Writer) {
    91  	z.w = w
    92  	// z.level and z.dict left unchanged.
    93  	if z.compressor != nil {
    94  		z.compressor.Reset(w)
    95  	}
    96  	if z.digest != nil {
    97  		z.digest.Reset()
    98  	}
    99  	z.err = nil
   100  	z.scratch = [4]byte{}
   101  	z.wroteHeader = false
   102  }
   103  
   104  // writeHeader writes the ZLIB header.
   105  func (z *Writer) writeHeader() (err error) {
   106  	z.wroteHeader = true
   107  	// ZLIB has a two-byte header (as documented in RFC 1950).
   108  	// The first four bits is the CINFO (compression info), which is 7 for the default deflate window size.
   109  	// The next four bits is the CM (compression method), which is 8 for deflate.
   110  	z.scratch[0] = 0x78
   111  	// The next two bits is the FLEVEL (compression level). The four values are:
   112  	// 0=fastest, 1=fast, 2=default, 3=best.
   113  	// The next bit, FDICT, is set if a dictionary is given.
   114  	// The final five FCHECK bits form a mod-31 checksum.
   115  	switch z.level {
   116  	case -2, 0, 1:
   117  		z.scratch[1] = 0 << 6
   118  	case 2, 3, 4, 5:
   119  		z.scratch[1] = 1 << 6
   120  	case 6, -1:
   121  		z.scratch[1] = 2 << 6
   122  	case 7, 8, 9:
   123  		z.scratch[1] = 3 << 6
   124  	default:
   125  		panic("unreachable")
   126  	}
   127  	if z.dict != nil {
   128  		z.scratch[1] |= 1 << 5
   129  	}
   130  	z.scratch[1] += uint8(31 - binary.BigEndian.Uint16(z.scratch[:2])%31)
   131  	if _, err = z.w.Write(z.scratch[0:2]); err != nil {
   132  		return err
   133  	}
   134  	if z.dict != nil {
   135  		// The next four bytes are the Adler-32 checksum of the dictionary.
   136  		binary.BigEndian.PutUint32(z.scratch[:], adler32.Checksum(z.dict))
   137  		if _, err = z.w.Write(z.scratch[0:4]); err != nil {
   138  			return err
   139  		}
   140  	}
   141  	if z.compressor == nil {
   142  		// Initialize deflater unless the Writer is being reused
   143  		// after a Reset call.
   144  		z.compressor, err = flate.NewWriterDict(z.w, z.level, z.dict)
   145  		if err != nil {
   146  			return err
   147  		}
   148  		z.digest = adler32.New()
   149  	}
   150  	return nil
   151  }
   152  
   153  // Write writes a compressed form of p to the underlying [io.Writer]. The
   154  // compressed bytes are not necessarily flushed until the [Writer] is closed or
   155  // explicitly flushed.
   156  func (z *Writer) Write(p []byte) (n int, err error) {
   157  	if !z.wroteHeader {
   158  		z.err = z.writeHeader()
   159  	}
   160  	if z.err != nil {
   161  		return 0, z.err
   162  	}
   163  	if len(p) == 0 {
   164  		return 0, nil
   165  	}
   166  	n, err = z.compressor.Write(p)
   167  	if err != nil {
   168  		z.err = err
   169  		return
   170  	}
   171  	z.digest.Write(p)
   172  	return
   173  }
   174  
   175  // Flush flushes the Writer to its underlying [io.Writer].
   176  func (z *Writer) Flush() error {
   177  	if !z.wroteHeader {
   178  		z.err = z.writeHeader()
   179  	}
   180  	if z.err != nil {
   181  		return z.err
   182  	}
   183  	z.err = z.compressor.Flush()
   184  	return z.err
   185  }
   186  
   187  // Close closes the Writer, flushing any unwritten data to the underlying
   188  // [io.Writer], but does not close the underlying io.Writer.
   189  func (z *Writer) Close() error {
   190  	if !z.wroteHeader {
   191  		z.err = z.writeHeader()
   192  	}
   193  	if z.err != nil {
   194  		return z.err
   195  	}
   196  	z.err = z.compressor.Close()
   197  	if z.err != nil {
   198  		return z.err
   199  	}
   200  	checksum := z.digest.Sum32()
   201  	// ZLIB (RFC 1950) is big-endian, unlike GZIP (RFC 1952).
   202  	binary.BigEndian.PutUint32(z.scratch[:], checksum)
   203  	_, z.err = z.w.Write(z.scratch[0:4])
   204  	return z.err
   205  }
   206
View as plain text