response.go 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229
  1. // Copyright 2011 Andy Balholm. 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. // Responding to ICAP requests.
  5. package icap
  6. import (
  7. "bytes"
  8. "errors"
  9. "fmt"
  10. "io"
  11. "log"
  12. "net/http"
  13. "net/http/httputil"
  14. "strconv"
  15. "strings"
  16. "time"
  17. )
  18. // ResponseWriter ---
  19. type ResponseWriter interface {
  20. // Header returns the header map that will be sent by WriteHeader.
  21. // Changing the header after a call to WriteHeader (or Write) has
  22. // no effect.
  23. Header() http.Header
  24. // Write writes the data to the connection as part of an ICAP reply.
  25. // If WriteHeader has not yet been called, Write calls WriteHeader(http.StatusOK, nil)
  26. // before writing the data.
  27. Write([]byte) (int, error)
  28. // Write raw data to the connection.
  29. WriteRaw(string)
  30. // WriteHeader sends an ICAP response header with status code.
  31. // Then it sends an HTTP header if httpMessage is not nil.
  32. // httpMessage may be an *http.Request or an *http.Response.
  33. // hasBody should be true if there will be calls to Write(), generating a message body.
  34. WriteHeader(code int, httpMessage interface{}, hasBody bool)
  35. }
  36. type respWriter struct {
  37. conn *conn // information on the connection
  38. req *Request // the request that is being responded to
  39. header http.Header // the ICAP header to write for the response
  40. wroteHeader bool // true if the headers have already been written
  41. wroteRaw bool // true if raw data was written to the connection
  42. cw io.WriteCloser // the chunked writer used to write the body
  43. }
  44. func (w *respWriter) Header() http.Header {
  45. return w.header
  46. }
  47. func (w *respWriter) Write(p []byte) (n int, err error) {
  48. if !w.wroteHeader {
  49. w.WriteHeader(http.StatusOK, nil, true)
  50. }
  51. if w.cw == nil {
  52. return 0, errors.New("called Write() on an icap.ResponseWriter that should not have a body")
  53. }
  54. return w.cw.Write(p)
  55. }
  56. func (w *respWriter) WriteRaw(p string) {
  57. bw := w.conn.buf.Writer
  58. io.WriteString(bw, p)
  59. w.wroteRaw = true
  60. }
  61. func (w *respWriter) WriteHeader(code int, httpMessage interface{}, hasBody bool) {
  62. if w.wroteHeader {
  63. log.Println("Called WriteHeader twice on the same connection")
  64. return
  65. }
  66. // Make the HTTP header and the Encapsulated: header.
  67. var header []byte
  68. var encap string
  69. var err error
  70. switch msg := httpMessage.(type) {
  71. case *http.Request:
  72. header, err = httpRequestHeader(msg)
  73. if err != nil {
  74. break
  75. }
  76. if hasBody {
  77. encap = fmt.Sprintf("req-hdr=0, req-body=%d", len(header))
  78. } else {
  79. encap = fmt.Sprintf("req-hdr=0, null-body=%d", len(header))
  80. }
  81. case *http.Response:
  82. header, err = httpResponseHeader(msg)
  83. if err != nil {
  84. break
  85. }
  86. if hasBody {
  87. encap = fmt.Sprintf("res-hdr=0, res-body=%d", len(header))
  88. } else {
  89. encap = fmt.Sprintf("res-hdr=0, null-body=%d", len(header))
  90. }
  91. }
  92. if encap == "" {
  93. if hasBody {
  94. method := w.req.Method
  95. if len(method) > 3 {
  96. method = method[0:3]
  97. }
  98. method = strings.ToLower(method)
  99. encap = fmt.Sprintf("%s-body=0", method)
  100. } else {
  101. encap = "null-body=0"
  102. }
  103. }
  104. w.header.Set("Encapsulated", encap)
  105. if _, ok := w.header["Date"]; !ok {
  106. w.Header().Set("Date", time.Now().UTC().Format(http.TimeFormat))
  107. }
  108. w.header.Set("Connection", "close")
  109. bw := w.conn.buf.Writer
  110. status := StatusText(code)
  111. if status == "" {
  112. status = fmt.Sprintf("status code %d", code)
  113. }
  114. fmt.Fprintf(bw, "ICAP/1.0 %d %s\r\n", code, status)
  115. w.header.Write(bw)
  116. io.WriteString(bw, "\r\n")
  117. if header != nil {
  118. bw.Write(header)
  119. }
  120. w.wroteHeader = true
  121. if hasBody {
  122. w.cw = httputil.NewChunkedWriter(w.conn.buf.Writer)
  123. }
  124. }
  125. func (w *respWriter) finishRequest() {
  126. if !w.wroteHeader {
  127. w.WriteHeader(http.StatusOK, nil, false)
  128. }
  129. if w.cw != nil && !w.wroteRaw {
  130. w.cw.Close()
  131. w.cw = nil
  132. io.WriteString(w.conn.buf, "\r\n")
  133. }
  134. w.conn.buf.Flush()
  135. }
  136. // httpRequestHeader returns the headers for an HTTP request
  137. // as a slice of bytes in a form suitable for including in an ICAP message.
  138. func httpRequestHeader(req *http.Request) (hdr []byte, err error) {
  139. buf := new(bytes.Buffer)
  140. if req.URL == nil {
  141. if err != nil {
  142. return nil, errors.New("icap: httpRequestHeader called on Request with no URL")
  143. }
  144. }
  145. host := req.URL.Host
  146. if host == "" {
  147. host = req.Host
  148. }
  149. req.Header.Set("Host", host)
  150. uri := req.URL.String()
  151. fmt.Fprintf(buf, "%s %s %s\r\n", valueOrDefault(req.Method, "GET"), uri, valueOrDefault(req.Proto, "HTTP/1.1"))
  152. req.Header.WriteSubset(buf, map[string]bool{
  153. "Transfer-Encoding": true,
  154. "Content-Length": true,
  155. })
  156. io.WriteString(buf, "\r\n")
  157. return buf.Bytes(), nil
  158. }
  159. // httpResponseHeader returns the headers for an HTTP response
  160. // as a slice of bytes.
  161. func httpResponseHeader(resp *http.Response) (hdr []byte, err error) {
  162. buf := new(bytes.Buffer)
  163. // Status line
  164. text := resp.Status
  165. if text == "" {
  166. text = http.StatusText(resp.StatusCode)
  167. if text == "" {
  168. text = "status code " + strconv.Itoa(resp.StatusCode)
  169. }
  170. }
  171. proto := resp.Proto
  172. if proto == "" {
  173. proto = "HTTP/1.1"
  174. }
  175. fmt.Fprintf(buf, "%s %d %s\r\n", proto, resp.StatusCode, text)
  176. if _, xIcap206Exists := resp.Header["X-Icap-206"]; xIcap206Exists {
  177. resp.Header.Write(buf)
  178. } else {
  179. resp.Header.WriteSubset(buf, map[string]bool{
  180. "Transfer-Encoding": true,
  181. "Content-Length": false,
  182. })
  183. }
  184. io.WriteString(buf, "\r\n")
  185. return buf.Bytes(), nil
  186. }
  187. // Return value if nonempty, def otherwise.
  188. func valueOrDefault(value, def string) string {
  189. if value != "" {
  190. return value
  191. }
  192. return def
  193. }