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>

diff --git a/src/cmd/pprof/internal/fetch/fetch.go b/src/cmd/pprof/internal/fetch/fetch.go
index 3ed16bb50d34e6bf06154b819442509330d5112d..45e02f2cd744c15b3460a71be887349eff6df040 100644 (file)
--- a/src/cmd/pprof/internal/fetch/fetch.go
+++ b/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) {

diff --git a/src/net/http/pprof/pprof.go b/src/net/http/pprof/pprof.go
index 126e9eaaa7041f84617a5212c5f7216f327c559c..6930df531b42ce55e67d30baf7e8fe6bae3fa4f2 100644 (file)
--- a/src/net/http/pprof/pprof.go
+++ b/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,6 +137,14 @@ 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")
@@ -130,6 +152,7 @@ func Trace(w http.ResponseWriter, r *http.Request) {
                // 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