compress.go 3.7 KB

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