另客网go项目公用的代码库

response_writer.go 2.4KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  1. // Copyright 2014 Manu Martinez-Almeida. All rights reserved.
  2. // Use of this source code is governed by a MIT style
  3. // license that can be found in the LICENSE file.
  4. package gin
  5. import (
  6. "bufio"
  7. "io"
  8. "net"
  9. "net/http"
  10. )
  11. const (
  12. noWritten = -1
  13. defaultStatus = http.StatusOK
  14. )
  15. type responseWriterBase interface {
  16. http.ResponseWriter
  17. http.Hijacker
  18. http.Flusher
  19. http.CloseNotifier
  20. // Returns the HTTP response status code of the current request.
  21. Status() int
  22. // Returns the number of bytes already written into the response http body.
  23. // See Written()
  24. Size() int
  25. // Writes the string into the response body.
  26. WriteString(string) (int, error)
  27. // Returns true if the response body was already written.
  28. Written() bool
  29. // Forces to write the http header (status code + headers).
  30. WriteHeaderNow()
  31. }
  32. type responseWriter struct {
  33. http.ResponseWriter
  34. size int
  35. status int
  36. }
  37. var _ ResponseWriter = &responseWriter{}
  38. func (w *responseWriter) reset(writer http.ResponseWriter) {
  39. w.ResponseWriter = writer
  40. w.size = noWritten
  41. w.status = defaultStatus
  42. }
  43. func (w *responseWriter) WriteHeader(code int) {
  44. if code > 0 && w.status != code {
  45. if w.Written() {
  46. debugPrint("[WARNING] Headers were already written. Wanted to override status code %d with %d", w.status, code)
  47. }
  48. w.status = code
  49. }
  50. }
  51. func (w *responseWriter) WriteHeaderNow() {
  52. if !w.Written() {
  53. w.size = 0
  54. w.ResponseWriter.WriteHeader(w.status)
  55. }
  56. }
  57. func (w *responseWriter) Write(data []byte) (n int, err error) {
  58. w.WriteHeaderNow()
  59. n, err = w.ResponseWriter.Write(data)
  60. w.size += n
  61. return
  62. }
  63. func (w *responseWriter) WriteString(s string) (n int, err error) {
  64. w.WriteHeaderNow()
  65. n, err = io.WriteString(w.ResponseWriter, s)
  66. w.size += n
  67. return
  68. }
  69. func (w *responseWriter) Status() int {
  70. return w.status
  71. }
  72. func (w *responseWriter) Size() int {
  73. return w.size
  74. }
  75. func (w *responseWriter) Written() bool {
  76. return w.size != noWritten
  77. }
  78. // Hijack implements the http.Hijacker interface.
  79. func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
  80. if w.size < 0 {
  81. w.size = 0
  82. }
  83. return w.ResponseWriter.(http.Hijacker).Hijack()
  84. }
  85. // CloseNotify implements the http.CloseNotify interface.
  86. func (w *responseWriter) CloseNotify() <-chan bool {
  87. return w.ResponseWriter.(http.CloseNotifier).CloseNotify()
  88. }
  89. // Flush implements the http.Flush interface.
  90. func (w *responseWriter) Flush() {
  91. w.WriteHeaderNow()
  92. w.ResponseWriter.(http.Flusher).Flush()
  93. }