Micro & pluggable web framework for Go
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

157 lines
3.4KB

  1. // Copyright 2015 The Tango 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. package tango
  5. import (
  6. "bufio"
  7. "compress/flate"
  8. "compress/gzip"
  9. "fmt"
  10. "io"
  11. "net"
  12. "net/http"
  13. "path"
  14. "strings"
  15. )
  16. // some http headers
  17. const (
  18. HeaderAcceptEncoding = "Accept-Encoding"
  19. HeaderContentEncoding = "Content-Encoding"
  20. HeaderContentLength = "Content-Length"
  21. HeaderContentType = "Content-Type"
  22. HeaderVary = "Vary"
  23. )
  24. // Compresser defines the interface return compress type
  25. type Compresser interface {
  26. CompressType() string
  27. }
  28. // GZip implements gzip Compresser
  29. type GZip struct{}
  30. // CompressType returns compress type
  31. func (GZip) CompressType() string {
  32. return "gzip"
  33. }
  34. // Deflate implements deflate Compresser
  35. type Deflate struct{}
  36. // CompressType returns compress type
  37. func (Deflate) CompressType() string {
  38. return "deflate"
  39. }
  40. // Compress implements auto Compresser
  41. type Compress struct{}
  42. // CompressType returns compress type
  43. func (Compress) CompressType() string {
  44. return "auto"
  45. }
  46. // Compresses defines a middleware to compress HTTP response
  47. func Compresses(exts []string) HandlerFunc {
  48. extsmap := make(map[string]bool)
  49. for _, ext := range exts {
  50. extsmap[strings.ToLower(ext)] = true
  51. }
  52. return func(ctx *Context) {
  53. ae := ctx.Req().Header.Get("Accept-Encoding")
  54. if ae == "" {
  55. ctx.Next()
  56. return
  57. }
  58. if len(extsmap) > 0 {
  59. ext := strings.ToLower(path.Ext(ctx.Req().URL.Path))
  60. if _, ok := extsmap[ext]; ok {
  61. compress(ctx, "auto")
  62. return
  63. }
  64. }
  65. if action := ctx.Action(); action != nil {
  66. if c, ok := action.(Compresser); ok {
  67. compress(ctx, c.CompressType())
  68. return
  69. }
  70. }
  71. // if blank, then no compress
  72. ctx.Next()
  73. }
  74. }
  75. func compress(ctx *Context, compressType string) {
  76. ae := ctx.Req().Header.Get("Accept-Encoding")
  77. acceptCompress := strings.SplitN(ae, ",", -1)
  78. var writer io.Writer
  79. var val string
  80. for _, val = range acceptCompress {
  81. val = strings.TrimSpace(val)
  82. if compressType == "auto" || val == compressType {
  83. if val == "gzip" {
  84. ctx.Header().Set("Content-Encoding", "gzip")
  85. writer = gzip.NewWriter(ctx.ResponseWriter)
  86. break
  87. } else if val == "deflate" {
  88. ctx.Header().Set("Content-Encoding", "deflate")
  89. writer, _ = flate.NewWriter(ctx.ResponseWriter, flate.BestSpeed)
  90. break
  91. }
  92. }
  93. }
  94. // not supported compress method, then ignore
  95. if writer == nil {
  96. ctx.Next()
  97. return
  98. }
  99. // for cache server
  100. ctx.Header().Add(HeaderVary, "Accept-Encoding")
  101. gzw := &compressWriter{writer, ctx.ResponseWriter}
  102. ctx.ResponseWriter = gzw
  103. ctx.Next()
  104. // delete content length after we know we have been written to
  105. gzw.Header().Del(HeaderContentLength)
  106. ctx.ResponseWriter = gzw.ResponseWriter
  107. switch writer.(type) {
  108. case *gzip.Writer:
  109. writer.(*gzip.Writer).Close()
  110. case *flate.Writer:
  111. writer.(*flate.Writer).Close()
  112. }
  113. }
  114. type compressWriter struct {
  115. w io.Writer
  116. ResponseWriter
  117. }
  118. func (grw *compressWriter) Write(p []byte) (int, error) {
  119. if len(grw.Header().Get(HeaderContentType)) == 0 {
  120. grw.Header().Set(HeaderContentType, http.DetectContentType(p))
  121. }
  122. return grw.w.Write(p)
  123. }
  124. func (grw *compressWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
  125. hijacker, ok := grw.ResponseWriter.(http.Hijacker)
  126. if !ok {
  127. return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
  128. }
  129. return hijacker.Hijack()
  130. }