src/pkg/http/serve_test.go | 7 +++++++ src/pkg/http/server.go | 15 +++++++++++---- diff --git a/src/pkg/http/serve_test.go b/src/pkg/http/serve_test.go index 1054d4797ca9d66865cdc62a7aceb76c040013c1..dc4594a7908aec8aa6585d675c44c85bdbc87ae2 100644 --- a/src/pkg/http/serve_test.go +++ b/src/pkg/http/serve_test.go @@ -12,6 +12,7 @@ "bytes" "fmt" . "http" "http/httptest" + "io" "io/ioutil" "log" "os" @@ -494,6 +495,12 @@ ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) { _, err := w.Write([]byte("Ignored body")) if err != ErrBodyNotAllowed { t.Errorf("on Write, expected ErrBodyNotAllowed, got %v", err) + } + + // Also exercise the ReaderFrom path + _, err = io.Copy(w, strings.NewReader("Ignored body")) + if err != ErrBodyNotAllowed { + t.Errorf("on Copy, expected ErrBodyNotAllowed, got %v", err) } })) defer ts.Close() diff --git a/src/pkg/http/server.go b/src/pkg/http/server.go index 4063fad224cf07b513d4e59c2f5a138fd7d410d7..d4638f127c7238fffb455bc64b081237352cbf31 100644 --- a/src/pkg/http/server.go +++ b/src/pkg/http/server.go @@ -129,7 +129,7 @@ // Flush before checking r.chunking, as Flush will call // WriteHeader if it hasn't been called yet, and WriteHeader // is what sets r.chunking. r.Flush() - if !r.chunking { + if !r.chunking && r.bodyAllowed() { if rf, ok := r.conn.rwc.(io.ReaderFrom); ok { n, err = rf.ReadFrom(src) r.written += n @@ -335,6 +335,15 @@ w.header.Write(w.conn.buf) io.WriteString(w.conn.buf, "\r\n") } +// bodyAllowed returns true if a Write is allowed for this response type. +// It's illegal to call this before the header has been flushed. +func (w *response) bodyAllowed() bool { + if !w.wroteHeader { + panic("") + } + return w.status != StatusNotModified && w.req.Method != "HEAD" +} + func (w *response) Write(data []byte) (n int, err os.Error) { if w.conn.hijacked { log.Print("http: response.Write on hijacked connection") @@ -346,9 +355,7 @@ } if len(data) == 0 { return 0, nil } - - if w.status == StatusNotModified || w.req.Method == "HEAD" { - // Must not have body. + if !w.bodyAllowed() { return 0, ErrBodyNotAllowed }