github.com/simonmittag/ws@v1.1.0-rc.5.0.20210419231947-82b846128245/wsflate/writer.go (about)

     1  package wsflate
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  )
     7  
     8  var (
     9  	compressionTail = [4]byte{
    10  		0, 0, 0xff, 0xff,
    11  	}
    12  	compressionReadTail = [9]byte{
    13  		0, 0, 0xff, 0xff,
    14  		1,
    15  		0, 0, 0xff, 0xff,
    16  	}
    17  )
    18  
    19  // Compressor is an interface holding deflate compression implementation.
    20  type Compressor interface {
    21  	io.Writer
    22  	Flush() error
    23  }
    24  
    25  // WriteResetter is an optional interface that Compressor can implement.
    26  type WriteResetter interface {
    27  	Reset(io.Writer)
    28  }
    29  
    30  // Writer implements compression for an io.Writer object using Compressor.
    31  // Essentially Writer is a thin wrapper around Compressor interface to meet
    32  // PMCE specs.
    33  //
    34  // After all data has been written client should call Flush() method.
    35  // If any error occurs after writing to or flushing a Writer, all subsequent
    36  // calls to Write(), Flush() or Close() will return the error.
    37  //
    38  // Writer might be reused for different io.Writer objects after its Reset()
    39  // method has been called.
    40  type Writer struct {
    41  	// NOTE: Writer uses compressor constructor function instead of field to
    42  	// reach these goals:
    43  	// 	1. To shrink Compressor interface and make it easier to be implemented.
    44  	//	2. If used as a field (and argument to the NewWriter()), Compressor object
    45  	//	will probably be initialized twice - first time to pass into Writer, and
    46  	//	second time during Writer initialization (which does Reset() internally).
    47  	// 	3. To get rid of wrappers if Reset() would be a part of	Compressor.
    48  	// 	E.g. non conformant implementations would have to provide it somehow,
    49  	// 	probably making a wrapper with the same constructor function.
    50  	// 	4. To make Reader and Writer API the same. That is, there is no Reset()
    51  	// 	method for flate.Reader already, so we need to provide it as a wrapper
    52  	// 	(see point #3), or drop the Reader.Reset() method.
    53  	dest io.Writer
    54  	ctor func(io.Writer) Compressor
    55  	c    Compressor
    56  	cbuf cbuf
    57  	err  error
    58  }
    59  
    60  // NewWriter returns a new Writer.
    61  func NewWriter(w io.Writer, ctor func(io.Writer) Compressor) *Writer {
    62  	// NOTE: NewWriter() is chosen against structure with exported fields here
    63  	// due its Reset() method, which in case of structure, would change
    64  	// exported field.
    65  	ret := &Writer{
    66  		dest: w,
    67  		ctor: ctor,
    68  	}
    69  	ret.Reset(w)
    70  	return ret
    71  }
    72  
    73  // Reset resets Writer to compress data into dest.
    74  // Any not flushed data will be lost.
    75  func (w *Writer) Reset(dest io.Writer) {
    76  	w.err = nil
    77  	w.cbuf.reset(dest)
    78  	if x, ok := w.c.(WriteResetter); ok {
    79  		x.Reset(&w.cbuf)
    80  	} else {
    81  		w.c = w.ctor(&w.cbuf)
    82  	}
    83  }
    84  
    85  // Write implements io.Writer.
    86  func (w *Writer) Write(p []byte) (n int, err error) {
    87  	if w.err != nil {
    88  		return 0, w.err
    89  	}
    90  	n, w.err = w.c.Write(p)
    91  	return n, w.err
    92  }
    93  
    94  // Flush writes any pending data into w.Dest.
    95  func (w *Writer) Flush() error {
    96  	if w.err != nil {
    97  		return w.err
    98  	}
    99  	w.err = w.c.Flush()
   100  	w.checkTail()
   101  	return w.err
   102  }
   103  
   104  // Close closes Writer and a Compressor instance used under the hood (if it
   105  // implements io.Closer interface).
   106  func (w *Writer) Close() error {
   107  	if w.err != nil {
   108  		return w.err
   109  	}
   110  	if c, ok := w.c.(io.Closer); ok {
   111  		w.err = c.Close()
   112  	}
   113  	w.checkTail()
   114  	return w.err
   115  }
   116  
   117  // Err returns an error happened during any operation.
   118  func (w *Writer) Err() error {
   119  	return w.err
   120  }
   121  
   122  func (w *Writer) checkTail() {
   123  	if w.err == nil && w.cbuf.buf != compressionTail {
   124  		w.err = fmt.Errorf(
   125  			"wsflate: bad compressor: unexpected stream tail: %#x vs %#x",
   126  			w.cbuf.buf, compressionTail,
   127  		)
   128  	}
   129  }