github.com/xushiwei/go@v0.0.0-20130601165731-2b9d83f45bc9/src/pkg/net/http/httputil/chunked.go (about)

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