net/http: panic on bogus use of CloseNotifier or Hijacker

Fixes #14001

Change-Id: I6f9bc3028345081758d8f537c3aaddb2e254e69e
Reviewed-on: https://go-review.googlesource.com/18708
Reviewed-by: Ian Lance Taylor <iant@golang.org>
Run-TryBot: Ian Lance Taylor <iant@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

1 files changed, 14 insertions, 3 deletions

diff --git a/src/net/http/server.go b/src/net/http/server.go
index 8d50ffc599..2ec106927b 100644
--- a/src/net/http/server.go
+++ b/src/net/http/server.go
@@ -347,7 +347,7 @@ type response struct {
 	// written.
 	trailers []string
 
-	handlerDone bool // set true when the handler exits
+	handlerDone atomicBool // set true when the handler exits
 
 	// Buffers for Date and Content-Length
 	dateBuf [len(TimeFormat)]byte
@@ -358,6 +358,11 @@ type response struct {
 	closeNotifyCh <-chan bool
 }
 
+type atomicBool int32
+
+func (b *atomicBool) isSet() bool { return atomic.LoadInt32((*int32)(b)) != 0 }
+func (b *atomicBool) setTrue()    { atomic.StoreInt32((*int32)(b), 1) }
+
 // declareTrailer is called for each Trailer header when the
 // response header is written. It notes that a header will need to be
 // written in the trailers at the end of the response.
@@ -911,7 +916,7 @@ func (cw *chunkWriter) writeHeader(p []byte) {
 	// send a Content-Length header.
 	// Further, we don't send an automatic Content-Length if they
 	// set a Transfer-Encoding, because they're generally incompatible.
-	if w.handlerDone && !trailers && !hasTE && bodyAllowedForStatus(w.status) && header.get("Content-Length") == "" && (!isHEAD || len(p) > 0) {
+	if w.handlerDone.isSet() && !trailers && !hasTE && bodyAllowedForStatus(w.status) && header.get("Content-Length") == "" && (!isHEAD || len(p) > 0) {
 		w.contentLength = int64(len(p))
 		setHeader.contentLength = strconv.AppendInt(cw.res.clenBuf[:0], int64(len(p)), 10)
 	}
@@ -1234,7 +1239,7 @@ func (w *response) write(lenData int, dataB []byte, dataS string) (n int, err er
 }
 
 func (w *response) finishRequest() {
-	w.handlerDone = true
+	w.handlerDone.setTrue()
 
 	if !w.wroteHeader {
 		w.WriteHeader(StatusOK)
@@ -1498,6 +1503,9 @@ func (w *response) sendExpectationFailed() {
 // Hijack implements the Hijacker.Hijack method. Our response is both a ResponseWriter
 // and a Hijacker.
 func (w *response) Hijack() (rwc net.Conn, buf *bufio.ReadWriter, err error) {
+	if w.handlerDone.isSet() {
+		panic("net/http: Hijack called after ServeHTTP finished")
+	}
 	if w.wroteHeader {
 		w.cw.flush()
 	}
@@ -1521,6 +1529,9 @@ func (w *response) Hijack() (rwc net.Conn, buf *bufio.ReadWriter, err error) {
 }
 
 func (w *response) CloseNotify() <-chan bool {
+	if w.handlerDone.isSet() {
+		panic("net/http: CloseNotify called after ServeHTTP finished")
+	}
 	c := w.conn
 	c.mu.Lock()
 	defer c.mu.Unlock()