chunked.go 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181
  1. // Copyright 2009 The Go Authors. 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. // The wire protocol for HTTP's "chunked" Transfer-Encoding.
  5. // This code is derived from the standard library's http/httputil/chunked.go,
  6. package icap
  7. import (
  8. "bufio"
  9. "errors"
  10. "fmt"
  11. "io"
  12. )
  13. const maxLineLength = 4096 // assumed <= bufio.defaultBufSize
  14. var errLineTooLong = errors.New("header line too long")
  15. // NewChunkedReader returns a new chunkedReader that translates the data read from r
  16. // out of HTTP "chunked" format before returning it.
  17. // The chunkedReader returns io.EOF when the final 0-length chunk is read.
  18. //
  19. // NewChunkedReader is not needed by normal applications. The http package
  20. // automatically decodes chunking when reading response bodies.
  21. func newChunkedReader(r io.Reader) io.Reader {
  22. br, ok := r.(*bufio.Reader)
  23. if !ok {
  24. br = bufio.NewReader(r)
  25. }
  26. return &chunkedReader{r: br}
  27. }
  28. type chunkedReader struct {
  29. r *bufio.Reader
  30. n uint64 // unread bytes in chunk
  31. err error
  32. buf [2]byte
  33. }
  34. func (cr *chunkedReader) beginChunk() {
  35. // chunk-size CRLF
  36. var line []byte
  37. line, cr.err = readLine(cr.r)
  38. if cr.err != nil {
  39. return
  40. }
  41. cr.n, cr.err = parseHexUint(line)
  42. if cr.err != nil {
  43. return
  44. }
  45. if cr.n == 0 {
  46. cr.err = io.EOF
  47. }
  48. }
  49. func (cr *chunkedReader) Read(b []uint8) (n int, err error) {
  50. if cr.err != nil {
  51. return 0, cr.err
  52. }
  53. if cr.n == 0 {
  54. cr.beginChunk()
  55. if cr.err != nil {
  56. return 0, cr.err
  57. }
  58. }
  59. if uint64(len(b)) > cr.n {
  60. b = b[0:cr.n]
  61. }
  62. n, cr.err = cr.r.Read(b)
  63. cr.n -= uint64(n)
  64. if cr.n == 0 && cr.err == nil {
  65. // end of chunk (CRLF)
  66. if _, cr.err = io.ReadFull(cr.r, cr.buf[:]); cr.err == nil {
  67. if cr.buf[0] != '\r' || cr.buf[1] != '\n' {
  68. cr.err = errors.New("malformed chunked encoding")
  69. }
  70. }
  71. }
  72. return n, cr.err
  73. }
  74. // Read a line of bytes (up to \n) from b.
  75. // Give up if the line exceeds maxLineLength.
  76. // The returned bytes are a pointer into storage in
  77. // the bufio, so they are only valid until the next bufio read.
  78. func readLine(b *bufio.Reader) (p []byte, err error) {
  79. if p, err = b.ReadSlice('\n'); err != nil {
  80. // We always know when EOF is coming.
  81. // If the caller asked for a line, there should be a line.
  82. if err == io.EOF {
  83. err = io.ErrUnexpectedEOF
  84. } else if err == bufio.ErrBufferFull {
  85. err = errLineTooLong
  86. }
  87. return nil, err
  88. }
  89. if len(p) >= maxLineLength {
  90. return nil, errLineTooLong
  91. }
  92. return trimTrailingWhitespace(p), nil
  93. }
  94. func trimTrailingWhitespace(b []byte) []byte {
  95. for len(b) > 0 && isASCIISpace(b[len(b)-1]) {
  96. b = b[:len(b)-1]
  97. }
  98. return b
  99. }
  100. func isASCIISpace(b byte) bool {
  101. return b == ' ' || b == '\t' || b == '\n' || b == '\r'
  102. }
  103. // NewChunkedWriter returns a new chunkedWriter that translates writes into HTTP
  104. // "chunked" format before writing them to w. Closing the returned chunkedWriter
  105. // sends the final 0-length chunk that marks the end of the stream.
  106. //
  107. // NewChunkedWriter is not needed by normal applications. The http
  108. // package adds chunking automatically if handlers don't set a
  109. // Content-Length header. Using NewChunkedWriter inside a handler
  110. // would result in double chunking or chunking with a Content-Length
  111. // length, both of which are wrong.
  112. func NewChunkedWriter(w io.Writer) io.WriteCloser {
  113. return &chunkedWriter{w}
  114. }
  115. // Writing to chunkedWriter translates to writing in HTTP chunked Transfer
  116. // Encoding wire format to the underlying Wire chunkedWriter.
  117. type chunkedWriter struct {
  118. Wire io.Writer
  119. }
  120. // Write the contents of data as one chunk to Wire.
  121. // NOTE: Note that the corresponding chunk-writing procedure in Conn.Write has
  122. // a bug since it does not check for success of io.WriteString
  123. func (cw *chunkedWriter) Write(data []byte) (n int, err error) {
  124. // Don't send 0-length data. It looks like EOF for chunked encoding.
  125. if len(data) == 0 {
  126. return 0, nil
  127. }
  128. if _, err = fmt.Fprintf(cw.Wire, "%x\r\n", len(data)); err != nil {
  129. return 0, err
  130. }
  131. if n, err = cw.Wire.Write(data); err != nil {
  132. return
  133. }
  134. if n != len(data) {
  135. err = io.ErrShortWrite
  136. return
  137. }
  138. _, err = io.WriteString(cw.Wire, "\r\n")
  139. return
  140. }
  141. func (cw *chunkedWriter) Close() error {
  142. _, err := io.WriteString(cw.Wire, "0\r\n")
  143. return err
  144. }
  145. func parseHexUint(v []byte) (n uint64, err error) {
  146. for _, b := range v {
  147. n <<= 4
  148. switch {
  149. case '0' <= b && b <= '9':
  150. b = b - '0'
  151. case 'a' <= b && b <= 'f':
  152. b = b - 'a' + 10
  153. case 'A' <= b && b <= 'F':
  154. b = b - 'A' + 10
  155. default:
  156. return 0, fmt.Errorf("invalid chunk length: '%s'", v)
  157. }
  158. n |= uint64(b)
  159. }
  160. return
  161. }