net/http: don't send an automatic Content-Length on a 304 Not Modified

Also start of some test helper unification, long overdue.
I refrained from cleaning up the rest in this CL.

Fixes #6157

R=golang-dev, adg
CC=golang-dev
https://golang.org/cl/13030043

src/pkg/net/http/serve_test.go

@@ -122,6 +122,28 @@ func reqBytes(req string) []byte {
 	return []byte(strings.Replace(strings.TrimSpace(req), "\n", "\r\n", -1) + "\r\n\r\n")
 }
 
+type handlerTest struct {
+	handler Handler
+}
+
+func newHandlerTest(h Handler) handlerTest {
+	return handlerTest{h}
+}
+
+func (ht handlerTest) rawResponse(req string) string {
+	reqb := reqBytes(req)
+	var output bytes.Buffer
+	conn := &rwTestConn{
+		Reader: bytes.NewReader(reqb),
+		Writer: &output,
+		closec: make(chan bool, 1),
+	}
+	ln := &oneConnListener{conn: conn}
+	go Serve(ln, ht.handler)
+	<-conn.closec
+	return output.String()
+}
+
 func TestConsumingBodyOnNextConn(t *testing.T) {
 	conn := new(testConn)
 	for i := 0; i < 2; i++ {
@@ -1588,7 +1610,6 @@ func TestOptions(t *testing.T) {
 // ones, even if the handler modifies them (~erroneously) after the
 // first Write.
 func TestHeaderToWire(t *testing.T) {
-	req := reqBytes("GET / HTTP/1.1\nHost: golang.org")
 	tests := []struct {
 		name    string
 		handler func(ResponseWriter, *Request)
@@ -1751,17 +1772,10 @@ func TestHeaderToWire(t *testing.T) {
 		},
 	}
 	for _, tc := range tests {
-		var output bytes.Buffer
+		ht := newHandlerTest(HandlerFunc(tc.handler))
-		conn := &rwTestConn{
+		got := ht.rawResponse("GET / HTTP/1.1\nHost: golang.org")
-			Reader: bytes.NewReader(req),
+		if err := tc.check(got); err != nil {
-			Writer: &output,
+			t.Errorf("%s: %v\nGot response:\n%s", tc.name, err, got)
-			closec: make(chan bool, 1),
-		}
-		ln := &oneConnListener{conn: conn}
-		go Serve(ln, HandlerFunc(tc.handler))
-		<-conn.closec
-		if err := tc.check(output.String()); err != nil {
-			t.Errorf("%s: %v\nGot response:\n%s", tc.name, err, output.Bytes())
 		}
 	}
 }
@@ -1952,6 +1966,34 @@ func TestServerReaderFromOrder(t *testing.T) {
 	}
 }
 
+// Issue 6157
+func TestNoContentTypeOnNotModified(t *testing.T) {
+	ht := newHandlerTest(HandlerFunc(func(w ResponseWriter, r *Request) {
+		if r.URL.Path == "/header" {
+			w.Header().Set("Content-Length", "123")
+		}
+		w.WriteHeader(StatusNotModified)
+		if r.URL.Path == "/more" {
+			w.Write([]byte("stuff"))
+		}
+	}))
+	for _, req := range []string{
+		"GET / HTTP/1.0",
+		"GET /header HTTP/1.0",
+		"GET /more HTTP/1.0",
+		"GET / HTTP/1.1",
+		"GET /header HTTP/1.1",
+		"GET /more HTTP/1.1",
+	} {
+		got := ht.rawResponse(req)
+		if !strings.Contains(got, "304 Not Modified") {
+			t.Errorf("Non-304 Not Modified for %q: %s", req, got)
+		} else if strings.Contains(got, "Content-Length") {
+			t.Errorf("Got a Content-Length from %q: %s", req, got)
+		}
+	}
+}
+
 func BenchmarkClientServer(b *testing.B) {
 	b.ReportAllocs()
 	b.StopTimer()

src/pkg/net/http/server.go

@@ -737,7 +737,15 @@ func (cw *chunkWriter) writeHeader(p []byte) {
 	// response header and this is our first (and last) write, set
 	// it, even to zero. This helps HTTP/1.0 clients keep their
 	// "keep-alive" connections alive.
-	if w.handlerDone && header.get("Content-Length") == "" && (!isHEAD || len(p) > 0) {
+	// Exceptions: 304 responses never get Content-Length, and if
+	// it was a HEAD request, we don't know the difference between
+	// 0 actual bytes and 0 bytes because the handler noticed it
+	// was a HEAD request and chose not to write anything.  So for
+	// HEAD, the handler should either write the Content-Length or
+	// write non-zero bytes.  If it's actually 0 bytes and the
+	// handler never looked at the Request.Method, we just don't
+	// send a Content-Length header.
+	if w.handlerDone && w.status != StatusNotModified && 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)
 	}