compress.go

Documentation: github.com/emicklei/go-restful/v3

     1  package restful
     2  
     3  // Copyright 2013 Ernest Micklei. All rights reserved.
     4  // Use of this source code is governed by a license
     5  // that can be found in the LICENSE file.
     6  
     7  import (
     8  	"bufio"
     9  	"compress/gzip"
    10  	"compress/zlib"
    11  	"errors"
    12  	"io"
    13  	"net"
    14  	"net/http"
    15  	"strings"
    16  )
    17  
    18  // OBSOLETE : use restful.DefaultContainer.EnableContentEncoding(true) to change this setting.
    19  var EnableContentEncoding = false
    20  
    21  // CompressingResponseWriter is a http.ResponseWriter that can perform content encoding (gzip and zlib)
    22  type CompressingResponseWriter struct {
    23  	writer     http.ResponseWriter
    24  	compressor io.WriteCloser
    25  	encoding   string
    26  }
    27  
    28  // Header is part of http.ResponseWriter interface
    29  func (c *CompressingResponseWriter) Header() http.Header {
    30  	return c.writer.Header()
    31  }
    32  
    33  // WriteHeader is part of http.ResponseWriter interface
    34  func (c *CompressingResponseWriter) WriteHeader(status int) {
    35  	c.writer.WriteHeader(status)
    36  }
    37  
    38  // Write is part of http.ResponseWriter interface
    39  // It is passed through the compressor
    40  func (c *CompressingResponseWriter) Write(bytes []byte) (int, error) {
    41  	if c.isCompressorClosed() {
    42  		return -1, errors.New("Compressing error: tried to write data using closed compressor")
    43  	}
    44  	return c.compressor.Write(bytes)
    45  }
    46  
    47  // CloseNotify is part of http.CloseNotifier interface
    48  func (c *CompressingResponseWriter) CloseNotify() <-chan bool {
    49  	return c.writer.(http.CloseNotifier).CloseNotify()
    50  }
    51  
    52  // Flush is part of http.Flusher interface. Noop if the underlying writer doesn't support it.
    53  func (c *CompressingResponseWriter) Flush() {
    54  	flusher, ok := c.writer.(http.Flusher)
    55  	if !ok {
    56  		// writer doesn't support http.Flusher interface
    57  		return
    58  	}
    59  	flusher.Flush()
    60  }
    61  
    62  // Close the underlying compressor
    63  func (c *CompressingResponseWriter) Close() error {
    64  	if c.isCompressorClosed() {
    65  		return errors.New("Compressing error: tried to close already closed compressor")
    66  	}
    67  
    68  	c.compressor.Close()
    69  	if ENCODING_GZIP == c.encoding {
    70  		currentCompressorProvider.ReleaseGzipWriter(c.compressor.(*gzip.Writer))
    71  	}
    72  	if ENCODING_DEFLATE == c.encoding {
    73  		currentCompressorProvider.ReleaseZlibWriter(c.compressor.(*zlib.Writer))
    74  	}
    75  	// gc hint needed?
    76  	c.compressor = nil
    77  	return nil
    78  }
    79  
    80  func (c *CompressingResponseWriter) isCompressorClosed() bool {
    81  	return nil == c.compressor
    82  }
    83  
    84  // Hijack implements the Hijacker interface
    85  // This is especially useful when combining Container.EnabledContentEncoding
    86  // in combination with websockets (for instance gorilla/websocket)
    87  func (c *CompressingResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    88  	hijacker, ok := c.writer.(http.Hijacker)
    89  	if !ok {
    90  		return nil, nil, errors.New("ResponseWriter doesn't support Hijacker interface")
    91  	}
    92  	return hijacker.Hijack()
    93  }
    94  
    95  // WantsCompressedResponse reads the Accept-Encoding header to see if and which encoding is requested.
    96  // It also inspects the httpWriter whether its content-encoding is already set (non-empty).
    97  func wantsCompressedResponse(httpRequest *http.Request, httpWriter http.ResponseWriter) (bool, string) {
    98  	if contentEncoding := httpWriter.Header().Get(HEADER_ContentEncoding); contentEncoding != "" {
    99  		return false, ""
   100  	}
   101  	header := httpRequest.Header.Get(HEADER_AcceptEncoding)
   102  	gi := strings.Index(header, ENCODING_GZIP)
   103  	zi := strings.Index(header, ENCODING_DEFLATE)
   104  	// use in order of appearance
   105  	if gi == -1 {
   106  		return zi != -1, ENCODING_DEFLATE
   107  	} else if zi == -1 {
   108  		return gi != -1, ENCODING_GZIP
   109  	} else {
   110  		if gi < zi {
   111  			return true, ENCODING_GZIP
   112  		}
   113  		return true, ENCODING_DEFLATE
   114  	}
   115  }
   116  
   117  // NewCompressingResponseWriter create a CompressingResponseWriter for a known encoding = {gzip,deflate}
   118  func NewCompressingResponseWriter(httpWriter http.ResponseWriter, encoding string) (*CompressingResponseWriter, error) {
   119  	httpWriter.Header().Set(HEADER_ContentEncoding, encoding)
   120  	c := new(CompressingResponseWriter)
   121  	c.writer = httpWriter
   122  	var err error
   123  	if ENCODING_GZIP == encoding {
   124  		w := currentCompressorProvider.AcquireGzipWriter()
   125  		w.Reset(httpWriter)
   126  		c.compressor = w
   127  		c.encoding = ENCODING_GZIP
   128  	} else if ENCODING_DEFLATE == encoding {
   129  		w := currentCompressorProvider.AcquireZlibWriter()
   130  		w.Reset(httpWriter)
   131  		c.compressor = w
   132  		c.encoding = ENCODING_DEFLATE
   133  	} else {
   134  		return nil, errors.New("Unknown encoding:" + encoding)
   135  	}
   136  	return c, err
   137  }
   138
View as plain text