net/http/pprof: return error when requested profile duration exceeds …

…WriteTimeout

Updates Profile and Trace handlers to reject requests for durations >=
WriteTimeout.

Modifies go tool pprof to print the body of the http response when
status != 200.

Fixes #18755

Change-Id: I6faed21685693caf39f315f003039538114937b0
Reviewed-on: https://go-review.googlesource.com/35564
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

src/cmd/pprof/internal/fetch/fetch.go

@@ -52,7 +52,8 @@ func FetchURL(source string, timeout time.Duration) (io.ReadCloser, error) {
 		return nil, fmt.Errorf("http fetch: %v", err)
 	}
 	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("server response: %s", resp.Status)
+		defer resp.Body.Close()
+		return nil, statusCodeError(resp)
 	}
 
 	return resp.Body, nil
@@ -64,13 +65,24 @@ func PostURL(source, post string) ([]byte, error) {
 	if err != nil {
 		return nil, fmt.Errorf("http post %s: %v", source, err)
 	}
+	defer resp.Body.Close()
 	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("server response: %s", resp.Status)
+		return nil, statusCodeError(resp)
 	}
-	defer resp.Body.Close()
 	return ioutil.ReadAll(resp.Body)
 }
 
+func statusCodeError(resp *http.Response) error {
+	if resp.Header.Get("X-Go-Pprof") != "" && strings.Contains(resp.Header.Get("Content-Type"), "text/plain") {
+		// error is from pprof endpoint
+		body, err := ioutil.ReadAll(resp.Body)
+		if err == nil {
+			return fmt.Errorf("server response: %s - %s", resp.Status, body)
+		}
+	}
+	return fmt.Errorf("server response: %s", resp.Status)
+}
+
 // httpGet is a wrapper around http.Get; it is defined as a variable
 // so it can be redefined during for testing.
 var httpGet = func(source string, timeout time.Duration) (*http.Response, error) {

src/net/http/pprof/pprof.go

@@ -90,6 +90,11 @@ func sleep(w http.ResponseWriter, d time.Duration) {
 	}
 }
 
+func durationExceedsWriteTimeout(r *http.Request, seconds float64) bool {
+	srv, ok := r.Context().Value(http.ServerContextKey).(*http.Server)
+	return ok && srv.WriteTimeout != 0 && seconds >= srv.WriteTimeout.Seconds()
+}
+
 // Profile responds with the pprof-formatted cpu profile.
 // The package initialization registers it as /debug/pprof/profile.
 func Profile(w http.ResponseWriter, r *http.Request) {
@@ -98,6 +103,14 @@ func Profile(w http.ResponseWriter, r *http.Request) {
 		sec = 30
 	}
 
+	if durationExceedsWriteTimeout(r, float64(sec)) {
+		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+		w.Header().Set("X-Go-Pprof", "1")
+		w.WriteHeader(http.StatusBadRequest)
+		fmt.Fprintln(w, "profile duration exceeds server's WriteTimeout")
+		return
+	}
+
 	// Set Content Type assuming StartCPUProfile will work,
 	// because if it does it starts writing.
 	w.Header().Set("Content-Type", "application/octet-stream")
@@ -106,6 +119,7 @@ func Profile(w http.ResponseWriter, r *http.Request) {
 		// Can change header back to text content
 		// and send error code.
 		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+		w.Header().Set("X-Go-Pprof", "1")
 		w.WriteHeader(http.StatusInternalServerError)
 		fmt.Fprintf(w, "Could not enable CPU profiling: %s\n", err)
 		return
@@ -123,13 +137,22 @@ func Trace(w http.ResponseWriter, r *http.Request) {
 		sec = 1
 	}
 
+	if durationExceedsWriteTimeout(r, sec) {
+		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+		w.Header().Set("X-Go-Pprof", "1")
+		w.WriteHeader(http.StatusBadRequest)
+		fmt.Fprintln(w, "profile duration exceeds server's WriteTimeout")
+		return
+	}
+
 	// Set Content Type assuming trace.Start will work,
 	// because if it does it starts writing.
 	w.Header().Set("Content-Type", "application/octet-stream")
 	if err := trace.Start(w); err != nil {
 		// trace.Start failed, so no writes yet.
 		// Can change header back to text content and send error code.
 		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+		w.Header().Set("X-Go-Pprof", "1")
 		w.WriteHeader(http.StatusInternalServerError)
 		fmt.Fprintf(w, "Could not enable tracing: %s\n", err)
 		return