response_writer.go 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  1. // Copyright 2013 Martini Authors
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License"): you may
  4. // not use this file except in compliance with the License. You may obtain
  5. // a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  11. // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  12. // License for the specific language governing permissions and limitations
  13. // under the License.
  14. package macaron
  15. import (
  16. "bufio"
  17. "fmt"
  18. "net"
  19. "net/http"
  20. )
  21. // ResponseWriter is a wrapper around http.ResponseWriter that provides extra information about
  22. // the response. It is recommended that middleware handlers use this construct to wrap a responsewriter
  23. // if the functionality calls for it.
  24. type ResponseWriter interface {
  25. http.ResponseWriter
  26. http.Flusher
  27. // Status returns the status code of the response or 0 if the response has not been written.
  28. Status() int
  29. // Written returns whether or not the ResponseWriter has been written.
  30. Written() bool
  31. // Size returns the size of the response body.
  32. Size() int
  33. // Before allows for a function to be called before the ResponseWriter has been written to. This is
  34. // useful for setting headers or any other operations that must happen before a response has been written.
  35. Before(BeforeFunc)
  36. }
  37. // BeforeFunc is a function that is called before the ResponseWriter has been written to.
  38. type BeforeFunc func(ResponseWriter)
  39. // NewResponseWriter creates a ResponseWriter that wraps an http.ResponseWriter
  40. func NewResponseWriter(method string, rw http.ResponseWriter) ResponseWriter {
  41. return &responseWriter{method, rw, 0, 0, nil}
  42. }
  43. type responseWriter struct {
  44. method string
  45. http.ResponseWriter
  46. status int
  47. size int
  48. beforeFuncs []BeforeFunc
  49. }
  50. func (rw *responseWriter) WriteHeader(s int) {
  51. rw.callBefore()
  52. rw.ResponseWriter.WriteHeader(s)
  53. rw.status = s
  54. }
  55. func (rw *responseWriter) Write(b []byte) (size int, err error) {
  56. if !rw.Written() {
  57. // The status will be StatusOK if WriteHeader has not been called yet
  58. rw.WriteHeader(http.StatusOK)
  59. }
  60. if rw.method != "HEAD" {
  61. size, err = rw.ResponseWriter.Write(b)
  62. rw.size += size
  63. }
  64. return size, err
  65. }
  66. func (rw *responseWriter) Status() int {
  67. return rw.status
  68. }
  69. func (rw *responseWriter) Size() int {
  70. return rw.size
  71. }
  72. func (rw *responseWriter) Written() bool {
  73. return rw.status != 0
  74. }
  75. func (rw *responseWriter) Before(before BeforeFunc) {
  76. rw.beforeFuncs = append(rw.beforeFuncs, before)
  77. }
  78. func (rw *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
  79. hijacker, ok := rw.ResponseWriter.(http.Hijacker)
  80. if !ok {
  81. return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
  82. }
  83. return hijacker.Hijack()
  84. }
  85. func (rw *responseWriter) CloseNotify() <-chan bool {
  86. return rw.ResponseWriter.(http.CloseNotifier).CloseNotify()
  87. }
  88. func (rw *responseWriter) callBefore() {
  89. for i := len(rw.beforeFuncs) - 1; i >= 0; i-- {
  90. rw.beforeFuncs[i](rw)
  91. }
  92. }
  93. func (rw *responseWriter) Flush() {
  94. flusher, ok := rw.ResponseWriter.(http.Flusher)
  95. if ok {
  96. flusher.Flush()
  97. }
  98. }