From e38fa9164894b2610a41b335a56e3b6494a3cef6 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@golang.org>
Date: Tue, 19 Aug 2014 18:45:05 -0700
Subject: [PATCH] net/http: fix TimeoutHandler data races; hold lock longer

The existing lock needed to be held longer. If a timeout occured
while writing (but after the guarded timeout check), the writes
would clobber a future connection's buffer.

Also remove a harmless warning by making Write also set the
flag that headers were sent (implicitly), so we don't try to
write headers later (a no-op + warning) on timeout after we've
started writing.

Fixes #8414
Fixes #8209

LGTM=ruiu, adg
R=adg, ruiu
CC=golang-codereviews
https://golang.org/cl/123610043
---
 src/pkg/net/http/serve_test.go | 77 ++++++++++++++++++++++++++++++++++
 src/pkg/net/http/server.go     |  9 ++--
 2 files changed, 81 insertions(+), 5 deletions(-)

diff --git a/src/pkg/net/http/serve_test.go b/src/pkg/net/http/serve_test.go
index 2a3fc307be8..ee4f2049953 100644
--- a/src/pkg/net/http/serve_test.go
+++ b/src/pkg/net/http/serve_test.go
@@ -15,6 +15,7 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
+	"math/rand"
 	"net"
 	. "net/http"
 	"net/http/httptest"
@@ -1188,6 +1189,82 @@ func TestTimeoutHandler(t *testing.T) {
 	}
 }
 
+// See issues 8209 and 8414.
+func TestTimeoutHandlerRace(t *testing.T) {
+	defer afterTest(t)
+
+	delayHi := HandlerFunc(func(w ResponseWriter, r *Request) {
+		ms, _ := strconv.Atoi(r.URL.Path[1:])
+		if ms == 0 {
+			ms = 1
+		}
+		for i := 0; i < ms; i++ {
+			w.Write([]byte("hi"))
+			time.Sleep(time.Millisecond)
+		}
+	})
+
+	ts := httptest.NewServer(TimeoutHandler(delayHi, 20*time.Millisecond, ""))
+	defer ts.Close()
+
+	var wg sync.WaitGroup
+	gate := make(chan bool, 10)
+	n := 50
+	if testing.Short() {
+		n = 10
+		gate = make(chan bool, 3)
+	}
+	for i := 0; i < n; i++ {
+		gate <- true
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			defer func() { <-gate }()
+			res, err := Get(fmt.Sprintf("%s/%d", ts.URL, rand.Intn(50)))
+			if err == nil {
+				io.Copy(ioutil.Discard, res.Body)
+				res.Body.Close()
+			}
+		}()
+	}
+	wg.Wait()
+}
+
+// See issues 8209 and 8414.
+func TestTimeoutHandlerRaceHeader(t *testing.T) {
+	defer afterTest(t)
+
+	delay204 := HandlerFunc(func(w ResponseWriter, r *Request) {
+		w.WriteHeader(204)
+	})
+
+	ts := httptest.NewServer(TimeoutHandler(delay204, time.Nanosecond, ""))
+	defer ts.Close()
+
+	var wg sync.WaitGroup
+	gate := make(chan bool, 50)
+	n := 500
+	if testing.Short() {
+		n = 10
+	}
+	for i := 0; i < n; i++ {
+		gate <- true
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			defer func() { <-gate }()
+			res, err := Get(ts.URL)
+			if err != nil {
+				t.Error(err)
+				return
+			}
+			defer res.Body.Close()
+			io.Copy(ioutil.Discard, res.Body)
+		}()
+	}
+	wg.Wait()
+}
+
 // Verifies we don't path.Clean() on the wrong parts in redirects.
 func TestRedirectMunging(t *testing.T) {
 	req, _ := NewRequest("GET", "http://example.com/", nil)
diff --git a/src/pkg/net/http/server.go b/src/pkg/net/http/server.go
index eae097eb8e9..203037e9f5e 100644
--- a/src/pkg/net/http/server.go
+++ b/src/pkg/net/http/server.go
@@ -1916,9 +1916,9 @@ func (tw *timeoutWriter) Header() Header {
 
 func (tw *timeoutWriter) Write(p []byte) (int, error) {
 	tw.mu.Lock()
-	timedOut := tw.timedOut
-	tw.mu.Unlock()
-	if timedOut {
+	defer tw.mu.Unlock()
+	tw.wroteHeader = true // implicitly at least
+	if tw.timedOut {
 		return 0, ErrHandlerTimeout
 	}
 	return tw.w.Write(p)
@@ -1926,12 +1926,11 @@ func (tw *timeoutWriter) Write(p []byte) (int, error) {
 
 func (tw *timeoutWriter) WriteHeader(code int) {
 	tw.mu.Lock()
+	defer tw.mu.Unlock()
 	if tw.timedOut || tw.wroteHeader {
-		tw.mu.Unlock()
 		return
 	}
 	tw.wroteHeader = true
-	tw.mu.Unlock()
 	tw.w.WriteHeader(code)
 }