Add HTTP timeouts for flamegraph web servers

3 files changed, 74 insertions, 4 deletions

diff --git a/internal/flamegraph/liveserver.go b/internal/flamegraph/liveserver.go
index 95063d0..801b841 100644
--- a/internal/flamegraph/liveserver.go
+++ b/internal/flamegraph/liveserver.go
@@ -8,6 +8,12 @@ import (
 	"time"
 )
 
+var liveServerTimeouts = serverTimeouts{
+	readTimeout:  10 * time.Second,
+	writeTimeout: 5 * time.Minute,
+	idleTimeout:  60 * time.Second,
+}
+
 // ServeLive starts the live flamegraph HTTP server and blocks until ctx is canceled.
 func ServeLive(ctx context.Context, lt *LiveTrie, interval time.Duration) error {
 	mux := http.NewServeMux()
@@ -15,7 +21,7 @@ func ServeLive(ctx context.Context, lt *LiveTrie, interval time.Duration) error
 	mux.HandleFunc("/events", handleSSE(lt, interval))
 	mux.HandleFunc("/reset", handleReset(lt))
 	mux.HandleFunc("/order", handleOrder(lt))
-	return runServer(ctx, mux, func(hostname string, port int) {
+	return runServer(ctx, mux, liveServerTimeouts, func(hostname string, port int) {
 		fmt.Printf("Live flamegraph available at http://%s:%d/\n", hostname, port)
 	})
 }
diff --git a/internal/flamegraph/webserver.go b/internal/flamegraph/webserver.go
index 49eeaba..2bf6286 100644
--- a/internal/flamegraph/webserver.go
+++ b/internal/flamegraph/webserver.go
@@ -13,6 +13,18 @@ import (
 	"time"
 )
 
+type serverTimeouts struct {
+	readTimeout  time.Duration
+	writeTimeout time.Duration
+	idleTimeout  time.Duration
+}
+
+var defaultServerTimeouts = serverTimeouts{
+	readTimeout:  10 * time.Second,
+	writeTimeout: 30 * time.Second,
+	idleTimeout:  60 * time.Second,
+}
+
 // ServeSVG starts a small HTTP server that serves a single flamegraph SVG.
 //
 // It prints a URL of the form http://HOSTNAME:PORT/abs/path/to.svg and blocks until
@@ -26,7 +38,7 @@ func ServeSVG(svgFile string) error {
 	urlPath := buildURLPath(absPath)
 	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
 	defer stop()
-	return runServer(ctx, buildSVGHandler(absPath, urlPath), func(hostname string, port int) {
+	return runServer(ctx, buildSVGHandler(absPath, urlPath), defaultServerTimeouts, func(hostname string, port int) {
 		printServerURL(hostname, port, urlPath)
 	})
 }
@@ -72,8 +84,17 @@ func printServerURL(hostname string, port int, urlPath string) {
 	fmt.Println("Press Ctrl+C to stop the web server.")
 }
 
-func runServer(ctx context.Context, mux *http.ServeMux, printURL func(hostname string, port int)) error {
-	srv := &http.Server{Handler: mux}
+func newHTTPServer(mux *http.ServeMux, timeouts serverTimeouts) *http.Server {
+	return &http.Server{
+		Handler:      mux,
+		ReadTimeout:  timeouts.readTimeout,
+		WriteTimeout: timeouts.writeTimeout,
+		IdleTimeout:  timeouts.idleTimeout,
+	}
+}
+
+func runServer(ctx context.Context, mux *http.ServeMux, timeouts serverTimeouts, printURL func(hostname string, port int)) error {
+	srv := newHTTPServer(mux, timeouts)
 
 	listener, err := listenRandomPort()
 	if err != nil {
diff --git a/internal/flamegraph/webserver_timeout_test.go b/internal/flamegraph/webserver_timeout_test.go
new file mode 100644
index 0000000..c1df7e5
--- /dev/null
+++ b/internal/flamegraph/webserver_timeout_test.go
@@ -0,0 +1,43 @@
+package flamegraph
+
+import (
+	"net/http"
+	"testing"
+	"time"
+)
+
+func TestNewHTTPServerUsesConfiguredTimeouts(t *testing.T) {
+	mux := http.NewServeMux()
+	timeouts := serverTimeouts{
+		readTimeout:  11 * time.Second,
+		writeTimeout: 44 * time.Second,
+		idleTimeout:  66 * time.Second,
+	}
+
+	srv := newHTTPServer(mux, timeouts)
+
+	if srv.Handler != mux {
+		t.Fatalf("Handler not set from mux")
+	}
+	if srv.ReadTimeout != timeouts.readTimeout {
+		t.Fatalf("ReadTimeout = %v, want %v", srv.ReadTimeout, timeouts.readTimeout)
+	}
+	if srv.WriteTimeout != timeouts.writeTimeout {
+		t.Fatalf("WriteTimeout = %v, want %v", srv.WriteTimeout, timeouts.writeTimeout)
+	}
+	if srv.IdleTimeout != timeouts.idleTimeout {
+		t.Fatalf("IdleTimeout = %v, want %v", srv.IdleTimeout, timeouts.idleTimeout)
+	}
+}
+
+func TestLiveServerWriteTimeoutIsLongerThanDefault(t *testing.T) {
+	if liveServerTimeouts.readTimeout != defaultServerTimeouts.readTimeout {
+		t.Fatalf("read timeout mismatch: live=%v default=%v", liveServerTimeouts.readTimeout, defaultServerTimeouts.readTimeout)
+	}
+	if liveServerTimeouts.idleTimeout != defaultServerTimeouts.idleTimeout {
+		t.Fatalf("idle timeout mismatch: live=%v default=%v", liveServerTimeouts.idleTimeout, defaultServerTimeouts.idleTimeout)
+	}
+	if liveServerTimeouts.writeTimeout <= defaultServerTimeouts.writeTimeout {
+		t.Fatalf("expected live write timeout > default write timeout, got live=%v default=%v", liveServerTimeouts.writeTimeout, defaultServerTimeouts.writeTimeout)
+	}
+}