func (t *transferWriter) WriteBody(w io.Writer) error {
var err error
var ncopy int64
// Write body
if t.Body != nil {
if chunked(t.TransferEncoding) {
if bw, ok := w.(*bufio.Writer); ok && !t.IsResponse {
w = &internal.FlushAfterChunkWriter{Writer: bw}
}
cw := internal.NewChunkedWriter(w)
_, err = io.Copy(cw, t.Body)
if err == nil {
err = cw.Close()
}
} else if t.ContentLength == -1 {
ncopy, err = io.Copy(w, t.Body)
} else {
ncopy, err = io.Copy(w, io.LimitReader(t.Body, t.ContentLength))
if err != nil {
return err
}
var nextra int64
nextra, err = io.Copy(ioutil.Discard, t.Body)
ncopy += nextra
}
if err != nil {
return err
}
if err = t.BodyCloser.Close(); err != nil {
return err
}
}
if !t.ResponseToHEAD && t.ContentLength != -1 && t.ContentLength != ncopy {
return fmt.Errorf("http: ContentLength=%d with Body length %d",
t.ContentLength, ncopy)
}
if chunked(t.TransferEncoding) {
// Write Trailer header
if t.Trailer != nil {
if err := t.Trailer.Write(w); err != nil {
return err
}
}
// Last chunk, empty trailer
_, err = io.WriteString(w, "\r\n")
}
return err
}
// TestReadResponseCloseInMiddle tests that closing a body after
// reading only part of its contents advances the read to the end of
// the request, right up until the next request.
func TestReadResponseCloseInMiddle(t *testing.T) {
t.Parallel()
for _, test := range readResponseCloseInMiddleTests {
fatalf := func(format string, args ...interface{}) {
args = append([]interface{}{test.chunked, test.compressed}, args...)
t.Fatalf("on test chunked=%v, compressed=%v: "+format, args...)
}
checkErr := func(err error, msg string) {
if err == nil {
return
}
fatalf(msg+": %v", err)
}
var buf bytes.Buffer
buf.WriteString("HTTP/1.1 200 OK\r\n")
if test.chunked {
buf.WriteString("Transfer-Encoding: chunked\r\n")
} else {
buf.WriteString("Content-Length: 1000000\r\n")
}
var wr io.Writer = &buf
if test.chunked {
wr = internal.NewChunkedWriter(wr)
}
if test.compressed {
buf.WriteString("Content-Encoding: gzip\r\n")
wr = gzip.NewWriter(wr)
}
buf.WriteString("\r\n")
chunk := bytes.Repeat([]byte{'x'}, 1000)
for i := 0; i < 1000; i++ {
if test.compressed {
// Otherwise this compresses too well.
_, err := io.ReadFull(rand.Reader, chunk)
checkErr(err, "rand.Reader ReadFull")
}
wr.Write(chunk)
}
if test.compressed {
err := wr.(*gzip.Writer).Close()
checkErr(err, "compressor close")
}
if test.chunked {
buf.WriteString("0\r\n\r\n")
}
buf.WriteString("Next Request Here")
bufr := bufio.NewReader(&buf)
resp, err := ReadResponse(bufr, dummyReq("GET"))
checkErr(err, "ReadResponse")
expectedLength := int64(-1)
if !test.chunked {
expectedLength = 1000000
}
if resp.ContentLength != expectedLength {
fatalf("expected response length %d, got %d", expectedLength, resp.ContentLength)
}
if resp.Body == nil {
fatalf("nil body")
}
if test.compressed {
gzReader, err := gzip.NewReader(resp.Body)
checkErr(err, "gzip.NewReader")
resp.Body = &readerAndCloser{gzReader, resp.Body}
}
rbuf := make([]byte, 2500)
n, err := io.ReadFull(resp.Body, rbuf)
checkErr(err, "2500 byte ReadFull")
if n != 2500 {
fatalf("ReadFull only read %d bytes", n)
}
if test.compressed == false && !bytes.Equal(bytes.Repeat([]byte{'x'}, 2500), rbuf) {
fatalf("ReadFull didn't read 2500 'x'; got %q", string(rbuf))
}
resp.Body.Close()
rest, err := ioutil.ReadAll(bufr)
checkErr(err, "ReadAll on remainder")
if e, g := "Next Request Here", string(rest); e != g {
g = regexp.MustCompile(`(xx+)`).ReplaceAllStringFunc(g, func(match string) string {
return fmt.Sprintf("x(repeated x%d)", len(match))
})
fatalf("remainder = %q, expected %q", g, e)
}
}
}
// NewChunkedWriter returns a new chunkedWriter that translates writes into HTTP
// "chunked" format before writing them to w. Closing the returned chunkedWriter
// sends the final 0-length chunk that marks the end of the stream.
//
// NewChunkedWriter is not needed by normal applications. The http
// package adds chunking automatically if handlers don't set a
// Content-Length header. Using NewChunkedWriter inside a handler
// would result in double chunking or chunking with a Content-Length
// length, both of which are wrong.
func NewChunkedWriter(w io.Writer) io.WriteCloser {
return internal.NewChunkedWriter(w)
}
