builder.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. // Copyright 2018 Frédéric Guillot. All rights reserved.
  2. // Use of this source code is governed by the Apache 2.0
  3. // license that can be found in the LICENSE file.
  4. package response // import "miniflux.app/http/response"
  5. import (
  6. "compress/flate"
  7. "compress/gzip"
  8. "fmt"
  9. "io"
  10. "net/http"
  11. "strings"
  12. "time"
  13. "miniflux.app/logger"
  14. )
  15. const compressionThreshold = 1024
  16. // Builder generates HTTP responses.
  17. type Builder struct {
  18. w http.ResponseWriter
  19. r *http.Request
  20. statusCode int
  21. headers map[string]string
  22. enableCompression bool
  23. body interface{}
  24. }
  25. // WithStatus uses the given status code to build the response.
  26. func (b *Builder) WithStatus(statusCode int) *Builder {
  27. b.statusCode = statusCode
  28. return b
  29. }
  30. // WithHeader adds the given HTTP header to the response.
  31. func (b *Builder) WithHeader(key, value string) *Builder {
  32. b.headers[key] = value
  33. return b
  34. }
  35. // WithBody uses the given body to build the response.
  36. func (b *Builder) WithBody(body interface{}) *Builder {
  37. b.body = body
  38. return b
  39. }
  40. // WithAttachment forces the document to be downloaded by the web browser.
  41. func (b *Builder) WithAttachment(filename string) *Builder {
  42. b.headers["Content-Disposition"] = fmt.Sprintf("attachment; filename=%s", filename)
  43. return b
  44. }
  45. // WithoutCompression disables HTTP compression.
  46. func (b *Builder) WithoutCompression() *Builder {
  47. b.enableCompression = false
  48. return b
  49. }
  50. // WithCaching adds caching headers to the response.
  51. func (b *Builder) WithCaching(etag string, duration time.Duration, callback func(*Builder)) {
  52. b.headers["ETag"] = etag
  53. b.headers["Cache-Control"] = "public"
  54. b.headers["Expires"] = time.Now().Add(duration).Format(time.RFC1123)
  55. if etag == b.r.Header.Get("If-None-Match") {
  56. b.statusCode = http.StatusNotModified
  57. b.body = nil
  58. b.Write()
  59. } else {
  60. callback(b)
  61. }
  62. }
  63. // Write generates the HTTP response.
  64. func (b *Builder) Write() {
  65. if b.body == nil {
  66. b.writeHeaders()
  67. return
  68. }
  69. switch v := b.body.(type) {
  70. case []byte:
  71. b.compress(v)
  72. case string:
  73. b.compress([]byte(v))
  74. case error:
  75. b.compress([]byte(v.Error()))
  76. case io.Reader:
  77. // Compression not implemented in this case
  78. b.writeHeaders()
  79. _, err := io.Copy(b.w, v)
  80. if err != nil {
  81. logger.Error("%v", err)
  82. }
  83. }
  84. }
  85. func (b *Builder) writeHeaders() {
  86. b.headers["X-XSS-Protection"] = "1; mode=block"
  87. b.headers["X-Content-Type-Options"] = "nosniff"
  88. b.headers["X-Frame-Options"] = "DENY"
  89. b.headers["Referrer-Policy"] = "no-referrer"
  90. for key, value := range b.headers {
  91. b.w.Header().Set(key, value)
  92. }
  93. b.w.WriteHeader(b.statusCode)
  94. }
  95. func (b *Builder) compress(data []byte) {
  96. if b.enableCompression && len(data) > compressionThreshold {
  97. acceptEncoding := b.r.Header.Get("Accept-Encoding")
  98. switch {
  99. case strings.Contains(acceptEncoding, "gzip"):
  100. b.headers["Content-Encoding"] = "gzip"
  101. b.writeHeaders()
  102. gzipWriter := gzip.NewWriter(b.w)
  103. defer gzipWriter.Close()
  104. gzipWriter.Write(data)
  105. return
  106. case strings.Contains(acceptEncoding, "deflate"):
  107. b.headers["Content-Encoding"] = "deflate"
  108. b.writeHeaders()
  109. flateWriter, _ := flate.NewWriter(b.w, -1)
  110. defer flateWriter.Close()
  111. flateWriter.Write(data)
  112. return
  113. }
  114. }
  115. b.writeHeaders()
  116. b.w.Write(data)
  117. }
  118. // New creates a new response builder.
  119. func New(w http.ResponseWriter, r *http.Request) *Builder {
  120. return &Builder{w: w, r: r, statusCode: http.StatusOK, headers: make(map[string]string), enableCompression: true}
  121. }