response.go 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203
  1. package authorization
  2. import (
  3. "bufio"
  4. "bytes"
  5. "encoding/json"
  6. "fmt"
  7. "net"
  8. "net/http"
  9. "github.com/Sirupsen/logrus"
  10. )
  11. // ResponseModifier allows authorization plugins to read and modify the content of the http.response
  12. type ResponseModifier interface {
  13. http.ResponseWriter
  14. http.Flusher
  15. http.CloseNotifier
  16. // RawBody returns the current http content
  17. RawBody() []byte
  18. // RawHeaders returns the current content of the http headers
  19. RawHeaders() ([]byte, error)
  20. // StatusCode returns the current status code
  21. StatusCode() int
  22. // OverrideBody replaces the body of the HTTP reply
  23. OverrideBody(b []byte)
  24. // OverrideHeader replaces the headers of the HTTP reply
  25. OverrideHeader(b []byte) error
  26. // OverrideStatusCode replaces the status code of the HTTP reply
  27. OverrideStatusCode(statusCode int)
  28. // FlushAll flushes all data to the HTTP response
  29. FlushAll() error
  30. // Hijacked indicates the response has been hijacked by the Docker daemon
  31. Hijacked() bool
  32. }
  33. // NewResponseModifier creates a wrapper to an http.ResponseWriter to allow inspecting and modifying the content
  34. func NewResponseModifier(rw http.ResponseWriter) ResponseModifier {
  35. return &responseModifier{rw: rw, header: make(http.Header)}
  36. }
  37. // responseModifier is used as an adapter to http.ResponseWriter in order to manipulate and explore
  38. // the http request/response from docker daemon
  39. type responseModifier struct {
  40. // The original response writer
  41. rw http.ResponseWriter
  42. // body holds the response body
  43. body []byte
  44. // header holds the response header
  45. header http.Header
  46. // statusCode holds the response status code
  47. statusCode int
  48. // hijacked indicates the request has been hijacked
  49. hijacked bool
  50. }
  51. func (rm *responseModifier) Hijacked() bool {
  52. return rm.hijacked
  53. }
  54. // WriterHeader stores the http status code
  55. func (rm *responseModifier) WriteHeader(s int) {
  56. // Use original request if hijacked
  57. if rm.hijacked {
  58. rm.rw.WriteHeader(s)
  59. return
  60. }
  61. rm.statusCode = s
  62. }
  63. // Header returns the internal http header
  64. func (rm *responseModifier) Header() http.Header {
  65. // Use original header if hijacked
  66. if rm.hijacked {
  67. return rm.rw.Header()
  68. }
  69. return rm.header
  70. }
  71. // StatusCode returns the http status code
  72. func (rm *responseModifier) StatusCode() int {
  73. return rm.statusCode
  74. }
  75. // OverrideBody replaces the body of the HTTP response
  76. func (rm *responseModifier) OverrideBody(b []byte) {
  77. rm.body = b
  78. }
  79. // OverrideStatusCode replaces the status code of the HTTP response
  80. func (rm *responseModifier) OverrideStatusCode(statusCode int) {
  81. rm.statusCode = statusCode
  82. }
  83. // OverrideHeader replaces the headers of the HTTP response
  84. func (rm *responseModifier) OverrideHeader(b []byte) error {
  85. header := http.Header{}
  86. if err := json.Unmarshal(b, &header); err != nil {
  87. return err
  88. }
  89. rm.header = header
  90. return nil
  91. }
  92. // Write stores the byte array inside content
  93. func (rm *responseModifier) Write(b []byte) (int, error) {
  94. if rm.hijacked {
  95. return rm.rw.Write(b)
  96. }
  97. rm.body = append(rm.body, b...)
  98. return len(b), nil
  99. }
  100. // Body returns the response body
  101. func (rm *responseModifier) RawBody() []byte {
  102. return rm.body
  103. }
  104. func (rm *responseModifier) RawHeaders() ([]byte, error) {
  105. var b bytes.Buffer
  106. if err := rm.header.Write(&b); err != nil {
  107. return nil, err
  108. }
  109. return b.Bytes(), nil
  110. }
  111. // Hijack returns the internal connection of the wrapped http.ResponseWriter
  112. func (rm *responseModifier) Hijack() (net.Conn, *bufio.ReadWriter, error) {
  113. rm.hijacked = true
  114. rm.FlushAll()
  115. hijacker, ok := rm.rw.(http.Hijacker)
  116. if !ok {
  117. return nil, nil, fmt.Errorf("Internal response writer doesn't support the Hijacker interface")
  118. }
  119. return hijacker.Hijack()
  120. }
  121. // CloseNotify uses the internal close notify API of the wrapped http.ResponseWriter
  122. func (rm *responseModifier) CloseNotify() <-chan bool {
  123. closeNotifier, ok := rm.rw.(http.CloseNotifier)
  124. if !ok {
  125. logrus.Error("Internal response writer doesn't support the CloseNotifier interface")
  126. return nil
  127. }
  128. return closeNotifier.CloseNotify()
  129. }
  130. // Flush uses the internal flush API of the wrapped http.ResponseWriter
  131. func (rm *responseModifier) Flush() {
  132. flusher, ok := rm.rw.(http.Flusher)
  133. if !ok {
  134. logrus.Error("Internal response writer doesn't support the Flusher interface")
  135. return
  136. }
  137. rm.FlushAll()
  138. flusher.Flush()
  139. }
  140. // FlushAll flushes all data to the HTTP response
  141. func (rm *responseModifier) FlushAll() error {
  142. // Copy the header
  143. for k, vv := range rm.header {
  144. for _, v := range vv {
  145. rm.rw.Header().Add(k, v)
  146. }
  147. }
  148. // Copy the status code
  149. // Also WriteHeader needs to be done after all the headers
  150. // have been copied (above).
  151. if rm.statusCode > 0 {
  152. rm.rw.WriteHeader(rm.statusCode)
  153. }
  154. var err error
  155. if len(rm.body) > 0 {
  156. // Write body
  157. _, err = rm.rw.Write(rm.body)
  158. }
  159. // Clean previous data
  160. rm.body = nil
  161. rm.statusCode = 0
  162. rm.header = http.Header{}
  163. return err
  164. }